Re-ranking Retrieved Contexts

Hello! In our previous lesson, we explored Query Decomposition, a technique to improve retrieval coverage by breaking down complex questions into simpler sub-queries. This ensures we gather a broad set of potentially relevant documents. However, this "wide net" approach often brings in noise—documents that are topically related but not the most relevant for answering the specific query.

This brings us to a fundamental trade-off in information retrieval: recall vs. precision. Our initial retrieval, using vector search, is optimized for recall (don't miss anything important). Now, we need a step to optimize for precision (make sure the top results are the most important).

In this lesson, you will learn how to apply a re-ranking model to re-order retrieved contexts for maximum relevance. This is the second crucial stage in a modern, production-grade retrieval pipeline, turning a large set of candidate documents into a small, highly-relevant context for the LLM.

The Two-Stage Retrieval Funnel

A standard RAG pipeline can be thought of as a funnel. The first stage, which we've implemented, casts a wide net to retrieve a large set of candidate documents. The second stage, which we're building today, applies a more computationally intensive and sophisticated model to re-rank this smaller set, ensuring only the most relevant documents proceed to the final generation step.

This multi-stage process is essential for building scalable and accurate AI systems.

Why Vector Search Isn't Enough: Bi-Encoders vs. Cross-Encoders

To understand why a re-ranking step is necessary, we need to revisit how our initial retrieval works and contrast it with the re-ranking approach.

• Bi-Encoders (Initial Retrieval): Our vector search relies on a bi-encoder architecture. The query and the documents are encoded into vectors independently. The system then finds documents whose vectors are closest to the query vector (e.g., using cosine similarity). This is very fast and scalable because all document vectors can be pre-computed and stored in a vector database. However, because the model never sees the query and document together, it can miss subtle relationships and nuances.

• Cross-Encoders (Re-ranking): Re-ranking models typically use a cross-encoder architecture. A cross-encoder takes both the user's query and a single document as a combined input. It processes them together, allowing the model to perform a deep, token-by-token comparison. The output is not a vector, but a single relevance score (e.g., from 0 to 1). This is far more accurate but also much slower, as it cannot be pre-computed.

The two-stage pattern leverages the best of both worlds: use a fast bi-encoder to retrieve, say, the top 50 candidates, and then use a slow but powerful cross-encoder to re-rank those 50 candidates and pick the top 3-5.

The following article provides an excellent deep dive into this distinction.

A Spectrum of Re-ranking Strategies

Cross-encoders are not the only way to re-rank. As an AI Engineer, you need to be aware of the different strategies and their trade-offs in terms of performance, cost, and latency.

The following article provides a superb overview of five different re-ranking strategies, which is invaluable for system design discussions.

For our hands-on work, we'll focus on the Cross-Encoder strategy. It offers a great balance of high accuracy and gives you full control, as you run the model yourself. It doesn't require a paid API, making it perfect for development and many production use cases.

Hands-On: Implementing Cross-Encoder Re-ranking

Let's add a re-ranking stage to our RAG pipeline. We will use the sentence-transformers library, which makes it incredibly easy to download and use pre-trained cross-encoder models from Hugging Face.

Our new workflow will be:

1. Retrieve: Fetch a larger number of documents (e.g., 20) using our existing ChromaDB setup.
2. Re-rank: Use a CrossEncoder model to score all 20 documents against the query.
3. Select: Pick the top 3-5 re-ranked documents to build the final context.
4. Synthesize: Pass the refined context to the LLM for answer generation.

Here is the updated Python code. Notice how we've added a rerank_with_cross_encoder function and modified the main pipeline logic to include this new step.

import chromadb
from openai import OpenAI
from sentence_transformers.cross_encoder import CrossEncoder
import json




# --- Setup (Same as previous lesson) ---
client = OpenAI() # Assumes OPENAI_API_KEY is set
chroma_client = chromadb.PersistentClient(path="./my_vector_db")
collection_name = "rag_system_docs"
collection = chroma_client.get_collection(name=collection_name)




# --- 1. REUSED: Initial Retrieval Step ---
def retrieve_initial_documents(query: str, n_results: int = 20):
    """
    Retrieves a larger set of initial candidate documents from ChromaDB.
    """
    print(f"--- Retrieving top {n_results} initial documents ---")
    results = collection.query(
        query_texts=[query],
        n_results=n_results
    )
    



    # Store documents as a list of strings for the reranker
    return results['documents'][0] if results['documents'] else []




# --- 2. NEW: Re-ranking Step with Cross-Encoder ---
def rerank_with_cross_encoder(query: str, documents: list[str], top_n: int = 3) -> list[str]:
    """
    Re-ranks a list of documents using a Cross-Encoder model and returns the top_n.
    """
    if not documents:
        return []

    print(f"--- Re-ranking {len(documents)} documents with Cross-Encoder ---")
    



    # Instantiate the Cross-Encoder model
    # ms-marco-MiniLM-L-6-v2 is a lightweight but powerful model trained for passage ranking.
    cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
    



    # The model expects pairs of [query, document]
    model_inputs = [[query, doc] for doc in documents]
    



    # Predict the relevance scores
    scores = cross_encoder.predict(model_inputs)
    



    # Combine documents with their scores and sort
    scored_docs = list(zip(documents, scores))
    scored_docs.sort(key=lambda x: x[1], reverse=True)
    



    # Return the content of the top N documents
    reranked_docs = [doc for doc, score in scored_docs[:top_n]]
    
    print(f"--- Top {top_n} re-ranked documents selected ---")
    return reranked_docs




# --- 3. REUSED: Synthesis Step ---
def synthesize_answer(query: str, context_docs: list[str]):
    """
    Generates a final answer based on the query and re-ranked context.
    """
    if not context_docs:
        return "I could not find any relevant information after re-ranking."
    
    context_str = "\n\n---\n\n".join(context_docs)
    
    prompt = f"""
    You are a helpful AI assistant. Your task is to provide a comprehensive answer to the user's question based on the provided, highly-relevant context.

    Context:
    {context_str}
    
    Question: {query}
    
    Answer:
    """
    
    print("\n--- Synthesizing final answer from re-ranked context ---")
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are an expert Q&A assistant."},
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message.content




# --- Run the full pipeline ---
# Make sure to populate your ChromaDB collection with some documents first!
user_query = "What is the difference between recursive and fixed-size chunking strategies?"




# 1. Retrieve a large set of candidate documents
initial_docs = retrieve_initial_documents(user_query, n_results=20)




# 2. Re-rank the documents and select the best ones
reranked_context = rerank_with_cross_encoder(initial_docs, user_query, top_n=3)




# 3. Synthesize the final answer using the refined context
final_answer = synthesize_answer(user_query, reranked_context)

print("\n\n--- FINAL ANSWER ---")
print(final_answer)

Conclusion

You've now implemented one of the most impactful optimizations for a RAG system. By adding a re-ranking stage, you've shifted your pipeline from a simple, single-stage process to a more sophisticated, two-stage funnel that balances speed with accuracy. This is a core pattern in modern information retrieval and a critical concept for any AI Engineer to master.

Key Takeaways:

• Recall vs. Precision: Initial retrieval is for recall (get everything), while re-ranking is for precision (pick the best).
• Bi- vs. Cross-Encoders: Bi-encoders are fast and scalable for initial retrieval because they process queries and documents independently. Cross-encoders are slower but more accurate for re-ranking because they analyze the query and document together.
• The Two-Stage Funnel: The standard production pattern is to use a bi-encoder to fetch a large candidate set (e.g., top 50-100) and a cross-encoder to re-rank that set to find the final top-k (e.g., 3-5).
• A Spectrum of Options: For a system design interview, you should be prepared to discuss the trade-offs between different re-ranking methods, including local cross-encoders, LLM-based scorers, and managed APIs like Cohere.

We have now added several powerful optimization techniques to our RAG pipeline. But a crucial question remains: how do we know if these changes are actually improving performance? In our next lesson, we will dive into the world of RAG evaluation, learning how to measure the quality of our retrieval system with metrics like context precision and context recall.