How to fine-tune an embedding model with Sentence Transformers

Embedding models inherit the vocabulary and similarity judgments of their pretraining data, so phrases that mean the same thing inside one support product can still land too far apart. A short fine-tuning run can shift those domain-specific relationships without replacing the base architecture.

The SentenceTransformerTrainer class accepts a datasets dataset, training arguments, and a loss matched to the row format. CoSENTLoss fits sentence pairs with numeric similarity scores and supplies a stronger ranking signal than the older CosineSimilarityLoss path for the same data shape.

Eight labeled pairs and eight CPU update steps keep the sample small enough to inspect while exercising the whole model lifecycle. A production run needs a larger reviewed corpus and held-out evaluation; this run instead confirms that embeddings changed, the saved model reloads without changing its output, and related support wording outranks unrelated billing text.

1. Install Sentence Transformers with its training dependencies in the active Python environment.

   $ python -m pip install --upgrade "sentence-transformers[train]"

   The train extra installs datasets and the trainer dependencies used by SentenceTransformerTrainer.
   Related: How to install Sentence Transformers with pip

2. Create train_support_embedding.py with the imports, runtime settings, and output paths.

   train_support_embedding.py

   import os
   from pathlib import Path
    
   import numpy as np
   from datasets import Dataset
   from sentence_transformers import (
       SentenceTransformer,
       SentenceTransformerTrainer,
       SentenceTransformerTrainingArguments,
   )
   from sentence_transformers.sentence_transformer.losses import CoSENTLoss
    
   os.environ["TOKENIZERS_PARALLELISM"] = "false"
   os.environ["WANDB_DISABLED"] = "true"
    
   model_dir = Path("models/support-embedding/final")
   training_dir = Path("training-output/support-embedding")

3. Add the labeled support pairs below the path variables in train_support_embedding.py.

   train_dataset = Dataset.from_dict(
       {
           "sentence1": [
               "reset a forgotten password",
               "reset a forgotten password",
               "replace an expired API token",
               "replace an expired API token",
               "restore a deleted project",
               "restore a deleted project",
               "export audit logs",
               "export audit logs",
           ],
           "sentence2": [
               "send a password recovery email",
               "download the latest invoice",
               "create a replacement access token",
               "change the billing contact",
               "recover a project from deleted items",
               "invite a new support agent",
               "download the audit log CSV file",
               "enable two-factor authentication",
           ],
           "score": [0.95, 0.05, 0.94, 0.08, 0.93, 0.04, 0.92, 0.06],
       }
   )

   CoSENTLoss compares pair relationships during each update, so the sample mixes high- and low-score pairs. Production training requires reviewed domain labels instead of these sample rows.

4. Add the pretraining probe block below the dataset block.

   probe_texts = [
       "replace an expired API token",
       "create a replacement access token",
       "download the latest invoice",
   ]
   model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
   before = model.encode(probe_texts, normalize_embeddings=True)
   loss = CoSENTLoss(model)

   The pre-training embeddings make the later weight-change check independent of the trainer's status message.

5. Add the training block below the loss definition.

   args = SentenceTransformerTrainingArguments(
       output_dir=str(training_dir),
       max_steps=8,
       per_device_train_batch_size=4,
       learning_rate=2e-5,
       warmup_steps=0,
       save_strategy="no",
       logging_strategy="no",
       disable_tqdm=True,
       report_to=[],
       use_cpu=True,
       seed=7,
   )
   trainer = SentenceTransformerTrainer(
       model=model,
       args=args,
       train_dataset=train_dataset,
       loss=loss,
   )
    
   train_result = trainer.train()
   after = model.encode(probe_texts, normalize_embeddings=True)
   weight_change = float(np.max(np.abs(after - before)))
   model.save_pretrained(model_dir)

   On an environment with a supported accelerator, omitting use_cpu=True lets the trainer select it. Longer training should follow only after held-out evaluation shows that the model is not overfitting.

6. Append the saved-model verification block to train_support_embedding.py.

   reloaded = SentenceTransformer(str(model_dir))
   reloaded_embeddings = reloaded.encode(probe_texts, normalize_embeddings=True)
   reload_delta = float(np.max(np.abs(reloaded_embeddings - after)))
   scores = reloaded_embeddings[:1] @ reloaded_embeddings[1:].T
    
   print(f"Training rows: {len(train_dataset)}")
   print(f"Loss: {loss.__class__.__name__}")
   print(f"Training steps: {train_result.global_step}")
   print(f"Maximum embedding change: {weight_change:.6f}")
   print(f"Saved model: {model_dir}")
   print(f"Reload delta: {reload_delta:.8f}")
   print(f"Related score: {scores[0, 0]:.4f}")
   print(f"Unrelated score: {scores[0, 1]:.4f}")
    
   if train_result.global_step != 8:
       raise SystemExit("trainer did not complete eight update steps")
   if weight_change <= 0:
       raise SystemExit("model embeddings did not change during training")
   if reload_delta >= 1e-6:
       raise SystemExit("reloaded model does not match the trained model")
   if scores[0, 0] <= scores[0, 1]:
       raise SystemExit("related support text did not score higher")

   The saved directory remains at models/support-embedding/final for held-out evaluation or application loading.
   Related: How to save and reload a Sentence Transformers model

7. Run train_support_embedding.py with its saved-model checks.

   $ python train_support_embedding.py
   ##### snipped #####
   Training rows: 8
   Loss: CoSENTLoss
   Training steps: 8
   Maximum embedding change: 0.023853
   Saved model: models/support-embedding/final
   Reload delta: 0.00000000
   Related score: 0.7047
   Unrelated score: 0.1679