RAG优化必备：如何微调自己的嵌入模型

前言

在大模型时代中，嵌入（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}

可以看到，进行少量训练后，指标是有所提高的