Reihnnnn / RAG-Retrieval

Unify Efficient Fine-tuning of RAG retrieval ,including embedding、colbert、cross encoder

RAG-Retrieval

创建环境 | 微调模型 | 实验结果 | License

RAG-Retrieval 提供了全链路的RAG检索微调代码，支持微调任意开源的RAG检索模型，包括向量（embedding、图a）、迟交互式模型（colbert、图d）、交互式模型（cross encoder、图c）。

最新更新

• 3/18/2024:发布RAG-Retrieval RAG-Retrieval知乎介绍

创建环境

conda create -n rag-retrieval python=3.8 && conda activate rag-retrieval
#为了避免自动安装的torch与本地的cuda不兼容，建议进行下一步之前先手动安装本地cuda版本兼容的torch。
pip install -r requirements.txt 

微调模型

向量（embedding）模型

• 支持微调任意开源的embedding模型（bge,m3e等等）

• 支持对两种数据进行微调：

  • query和正例（负例采用batch内随机负例），
  • query和正例以及难负例。（负例为对应的难负例，以及batch内随机负例）

微调embedding模型流程

cd ./rag-retrieval/embedding
bash train_embedding.sh

迟交互式（colbert）模型

• 支持微调开源的bge-m3e模型中的colbert。
• 支持query和正例以及难负例。（负例为对应的难负例，以及batch内随机负例）

微调colbert模型流程

cd ./rag-retrieval/colbert
bash train_colbert.sh

排序（reranker,cross encoder）模型

• 支持微调任意开源的reranker模型（例如，bge-rerank、bce-rerank等）
• 支持两种数据进行微调：
  • query和doc的相关性为二分类（1代表相关、0代表不相关）
  • query和doc的相关性为四分类。（3，2，1，0，相关性依次降低。）

微调reranker模型流程

cd ./rag-retrieval/reranker
bash train_reranker.sh

实验结果

reranker模型在 MTEB Reranking 任务的结果

Model	Model Size(GB)	T2Reranking	MMarcoReranking	CMedQAv1	CMedQAv2	Avg
bge-reranker-base	1.11	67.28	35.46	81.27	84.10	67.03
bce-reranker-base_v1	1.11	70.25	34.13	79.64	81.31	66.33
rag-retrieval-reranker	0.41	67.33	31.57	83.54	86.03	67.12

其中，rag-retrieval-reranker是我们使用RAG-Retrieval代码在hfl/chinese-roberta-wwm-ext模型上训练所得，训练数据使用bge-rerank模型的训练数据.

colbert模型在 MTEB Reranking 任务的结果

Model	Model Size(GB)	Dim	T2Reranking	MMarcoReranking	CMedQAv1	CMedQAv2	Avg
bge-m3-colbert	2.24	1024	66.82	26.71	75.88	76.83	61.56
rag-retrieval-colbert	0.41	1024	66.85	31.46	81.05	84.22	65.90

其中，rag-retrieval-colbert是我们使用RAG-Retrieval代码在hfl/chinese-roberta-wwm-ext模型上训练所得，训练数据使用bge-rerank模型的训练数据.

用领域内数据微调开源的BGE系列模型

Model	T2ranking
bge-v1.5-embedding	66.49
bge-v1.5-embedding finetune	67.15	+0.66
bge-m3-colbert	66.82
bge-m3-colbert finetune	67.22	+0.40
bge-reranker-base	67.28
bge-reranker-base finetune	67.57	+0.29

后面带有finetune的代表我们使用RAG-Retrieval在对应开源模型的基础上继续微调所得，训练数据使用T2-Reranking的训练集。

License

RAG-Retrieval is licensed under the MIT License.