前言
在大模型时代中，嵌入（embedding）模型是非常重要且基础的能力，例如在RAG中，匹配问题和知识条文时，好的嵌入模型会让检索更精准。但是通用的嵌入模型，无论其宣传效果有多好，在自己的使用场景中往往会出现水土不服。所以，如何微调嵌入模型，使得匹配自己的业务场景使得尤为重要。本文介绍了使用Sentence Transformers框架对开源的嵌入模型进行微调，使用的Sentence Transformers版本为3.0.0版本。3版本相对2版本发生了重大改变，个人感觉最主要的变化是继承了transformers的很多能力，使得在多卡训练上更方便，这在动辄数百万级的训练样本情况下尤为重要。网上很多教程都是2版本的，在3版本上不适用，请注意。更多的内容请关注官方文档：https://sbert.net/
1.加载库
from datasets import load_dataset
from sentence_transformers import (
SentenceTransformer,
SentenceTransformerTrainer,
SentenceTransformerTrainingArguments,
SentenceTransformerModelCardData,
)
from sentence_transformers.losses import CoSENTLoss
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers.evaluation import BinaryClassificationEvaluator, EmbeddingSimilarityEvaluator, TripletEvaluator
import pandas as pd
2.加载模型，设置损失函数
# 模型地址：https://huggingface.co/BAAI/bge-small-en-v1.5
model = SentenceTransformer("models/bge-small-en-v1.5/")
# 根据不同的数据类型选择不同的损失函数，详见：https://sbert.net/docs/sentence_transformer/loss_overview.html?highlight=loss
loss = CoSENTLoss(model)
根据不同的数据类型选择不同的损失函数，详见：https://sbert.net/docs/sentence_transformer/loss_overview.html?highlight=loss
，本案例的场景可选用CoSENTLoss
3.读取数据
# 数据集地址：https://huggingface.co/datasets/sentence-transformers/all-nli,请下载后再加载
train_path = 'data/all-nli/pair-score/train-00000-of-00001.parquet'
dev_path = 'data/all-nli/pair-score/dev-00000-of-00001.parquet'
test_path = 'data/all-nli/pair-score/test-00000-of-00001.parquet'
train_dataset = load_dataset("parquet", data_files=train_path)
eval_dataset = load_dataset("parquet", data_files=dev_path)
test_dataset = load_dataset("parquet", data_files=test_path)
print(train_dataset) # train_dataset['train']['premise']才是真正的样本
DatasetDict({
train: Dataset({
features: ['sentence1', 'sentence2', 'score'],
num_rows: 942069
})
})我们也可以读取成dataframe的形式，更方便观察数据
df = pd.read_parquet(train_path)
df
可以看见数据为三列的方式，前两列是句子，最后一列是两个句子的相似度评分。如果觉得评分太麻烦，最后一列也可以是0或1的label，但对应的代码需要修改。我们在制作训练集的时候，按照这三列的形式是更方便的。
4.设置训练参数
这里是继承了transformers的参数设置，运行的时候会默认机器的卡全部都用于训练。如果要指定某些卡进行训练，需要使用启动命令设置该环境变量，例如：CUDA_VISIBLE_DEVICES="0" python main.py
args = SentenceTransformerTrainingArguments(
# 必须的参数:
output_dir="model_save/pair-class",
# 训练参数:
num_train_epochs=1,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
learning_rate=2e-5,
warmup_ratio=0.1,
fp16=True,
bf16=False,
batch_sampler=BatchSamplers.NO_DUPLICATES, # 确保同一批批训练样本中没有重复的样本
# 验证/保存参数:
eval_strategy="steps",
eval_steps=100,
save_strategy="steps",
save_steps=100,
save_total_limit=2,
logging_steps=100,
run_name="pair-score", # 设置好名称，可在`wandb`上记录
)
5.设置验证方式
BinaryClassificationEvaluator用于分类标签，EmbeddingSimilarityEvaluator用于分数标签，TripletEvaluator用于句子-正例-负例
本案例适用EmbeddingSimilarityEvaluator
不同的验证方式对应的超参不同，具体参考：https://sbert.net/docs/package_reference/sentence_transformer/evaluation.html?highlight=sentence_transformers%20evaluation
dev_evaluator = EmbeddingSimilarityEvaluator(
sentences1=eval_dataset['train']["sentence1"],
sentences2=eval_dataset['train']["sentence2"],
scores=eval_dataset['train']["score"],
name="pair-class-dev",
)
# 训练前先测一下
dev_evaluator(model)
{'pair-class-dev_pearson_cosine': 0.5586668614994776,
'pair-class-dev_spearman_cosine': 0.5591719112836276,
'pair-class-dev_pearson_manhattan': 0.548648858854572,
'pair-class-dev_spearman_manhattan': 0.555145545980968,
'pair-class-dev_pearson_euclidean': 0.5517466850163146,
'pair-class-dev_spearman_euclidean': 0.5591718680615544,
'pair-class-dev_pearson_dot': 0.5586668602820397,
'pair-class-dev_spearman_dot': 0.5591718350908823,
'pair-class-dev_pearson_max': 0.5586668614994776,
'pair-class-dev_spearman_max': 0.5591719112836276}以上指标都是衡量向量相似度的，等会训练几个周期，看看有没有提升
6.开始训练
trainer = SentenceTransformerTrainer(
model=model,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
loss=loss,
evaluator=dev_evaluator,
)
trainer.train()
这里只训练500步，可以看到评价指标是逐渐上升的，训练集的损失函数逐渐下降。验证集的损失函数显示No log，可能是验证方式的原因。
7.测试集效果
test_evaluator = EmbeddingSimilarityEvaluator(
sentences1=test_dataset['train']["sentence1"],
sentences2=test_dataset['train']["sentence2"],
scores=test_dataset['train']["score"],
name="pair-class-dev",
)
test_evaluator(model)
# 保存最终模型
model.save_pretrained("model_save/pair-class/final")
{'pair-class-dev_pearson_cosine': 0.6001533384580744,
'pair-class-dev_spearman_cosine': 0.6145674695662575,
'pair-class-dev_pearson_manhattan': 0.5995323301973094,
'pair-class-dev_spearman_manhattan': 0.6124315808477351,
'pair-class-dev_pearson_euclidean': 0.6013774845569113,
'pair-class-dev_spearman_euclidean': 0.6145674696825638,
'pair-class-dev_pearson_dot': 0.6001533367553711,
'pair-class-dev_spearman_dot': 0.6145674475881859,
'pair-class-dev_pearson_max': 0.6013774845569113,
'pair-class-dev_spearman_max': 0.6145674696825638}可以看到，进行少量训练后，指标是有所提高的
