如何用 Sentence Transformers 训练和微调嵌入模型

Image 1: Hugging Face's logo Hugging Face

模型
数据集
Spaces
Buckets 新
文档
企业版
定价
- 网站
  - 任务
  - HuggingChat
  - 合集
  - 语言
  - 组织
- 社区
  - 博客
  - 帖子
  - 每日论文
  - 硬件
  - 学习
  - Discord
  - 论坛
  - GitHub
- 解决方案

返回文章

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

发布于 2024年5月28日

在 GitHub 上更新

- [x] 点赞 279

+273

Tom Aarsen tomaarsen 关注

Sentence Transformers 是一个 Python 库，用于使用和训练嵌入模型，适用于检索增强生成、语义搜索、语义文本相似度、复述挖掘等广泛的应用场景。在这篇博客文章中，我将向你展示如何使用它来微调 Sentence Transformer 模型，以提升它们在特定任务上的表现。你也可以使用这种方法从头训练新的 Sentence Transformer 模型。微调 Sentence Transformers 涉及多个组件，包括数据集、损失函数、训练参数、评估器和训练器本身。我将逐一详细介绍这些组件，并提供如何使用它们来训练有效模型的示例。

为什么要微调？

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

"Apple 发布新款 iPad"
"NVIDIA 正在为下一代 GPU 做准备"

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

训练组件

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

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

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

数据集

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

注意：许多可直接与 Sentence Transformers 配合使用的 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,
})

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

数据集格式

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

如果你的损失函数需要标签（如损失函数概览表中所示），你的数据集必须有一个名为 "label" 或 "score" 的列。
除 "label" 或 "score" 之外的所有列都被视为输入（如损失函数概览表中所示）。这些列的数量必须与你所选损失函数的有效输入数量相匹配。列的名称无关紧要，只有它们的顺序重要。

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

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

该数据集有一个 "label" 列，这些损失函数需要该列。
该数据集有 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 自带的已实现评估器：

评估器	所需数据
`BinaryClassificationEvaluator`	带有类别标签的句子对
`EmbeddingSimilarityEvaluator`	带有相似度分数的句子对
`InformationRetrievalEvaluator`	查询（qid => question）、语料库（cid => document）和相关文档（qid => set[cid]）
`MSEEvaluator`	用教师模型嵌入的源句子和用学生模型嵌入的目标句子。可以是相同的文本。
`ParaphraseMiningEvaluator`	ID 到句子的映射，以及带有重复句子 ID 的句子对。
`RerankingEvaluator`	{'query': '..', 'positive': [...], 'negative': [...]} 字典的列表。
`TranslationEvaluator`	两种不同语言的句子对。
`TripletEvaluator`	(anchor, positive, negative) 三元组。

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

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

使用 STSb 的 EmbeddingSimilarityEvaluator

STS Benchmark（又称 STSb）是一个常用的基准数据集，用于衡量模型对短文本（如 "A man is feeding a mouse to a snake."）语义文本相似度的理解。

欢迎在 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) 三元组：这是训练嵌入模型的一种常见格式。

在此片段中，它用于评估模型认为锚文本与蕴含文本比锚文本与矛盾文本更相似的频率。一个示例文本是“An older man is drinking orange juice at a restaurant.”。

欢迎在 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

Trainer

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%。

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

Callbacks

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

WandbCallback 用于在安装了 wandb 时将训练指标记录到 W&B
TensorBoardCallback 用于在可访问 tensorboard 时将训练指标记录到 TensorBoard
CodeCarbonCallback 用于在安装了 codecarbon 时跟踪训练期间的碳排放

只要安装了所需的依赖项，这些都会自动使用，无需你指定任何内容。

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

Multi-Dataset Training

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

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

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

MultiDatasetBatchSamplers.ROUND_ROBIN：以轮询方式从每个数据集中采样，直到一个数据集被耗尽。此策略可能不会使用每个数据集中的所有样本，但确保从每个数据集中等量采样。
MultiDatasetBatchSamplers.PROPORTIONAL（默认）：按数据集大小比例从每个数据集中采样。此策略确保使用每个数据集中的所有样本，并且较大的数据集被更频繁地采样。

多任务训练已被证明非常有效。例如，Huang 等人 2024 使用了 MultipleNegativesRankingLoss、CoSENTLoss 以及 MultipleNegativesRankingLoss 的变体（不使用批内负样本，仅使用难负样本），在中文上达到了最先进的性能。他们还应用了 MatryoshkaLoss 使模型能够生成 Matryoshka 嵌入。

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

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 训练、损失日志记录等新功能。尽管如此，新的训练方法要强大得多，因此推荐使用新方法编写新的训练脚本。