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 models, 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, Distillation of support embedding models and reranker models, support distill from a larger model to a smaller model (0.5b llm or bert-base).

Communication between communities

Join our WeChat group chat

News

• 🔥 22/05/2025: RAG-Retrieval released Myopic Trap, an empirical study of positional bias across the full IR pipeline. We systematically evaluate a range of SOTA retrieval models—including BM25, dense embeddings, ColBERT-style models, and rerankers—on two carefully designed position-aware benchmarks: SQuAD-PosQ and FineWeb-PosQ. Learn more

• 29/12/2024: RAG-Retrieval released the core training code (stage3) of Stella and Jasper embedding model Jasper and Stella: distillation of SOTA embedding models.

• 21/10/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)

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

• 02/06/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

• 05/05/2024: Released a lightweight Python library for RAG-Retrieval. RAG-Retrieval: Your RAG Application Deserves a better infer framework

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

Features

• Simple yet Elegant: Rejects complex, with a simple and understandable code structure for easy modifications.
• 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 larger models into smaller models: Enables the distillation of larger LLM-based reranker or embedding models into smaller ones (e.g., a 0.5B-parameter LLM or BERT-base).
• Advanced Algorithms: For embedding models, supports the MRL algorithm to reduce the dimensionality of output vectors and Stella distillation method.
• Multi-gpu training strategy: Includes deepspeed, fsdp.

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

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

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

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.

Citation

If you find this repository helpful, please cite our work:

@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}, 
}

Acknowledge

During the development process, we borrowed or based on the implementation of the following projects. We sincerely appreciate the efforts of these teams for their contributions to open-source research and development.

• FlagEmbedding
• uniem
• sentence-transformers
• rerankers

Star History

License

RAG-Retrieval is licensed under the MIT License.