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

发布日期：2024-07-19 07:19:15 浏览次数： 4556

前言

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

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