DEV Community

Ayi NEDJIMI
Ayi NEDJIMI

Posted on

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

RAG solves a real problem: language models hallucinate when they don't have relevant context, and they can't access data beyond their training cutoff. But naive RAG — embed a query, retrieve the top-K chunks, feed them to the model — has a precision problem. The chunks you retrieve are often related to the query, not relevant. Re-ranking fixes this by scoring retrieved passages a second time, with a model that understands semantic relationships at a finer grain.

This article walks through building a RAG pipeline with re-ranking in Python, from embedding to final answer generation.

Why Naive RAG Falls Short

The standard retrieval step uses a bi-encoder: embed the query and all documents into the same vector space, then find the nearest neighbors by cosine similarity. This is fast enough to run at scale, but bi-encoders compress meaning into a fixed-size vector — they can't compare query and passage directly.

The result: you retrieve chunks that are topically similar but not actually useful. Ask "what are the access control requirements for SOC 2?" and you might get chunks that discuss HIPAA or ISO 27001 instead. Nearby in vector space, useless in practice.

Re-ranking addresses this with a cross-encoder: a model that takes the query and each candidate passage together as input and produces a relevance score. Cross-encoders are slower (you can't pre-compute embeddings), but they're significantly more accurate.

The standard pipeline becomes:

  1. Retrieve top-N candidates using a fast bi-encoder (N = 20–50)
  2. Re-rank those candidates with a cross-encoder
  3. Select top-K re-ranked passages (K = 3–5) and pass them to the language model

Setting Up the Environment

pip install sentence-transformers faiss-cpu openai
Enter fullscreen mode Exit fullscreen mode

We'll use sentence-transformers for both embedding and re-ranking, FAISS for the vector index, and any language model API for generation.

Building the Retrieval Layer

Index your documents with a bi-encoder. For this example, we're indexing a collection of security policy documents — the kind of structured knowledge base where precision matters.

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

bi_encoder = SentenceTransformer("BAAI/bge-small-en-v1.5")

documents = [
    "Access control policies must enforce least privilege for all users.",
    "Encryption at rest is required for all data classified as sensitive.",
    "Multi-factor authentication is mandatory for administrative accounts.",
    "Log retention must cover a minimum of 12 months for SOC 2 compliance.",
    "Vulnerability scans must be conducted at least quarterly.",
    # ... your actual documents
]

doc_embeddings = bi_encoder.encode(documents, normalize_embeddings=True)
dimension = doc_embeddings.shape[1]

index = faiss.IndexFlatIP(dimension)  # Inner product = cosine sim when normalized
index.add(doc_embeddings.astype(np.float32))

def retrieve(query: str, top_n: int = 20) -> list[tuple[str, float]]:
    q_emb = bi_encoder.encode([query], normalize_embeddings=True)
    scores, indices = index.search(q_emb.astype(np.float32), top_n)
    return [(documents[i], float(scores[0][j])) for j, i in enumerate(indices[0])]
Enter fullscreen mode Exit fullscreen mode

This retrieves 20 candidates quickly. Now we re-rank them.

Adding the Cross-Encoder Re-Ranker

The cross-encoder takes (query, passage) pairs and scores each one directly:

from sentence_transformers import CrossEncoder

cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

def rerank(
    query: str,
    candidates: list[tuple[str, float]],
    top_k: int = 5
) -> list[str]:
    pairs = [[query, doc] for doc, _ in candidates]
    scores = cross_encoder.predict(pairs)

    ranked = sorted(
        zip([doc for doc, _ in candidates], scores),
        key=lambda x: x[1],
        reverse=True,
    )
    return [doc for doc, _ in ranked[:top_k]]
Enter fullscreen mode Exit fullscreen mode

ms-marco-MiniLM-L-6-v2 is a solid default: 22M parameters, fast on CPU, and trained on the MS MARCO passage ranking dataset.

Full Pipeline: Retrieve → Re-Rank → Generate

import os
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

def rag_with_reranking(query: str) -> str:
    # Step 1: retrieve broad candidates
    candidates = retrieve(query, top_n=20)

    # Step 2: re-rank to a precise top-K
    top_passages = rerank(query, candidates, top_k=4)

    # Step 3: build grounded context and generate
    context = "\n\n".join(
        f"[{i+1}] {passage}" for i, passage in enumerate(top_passages)
    )
    prompt = (
        "Answer using only the provided context. "
        "If the answer is not there, say so explicitly.\n\n"
        f"Context:\n{context}\n\nQuestion: {query}"
    )
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1,
    )
    return response.choices[0].message.content

print(rag_with_reranking("What are the log retention requirements?"))
Enter fullscreen mode Exit fullscreen mode

Temperature is set to 0.1 deliberately — low creativity keeps the model grounded in retrieved content rather than extrapolating.

Measuring the Gain

Re-ranking adds latency: roughly 50–80ms on CPU for 20 candidates with the MiniLM model. Whether that trade-off is worth it depends on your context. In practice, re-ranking improves hit rate by 15–30% on most retrieval benchmarks.

A minimal evaluation harness:

def evaluate_pipeline(test_pairs: list[dict]) -> dict:
    """test_pairs: [{"query": str, "expected": str}]"""
    hits_base, hits_reranked = 0, 0

    for pair in test_pairs:
        query, expected = pair["query"], pair["expected"]

        # Baseline: top-5 from bi-encoder only
        base = [doc for doc, _ in retrieve(query, top_n=5)]
        hits_base += any(expected in doc for doc in base)

        # With re-ranking: retrieve 20, re-rank to 5
        candidates = retrieve(query, top_n=20)
        reranked = rerank(query, candidates, top_k=5)
        hits_reranked += any(expected in doc for doc in reranked)

    n = len(test_pairs)
    return {
        "hit_rate_baseline": hits_base / n,
        "hit_rate_reranked": hits_reranked / n,
    }
Enter fullscreen mode Exit fullscreen mode

Build a test set of 50–100 query/answer pairs from your actual data. Even a small evaluation corpus will surface where your chunking strategy breaks down.

Practical Notes

Chunk size affects re-ranking quality. Chunks shorter than 100 tokens lose context; longer than 512 and the cross-encoder degrades. 200–300 tokens with 20% overlap is a reasonable starting point.

Model alternatives. For multilingual corpora, swap bge-small-en-v1.5 for BAAI/bge-m3. For higher re-ranking accuracy at the cost of 4× latency, try cross-encoder/ms-marco-electra-base.

Don't strip metadata. If your documents have categories, dates, or source identifiers, include them in the chunk. Cross-encoders can use structural signals to break ties between semantically equivalent passages.

Where this pattern really pays off. Precision-critical retrieval — compliance documentation, access control policies, medical protocols — is where re-ranking earns its latency cost. If your RAG pipeline is querying a corpus of security hardening guidelines (the kind you'd find in a structured security checklist library), a hallucinated answer is actively dangerous. Get retrieval right before you trust the model's output.

The Takeaway

Retrieve many, re-rank to precision, generate from the best. The bi-encoder handles scale; the cross-encoder handles accuracy. You don't have to choose.

If your RAG pipeline produces plausible-but-imprecise answers, add a cross-encoder re-ranker before you try anything else — it's the highest-leverage improvement you can make without changing your architecture.


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

Top comments (0)