Model or dataset
NovaSearch-Team/RAG-Retrieval avatar
NovaSearch-Team/RAG-Retrieval

RAG-Retrieval: One Codebase for Embedding, ColBERT and Reranker Fine-Tuning

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

1,129 stars87 forksPythonMIT

At a glance

What is it?
NovaSearch-Team's RAG-Retrieval bundles training, distillation and a separate PyPI inference package for three retrieval model families. The training side is a set of per-task directories rather than a single CLI, and the inference side only covers rerankers.
Who is it for?
Adopt RAG-Retrieval if you already have a retrieval fine-tuning task and want embedding, ColBERT and reranker training plus distillation under one MIT-licensed tree, and if you are willing to work inside per-model subdirectories rather than a single command. Do not adopt it if you need a served, unified inference path across embedding and reranker models, because the PyPI package covers rerankers only.
Can I use it commercially?
Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
Is it still maintained?
Yes. The repository last received commits 18 days ago.
What is it written in?
Mainly Python, according to GitHub's language statistics.

Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.

DEEP OPEN-SOURCE ANALYSIS

The training-side problem RAG-Retrieval is built around

Retrieval fine-tuning in RAG stacks is usually split across unrelated codebases. An embedding model comes from one training script, a ColBERT-style late-interaction model from another, and a cross-encoder reranker from a third, each with its own data format, loss implementation and launch convention. RAG-Retrieval's stated goal is to collapse that into one repository that covers embedding models (BERT-based and LLM-based), late-interaction models (ColBERT), and rerankers (BERT-based and LLM-based).

The audience is narrow but real: engineers who already have query-document pairs and want to fine-tune an open-source retrieval model rather than call a hosted API. The README lists bge, bce and gte families as compatible targets, including bge-m3 and gte-multilingual-reranker-base. If your retrieval stack is built on one of those checkpoints and you want to keep it, this repository is aimed at you. If you are looking for a retrieval framework that also handles indexing, chunking and serving, this is not that.

Three model families, one repository layout

The architecture visible in the README is directory-based rather than abstraction-based. Training code lives under rag_retrieval/train/, with a subdirectory per model type: embedding, and by implication colbert and reranker. Each subdirectory carries its own README and its own launch script, such as train_embedding.sh. There is no single unified training entry point described in the material.

That layout has a consequence worth stating plainly. Shared logic exists, but the mental model is copy-and-adapt: you enter the subdirectory for your model family, read its README, and adjust its shell script. The README frames this as a deliberate choice, describing the code as having a "simple and understandable code structure for easy modifications," and explicitly says the project "rejects complex." For a research-oriented fine-tuning repo that is a defensible trade: fewer layers between you and the loss function. It does mean that switching from embedding to reranker training is a change of directory and script, not a change of flag.

The algorithms named in the README are specific. For embedding models it lists the MRL algorithm for reducing output vector dimensionality, and the Stella distillation method from the Jasper and Stella work. It also lists LLM preference-based supervised fine-tuning of the retriever, added in June 2024. Multi-GPU training is handled through deepspeed and fsdp, which are named as supported strategies rather than described in detail.

Training and inference are two separate products

The most important structural fact about this project is that the training repository and the inference library are not the same artifact. The README states that for inference, RAG-Retrieval "focuses reranker" and ships a lightweight Python library on PyPI called rag-retrieval. That library provides a unified interface for calling ranking models, and the README describes support for Cross Encoder Reranker and Decoder-Only LLM Reranker, with two long-document handling modes: maximum length truncation, and splitting the document and taking the maximum score.

Extensibility is defined by a base class. According to the README, adding a new ranking model means inheriting from BaseReranker and implementing the rank and compute_score functions. That is a small, legible surface, and it tells you the intended extension path is subclassing rather than configuration.

What the material does not describe is a unified inference path for embedding or ColBERT models through that same package. The stated inference scope is rerankers. If you fine-tune an embedding model here, the README does not document a matching rag-retrieval serving interface for it. Plan for the two halves separately: training from the repository, reranker inference from PyPI.

Getting it running: the commands the README gives

Installation is split by purpose. For training, the README gives a conda environment and a requirements install:

conda create -n rag-retrieval python=3.8 && conda activate rag-retrieval pip install -r requirements.txt

The README adds a warning directly above that second line: to avoid incompatibility between the automatically installed torch and the local CUDA, it recommends manually installing a compatible torch version before proceeding. Take that seriously. A requirements file that pulls torch will otherwise resolve against whatever index default is active, and a CUDA mismatch surfaces later as a runtime error rather than an install error.

For prediction, installation is a single package: pip install rag-retrieval, with the same torch caveat repeated. Note the name collision: the PyPI distribution is rag-retrieval (hyphen), while the training code lives under the rag_retrieval (underscore) package directory in the repository.

Training itself is a directory change plus a script. The README's example is embedding:

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

For other model types the README directs you into the corresponding subdirectory and its own README. There is no config key list in the top-level README; the documented procedure is that detailed steps live in each subdirectory's README file. That is the first place to look, and the material does not enumerate the arguments those scripts accept.

Where the project stops short

The clearest limitation is the asymmetry between training and inference. Three model families are trainable; the packaged inference library is described as covering rerankers. Anyone expecting a single rag-retrieval call to serve an embedding model will not find it documented here.

A second limitation is the absence of a unified training interface. The per-subdirectory structure means duplicated conventions across model families. If your workflow needs to sweep hyperparameters across embedding, ColBERT and reranker runs from one harness, you are writing that glue yourself, because the repository's documented entry point is a shell script inside each subdirectory.

The third is environmental. The README pins the training environment to Python 3.8 and explicitly flags torch/CUDA compatibility as a manual step. That is not unusual for fine-tuning code, but it does mean the install is not a one-liner you can trust on an arbitrary machine.

Finally, the published release list is thin. The single release shown is rag_retrieval_only_train (RAG-Retrieval v0.1) from May 2024, while the README's news entries run through May 2025 and the repository shows later activity. The material does not establish a versioning cadence, so treat the master branch, not the release tag, as the reference point. The README's own links also point to NLPJCL/RAG-Retrieval rather than the NovaSearch-Team path, which is worth noting if you are resolving URLs by hand.

RAG-Retrieval against a general fine-tuning stack such as Sentence Transformers

The obvious alternative for the embedding half is a general sentence-embedding fine-tuning library. Sentence Transformers, for example, trains embedding models from query-document or triplet data and is widely used for exactly that. The difference in approach is scope, not quality. A general embedding trainer is built around one model family and one loss interface. RAG-Retrieval is built around three families, and its distinguishing features are the ones a general trainer would not carry by default: the MRL loss for truncatable embeddings, the Stella distillation path for shrinking a large teacher into a 0.5B LLM or BERT-base student, and LLM-based reranker training with distillation into BERT.

If your task is embedding-only and the loss you need is already in a general library, RAG-Retrieval's breadth buys you nothing and its per-directory structure costs you time. The trade flips if you need to move between families, or if distillation from a large model into a small one is the point of the project. The README's own experimental table is a useful signal here: it lists rag-retrieval-reranker at 0.41 GB against bge-reranker-base and bce-reranker-base_v1 at 1.11 GB each, with average scores of 67.03 and 66.33 respectively for the baselines. The rag-retrieval-reranker average column is not filled in the supplied material, so the size comparison is the only part of that table that can be read directly. Do not treat the table as a claim about quality without reproducing it.

Licence, maintenance and what to verify before you commit

The repository is MIT-licensed, and the README links the LICENSE file on the master branch. MIT is permissive: it allows commercial use, modification and redistribution provided the copyright notice and licence text are retained. That applies to the code in this repository. It does not automatically extend to model weights you download separately, whose licences are set by their own publishers, and it does not cover the terms of any dataset you train on. Check those independently; this is a description of the licence, not legal advice.

Maintenance cost is dominated by the environment rather than the code. Python 3.8 plus a manually pinned torch version means upgrades to newer CUDA or Python releases require you to re-verify the training scripts, and the README does not document a tested compatibility matrix. The deepspeed and fsdp options add their own version sensitivity on top of that.

The upgrade path is also worth checking before you depend on it. The only release listed is v0.1 from May 2024, tagged rag_retrieval_only_train, which suggests the released artifact was scoped to training. The README's news entries through May 2025 and the later repository activity indicate the master branch has moved on. Pin a commit if you need reproducibility, and verify whether the rag-retrieval package you install from PyPI corresponds to the training code you are reading.

Editorial conclusion

Adopt RAG-Retrieval if you already have a retrieval fine-tuning task and want embedding, ColBERT and reranker training plus distillation under one MIT-licensed tree, and if you are willing to work inside per-model subdirectories rather than a single command. Do not adopt it if you need a served, unified inference path across embedding and reranker models, because the PyPI package covers rerankers only. Before committing, read the README in the specific subdirectory you intend to use (for example rag_retrieval/train/embedding), confirm the torch version your CUDA build needs before running pip install -r requirements.txt, and check whether the released rag-retrieval package on PyPI matches the master branch you are reading.

Official sources

  1. Issues
  2. License: MIT
  3. NovaSearch-Team/RAG-Retrieval on GitHub
  4. README
  5. Releases
Community notes

Community notes