Advanced RAG: Increase RAG Quality with ColBERT Reranker and llamaindex

On this page

Reranking improves Retrieval Augmented Generation (RAG) by sorting the documents from semantic search a second time, with a model trained to judge how relevant a document is to a question. This post explains why plain vector search often puts the wrong chunks first and shows how to add the ColBERT reranker to a LlamaIndex query engine in a few lines of Python.

RAG is one of the most useful ways of integrating your own information and data with LLMs. It is also fiddly to set up and tune. There are many knobs to turn, and each one has a large effect on the final result: the LLM's answer to your question.

One of the main factors behind answer quality is context precision, a metric for how well the retrieved documents match the question (you can measure it with RAGAS). The better the documents match, the better the answer.

In one of our last posts, we showed how recursive retrieval can improve RAG quality. Reranking is a separate technique, and you can (and should) combine the two. We start with what reranking is and why you need it, then walk through a hands-on implementation with ColBERT and LlamaIndex.

How does reranking work and why do we need reranking?

In a default RAG setup, documents are retrieved with semantic search, which ranks them by their similarity to the question. If the question is "How much does a house cost?", the similarity search uses vector embeddings to find documents that are similar to the question. These documents are then sent to the LLM, together with the question, to generate the answer.

A user question and context documents are sent together to the LLM (GPT-4)

With a perfect embedding model, the retrieved documents would always match the question. That is not the case, for two reasons. First, even state-of-the-art embedding models can't fully capture the meaning of both the document texts and the users' questions. Second, documents are usually split into chunks automatically, and a chunk often covers several topics instead of one. There are strategies to handle this (such as late chunking), but the problem remains.

Reranking targets exactly this problem. It takes the documents returned by semantic search and sorts them again with another model. That model is often a neural network trained to judge how relevant a document is to a question. Because it is trained on large datasets of questions and documents, it can capture relevance better than a plain embedding model. And even when it doesn't, it gives you a "second opinion" on each document.

The second big issue for semantic search is that embedding models are mostly trained on general texts from the internet. They have seen little text from your specific domain, so semantic search often performs poorly on specialized content. We call this the "out of domain" problem. A reranking model trained on domain-specific texts can reduce it, and since reranking is a simpler task, reranking models are easier to train on domain data than embedding models.

Semantic search results pass through a reranking model and reach the LLM reordered

Enough theory. Let's implement reranking with ColBERT and LlamaIndex.

How to implement RAG document reranking with ColBERT and LlamaIndex

We introduced LlamaIndex in one of our last posts; see that post for more about the library. Besides being a good general retrieval library, LlamaIndex makes it simple to plug reranking models into a RAG pipeline. In this guide, we use ColBERT as the reranking model.

ColBERT is a fast and accurate retrieval model that enables scalable BERT-based search over large text collections in tens of milliseconds. The speed matters. RAG is not a fast technology to begin with, since every LLM call adds latency, and a reranker adds another model call on top. ColBERT is one of the fastest reranking models available, so it keeps that extra cost small.

Start by installing the required libraries:

1pip install llama-index
2pip install llama-index-core
3pip install --quiet transformers torch
4pip install llama-index-embeddings-openai
5pip install llama-index-llms-openai
6pip install llama-index-postprocessor-colbert-rerank

If you want to follow along, download the sample data from our website.

Next, import the libraries and define which OpenAI models to use. Replace the OPENAI_API_KEY value with your own API key.

1import os
2from llama_index.core import VectorStoreIndex
3from llama_index.llms.openai import OpenAI
4from llama_index.embeddings.openai import OpenAIEmbedding
5from llama_index.readers.file import PyMuPDFReader
6from llama_index.core import Settings
7from llama_index.postprocessor.colbert_rerank import ColbertRerank
8
9os.environ["OPENAI_API_KEY"] = "Your API Key"
10
11# Setup the OpenAI LLM
12Settings.llm = OpenAI(model="gpt-4")
13Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

Optionally, enable debug logging to see what LlamaIndex does under the hood:

1
2# Optional: Set up debug logging to see exactly what llamaindex is doing
3import logging
4import sys
5logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
6logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

Now load the document, split it into chunks and create a VectorStoreIndex. LlamaIndex does all of this in two calls:

1# Load the document
2file = "public/blog/recursiveRetrieval/world_billionaires.pdf"
3reader = PyMuPDFReader()
4docs = reader.load(file)
5
6index = VectorStoreIndex.from_documents(documents=docs)

To add the ColBERT reranker to the LlamaIndex pipeline, pass it as a node postprocessor to the query engine. The query engine then runs these steps:

1. Creates a query vector from the question
2. Retrieves the most similar documents from the index
3. Reranks the documents using the ColBERT model
4. Sends the reranked documents to the LLM

1colbert_reranker = ColbertRerank(
2    top_n=5,
3    model="colbert-ir/colbertv2.0",
4    tokenizer="colbert-ir/colbertv2.0",
5    keep_retrieval_score=True,
6)
7
8query_engine = index.as_query_engine(
9    similarity_top_k=10,
10    node_postprocessors=[colbert_reranker],
11)

That's all the setup. You can now use the query_engine to ask questions and get answers based on the indexed documents, with the ColBERT reranker reordering the documents before they are sent to the LLM.

1response = query_engine.query(
2    "Who was the richest person in the world in 2017?",
3)
4
5print(response)

---

Interested in building high-quality AI agent systems?

We prepared a comprehensive guide based on cutting-edge research for how to build robust, reliable AI agent systems that actually work in production. This guide covers:

• Understanding the 14 systematic failure modes in multi-agent systems
• Evidence-based best practices for agent design
• Structured communication protocols and verification mechanisms

---

See the reranker in action

So far we can't see what the reranker changed. Print the retrieved document nodes to compare the scores:

1for node in response.source_nodes:
2    print(node.id_)
3    print(node.node.get_content())
4    print("reranking score: ", node.score)
5    print("retrieval score: ", node.node.metadata["retrieval_score"])
6    print("=====================================")

This snippet prints the retrieved documents with two scores:

• The retrieval score is the score the semantic search assigned to the document. The higher the score, the more similar the document is to the question, according to semantic search.

• The reranking score is the score the ColBERT reranker assigned to the document. The higher the score, the more relevant the document is to the question, according to ColBERT.

In the printed output, the order by reranking score is often different from the order by retrieval score. More often than not, the new order is a better match for the question. Keep in mind that the first document should be the most relevant one.

Conclusion

Answer quality in RAG depends on how relevant the retrieved documents are, and semantic search alone often gets the order wrong. Embedding models don't capture meaning perfectly, chunks mix topics, and general-purpose models struggle with domain-specific texts.

A reranker such as ColBERT adds a second opinion on relevance and helps with the out-of-domain problem. With LlamaIndex, adding it takes one ColbertRerank object passed as a node postprocessor: retrieve more candidates than you need (here 10), rerank them, and send only the best ones (here 5) to the LLM. Because ColBERT is fast, the extra latency stays low. If you run on Azure, Azure AI Search has a managed reranking step built in, the semantic ranker, which we cover in our Azure AI Search RAG tutorial.

A good next step is to combine reranking with recursive retrieval and measure the effect on context precision with an evaluation framework such as RAGAS.

Further reading

• How to test your RAG pipeline?
• Advanced RAG: Recursive Retrieval with llamaindex
• Advanced RAG: Hypothetical Document Embeddings (HyDE)
• RAGFlow: an open source RAG engine with deep document understanding
• Langfuse: tracing retrieval latency and answer quality of LLM applications

About the author

Andreas Nigg

Co-Founder, Pondhouse Data

Andreas has spent more than eight years building data applications. At Pondhouse Data he designs and ships AI systems and integrates them into complex, real-world environments.

Talk to Andreas about your projectFollow on LinkedIn