使用 Sentence Transformers v3 訓練和微調嵌入模型

釋出於 2024 年 5 月 28 日

Sentence Transformers 是一個用於使用和訓練嵌入模型的 Python 庫，適用於廣泛的應用，例如檢索增強生成、語義搜尋、語義文字相似度、釋義挖掘等。其 v3.0 更新是該專案自成立以來最大的一次更新，引入了一種新的訓練方法。在這篇博文中，我將向您展示如何使用它來微調 Sentence Transformer 模型，以提高它們在特定任務上的效能。您也可以使用此方法從頭開始訓練新的 Sentence Transformer 模型。

現在，微調 Sentence Transformers 涉及多個元件，包括資料集、損失函式、訓練引數、評估器以及新的訓練器本身。我將詳細介紹每個元件，並提供如何使用它們來訓練有效模型的示例。

為什麼要微調？

微調 Sentence Transformer 模型可以顯著提升它們在特定任務上的效能。這是因為每個任務都需要獨特的相似性概念。我們以幾個新聞標題為例：

“蘋果釋出新款 iPad”
“英偉達正為下一代 GPU 做準備”

根據用例，我們可能希望這些文字的嵌入相似或不相似。例如，新聞文章的分類模型可以將這些文字視為相似，因為它們都屬於“科技”類別。另一方面，語義文字相似度或檢索模型應將它們視為不相似，因為它們的含義不同。

訓練元件

訓練 Sentence Transformer 模型涉及以下元件：

資料集：用於訓練和評估的資料。
損失函式：衡量模型效能並指導最佳化過程的函式。
訓練引數（可選）：影響訓練效能和跟蹤/除錯的引數。
評估器（可選）：用於在訓練前、訓練中或訓練後評估模型的工具。
訓練器：將模型、資料集、損失函式和其他元件整合起來進行訓練。

現在，讓我們更詳細地深入瞭解這些元件。

資料集

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,
})

字典中的每個鍵都將成為結果資料集中的一列。

資料集格式

確保資料集格式與您選擇的損失函式匹配至關重要。這涉及檢查兩件事：

如果您的損失函式需要一個 Label（如損失概述表中所示），您的資料集必須有一個名為 "label" 或 "score" 的列。
除了 "label" 或 "score" 之外的所有列都被視為 Inputs（如損失概述表中所示）。這些列的數量必須與您所選損失函式的有效輸入數量相匹配。列的名稱不重要，只有它們的順序重要。

例如，如果您的損失函式接受 (anchor, positive, negative) 三元組，那麼您的第一個、第二個和第三個資料集列分別對應於 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 => 問題），語料庫（cid => 文件），以及相關文件（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) 是一個常用的基準資料集，用於衡量模型對“一個人正在給蛇喂老鼠。”等短文字語義相似度的理解。

歡迎瀏覽 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 透過允許您使用多個數據集進行訓練而無需將它們轉換為相同格式來簡化此過程。您甚至可以對每個資料集應用不同的損失函式。以下是多資料集訓練的步驟：

使用 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，使模型能夠生成套娃嵌入。

以下是多資料集訓練的示例：

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 訓練、損失日誌記錄等。儘管如此，新的訓練方法功能更強大，因此建議使用新方法編寫新的訓練指令碼。