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

最新推荐文章于 2024-09-13 10:16:57 发布

Python编程杰哥

最新推荐文章于 2024-09-13 10:16:57 发布

阅读量868

点赞数 17

文章标签：人工智能 flask python 后端语言模型

本文链接：https://blog.csdn.net/xx_nm98/article/details/140672782

版权

前言

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

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

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述