使用统一的方式推理不同的RAG Reranker模型 | 微调全链路的RAG检索模型 | 实验结果 | License
RAG-Retrieval 提供了全链路的RAG检索微调(train)和推理(infer)以及蒸馏(distill)代码。
加入我们微信群聊
10/21/2024: RAG-Retrieval发布基于LLM做Reranker任务的两种不同方法,以及将其蒸馏到bert中的方法。LLM在Reranker任务上的最佳实践?A simple experiment report(with code)
6/5/2024: RAG-Retrieval的Embedding模型的MRL loss实现。RAG-Retrieval:让MRL loss成为训练向量(embedding)模型的标配
6/2/2024: RAG-Retrieval实现基于LLM偏好监督RAG检索器微调。RAG-Retrieval实现基于LLM偏好监督RAG检索器微调
5/5/2024:发布RAG-Retrieval的轻量级的python库RAG-Retrieval:你的RAG应用值得更好的排序推理框架
3/18/2024:发布RAG-Retrieval RAG-Retrieval知乎介绍
排序模型是任何检索架构的重要组成部分,也是 RAG 的重要组成部分,但目前的现状是:
因此,RAG-Retrieval开发了一个轻量级的python库rag-retrieval,提供统一的方式调用任意不同的RAG排序模型,其有以下的特点。
支持多种排序模型:支持常见的开源排序模型(Cross Encoder Reranker,Decoder-Only 的LLM Reranker)
长doc友好:支持两种不同的对于长doc的处理逻辑(最大长度截断,切分取最大分值)。
益于扩展:如果有新的排序模型,用户只需要继承basereranker,并且实现rank以及comput_score函数即可。
#为了避免自动安装的torch与本地的cuda不兼容,建议进行下一步之前先手动安装本地cuda版本兼容的torch。
pip install rag-retrieval
对于Cross Encoder Reranker,只要其使用transformers的AutoModelForSequenceClassification,那么就可以使用rag_retrieval的Reranker来进行推理。举例如下。
bge系列的Cross Encoder模型,例如(BAAI/bge-reranker-base, BAAI/bge-reranker-large, BAAI/bge-reranker-v2-m3)
bce的Cross Encoder模型,例如(maidalun1020/bce-reranker-base_v1)
对于LLM Reranker,rag_retrieval的Reranker支持多种强大的LLM排序模型。也支持使用任意的LLM的chat模型来进行zero shot排序。举例如下。
bge系列的LLM Reranker模型,例如(BAAI/bge-reranker-v2-gemma, BAAI/bge-reranker-v2-minicpm-layerwise, BAAI/bge-reranker-v2-m3 )
也支持使用任意LLM的chat模型来进行zero shot排序
rag-retrieval包详细的使用方法和注意事项可以参考Tutorial
我们做了大量的测试与下面原有推理框架对齐,详见Tests,他们需要不同的模块来执行,rag_retrieval则使用了统一的接口。
如 FlagEmbedding的 FlagReranker,FlagLLMReranker,LayerWiseFlagLLMReranker。
如 BCEmbedding 的 RerankerModel
conda create -n rag-retrieval python=3.8 && conda activate rag-retrieval
#为了避免自动安装的torch与本地的cuda不兼容,建议进行下一步之前先手动安装本地cuda版本兼容的torch。
pip install -r requirements.txt
支持微调任意开源的embedding模型(bge,m3e等等)
支持对两种数据进行微调:
微调embedding模型流程,详细的流程可参考模型目录下的Tutorial。
cd ./rag_retrieval/train/embedding
bash train_embedding.sh
微调colbert模型流程,详细的流程可参考模型目录下的Tutorial。
cd ./rag_retrieval/train/colbert
bash train_colbert.sh
微调reranker模型流程,详细的流程可参考模型目录下的Tutorial。
cd ./rag_retrieval/train/reranker
bash train_reranker.sh
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模型的训练数据.
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模型的训练数据.
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的训练集。
值得注意的是bge的三种开源模型,训练集中已经包含了T2-Reranking,并且该数据较为通用,因此使用该数据继续微调的性能提升效果不大,但是如果使用垂直领域的数据集继续微调开源模型,性能提升会更大。
RAG-Retrieval is licensed under the MIT License.