RAG-检索

English | 中文

RAG-检索提供RAG检索模型的训练、推理和蒸馏的端到端代码。

• 在训练方面，RAG-检索支持任何开源RAG检索模型的微调，包括嵌入模型（图a，基于BERT、基于LLM）、晚期交互模型（图d，ColBERT）以及重排序模型（图c，基于BERT、基于LLM）。
• 在推理方面，RAG-检索专注于重排序，并开发了一个轻量级的Python库rag-retrieval，该库提供了一种统一的方式来调用任何不同的RAG重排序模型。
• 在蒸馏方面，支持嵌入模型和重排序模型的蒸馏，可以将较大的模型蒸馏为较小的模型（0.5B参数的LLM或BERT-base）。

社区交流

加入我们的微信群聊

新闻

• 🔥 22/05/2025: RAG-检索发布了《短视陷阱》，这是一项关于整个IR流程中位置偏差的实证研究。我们系统地评估了一系列SOTA检索模型——包括BM25、密集嵌入、ColBERT风格模型和重排序器——在两个精心设计的位置感知基准测试集上：SQuAD-PosQ和FineWeb-PosQ。了解更多

• 29/12/2024: RAG-检索发布了Stella和Jasper嵌入模型的核心训练代码（stage3）Jasper和Stella：SOTA嵌入模型的蒸馏。

• 21/10/2024: RAG-检索发布了两种基于LLM的重排序任务方法，以及一种将其蒸馏为BERT的方法。LLM在重排序任务中的最佳实践？一份简单的实验报告（附代码）

• 05/06/2024: RAG-检索实现了用于嵌入模型的MRL损失函数。RAG-检索：让MRL损失成为向量（嵌入）模型训练的标准

• 02/06/2024: RAG-检索实现了基于LLM偏好的监督微调的RAG检索器。RAG-检索实现基于LLM偏好的监督微调的RAG检索器

• 05/05/2024: RAG-检索发布了一个轻量级的Python库。RAG-检索：你的RAG应用值得一个更好的推理框架

• 18/03/2024: RAG-检索发布知乎上的RAG-检索介绍

特性

• 简单而优雅：摒弃复杂性，采用简单易懂的代码结构，便于修改。
• 支持RAG检索模型的端到端微调：嵌入模型（基于BERT、基于LLM）、晚期交互模型（ColBERT）以及重排序模型（基于BERT、基于LLM）。
• 支持任何开源RAG检索模型的微调：兼容大多数开源嵌入和重排序模型，例如：bge（bge-embedding、bge-m3、bge-reranker）、bce（bce-embedding、bce-reranker）、gte（gte-embedding、gte-multilingual-reranker-base）。
• 支持将大模型蒸馏为小模型：能够将较大的基于LLM的重排序或嵌入模型蒸馏为较小的模型（如0.5B参数的LLM或BERT-base）。
• 先进算法：对于嵌入模型，支持使用MRL算法来降低输出向量的维度，以及Stella蒸馏方法。
• 多GPU训练策略：包含deepspeed、fsdp。

快速入门

安装

用于训练（全部）：

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

用于预测（重排序）：

# 为了避免自动安装的torch与本地cuda不兼容，建议手动安装兼容版本的torch后再进行下一步。
pip install rag-retrieval

训练

针对不同类型的模型，请进入不同的子目录。例如： 对于嵌入，其他类型也类似。详细步骤可在各子目录的README文件中找到。

cd ./rag_retrieval/train/embedding
bash train_embedding.sh

推理

RAG-检索开发了一个轻量级的Python库rag-retrieval，它提供了一个统一的接口来调用各种RAG重排序模型，具有以下特点：

• 支持多种重排序模型：兼容常见的开源重排序模型（交叉编码器重排序器、仅解码器LLM重排序器）。

• 对长文档友好：支持两种不同的处理逻辑来应对长文档（最大长度截断和拆分以取最高分）。

• 易于扩展：如果有新的重排序模型，用户只需继承BaseReranker并实现rank和compute_score函数即可。

有关rag-retrieval包的详细使用方法和注意事项，请参阅教程

实验结果

重排序模型在 MTEB 重排序任务上的结果

模型	模型大小(GB)	T2重排序	MMarco重排序	CMedQA v1	CMedQA v2	平均
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 是基于 hfl/chinese-roberta-wwm-ext 模型，使用 RAG-Retrieval 代码进行训练的结果，训练数据采用了 bge-rerank 模型的训练数据。

Colbert 模型在 MTEB 重排序任务中的结果

模型	模型大小(GB)	维度	T2重排序	MMarco重排序	CMedQA v1	CMedQA v2	平均
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 是基于 hfl/chinese-roberta-wwm-ext 模型，使用 RAG-Retrieval 代码进行训练的结果，训练数据采用了 bge-rerank 模型的训练数据。

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

模型	T2重排序
bge-v1.5-embedding	66.49
bge-v1.5-embedding 微调	67.15	+0.66
bge-m3-colbert	66.82
bge-m3-colbert 微调	67.22	+0.40
bge-reranker-base	67.28
bge-reranker-base 微调	67.57	+0.29

末尾标注“微调”的数字表示我们使用了 RAG-Retrieval 对相应开源模型进行了微调，训练数据为 T2-重排序的训练集。

值得注意的是，bge 的三个开源模型的训练集中已经包含了 T2-重排序的数据，且数据较为通用，因此使用该数据进行微调后的性能提升并不显著。然而，如果使用垂直领域的数据集对开源模型进行微调，则性能提升会更加明显。

引用

如果您觉得本仓库有所帮助，请引用我们的工作：

@misc{zhang2025jasperstelladistillationsota,
      title={Jasper and Stella: distillation of SOTA embedding models}, 
      author={Dun Zhang and Jiacheng Li and Ziyang Zeng and Fulong Wang},
      year={2025},
      eprint={2412.19048},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2412.19048}, 
}

致谢

在开发过程中，我们借鉴或基于以下项目的实现。我们衷心感谢这些团队为开源研究与开发所做出的贡献。

• FlagEmbedding
• uniem
• sentence-transformers
• rerankers

星标历史

许可证

RAG-Retrieval 采用 MIT 许可证 许可。

RAG-Retrieval 快速上手指南

RAG-Retrieval 是一个端到端的工具库，支持 RAG 检索模型（Embedding、ColBERT、Reranker）的训练、推理和蒸馏。它兼容主流开源模型（如 BGE、BCE、GTE 系列），并提供统一的推理接口。

1. 环境准备

• 操作系统: Linux / macOS / Windows
• Python 版本: 推荐 Python 3.8+
• 硬件要求:
  • 训练: 需要 NVIDIA GPU，建议安装与本地 CUDA 版本匹配的 PyTorch。
  • 推理: CPU 或 GPU 均可运行。
• 前置依赖: 建议先手动安装适配本地环境的 torch，以避免自动安装版本不兼容的问题。

2. 安装步骤

场景 A：模型训练 (Training)

如果你需要微调 Embedding、ColBERT 或 Reranker 模型，请执行以下命令：

# 创建并激活虚拟环境
conda create -n rag-retrieval python=3.8 && conda activate rag-retrieval

# 【重要】建议先手动安装与你本地 CUDA 版本匹配的 torch，例如：
# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 安装项目依赖
pip install -r requirements.txt

国内加速提示：如遇网络问题，可使用清华源或阿里源安装依赖： pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

场景 B：模型推理 (Inference)

如果你仅需使用现成的 Reranker 模型进行排序推理，可直接安装轻量级 Python 库：

# 【重要】同样建议先手动安装适配的 torch

# 安装推理库
pip install rag-retrieval

3. 基本使用

3.1 训练模型 (Training)

RAG-Retrieval 将不同模型的训练代码分置于子目录中。以下以 Embedding 模型 的微调为例：

# 进入 Embedding 训练目录
cd ./rag_retrieval/train/embedding

# 执行训练脚本 (请根据实际需求修改 train_embedding.sh 中的参数)
bash train_embedding.sh

注：ColBERT 和 Reranker 的训练方式类似，只需进入对应的 train/colbert 或 train/reranker 目录即可。具体参数配置请参考各子目录下的 README。

3.2 模型推理 (Inference)

安装 rag-retrieval 包后，你可以使用统一的接口调用各种开源 Reranker 模型（支持 Cross Encoder 和 LLM-based 模型）。

以下是一个最简单的 Python 使用示例：

from rag_retrieval import Reranker

# 初始化 Reranker，自动加载默认模型或指定本地/HF 模型路径
# 支持长文档处理策略（截断或分段取最大分）
reranker = Reranker('BAAI/bge-reranker-base', device='cuda') 

# 准备查询和候选文档列表
query = "什么是人工智能？"
documents = [
    "人工智能是计算机科学的一个分支。",
    "今天天气真不错。",
    "深度学习是人工智能的重要子领域。"
]

# 执行排序
results = reranker.rank(query=query, docs=documents)

# 输出结果
for res in results:
    print(f"文档: {res['text']}, 得分: {res['score']}")

特性说明：

• 多模型兼容: 支持 bge-reranker, bce-reranker, gte-reranker 等主流模型。
• 长文档友好: 内置处理超长文本的逻辑（最大长度截断 或 切片后取最高分）。
• 易于扩展: 若有新模型，只需继承 BaseReranker 类并实现 rank 和 compute_score 方法即可。

更多高级用法（如自定义打分策略、批量处理等）请参考官方 Reranker 教程。