第 10 章 · 构建文本嵌入模型

本章目标：学习训练与微调嵌入模型的方法——从 Softmax Loss 到 Cosine Similarity 与 MultipleNegativesRankingLoss，再到监督微调、Augmented SBERT 与无监督 TSDAE。

本章对应《Hands-On Large Language Models》第 10 章官方笔记本。运行示例需要 GPU：在 Google Colab 中选择 Runtime > Change runtime type > Hardware accelerator > GPU > GPU type > T4。如需安装依赖，取消注释并运行以下代码块：

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）

加载 GLUE 基准中的 MNLI 数据集，取前 5 万条作训练集。标签含义：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）

以 bert-base-uncased 作为底座模型：

python

from sentence_transformers import SentenceTransformer

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

10.1.3 损失函数（Loss Function）

定义 Softmax Loss——它需要显式指定标签数量：

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）

为 STS-B 验证集创建一个嵌入相似度评估器：

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）

定义训练参数并使用 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

用 MTEB（Massive Text Embedding Benchmark）在 Banking77 分类任务上评估刚训练好的模型：

python

from mteb import MTEB

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

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

⚠️ 显存清理（VRAM Clean-up）——你需要运行下面的代码来部分清空显存（GPU RAM）。如果不起作用，建议直接重启笔记本。（使用 Google Colab 时）可以在右侧资源面板确认已用显存确实降低了。也可以运行 !nvidia-smi 查看当前占用。

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）

⚠️ 显存清理：进入下一个训练示例前，先 Restart 笔记本以释放内存。

10.3.1 余弦相似度损失（Cosine Similarity Loss）

把 MNLI 三分类标签转换成「相似/不相似」的二值对：(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"]]
})

创建 STS-B 评估器：

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"
)

使用 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)

⚠️ 显存清理：进入下一个训练示例前，先 Restart 笔记本以释放内存。

python

import gc
import torch

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

10.3.2 多负例排序损失（Multiple Negatives Ranking Loss）

构造 (anchor, positive, negative) 三元组：positive 取蕴含关系的 hypothesis，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)

创建 STS-B 评估器：

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"
)

使用 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）

⚠️ 显存清理：进入下一个训练示例前，先 Restart 笔记本以释放内存。

python

import gc
import torch

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

10.4.1 监督学习（Supervised）

这次不再从 BERT 底座开始，而是对一个现成的预训练嵌入模型 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)

同时评估未微调的原始模型以便对比：

python

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

⚠️ 显存清理：进入下一个训练示例前，先 Restart 笔记本以释放内存。

python

import gc
import torch

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

10.4.2 Augmented SBERT

Step 1: 微调一个交叉编码器（cross-encoder）

准备 1 万条 MNLI 样本作为 gold 数据集并用 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']]
    }
)

在 gold 数据集上训练交叉编码器：

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: 创建新的句对

从第 1 万到第 5 万条样本构造 silver 待标注集：

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: 用微调后的交叉编码器给新句对打标签（silver 数据集）

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: 在扩充后的数据集（gold + silver）上训练双编码器（SBERT）

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)

创建 STS-B 评估器：

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: 不使用 silver 数据集再评估一次

只用 gold 数据集重复同样的训练流程：

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)

与 gold + silver 数据集的结果相比，只用 gold 数据集会降低模型的性能！

⚠️ 显存清理：进入下一个训练示例前，先 Restart 笔记本以释放内存。

python

import gc
import torch

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

10.4.3 无监督学习（Unsupervised Learning）

TSDAE（基于 Transformer 的去噪自编码器）

先下载 nltk 分词器：

python

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

把 MNLI 的 premise 与 hypothesis 拍平成句子列表，用 DenoisingAutoEncoderDataset 自动加噪：

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]

也可以自定义删除比例：

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)
# )

创建 STS-B 评估器：

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"
)

手动组装「Transformer + CLS 池化」的嵌入模型：

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])

使用去噪自编码损失：

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")

定义训练参数并训练：

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 本章小结

训练嵌入模型的四要素：底座模型（如 bert-base-uncased）、配对数据集（MNLI）、损失函数、STS-B 相似度评估器；
Softmax Loss 需要显式指定标签数；Cosine Similarity Loss 把任务转化为「相似/不相似」二值回归；
MultipleNegativesRankingLoss 使用 (anchor, positive, negative) 三元组并可加入随机软负例，是当前更常用的选择；
MTEB 是大规模文本嵌入基准，可用一行代码在 Banking77 等任务上评估嵌入模型；
Augmented SBERT 五步法：训交叉编码器 → 造 silver 句对 → 打标签 → gold+silver 合并训双编码器 → 对比仅 gold 的退化结果；
TSDAE 通过对输入句子加噪再用去噪自编码损失训练，实现无需标注数据的无监督嵌入学习。

🧪 随堂测验

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

1. SoftmaxLoss 与 CosineSimilarityLoss 在数据要求上的关键区别是？

2. MultipleNegativesRankingLoss 一节中「软负例」是如何构造的？

3. Augmented SBERT 中 silver 数据集指的是？

4. 关于 TSDAE，下列说法正确的是？

🛠️ 动手实践

把 10.1 节的训练数据量从 50 000 条分别改为 5 000 条与 100 000 条，保持其他参数不变，比较三者在 STS-B 评估器上的得分差异。
参考 10.3.2 节的软负例做法，把软负例数量从每条 1 个扩展到通过两次 shuffle 生成的两组负例，观察 MultipleNegativesRankingLoss 下模型得分是否变化。
在 10.4.2 节 Augmented SBERT 流程中，把交叉编码器的底座换成 all-MiniLM-L6-v2 并重跑 Step 1–5，对比 silver 标签质量与最终双编码器性能的变化。

第 10 章 · 构建文本嵌入模型 ​

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） ​

10.4.2 Augmented SBERT ​

10.4.3 无监督学习（Unsupervised Learning） ​

TSDAE（基于 Transformer 的去噪自编码器） ​

10.5 本章小结 ​

🧪 随堂测验

🛠️ 动手实践 ​

第 10 章 · 构建文本嵌入模型

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）

10.4.2 Augmented SBERT

10.4.3 无监督学习（Unsupervised Learning）

TSDAE（基于 Transformer 的去噪自编码器）

10.5 本章小结

🛠️ 动手实践