使用 Sentence Transformers v3 训练和微调嵌入模型

Sentence Transformers 是一个用于使用和训练嵌入模型的 Python 库，适用于广泛的应用，例如检索增强生成、语义搜索、语义文本相似度、释义挖掘等。其 v3.0 更新是该项目自成立以来最大的一次更新，引入了一种新的训练方法。在这篇博文中，我将向您展示如何使用它来微调 Sentence Transformer 模型，以提高它们在特定任务上的性能。您也可以使用此方法从头开始训练新的 Sentence Transformer 模型。

现在，微调 Sentence Transformers 涉及多个组件，包括数据集、损失函数、训练参数、评估器以及新的训练器本身。我将详细介绍每个组件，并提供如何使用它们来训练有效模型的示例。

目录

• 为什么要微调？
• 训练组件
• 数据集
  • Hugging Face Hub 上的数据
  • 本地数据（CSV、JSON、Parquet、Arrow、SQL）
  • 需要预处理的本地数据
  • 数据集格式
• 损失函数
• 训练参数
• 评估器
  • 使用 STSb 的 EmbeddingSimilarityEvaluator
  • 使用 AllNLI 的 TripletEvaluator
• 训练器
  • 回调
• 多数据集训练
• 弃用
• 其他资源
  • 训练示例
  • 文档

为什么要微调？

微调 Sentence Transformer 模型可以显著提升它们在特定任务上的性能。这是因为每个任务都需要独特的相似性概念。我们以几个新闻标题为例：

• “苹果发布新款 iPad”
• “英伟达正为下一代 GPU 做准备”

根据用例，我们可能希望这些文本的嵌入相似或不相似。例如，新闻文章的分类模型可以将这些文本视为相似，因为它们都属于“科技”类别。另一方面，语义文本相似度或检索模型应将它们视为不相似，因为它们的含义不同。

训练组件

训练 Sentence Transformer 模型涉及以下组件：

1. 数据集：用于训练和评估的数据。
2. 损失函数：衡量模型性能并指导优化过程的函数。
3. 训练参数（可选）：影响训练性能和跟踪/调试的参数。
4. 评估器（可选）：用于在训练前、训练中或训练后评估模型的工具。
5. 训练器：将模型、数据集、损失函数和其他组件整合起来进行训练。

现在，让我们更详细地深入了解这些组件。

数据集

SentenceTransformerTrainer 使用 datasets.Dataset 或 datasets.DatasetDict 实例进行训练和评估。您可以从 Hugging Face Datasets Hub 加载数据，或使用各种格式的本地数据，例如 CSV、JSON、Parquet、Arrow 或 SQL。

注意：许多开箱即用的 Hugging Face 数据集已用 sentence-transformers 标记，您可以轻松地通过浏览 https://huggingface.co/datasets?other=sentence-transformers 找到它们。我们强烈建议您浏览这些数据集，以找到可能对您的任务有用的训练数据集。

Hugging Face Hub 上的数据

要从 Hugging Face Hub 中的数据集加载数据，请使用 load_dataset 函数

from datasets import load_dataset

train_dataset = load_dataset("sentence-transformers/all-nli", "pair-class", split="train")
eval_dataset = load_dataset("sentence-transformers/all-nli", "pair-class", split="dev")

print(train_dataset)
"""
Dataset({
    features: ['premise', 'hypothesis', 'label'],
    num_rows: 942069
})
"""

一些数据集，例如 sentence-transformers/all-nli，具有多个不同数据格式的子集。您需要指定子集名称以及数据集名称。

本地数据（CSV、JSON、Parquet、Arrow、SQL）

如果您有常见文件格式的本地数据，您也可以使用 load_dataset 轻松加载

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")
# or
dataset = load_dataset("json", data_files="my_file.json")

需要预处理的本地数据

如果您的本地数据需要预处理，您可以使用 datasets.Dataset.from_dict 使用列表字典初始化数据集

from datasets import Dataset

anchors = []
positives = []
# Open a file, perform preprocessing, filtering, cleaning, etc.
# and append to the lists

dataset = Dataset.from_dict({
    "anchor": anchors,
    "positive": positives,
})

字典中的每个键都将成为结果数据集中的一列。

数据集格式

确保数据集格式与您选择的损失函数匹配至关重要。这涉及检查两件事：

1. 如果您的损失函数需要一个 Label（如 损失概述 表中所示），您的数据集必须有一个名为 "label" 或 "score" 的列。
2. 除了 "label" 或 "score" 之外的所有列都被视为 Inputs（如 损失概述 表中所示）。这些列的数量必须与您所选损失函数的有效输入数量相匹配。列的名称不重要，只有它们的顺序重要。

例如，如果您的损失函数接受 (anchor, positive, negative) 三元组，那么您的第一个、第二个和第三个数据集列分别对应于 anchor、positive 和 negative。这意味着您的第一个和第二个列必须包含应该紧密嵌入的文本，而您的第一个和第三个列必须包含应该相距很远的文本。这就是为什么根据您的损失函数，数据集列的顺序很重要。

考虑一个包含 ["text1", "text2", "label"] 列的数据集，其中 "label" 列包含浮点相似度分数。此数据集可与 CoSENTLoss、AnglELoss 和 CosineSimilarityLoss 一起使用，因为：

1. 数据集有一个“label”列，这是这些损失函数所必需的。
2. 数据集有 2 个非标签列，与这些损失函数所需的输入数量相匹配。

如果数据集中的列顺序不正确，请使用 Dataset.select_columns 重新排序。此外，使用 Dataset.remove_columns 删除任何多余的列（例如，sample_id、metadata、source、type），否则它们将被视为输入。

损失函数

损失函数衡量模型在给定批次数据上的性能，并指导优化过程。损失函数的选择取决于您可用的数据和目标任务。有关选项的完整列表，请参阅 损失概述。

大多数损失函数都可以通过仅使用您正在训练的 SentenceTransformer model 来初始化

from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sentence_transformers.losses import CoSENTLoss

# Load a model to train/finetune
model = SentenceTransformer("FacebookAI/xlm-roberta-base")

# Initialize the CoSENTLoss
# This loss requires pairs of text and a floating point similarity score as a label
loss = CoSENTLoss(model)

# Load an example training dataset that works with our loss function:
train_dataset = load_dataset("sentence-transformers/all-nli", "pair-score", split="train")
"""
Dataset({
    features: ['sentence1', 'sentence2', 'label'],
    num_rows: 942069
})
"""

训练参数

SentenceTransformersTrainingArguments 类允许您指定影响训练性能和跟踪/调试的参数。虽然是可选的，但尝试这些参数可以帮助提高训练效率并提供对训练过程的洞察。

在 Sentence Transformers 文档中，我概述了一些最有用的训练参数。我建议您阅读 训练概述 > 训练参数。

以下是初始化 SentenceTransformersTrainingArguments 的示例：

from sentence_transformers.training_args import SentenceTransformerTrainingArguments

args = SentenceTransformerTrainingArguments(
    # Required parameter:
    output_dir="models/mpnet-base-all-nli-triplet",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    warmup_ratio=0.1,
    fp16=True,  # Set to False if your GPU can't handle FP16
    bf16=False,  # Set to True if your GPU supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # Losses using "in-batch negatives" benefit from no duplicates
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=2,
    logging_steps=100,
    run_name="mpnet-base-all-nli-triplet",  # Used in W&B if `wandb` is installed
)

请注意，eval_strategy 是在 transformers 版本 4.41.0 中引入的。之前的版本应使用 evaluation_strategy。

评估器

您可以向 SentenceTransformerTrainer 提供一个 eval_dataset 以在训练期间获取评估损失，但在训练期间获取更具体的指标也可能很有用。为此，您可以使用评估器在训练前、训练期间或训练后使用有用的指标评估模型的性能。您可以同时使用 eval_dataset 和评估器，或两者之一，或两者都不用。它们根据 eval_strategy 和 eval_steps 训练参数 进行评估。

以下是 Sentence Transformers 提供的已实现的评估器：

此外，您可以使用 SequentialEvaluator 将多个评估器组合成一个，然后将其传递给 SentenceTransformerTrainer。

如果您没有必要的评估数据，但仍想跟踪模型在常见基准上的性能，则可以使用 Hugging Face 的数据与这些评估器一起使用：

使用 STSb 的 EmbeddingSimilarityEvaluator

STS Benchmark (又称 STSb) 是一个常用的基准数据集，用于衡量模型对“一个人正在给蛇喂老鼠。”等短文本语义相似度的理解。

欢迎浏览 Hugging Face 上的 sentence-transformers/stsb 数据集。

from datasets import load_dataset
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator, SimilarityFunction

# Load the STSB dataset
eval_dataset = load_dataset("sentence-transformers/stsb", split="validation")

# Initialize the evaluator
dev_evaluator = EmbeddingSimilarityEvaluator(
    sentences1=eval_dataset["sentence1"],
    sentences2=eval_dataset["sentence2"],
    scores=eval_dataset["score"],
    main_similarity=SimilarityFunction.COSINE,
    name="sts-dev",
)
# Run evaluation manually:
# print(dev_evaluator(model))

# Later, you can provide this evaluator to the trainer to get results during training

使用 AllNLI 的 TripletEvaluator

AllNLI 是 SNLI 和 MultiNLI 数据集的串联，两者都是自然语言推理数据集。此任务传统上用于确定两个文本是蕴涵、矛盾还是两者都不是。此后，它已被用于训练嵌入模型，因为蕴涵和矛盾的句子可以构成有用的 (anchor, positive, negative) 三元组：这是一种训练嵌入模型的常见格式。

在这段代码中，它用于评估模型将锚文本和蕴涵文本视为比锚文本和矛盾文本更相似的频率。一个示例文本是“一个老人正在一家餐馆喝橙汁。”

欢迎浏览 Hugging Face 上的 sentence-transformers/all-nli 数据集。

from datasets import load_dataset
from sentence_transformers.evaluation import TripletEvaluator, SimilarityFunction

# Load triplets from the AllNLI dataset
max_samples = 1000
eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split=f"dev[:{max_samples}]")

# Initialize the evaluator
dev_evaluator = TripletEvaluator(
    anchors=eval_dataset["anchor"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
    main_distance_function=SimilarityFunction.COSINE,
    name=f"all-nli-{max_samples}-dev",
)
# Run evaluation manually:
# print(dev_evaluator(model))

# Later, you can provide this evaluator to the trainer to get results during training

训练器

SentenceTransformerTrainer 将模型、数据集、损失函数和其他组件整合在一起进行训练

from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    SentenceTransformerModelCardData,
)
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers.evaluation import TripletEvaluator

# 1. Load a model to finetune with 2. (Optional) model card data
model = SentenceTransformer(
    "microsoft/mpnet-base",
    model_card_data=SentenceTransformerModelCardData(
        language="en",
        license="apache-2.0",
        model_name="MPNet base trained on AllNLI triplets",
    )
)

# 3. Load a dataset to finetune on
dataset = load_dataset("sentence-transformers/all-nli", "triplet")
train_dataset = dataset["train"].select(range(100_000))
eval_dataset = dataset["dev"]
test_dataset = dataset["test"]

# 4. Define a loss function
loss = MultipleNegativesRankingLoss(model)

# 5. (Optional) Specify training arguments
args = SentenceTransformerTrainingArguments(
    # Required parameter:
    output_dir="models/mpnet-base-all-nli-triplet",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    warmup_ratio=0.1,
    fp16=True,  # Set to False if GPU can't handle FP16
    bf16=False,  # Set to True if GPU supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicates
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=2,
    logging_steps=100,
    run_name="mpnet-base-all-nli-triplet",  # Used in W&B if `wandb` is installed
)

# 6. (Optional) Create an evaluator & evaluate the base model
dev_evaluator = TripletEvaluator(
    anchors=eval_dataset["anchor"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
    name="all-nli-dev",
)
dev_evaluator(model)

# 7. Create a trainer & train
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=loss,
    evaluator=dev_evaluator,
)
trainer.train()

# (Optional) Evaluate the trained model on the test set, after training completes
test_evaluator = TripletEvaluator(
    anchors=test_dataset["anchor"],
    positives=test_dataset["positive"],
    negatives=test_dataset["negative"],
    name="all-nli-test",
)
test_evaluator(model)

# 8. Save the trained model
model.save_pretrained("models/mpnet-base-all-nli-triplet/final")

# 9. (Optional) Push it to the Hugging Face Hub
model.push_to_hub("mpnet-base-all-nli-triplet")

在这个例子中，我正在从 microsoft/mpnet-base 进行微调，这是一个尚未成为 Sentence Transformer 模型的基础模型。这需要比微调现有 Sentence Transformer 模型（如 all-mpnet-base-v2）更多的训练数据。

运行此脚本后，tomaarsen/mpnet-base-all-nli-triplet 模型已为我上传。使用余弦相似度的三元组准确率，即 cosine_similarity(anchor, positive) > cosine_similarity(anchor, negative) 的时间百分比，对于开发集为 90.04%，对于测试集为 91.5%！作为参考，microsoft/mpnet-base 模型在训练前在开发集上的得分仅为 68.32%。

所有这些信息都存储在自动生成的模型卡中，包括基础模型、语言、许可证、评估结果、训练和评估数据集信息、超参数、训练日志等。无需任何努力，您上传的模型应该包含您的潜在用户确定您的模型是否适合他们所需的所有信息。

回调

Sentence Transformers 训练器支持各种 transformers.TrainerCallback 子类，包括：

• 如果安装了wandb，则使用WandbCallback将训练指标记录到W&B。
• 如果可以访问tensorboard，则使用TensorBoardCallback将训练指标记录到TensorBoard。
• 如果安装了codecarbon，则使用CodeCarbonCallback跟踪训练期间的碳排放。

只要安装了所需的依赖项，这些功能就会自动使用，您无需进行任何指定。

有关这些回调以及如何创建自己的回调的更多信息，请参阅Transformers回调文档。

多数据集训练

表现最佳的模型通常使用多个数据集同时进行训练。SentenceTransformerTrainer 通过允许您使用多个数据集进行训练而无需将它们转换为相同格式来简化此过程。您甚至可以对每个数据集应用不同的损失函数。以下是多数据集训练的步骤：

1. 使用 datasets.Dataset 实例的字典（或 datasets.DatasetDict）作为 train_dataset 和 eval_dataset。
2. （可选）如果您想对不同的数据集使用不同的损失函数，请使用一个将数据集名称映射到损失的损失函数字典。

每个训练/评估批次将只包含来自一个数据集的样本。从多个数据集中采样批次的顺序由 MultiDatasetBatchSamplers 枚举确定，该枚举可以通过 multi_dataset_batch_sampler 传递给 SentenceTransformersTrainingArguments。有效选项包括：

• MultiDatasetBatchSamplers.ROUND_ROBIN：以轮循方式从每个数据集中采样，直到其中一个耗尽。此策略可能不会使用每个数据集中的所有样本，但它确保了从每个数据集中进行相同的采样。
• MultiDatasetBatchSamplers.PROPORTIONAL（默认）：按数据集大小比例从每个数据集中采样。此策略确保使用每个数据集中的所有样本，并且从较大的数据集中采样的频率更高。

多任务训练已被证明非常有效。例如，Huang 等人（2024）采用 MultipleNegativesRankingLoss、CoSENTLoss 以及不带批内负样本且只有硬负样本的 MultipleNegativesRankingLoss 变体，在中国实现了最先进的性能。他们还应用了 MatryoshkaLoss，使模型能够生成套娃嵌入。

以下是多数据集训练的示例：

from datasets import load_dataset
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
from sentence_transformers.losses import CoSENTLoss, MultipleNegativesRankingLoss, SoftmaxLoss

# 1. Load a model to finetune
model = SentenceTransformer("bert-base-uncased")

# 2. Loadseveral Datasets to train with
# (anchor, positive)
all_nli_pair_train = load_dataset("sentence-transformers/all-nli", "pair", split="train[:10000]")
# (premise, hypothesis) + label
all_nli_pair_class_train = load_dataset("sentence-transformers/all-nli", "pair-class", split="train[:10000]")
# (sentence1, sentence2) + score
all_nli_pair_score_train = load_dataset("sentence-transformers/all-nli", "pair-score", split="train[:10000]")
# (anchor, positive, negative)
all_nli_triplet_train = load_dataset("sentence-transformers/all-nli", "triplet", split="train[:10000]")
# (sentence1, sentence2) + score
stsb_pair_score_train = load_dataset("sentence-transformers/stsb", split="train[:10000]")
# (anchor, positive)
quora_pair_train = load_dataset("sentence-transformers/quora-duplicates", "pair", split="train[:10000]")
# (query, answer)
natural_questions_train = load_dataset("sentence-transformers/natural-questions", split="train[:10000]")

# Combine all datasets into a dictionary with dataset names to datasets
train_dataset = {
    "all-nli-pair": all_nli_pair_train,
    "all-nli-pair-class": all_nli_pair_class_train,
    "all-nli-pair-score": all_nli_pair_score_train,
    "all-nli-triplet": all_nli_triplet_train,
    "stsb": stsb_pair_score_train,
    "quora": quora_pair_train,
    "natural-questions": natural_questions_train,
}

# 3. Load several Datasets to evaluate with
# (anchor, positive, negative)
all_nli_triplet_dev = load_dataset("sentence-transformers/all-nli", "triplet", split="dev")
# (sentence1, sentence2, score)
stsb_pair_score_dev = load_dataset("sentence-transformers/stsb", split="validation")
# (anchor, positive)
quora_pair_dev = load_dataset("sentence-transformers/quora-duplicates", "pair", split="train[10000:11000]")
# (query, answer)
natural_questions_dev = load_dataset("sentence-transformers/natural-questions", split="train[10000:11000]")

# Use a dictionary for the evaluation dataset too, or just use one dataset or none at all
eval_dataset = {
    "all-nli-triplet": all_nli_triplet_dev,
    "stsb": stsb_pair_score_dev,
    "quora": quora_pair_dev,
    "natural-questions": natural_questions_dev,
}

# 4. Load several loss functions to train with
# (anchor, positive), (anchor, positive, negative)
mnrl_loss = MultipleNegativesRankingLoss(model)
# (sentence_A, sentence_B) + class
softmax_loss = SoftmaxLoss(model)
# (sentence_A, sentence_B) + score
cosent_loss = CoSENTLoss(model)

# Create a mapping with dataset names to loss functions, so the trainer knows which loss to apply where
# Note: You can also just use one loss if all your training/evaluation datasets use the same loss
losses = {
    "all-nli-pair": mnrl_loss,
    "all-nli-pair-class": softmax_loss,
    "all-nli-pair-score": cosent_loss,
    "all-nli-triplet": mnrl_loss,
    "stsb": cosent_loss,
    "quora": mnrl_loss,
    "natural-questions": mnrl_loss,
}

# 5. Define a simple trainer, although it's recommended to use one with args & evaluators
trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=losses,
)
trainer.train()

# 6. Save the trained model and optionally push it to the Hugging Face Hub
model.save_pretrained("bert-base-all-nli-stsb-quora-nq")
model.push_to_hub("bert-base-all-nli-stsb-quora-nq")

弃用

在 Sentence Transformer v3 发布之前，所有模型都将使用 SentenceTransformer.fit 方法进行训练。从 v3.0 开始，此方法将不再被弃用，而是将使用 SentenceTransformerTrainer 在后台运行。这意味着您的旧训练代码仍然可以运行，甚至应该升级以获得新功能，例如多 GPU 训练、损失日志记录等。尽管如此，新的训练方法功能更强大，因此建议使用新方法编写新的训练脚本。

更多资源

训练示例

以下页面包含训练示例以及解释和代码链接。我们建议您浏览这些内容，以熟悉训练循环：

• 语义文本相似度
• 自然语言推理
• 释义
• Quora 重复问题
• 套娃嵌入
• 自适应层模型
• 多语言模型
• 模型蒸馏
• 增强版 Sentence Transformers

文档

此外，以下页面可能有助于您了解更多关于 Sentence Transformers 的信息：

• 安装
• 快速入门
• 用法
• 预训练模型
• 训练概述（本博文是训练概述文档的精髓）
• 数据集概述
• 损失概述
• API参考

最后，这里有一些您可能感兴趣的高级页面：

• 超参数优化
• 分布式训练