Chapter 10 · Building Text Embedding Models

Chapter goals: Learn methods for training and fine-tuning embedding models — from Softmax Loss to Cosine Similarity and MultipleNegativesRankingLoss, to supervised fine-tuning, Augmented SBERT, and unsupervised TSDAE.

This chapter corresponds to the official notebook for Chapter 10 of Hands-On Large Language Models. Running the examples requires a GPU: in Google Colab select Runtime > Change runtime type > Hardware accelerator > GPU > GPU type > T4. To install dependencies, uncomment and run the code block below:

python

# %%capture
# !pip install -q accelerate>=0.27.2 peft>=0.9.0 bitsandbytes>=0.43.0 transformers>=4.38.2 trl>=0.7.11 sentencepiece>=0.1.99
# !pip install -q sentence-transformers>=3.0.0 mteb>=1.1.2 datasets>=2.18.0

10.1 Creating an Embedding Model

10.1.1 Data

Load the MNLI dataset from the GLUE benchmark, taking the first 50,000 rows as the training set. Label meanings: 0 = entailment, 1 = neutral, 2 = contradiction:

python

from datasets import load_dataset

# Load MNLI dataset from GLUE
# 0 = entailment, 1 = neutral, 2 = contradiction
train_dataset = load_dataset("glue", "mnli", split="train").select(range(50_000))
train_dataset = train_dataset.remove_columns("idx")

python

train_dataset[2]

10.1.2 Model

Use bert-base-uncased as the base model:

python

from sentence_transformers import SentenceTransformer

# Use a base model
embedding_model = SentenceTransformer('bert-base-uncased')

10.1.3 Loss Function

Define Softmax Loss — it requires explicitly specifying the number of labels:

python

from sentence_transformers import losses

# Define the loss function. In soft-max loss, we will also need to explicitly set the number of labels.
train_loss = losses.SoftmaxLoss(
    model=embedding_model,
    sentence_embedding_dimension=embedding_model.get_sentence_embedding_dimension(),
    num_labels=3
)

10.1.4 Evaluation

Create an embedding similarity evaluator for the STS-B validation set:

python

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine",
)

10.1.5 Training

Define training arguments and start training with SentenceTransformerTrainer:

python

from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="base_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

python

from sentence_transformers.trainer import SentenceTransformerTrainer

# Train embedding model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()

python

# Evaluate our trained model
evaluator(embedding_model)

10.2 MTEB

Evaluate the freshly trained model on the Banking77 classification task using MTEB (Massive Text Embedding Benchmark):

python

from mteb import MTEB

# Choose evaluation task
evaluation = MTEB(tasks=["Banking77Classification"])

# Calculate results
results = evaluation.run(embedding_model)
results

⚠️ VRAM Clean-up — you need to run the code below to partially free GPU memory. If it doesn't work, it is recommended to restart the notebook directly. (When using Google Colab) you can confirm that used VRAM has indeed dropped in the right-side resource panel. You can also run !nvidia-smi to check current usage.

python

# # Empty and delete trainer/model
# trainer.accelerator.clear()
# del trainer, embedding_model

# # Garbage collection and empty cache
# import gc
# import torch

# gc.collect()
# torch.cuda.empty_cache()

python

import gc
import torch

gc.collect()
torch.cuda.empty_cache()

10.3 Loss Functions

⚠️ VRAM clean-up: before entering the next training example, Restart the notebook first to free memory.

10.3.1 Cosine Similarity Loss

Convert the MNLI three-class labels into binary similar/not-similar pairs: (neutral/contradiction) = 0, (entailment) = 1:

python

from datasets import Dataset, load_dataset

# Load MNLI dataset from GLUE
# 0 = entailment, 1 = neutral, 2 = contradiction
train_dataset = load_dataset("glue", "mnli", split="train").select(range(50_000))
train_dataset = train_dataset.remove_columns("idx")

# (neutral/contradiction)=0 and (entailment)=1
mapping = {2: 0, 1: 0, 0:1}
train_dataset = Dataset.from_dict({
    "sentence1": train_dataset["premise"],
    "sentence2": train_dataset["hypothesis"],
    "label": [float(mapping[label]) for label in train_dataset["label"]]
})

Create the STS-B evaluator:

python

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)

Train using CosineSimilarityLoss:

python

from sentence_transformers import losses, SentenceTransformer
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define model
embedding_model = SentenceTransformer('bert-base-uncased')

# Loss function
train_loss = losses.CosineSimilarityLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="cosineloss_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()

python

# Evaluate our trained model
evaluator(embedding_model)

⚠️ VRAM clean-up: before entering the next training example, Restart the notebook first to free memory.

python

import gc
import torch

gc.collect()
torch.cuda.empty_cache()

10.3.2 Multiple Negatives Ranking Loss

Construct (anchor, positive, negative) triplets: positive takes the entailment-relationship hypothesis; negative is randomly sampled from the dataset as a "soft negative":

python

import random
from tqdm import tqdm
from datasets import Dataset, load_dataset

# # Load MNLI dataset from GLUE
mnli = load_dataset("glue", "mnli", split="train").select(range(50_000))
mnli = mnli.remove_columns("idx")
mnli = mnli.filter(lambda x: True if x['label'] == 0 else False)

# Prepare data and add a soft negative
train_dataset = {"anchor": [], "positive": [], "negative": []}
soft_negatives = mnli["hypothesis"]
random.shuffle(soft_negatives)
for row, soft_negative in tqdm(zip(mnli, soft_negatives)):
    train_dataset["anchor"].append(row["premise"])
    train_dataset["positive"].append(row["hypothesis"])
    train_dataset["negative"].append(soft_negative)
train_dataset = Dataset.from_dict(train_dataset)
len(train_dataset)

Create the STS-B evaluator:

python

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)

Train using MultipleNegativesRankingLoss:

python

from sentence_transformers import losses, SentenceTransformer
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define model
embedding_model = SentenceTransformer('bert-base-uncased')

# Loss function
train_loss = losses.MultipleNegativesRankingLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="mnrloss_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()

python

# Evaluate our trained model
evaluator(embedding_model)

10.4 Fine-tuning

⚠️ VRAM clean-up: before entering the next training example, Restart the notebook first to free memory.

python

import gc
import torch

gc.collect()
torch.cuda.empty_cache()

10.4.1 Supervised Learning

This time instead of starting from a BERT base, fine-tune an existing pre-trained embedding model all-MiniLM-L6-v2:

python

from datasets import load_dataset
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Load MNLI dataset from GLUE
# 0 = entailment, 1 = neutral, 2 = contradiction
train_dataset = load_dataset("glue", "mnli", split="train").select(range(50_000))
train_dataset = train_dataset.remove_columns("idx")

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)

python

from sentence_transformers import losses, SentenceTransformer
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define model
embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

# Loss function
train_loss = losses.MultipleNegativesRankingLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="finetuned_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()

python

# Evaluate our trained model
evaluator(embedding_model)

Also evaluate the un-fine-tuned original model for comparison:

python

# Evaluate the pre-trained model
original_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
evaluator(original_model)

⚠️ VRAM clean-up: before entering the next training example, Restart the notebook first to free memory.

python

import gc
import torch

gc.collect()
torch.cuda.empty_cache()

10.4.2 Augmented SBERT

Step 1: Fine-tune a cross-encoder

Prepare 10,000 MNLI samples as the gold dataset and load them with NoDuplicatesDataLoader:

python

import pandas as pd
from tqdm import tqdm
from datasets import load_dataset, Dataset
from sentence_transformers import InputExample
from sentence_transformers.datasets import NoDuplicatesDataLoader

# Prepare a small set of 10000 documents for the cross-encoder
dataset = load_dataset("glue", "mnli", split="train").select(range(10_000))
mapping = {2: 0, 1: 0, 0:1}

# Data Loader
gold_examples = [
    InputExample(texts=[row["premise"], row["hypothesis"]], label=mapping[row["label"]])
    for row in tqdm(dataset)
]
gold_dataloader = NoDuplicatesDataLoader(gold_examples, batch_size=32)

# Pandas DataFrame for easier data handling
gold = pd.DataFrame(
    {
    'sentence1': dataset['premise'],
    'sentence2': dataset['hypothesis'],
    'label': [mapping[label] for label in dataset['label']]
    }
)

Train the cross-encoder on the gold dataset:

python

from sentence_transformers.cross_encoder import CrossEncoder

# Train a cross-encoder on the gold dataset
cross_encoder = CrossEncoder('bert-base-uncased', num_labels=2)
cross_encoder.fit(
    train_dataloader=gold_dataloader,
    epochs=1,
    show_progress_bar=True,
    warmup_steps=100,
    use_amp=False
)

Step 2: Create new sentence pairs

Construct the silver unlabeled set from samples 10,000 through 50,000:

python

# Prepare the silver dataset by predicting labels with the cross-encoder
silver = load_dataset("glue", "mnli", split="train").select(range(10_000, 50_000))
pairs = list(zip(silver['premise'], silver['hypothesis']))

Step 3: Label the new sentence pairs using the fine-tuned cross-encoder (silver dataset)

python

import numpy as np

# Label the sentence pairs using our fine-tuned cross-encoder
output = cross_encoder.predict(pairs, apply_softmax=True, show_progress_bar=True)
silver = pd.DataFrame(
    {
        "sentence1": silver["premise"],
        "sentence2": silver["hypothesis"],
        "label": np.argmax(output, axis=1)
    }
)

Step 4: Train the bi-encoder (SBERT) on the augmented dataset (gold + silver)

python

# Combine gold + silver
data = pd.concat([gold, silver], ignore_index=True, axis=0)
data = data.drop_duplicates(subset=['sentence1', 'sentence2'], keep="first")
train_dataset = Dataset.from_pandas(data, preserve_index=False)

Create the STS-B evaluator:

python

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)

python

from sentence_transformers import losses, SentenceTransformer
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define model
embedding_model = SentenceTransformer('bert-base-uncased')

# Loss function
train_loss = losses.CosineSimilarityLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="augmented_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()

python

# Evaluate our trained model
evaluator(embedding_model)

python

trainer.accelerator.clear()

Step 5: Rerun evaluation without the silver dataset

Repeat the same training流程 using only the gold dataset:

python

# Combine gold + silver
data = pd.concat([gold], ignore_index=True, axis=0)
data = data.drop_duplicates(subset=['sentence1', 'sentence2'], keep="first")
train_dataset = Dataset.from_pandas(data, preserve_index=False)

# Define model
embedding_model = SentenceTransformer('bert-base-uncased')

# Loss function
train_loss = losses.CosineSimilarityLoss(model=embedding_model)

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="gold_only_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()

python

# Evaluate our trained model
evaluator(embedding_model)

Compared with the gold + silver dataset result, using only the gold dataset degrades the model's performance!

⚠️ VRAM clean-up: before entering the next training example, Restart the notebook first to free memory.

python

import gc
import torch

gc.collect()
torch.cuda.empty_cache()

10.4.3 Unsupervised Learning

TSDAE (Transformer-based Denoising AutoEncoder)

First download the nltk tokenizer:

python

# Download additional tokenizer
import nltk
nltk.download('punkt')

Flatten the MNLI premises and hypotheses into a sentence list, and use DenoisingAutoEncoderDataset to automatically add noise:

python

from tqdm import tqdm
from datasets import Dataset, load_dataset
from sentence_transformers.datasets import DenoisingAutoEncoderDataset

# Create a flat list of sentences
mnli = load_dataset("glue", "mnli", split="train").select(range(25_000))
flat_sentences = mnli["premise"] + mnli["hypothesis"]

# Add noise to our input data
damaged_data = DenoisingAutoEncoderDataset(list(set(flat_sentences)))

# Create dataset
train_dataset = {"damaged_sentence": [], "original_sentence": []}
for data in tqdm(damaged_data):
    train_dataset["damaged_sentence"].append(data.texts[0])
    train_dataset["original_sentence"].append(data.texts[1])
train_dataset = Dataset.from_dict(train_dataset)

python

train_dataset[0]

You can also customize the deletion ratio:

python

# # Choose a different deletion ratio
# flat_sentences = list(set(flat_sentences))
# damaged_data = DenoisingAutoEncoderDataset(
#     flat_sentences,
#     noise_fn=lambda s: DenoisingAutoEncoderDataset.delete(s, del_ratio=0.6)
# )

Create the STS-B evaluator:

python

from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator

# Create an embedding similarity evaluator for stsb
val_sts = load_dataset('glue', 'stsb', split='validation')
evaluator = EmbeddingSimilarityEvaluator(
    sentences1=val_sts["sentence1"],
    sentences2=val_sts["sentence2"],
    scores=[score/5 for score in val_sts["label"]],
    main_similarity="cosine"
)

Manually assemble an embedding model of "Transformer + CLS pooling":

python

from sentence_transformers import models, SentenceTransformer

# Create your embedding model
word_embedding_model = models.Transformer('bert-base-uncased')
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), 'cls')
embedding_model = SentenceTransformer(modules=[word_embedding_model, pooling_model])

Use the denoising auto-encoder loss:

python

from sentence_transformers import losses

# Use the denoising auto-encoder loss
train_loss = losses.DenoisingAutoEncoderLoss(
    embedding_model, tie_encoder_decoder=True
)
train_loss.decoder = train_loss.decoder.to("cuda")

Define training parameters and train:

python

from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

# Define the training arguments
args = SentenceTransformerTrainingArguments(
    output_dir="tsdae_embedding_model",
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    warmup_steps=100,
    fp16=True,
    eval_steps=100,
    logging_steps=100,
)

# Train model
trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    train_dataset=train_dataset,
    loss=train_loss,
    evaluator=evaluator
)
trainer.train()

python

# Evaluate our trained model
evaluator(embedding_model)

python

import gc
import torch

gc.collect()
torch.cuda.empty_cache()

10.5 Chapter Summary

Four elements for training an embedding model: base model (e.g. bert-base-uncased), paired dataset (MNLI), loss function, and STS-B similarity evaluator;
Softmax Loss requires explicitly specifying the number of labels; Cosine Similarity Loss converts the task into a binary similar/not-similar regression;
MultipleNegativesRankingLoss uses (anchor, positive, negative) triplets and can incorporate random soft negatives — currently the more commonly used choice;
MTEB is a large-scale text embedding benchmark; embedding models can be evaluated on tasks like Banking77 in one line of code;
Augmented SBERT five-step method: train cross-encoder → create silver pairs → label them → train bi-encoder on gold+silver → compare against gold-only degraded results;
TSDAE learns unsupervised embeddings by adding noise to input sentences and using a denoising auto-encoder loss to reconstruct the original sentence, requiring no labels.

🧪 随堂测验

点击你认为正确的选项。答错时会展示正确答案与原因解析。

1. What is the key difference between SoftmaxLoss and CosineSimilarityLoss in terms of data requirements?

2. In the MultipleNegativesRankingLoss section, how is the "soft negative" constructed?

3. In Augmented SBERT, what does the silver dataset refer to?

4. Regarding TSDAE, which statement is correct?

🛠️ Hands-on Practice

Change the training data size in Section 10.1 from 50,000 rows to 5,000 and 100,000 respectively; keep all other parameters unchanged; compare the score differences across the three on the STS-B evaluator.
Following the soft-negative approach in Section 10.3.2, extend the number of soft negatives from 1 per row to two groups generated via two shuffles; observe whether the model score changes under MultipleNegativesRankingLoss.
In the Augmented SBERT flow from Section 10.4.2, swap the cross-encoder's base to all-MiniLM-L6-v2 and rerun Steps 1–5; compare changes in silver label quality and final bi-encoder performance.

Chapter 10 · Building Text Embedding Models ​

10.1 Creating an Embedding Model ​

10.1.1 Data ​

10.1.2 Model ​

10.1.3 Loss Function ​

10.1.4 Evaluation ​

10.1.5 Training ​

10.2 MTEB ​

10.3 Loss Functions ​

10.3.1 Cosine Similarity Loss ​

10.3.2 Multiple Negatives Ranking Loss ​

10.4 Fine-tuning ​

10.4.1 Supervised Learning ​

10.4.2 Augmented SBERT ​

10.4.3 Unsupervised Learning ​

TSDAE (Transformer-based Denoising AutoEncoder) ​

10.5 Chapter Summary ​

🧪 随堂测验

🛠️ Hands-on Practice ​

Chapter 10 · Building Text Embedding Models

10.1 Creating an Embedding Model

10.1.1 Data

10.1.2 Model

10.1.3 Loss Function

10.1.4 Evaluation

10.1.5 Training

10.2 MTEB

10.3 Loss Functions

10.3.1 Cosine Similarity Loss

10.3.2 Multiple Negatives Ranking Loss

10.4 Fine-tuning

10.4.1 Supervised Learning

10.4.2 Augmented SBERT

10.4.3 Unsupervised Learning

TSDAE (Transformer-based Denoising AutoEncoder)

10.5 Chapter Summary

🛠️ Hands-on Practice