NovaSearch-Team/RAG-Retrieval

Unify Efficient Fine-tuning of RAG Retrieval, including Embedding, ColBERT, ReRanker.

What it solves

RAG-Retrieval 提供了一个用于优化检索增强生成 (RAG) 流流线中的检索阶段的全面工具包。它解决了使用统一框架来训练、微调和部署不同类型的检索模型——例如嵌入模型、后期交互模型和重排序器——的困难。

How it works

该项目实施了检索模型的完整生命周期，通过三个主要功能：

  • Training and Fine-tuning: 支持 BERT-based 和 LLM-based 嵌入模型、ColBERT-style 后期交互模型以及重排序器的端到端微调。它与 BGE、BCE 和 GTE 等热门开源模型兼容。
  • Distillation: 允许用户将知识从较大的、高性能的模型（如大型 LLM-based 重排序器）转移到较小的、更高效的模型（如 BERT-base 或 0.5B LLMs）。
  • Inference: 包括一个轻量级的 Python 库 (rag-retrieval)，提供了一个统一的接口来调用各种重排序器模型，处理长文档-通过截断或分割来优化评分。

Who it’s for

开发者和研究人员，他们正在从事 RAG 系统的工作，需要通过自定义训练或模型压缩来提高检索组件的准确性与效率。

Highlights

  • Unified Framework: 一个代码库用于嵌入、后期交互和重排序器模型。
  • Model Distillation: 内置支持将大模型蒸馏成较小的、可部署的版本。
  • Advanced Algorithms: 实施了用于降低输出向量维度的 MRL 算法和 Stella 蒸馏方法。
  • Scalable Training: 集成了使用 DeepSpeed 和 FSDP 的多 GPU 策略支持。