🔀 Reranking

Concept

A bi-encoder (the sentence-transformer used for indexing) is fast but approximate — it independently encodes the query and each document. A cross-encoder scores each (question, candidate) pair together, which is much more accurate but too slow to run over an entire corpus.

The trick: retrieve many cheap candidates with the bi-encoder (initial_k=25), then rerank the top ones with the cross-encoder to get the best top_k=3. This two-stage approach gives near-cross-encoder accuracy at bi-encoder speed.

Theory & Concepts

When you submit a knowledge base and a question to this Reranking demo, you are witnessing a powerful two-stage retrieval architecture designed to solve the precision issues of traditional RAG. The goal is to find the most relevant chunks of your text to answer the question, balancing speed and pinpoint accuracy.

In the first stage, the demo uses a Bi-Encoder to rapidly sift through the knowledge base. The bi-encoder processes your question and each chunk of the knowledge base independently to create vector embeddings in Chroma. It then performs a fast similarity search to retrieve a broad initial set of candidates (in this code, initial_k=25). While fast, this method is approximate. Because the question and text chunks are encoded separately, the bi-encoder might miss subtle semantic connections, occasionally ranking a highly relevant document lower than it should.

To fix this, the demo introduces a second stage: the Cross-Encoder (specifically, ms-marco-MiniLM-L6-v2). The backend takes the 25 candidates and pairs each one with your original question. The cross-encoder evaluates these pairs together simultaneously, allowing its attention mechanisms to analyze the deep contextual interactions between the specific words in your question and the text chunk. This generates a highly accurate relevance score for each pair. Finally, the code sorts the candidates by these new scores and returns the absolute best 3 (top_k=3) to the browser. By applying the expensive cross-encoder only to a small subset of fast bi-encoder results, the demo achieves near-perfect accuracy without sacrificing performance.

Request flow

Code flow

Backend

from sentence_transformers import CrossEncoder
from langchain_chroma import Chroma

from config import get_embedder

# Lazy-loaded cross-encoder — initialized on first use to prevent blocking Flask startup
_reranker = None


def _get_reranker():
    global _reranker
    if _reranker is None:
        _reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")
    return _reranker


def retrieve_with_rerank(
    db: Chroma,
    question: str,
    top_k: int = 3,
    initial_k: int = 25,
) -> list:
    """Retrieve candidates with bi-encoder, rerank with cross-encoder."""
    # Stage 1: Fast bi-encoder vector search retrieving a broad candidate pool (top 25)
    candidates = db.similarity_search(question, k=initial_k)
    if not candidates:
        return []

    # Stage 2: Initialize cross-encoder and build (query, passage) pairs
    reranker = _get_reranker()
    pairs = [(question, c.page_content) for c in candidates]

    # Stage 3: Cross-encoder jointly attends over query and chunk tokens to produce accurate relevance scores
    scores = reranker.predict(pairs)

    # Stage 4: Sort candidate documents by cross-encoder score in descending order
    # (Sort by score only to prevent comparison errors on tie scores between Document instances)
    ranked = sorted(zip(scores, candidates), key=lambda pair: pair[0], reverse=True)

    # Stage 5: Return top-k highest scoring chunks
    return [c for _, c in ranked[:top_k]]

API route

@bp.route("/rerank", methods=["POST"])
def rerank_route():
    """Retrieve and rerank results from a user-provided knowledge base.

    Request body (JSON): ``{ "knowledge_base": "<text>", "question": "<text>" }``
    Response (JSON):     ``{ "result": { "results": ["...", ...] } }``
    Error response:      ``{ "error": "<message>" }`` with HTTP 400 or 500
    """
    # Step 1: Parse request JSON payload
    data = request.get_json(force=True)
    knowledge = (data.get("knowledge_base") or "").strip()
    question = (data.get("question") or "").strip()

    # Step 2: Validate knowledge base and question inputs
    validation = validate_textarea(knowledge, "Knowledge base")
    if validation:
        return validation
    if not question:
        return jsonify({"error": "A question is required."}), 400

    # Step 3: Run two-stage retrieval (bi-encoder retrieval followed by cross-encoder reranking)
    try:
        docs = [line.strip() for line in knowledge.splitlines() if line.strip()]
        db, _ = build_index(docs, persist_directory=None)
        reranked = retrieve_with_rerank(db, question, top_k=3)  # Select top 3 after cross-encoder reranking
        results = [doc.page_content for doc in reranked]
        return jsonify({"result": {"results": results}})
    except Exception as e:
        # Return HTTP 500 on unexpected errors
        return jsonify({"error": str(e)}), 500