How to Implement Retrieval-Augmented Generation (RAG) with Re-Ranking

Basic RAG retrieves the top-k documents and passes them all to the language model. The problem: those top-k documents are ranked by cosine similarity of embeddings, which is a blunt instrument. You end up sending irrelevant chunks to the LLM, filling the context window with noise. Re-ranking fixes this by adding a second pass with a more expensive but more precise scoring model.

This post walks through building a RAG pipeline with re-ranking from scratch, using sentence-transformers and a cross-encoder.

Why Vector Search Alone Isn't Enough

Embedding-based retrieval converts your query and documents into dense vectors, then finds the nearest neighbors. It's fast and works at scale — but it captures semantic similarity, not relevance in the strict sense.

Consider a query: "How do I prevent SQL injection in Django?"

A bi-encoder will happily return documents about:

• SQL injection in Flask (similar, but different framework)
• Django ORM performance tips (related topic, wrong answer)
• Parameterized queries in Java (correct concept, wrong language)

All of these might score a cosine similarity of 0.78–0.82. None of them directly answer the question.

A cross-encoder, on the other hand, takes the query and each document together as a pair and produces a single relevance score. It's slower (it can't pre-compute document embeddings), but it's dramatically more accurate.

The Two-Stage Architecture

The standard pattern is two stages:

1. Retrieval: retrieve top-n candidates with a bi-encoder (fast, scales to millions of docs)
2. Re-ranking: score each candidate against the query with a cross-encoder, keep top-k (slower, but precise)

Typically n = 20–50 and k = 3–5. The cross-encoder never touches the full corpus — only the candidates the bi-encoder pre-selected.

query → bi-encoder → top-50 candidates → cross-encoder → top-5 → LLM

This keeps latency manageable while dramatically improving the quality of context sent to the model.

Building the Retrieval and Re-Ranking Pipeline

Install the dependencies:

pip install sentence-transformers faiss-cpu

Here is a minimal, working implementation:

import numpy as np
import faiss
from sentence_transformers import SentenceTransformer, CrossEncoder

# Load models once at startup
bi_encoder = SentenceTransformer("BAAI/bge-small-en-v1.5")
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")


def build_index(documents: list[str]):
    """Encode documents and build a FAISS index."""
    embeddings = bi_encoder.encode(documents, normalize_embeddings=True)
    dim = embeddings.shape[1]

    # IndexFlatIP + normalized vectors = cosine similarity
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings.astype(np.float32))

    return index


def retrieve_and_rerank(
    query: str,
    documents: list[str],
    index,
    top_n: int = 20,
    top_k: int = 5,
) -> list[str]:
    """Two-stage retrieval: bi-encoder candidates, then cross-encoder re-ranking."""
    # Stage 1: fast approximate retrieval
    query_vec = bi_encoder.encode([query], normalize_embeddings=True)
    _, indices = index.search(query_vec.astype(np.float32), top_n)
    candidates = [documents[i] for i in indices[0]]

    # Stage 2: precise re-ranking
    pairs = [(query, doc) for doc in candidates]
    scores = cross_encoder.predict(pairs)

    ranked = sorted(zip(scores, candidates), key=lambda x: x[0], reverse=True)
    return [doc for _, doc in ranked[:top_k]]

A few implementation notes:

• normalize_embeddings=True is required for IndexFlatIP to behave as cosine similarity
• cross_encoder.predict(pairs) processes all pairs in one batch — never call it in a loop
• For corpora larger than ~500k chunks, replace IndexFlatIP with IndexIVFFlat and tune nlist and nprobe

Adding the Generation Step

Once you have the re-ranked chunks, the generation step is straightforward:

from openai import OpenAI

client = OpenAI()


def generate_answer(query: str, context_chunks: list[str]) -> str:
    """Generate an answer from re-ranked context."""
    context = "\n\n---\n\n".join(context_chunks)

    system_prompt = (
        "You are a precise assistant. Answer based only on the provided context. "
        "If the context is insufficient, say so explicitly."
    )
    user_prompt = f"Context:\n{context}\n\nQuestion: {query}"

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=0.1,
    )

    return response.choices[0].message.content


def rag_pipeline(query: str, documents: list[str], index) -> str:
    """Full RAG pipeline: retrieve → re-rank → generate."""
    top_chunks = retrieve_and_rerank(query, documents, index, top_n=20, top_k=4)
    return generate_answer(query, top_chunks)

Keep temperature low (0.0–0.2) for retrieval-grounded tasks. High temperature defeats the point of precise retrieval by introducing hallucination at generation time.

Choosing the Right Cross-Encoder

Model	Size	Latency (20 pairs, CPU)	Quality
ms-marco-MiniLM-L-6-v2	22 MB	~50 ms	Good
ms-marco-MiniLM-L-12-v2	33 MB	~100 ms	Better
bge-reranker-base	278 MB	~400 ms	High
bge-reranker-v2-m3	570 MB	~900 ms	Best, multilingual

For most production use cases, ms-marco-MiniLM-L-6-v2 is the right default. The quality jump from L-6 to L-12 is modest; the jump from MiniLM to BGE reranker is where you see meaningful precision gains on domain-specific corpora.

If you're building a security-focused RAG system — querying CVE advisories, policy documents, or hardening guides — the precision improvement from re-ranking matters more than in general-purpose assistants. Returning the wrong vulnerability context is a worse failure mode than returning nothing. The security hardening checklists we publish are a good example of the kind of structured reference material that rewards precise retrieval.

What to Measure

Don't tune by feel. Set up an evaluation harness before you ship:

• Recall@k: does the correct document appear in the top-k re-ranked results?
• MRR (Mean Reciprocal Rank): where does the correct document rank on average?
• End-to-end answer accuracy: does the generated answer correctly address a labeled set of questions?

Run the evaluation on both the bi-encoder output and the cross-encoder output. The delta tells you exactly how much re-ranking is helping on your specific data. If it's under 5%, you may not need it.

The Takeaway

Re-ranking is the most impactful single improvement you can make to a RAG pipeline after basic retrieval is working. The two-stage architecture adds 50–150ms on CPU, and the precision gains are typically large enough to show up immediately in a properly structured eval.

Both models are open-source and run locally with no hosted dependency. The bi-encoder handles scale; the cross-encoder handles accuracy. Neither one alone is enough for production-quality RAG.

---

I run AYI NEDJIMI Consultants, a cybersecurity consulting firm. We publish free security hardening checklists — PDF and Excel.