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 策略支持。