RAG-Retrieval

English | 中文

The RAG-Retrieval offers end-to-end code for training, inference, and distillation of the RAG retrieval model.

• For training, RAG-Retrieval supports fine-tuning of any open-source RAG retrieval model, including embedding models (figure a,bert-based, llm-based), late interactive models (figure d,colbert), and reranker models (figure c,bert-based, llm-based).
• For inference, RAG-Retrieval focuses reranker and has developed a lightweight Python library rag-retrieval, which provides a unified way to call any different RAG ranking models.
• For distillation, it supports distilling LLM-based reranker models into bert-based reranker models.

Communication between communities

Join our WeChat group chat

News

• 12/29/2024: We will release the training code for Jasper and Stella in the next few days. Thank you.

• 10/21/2024: RAG-Retrieval released two different methods for Reranker tasks based on LLM, as well as a method for distilling them into BERT. Best Practices for LLM in Reranker Tasks? A Simple Experiment Report (with code)

• 6/5/2024: Implementation of MRL loss for the Embedding model in RAG-Retrieval. RAG-Retrieval: Making MRL Loss a Standard for Training Vector (Embedding) Models

• 6/2/2024: RAG-Retrieval implements LLM preference-based supervised fine-tuning of the RAG retriever. RAG-Retrieval Implements LLM Preference-Based Supervised Fine-Tuning of the RAG Retriever

• 5/5/2024: Released a lightweight Python library for RAG-Retrieval. RAG-Retrieval: Your RAG Application Deserves a Better Ranking Reasoning Framework

• 3/18/2024: Released RAG-Retrieval Introduction to RAG-Retrieval on Zhihu

Features

• Supports end-to-end fine-tuning of RAG retrieval models: Embedding (bert-based, llm-based), late interaction models (colbert), and reranker models (bert-based, llm-based).
• Supports fine-tuning of any open-source RAG retrieval models: Compatible with most open-source embedding and reranker models, such as: bge (bge-embedding, bge-m3, bge-reranker), bce (bce-embedding, bce-reranker), gte (gte-embedding, gte-multilingual-reranker-base).
• Supports distillation of llm-based large models to bert-based smaller models: Currently supports the distillation of llm-based reranker models into bert-based reranker models (implementation of mean squared error and cross-entropy loss).
• Advanced Algorithms: For embedding models, supports the MRL algorithm to reduce the dimensionality of output vectors.
• Multi-gpu training strategy: Includes deepspeed, fsdp.
• Simple yet Elegant: Rejects complex, with a simple and understandable code structure for easy modifications.

Quick Start

Installation

For training (all):

conda create -n rag-retrieval python=3.8 && conda activate rag-retrieval
# To avoid incompatibility between the automatically installed torch and the local cuda, it is recommended to manually install the compatible version of torch before proceeding to the next step.
pip install -r requirements.txt 

For prediction (reranker):

# To avoid incompatibility between the automatically installed torch and the local cuda, it is recommended to manually install the compatible version of torch before proceeding to the next step.
pip install rag-retrieval

Training

For different model types, please go into different subdirectories. For example: For embedding, and similarly for others. Detailed procedures can be found in the README file in each subdirectories.

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

inference

RAG-Retrieval has developed a lightweight Python library, rag-retrieval, which provides a unified interface for calling various RAG reranker models with the following features:

• Supports multiple ranking models: Compatible with common open-source ranking models (Cross Encoder Reranker, Decoder-Only LLM Reranker).

• Long document friendly: Supports two different handling logics for long documents (maximum length truncation and splitting to take the maximum score).

• Easy to Extend: If there is a new ranking model, users only need to inherit from BaseReranker and implement the rank and compute_score functions.

For detailed usage and considerations of the rag-retrieval package, please refer to the Tutorial

Experimental Results

Results of the reranker model on the MTEB Reranking task

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

Among them, rag-retrieval-reranker is the result of training on the hfl/chinese-roberta-wwm-ext model using the RAG-Retrieval code, and the training data uses the training data of the bge-rerank model.

Results of the Colbert model in the MTEB Reranking task

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

Among them, rag-retrieval-colbert is the result of training on the hfl/chinese-roberta-wwm-ext model using the RAG-Retrieval code, and the training data uses the training data of the bge-rerank model.

Fine-tune the open source BGE series models with domain data

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

The number with finetune at the end means that we used RAG-Retrieval to fine-tune the corresponding open source model, and the training data used the training set of T2-Reranking.

It is worth noting that the training set of the three open source models of bge already includes T2-Reranking, and the data is relatively general, so the performance improvement of fine-tuning using this data is not significant. However, if the open source model is fine-tuned using a vertical field data set, the performance improvement will be greater.

Star History

License

RAG-Retrieval is licensed under the MIT License.