Module 4: Re-ranking — the second stage that turns mediocre retrieval into excellent retrieval
Capsule 08: Capstone project — a re-ranking system with A/B testing
Project description
This is the close of the module. You're going to build a production-ready re-ranking system that takes everything you learned in capsules 02-07 and integrates it into a working pipeline with observability. This isn't a tutorial — it's the project that goes into your portfolio or your real codebase.
The system implements the three techniques (cross-encoder, LLM, Cohere) behind a unified interface, runs an A/B test on your own eval set, measures precision/recall/latency, and produces a report that justifies with data which technique to deploy.
By the end of this project you'll have:
- ✅ A re-ranking system with a reusable
Rerankerinterface - ✅ Three implementations (local cross-encoder, LLM, Cohere) with fallback
- ✅ An eval set built on your corpus with ground truth
- ✅ An A/B testing script that compares the techniques
- ✅ A comparison report that justifies the final decision
- ✅ An updated RAG pipeline with the winning technique integrated
Estimated time: 2-3 hours for implementation + 1 hour for analysis.
Project architecture
reranking_project/
├── src/
│ ├── rerankers/
│ │ ├── __init__.py
│ │ ├── base.py # Reranker interface
│ │ ├── cross_encoder.py # Local cross-encoder
│ │ ├── llm_based.py # LLM with structured outputs
│ │ ├── cohere_managed.py # Cohere Rerank API
│ │ └── factory.py # Factory that builds rerankers
│ ├── retrieval/
│ │ ├── pipeline.py # Pipeline with re-ranking
│ │ └── chromadb_setup.py
│ └── evaluation/
│ ├── eval_set.py # Building and loading the eval set
│ ├── metrics.py # Precision, Recall, MRR, latency
│ └── ab_test.py # A/B testing across rerankers
├── eval_data/
│ └── golden_set.json # Eval set with ground truth
├── benchmarks/
│ ├── run_benchmark.py
│ └── reports/ # Generated reports
├── .env
└── requirements.txt
Step 1: the unified interface
Create an abstraction that lets you swap rerankers without touching the rest of the code.
# src/rerankers/base.py
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import List
@dataclass
class RerankResult:
document: str
score: float
original_index: int
class Reranker(ABC):
"""Unified interface for every reranker."""
@abstractmethod
def rerank(
self,
query: str,
documents: List[str],
top_k: int = 5,
) -> List[RerankResult]:
"""Re-ranks documents by relevance to the query."""
pass
@property
@abstractmethod
def name(self) -> str:
"""Identifying name of the reranker."""
passStep 2: the concrete implementations
Cross-encoder (capsule 03, reused)
# src/rerankers/cross_encoder.py
from sentence_transformers import CrossEncoder
from .base import Reranker, RerankResult
class CrossEncoderReranker(Reranker):
"""Local re-ranker with a sentence-transformers cross-encoder."""
def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-12-v2"):
self.model = CrossEncoder(model_name)
self._model_name = model_name
@property
def name(self) -> str:
return f"cross_encoder_{self._model_name.split('/')[-1]}"
def rerank(self, query, documents, top_k=5):
if not documents:
return []
pairs = [(query, doc) for doc in documents]
scores = self.model.predict(pairs, batch_size=32, show_progress_bar=False)
indexed = list(enumerate(zip(documents, scores)))
sorted_results = sorted(indexed, key=lambda x: -x[1][1])
return [
RerankResult(document=doc, score=float(score), original_index=idx)
for idx, (doc, score) in sorted_results[:top_k]
]LLM-based (capsule 04)
# src/rerankers/llm_based.py
from openai import OpenAI
from pydantic import BaseModel, Field
from concurrent.futures import ThreadPoolExecutor
import os
from .base import Reranker, RerankResult
class RelevanceScore(BaseModel):
score: float = Field(ge=0.0, le=10.0)
class LLMReranker(Reranker):
"""LLM-based re-ranker (premium, expensive)."""
SYSTEM_PROMPT = """You are an expert relevance evaluator.
Score how well the document answers the query (0-10).
Be strict. Most docs should score 4-7. Reserve 9-10 for excellent matches.
Return JSON with score field."""
def __init__(self, model: str = "gpt-4o-mini", workers: int = 5):
self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
self.model = model
self.workers = workers
@property
def name(self) -> str:
return f"llm_{self.model}"
def _score_pair(self, args):
idx, query, doc = args
try:
response = self.client.beta.chat.completions.parse(
model=self.model,
messages=[
{"role": "system", "content": self.SYSTEM_PROMPT},
{"role": "user", "content": f'Query: {query}\n\nDocument:\n{doc[:1500]}'},
],
response_format=RelevanceScore,
temperature=0.1,
)
return idx, response.choices[0].message.parsed.score
except Exception as e:
print(f"LLM rerank failed for doc {idx}: {e}")
return idx, 0.0
def rerank(self, query, documents, top_k=5):
args = [(i, query, doc) for i, doc in enumerate(documents)]
with ThreadPoolExecutor(max_workers=self.workers) as executor:
scored = list(executor.map(self._score_pair, args))
scored.sort(key=lambda x: -x[1])
return [
RerankResult(document=documents[idx], score=float(score), original_index=idx)
for idx, score in scored[:top_k]
]Cohere (capsule 05)
# src/rerankers/cohere_managed.py
import cohere
import os
from .base import Reranker, RerankResult
class CohereReranker(Reranker):
"""Managed re-ranker via the Cohere API."""
def __init__(self, model: str = "rerank-multilingual-v3"):
self.co = cohere.Client(api_key=os.getenv("COHERE_API_KEY"))
self.model = model
@property
def name(self) -> str:
return f"cohere_{self.model}"
def rerank(self, query, documents, top_k=5):
if not documents:
return []
response = self.co.rerank(
model=self.model,
query=query,
documents=documents,
top_n=top_k,
)
return [
RerankResult(
document=documents[r.index],
score=r.relevance_score,
original_index=r.index,
)
for r in response.results
]Factory
# src/rerankers/factory.py
from .cross_encoder import CrossEncoderReranker
from .llm_based import LLMReranker
from .cohere_managed import CohereReranker
def get_reranker(reranker_type: str):
"""Factory that builds a reranker from its type."""
if reranker_type == "cross_encoder":
return CrossEncoderReranker()
elif reranker_type == "cross_encoder_multilingual":
return CrossEncoderReranker(model_name="cross-encoder/mmarco-mMiniLMv2-L12-H384-v1")
elif reranker_type == "llm":
return LLMReranker()
elif reranker_type == "cohere":
return CohereReranker()
elif reranker_type == "none":
return None # baseline without reranking
else:
raise ValueError(f"Unknown reranker type: {reranker_type}")Step 3: the eval set with ground truth
The eval set is the most important part of the project. Without valid ground truth, the benchmarks are blind.
Structure of the eval set
# eval_data/golden_set.json (format)
[
{
"query": "How do I implement OAuth2 authentication in FastAPI?",
"expected_doc_ids": ["doc_42", "doc_87", "doc_103"],
"category": "auth",
"language": "en",
"difficulty": "medium"
},
{
"query": "¿Cómo configuro PostgreSQL con SQLAlchemy?",
"expected_doc_ids": ["doc_215", "doc_301"],
"category": "database",
"language": "es",
"difficulty": "easy"
},
...
]How to build it
# src/evaluation/eval_set.py
import json
from pathlib import Path
from dataclasses import dataclass, asdict
from typing import List
@dataclass
class EvalQuery:
query: str
expected_doc_ids: List[str]
category: str = "general"
language: str = "en"
difficulty: str = "medium"
def load_eval_set(path: str = "eval_data/golden_set.json") -> List[EvalQuery]:
"""Loads the eval set from JSON."""
with open(path) as f:
data = json.load(f)
return [EvalQuery(**item) for item in data]
def save_eval_set(queries: List[EvalQuery], path: str):
"""Saves the eval set to JSON."""
with open(path, "w") as f:
json.dump([asdict(q) for q in queries], f, indent=2)
def build_eval_set_interactive(corpus_collection, n_queries: int = 50):
"""
Build the eval set interactively.
For each candidate query, the operator validates which docs are relevant.
"""
eval_set = []
candidate_queries = [
# Use real queries from the production log
"How do I implement OAuth2 in FastAPI?",
# ...
][:n_queries]
for query in candidate_queries:
# Retrieve the top-20 with cosine
results = corpus_collection.query(query_texts=[query], n_results=20)
print(f"\n{'='*80}")
print(f"Query: {query}")
print(f"{'='*80}")
print("Mark the relevant ones with [r]:")
relevant_ids = []
for i, (doc, doc_id) in enumerate(zip(results['documents'][0], results['ids'][0])):
print(f"\n[{i}] {doc_id}")
print(f" {doc[:200]}...")
mark = input(" Relevant? (r/n/q to quit): ").strip().lower()
if mark == "r":
relevant_ids.append(doc_id)
elif mark == "q":
break
if relevant_ids:
eval_set.append(EvalQuery(
query=query,
expected_doc_ids=relevant_ids,
category="manual",
language="en",
))
save_eval_set(eval_set, "eval_data/golden_set.json")
return eval_setRecommendation: build 50-100 queries with manual ground truth. It's worth the 2-4 hours.
Step 4: metrics and the A/B test
# src/evaluation/metrics.py
from typing import List
from src.rerankers.base import Reranker, RerankResult
from src.evaluation.eval_set import EvalQuery
import time
import statistics
def precision_at_k(retrieved_ids: List[str], expected_ids: List[str], k: int) -> float:
"""% of the top-K that are relevant."""
top_k = retrieved_ids[:k]
relevant_in_top_k = sum(1 for doc_id in top_k if doc_id in expected_ids)
return relevant_in_top_k / k
def recall_at_k(retrieved_ids: List[str], expected_ids: List[str], k: int) -> float:
"""% of the relevant docs that show up in the top-K."""
if not expected_ids:
return 0.0
top_k = retrieved_ids[:k]
relevant_in_top_k = sum(1 for doc_id in top_k if doc_id in expected_ids)
return relevant_in_top_k / len(expected_ids)
def mean_reciprocal_rank(retrieved_ids: List[str], expected_ids: List[str]) -> float:
"""1 / position of the first relevant doc. 0 if none is there."""
for i, doc_id in enumerate(retrieved_ids, start=1):
if doc_id in expected_ids:
return 1.0 / i
return 0.0
def evaluate_reranker(
reranker: Reranker,
eval_set: List[EvalQuery],
collection,
initial_n: int = 20,
final_top_k: int = 5,
) -> dict:
"""Evaluates one reranker over the full eval set."""
precisions = []
recalls = []
mrrs = []
latencies = []
for item in eval_set:
# Initial retrieval
results = collection.query(query_texts=[item.query], n_results=initial_n)
candidates = results['documents'][0]
candidate_ids = results['ids'][0]
# Re-rank
start = time.perf_counter()
if reranker is None:
# Baseline without re-rank
reranked_ids = candidate_ids[:final_top_k]
else:
reranked = reranker.rerank(item.query, candidates, top_k=final_top_k)
reranked_ids = [candidate_ids[r.original_index] for r in reranked]
elapsed = (time.perf_counter() - start) * 1000
# Metrics
precisions.append(precision_at_k(reranked_ids, item.expected_doc_ids, k=final_top_k))
recalls.append(recall_at_k(reranked_ids, item.expected_doc_ids, k=final_top_k))
mrrs.append(mean_reciprocal_rank(reranked_ids, item.expected_doc_ids))
latencies.append(elapsed)
latencies.sort()
return {
"precision_at_k": statistics.mean(precisions),
"recall_at_k": statistics.mean(recalls),
"mrr": statistics.mean(mrrs),
"latency_p50_ms": latencies[len(latencies) // 2],
"latency_p95_ms": latencies[int(len(latencies) * 0.95)],
"n_queries": len(eval_set),
}# src/evaluation/ab_test.py
from src.rerankers.factory import get_reranker
from src.evaluation.metrics import evaluate_reranker
from src.evaluation.eval_set import load_eval_set
def ab_test_all_rerankers(collection):
"""Compares every reranker over the same eval set."""
eval_set = load_eval_set()
print(f"Eval set: {len(eval_set)} queries")
rerankers_to_test = [
("none (baseline)", "none"),
("cross_encoder", "cross_encoder"),
("llm", "llm"),
("cohere", "cohere"),
]
results = {}
for label, reranker_type in rerankers_to_test:
print(f"\nEvaluating {label}...")
reranker = get_reranker(reranker_type)
result = evaluate_reranker(reranker, eval_set, collection)
results[label] = result
print(f" Precision@5: {result['precision_at_k']:.2%}")
print(f" Recall@5: {result['recall_at_k']:.2%}")
print(f" MRR: {result['mrr']:.3f}")
print(f" Latency p95: {result['latency_p95_ms']:.0f}ms")
return results
def print_comparison_report(results: dict):
"""Prints the comparison table."""
print("\n" + "="*80)
print("AB TEST REPORT")
print("="*80)
headers = ["Reranker", "P@5", "R@5", "MRR", "p50 ms", "p95 ms"]
print(f"\n{headers[0]:<25} {headers[1]:>8} {headers[2]:>8} {headers[3]:>8} {headers[4]:>10} {headers[5]:>10}")
print("-"*80)
for label, metrics in results.items():
print(f"{label:<25} {metrics['precision_at_k']:>7.1%} {metrics['recall_at_k']:>7.1%} {metrics['mrr']:>8.3f} {metrics['latency_p50_ms']:>9.0f} {metrics['latency_p95_ms']:>9.0f}")
# Identify the winner on each metric
print("\nBest on each metric:")
metrics_to_check = ["precision_at_k", "recall_at_k", "mrr"]
for metric in metrics_to_check:
best = max(results.items(), key=lambda x: x[1][metric])
print(f" {metric}: {best[0]} ({best[1][metric]:.3f})")
fastest = min(results.items(), key=lambda x: x[1]["latency_p95_ms"])
print(f" fastest p95: {fastest[0]} ({fastest[1]['latency_p95_ms']:.0f}ms)")Step 5: run it and report
# benchmarks/run_benchmark.py
from src.evaluation.ab_test import ab_test_all_rerankers, print_comparison_report
from src.retrieval.chromadb_setup import get_collection
import json
from datetime import datetime
def main():
collection = get_collection()
results = ab_test_all_rerankers(collection)
print_comparison_report(results)
# Save the report for auditing
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
report_path = f"benchmarks/reports/ab_test_{timestamp}.json"
with open(report_path, "w") as f:
json.dump(results, f, indent=2)
print(f"\nReport saved: {report_path}")
if __name__ == "__main__":
main()Expected output:
Eval set: 80 queries
Evaluating none (baseline)...
Precision@5: 71.50%
Recall@5: 63.20%
MRR: 0.587
Latency p95: 195ms
Evaluating cross_encoder...
Precision@5: 89.20%
Recall@5: 78.40%
MRR: 0.812
Latency p95: 348ms
Evaluating llm...
Precision@5: 92.80%
Recall@5: 81.50%
MRR: 0.849
Latency p95: 1480ms
Evaluating cohere...
Precision@5: 90.40%
Recall@5: 79.10%
MRR: 0.823
Latency p95: 312ms
================================================================================
AB TEST REPORT
================================================================================
Reranker P@5 R@5 MRR p50 ms p95 ms
--------------------------------------------------------------------------------
none (baseline) 71.5% 63.2% 0.587 165 195
cross_encoder 89.2% 78.4% 0.812 285 348
llm 92.8% 81.5% 0.849 1180 1480
cohere 90.4% 79.1% 0.823 265 312
Best on each metric:
precision_at_k: llm (0.928)
recall_at_k: llm (0.815)
mrr: llm (0.849)
fastest p95: none (baseline) (195ms)
Step 6: analysis and the final decision
With the benchmark data in hand, apply the decision framework from capsule 07:
# Decision: Cross-encoder
## Key A/B test results
- **The cross-encoder improves precision by +17.7 points** over the baseline (71.5% → 89.2%).
- **The LLM rerank improves it by +21.3 points** but with +1100ms extra and +$200/month in cost.
- **Cohere improves it by +18.9 points** at a cost similar to the LLM.
## Why the cross-encoder wins in this case
1. **Cost:** $0 vs $200/month (LLM) or $30/month (Cohere).
2. **Latency:** 348ms p95 vs 1480ms (LLM) — decisive for a <500ms UX.
3. **The quality gap vs the LLM is only 3.6%** — it doesn't justify the extra cost.
4. **The corpus is mostly English** — the MS MARCO cross-encoder works well here.
## When I'd migrate to Cohere or the LLM
- **To Cohere:** if we add significant multilingual content (>20% non-English).
- **To the LLM:** if compliance demands precision >92% and the $200/month is acceptable.
- **To the LLM in a cascade:** for the 10% of queries flagged "high stakes".
## Risks identified
- The cross-encoder has a cold start of ~5s. Mitigation: load the model at startup.
- p95 latency goes from 195ms to 348ms. Within the current SLA (<500ms) but the margin is tight.Step 7: integration into the RAG pipeline
# src/retrieval/pipeline.py (updated with re-ranking)
from src.rerankers.factory import get_reranker
from src.retrieval.chromadb_setup import get_collection
class RAGPipeline:
def __init__(self, reranker_type: str = "cross_encoder"):
self.collection = get_collection()
self.reranker = get_reranker(reranker_type)
def retrieve(self, query: str, top_k: int = 5):
"""The full pipeline with re-ranking."""
# Stage 1: broad retrieval
results = self.collection.query(query_texts=[query], n_results=20)
candidates = results['documents'][0]
candidate_ids = results['ids'][0]
# Stage 2: re-rank (if enabled)
if self.reranker is None:
return candidates[:top_k], candidate_ids[:top_k]
reranked = self.reranker.rerank(query, candidates, top_k=top_k)
reranked_ids = [candidate_ids[r.original_index] for r in reranked]
reranked_docs = [r.document for r in reranked]
return reranked_docs, reranked_idsProject delivery checklist
Before you consider the project closed, verify:
- The 3 techniques implemented behind the
Rerankerinterface. - An eval set of 50+ queries with manually validated ground truth.
- A
run_benchmark.pyscript that runs the A/B test. - A comparison report with a metrics table and a justified decision.
- An updated RAG pipeline with the winning technique integrated.
- A README with setup, how to run the benchmark, and the final decision.
- Environment variables documented (
OPENAI_API_KEY,COHERE_API_KEY). - Tests for each reranker (sanity checks at minimum).
Optional extensions
Once you finish the base project, consider these:
Extension 1: a cross-encoder + LLM cascade
For "high stakes" queries:
class CascadeReranker(Reranker):
def __init__(self):
self.cross_encoder = CrossEncoderReranker()
self.llm = LLMReranker()
def rerank(self, query, documents, top_k=5):
# Stage A: the cross-encoder filters down to the top-15
cross_top_15 = self.cross_encoder.rerank(query, documents, top_k=15)
# Stage B: the LLM refines down to the top-K
cross_docs = [r.document for r in cross_top_15]
return self.llm.rerank(query, cross_docs, top_k=top_k)Extension 2: caching results
For repeated queries, cache the re-ranking results (covered in capsule 06).
Extension 3: routing by query type
Detect the query type and route it to the right reranker (simple queries → cross-encoder, critical queries → LLM).
Extension 4: monitoring in production
Log every rerank with a trace_id, the scores, the latency. Build a dashboard that shows the score distribution.
Traps and common mistakes in the project
Trap 1: a small or biased eval set
If you build the eval set with 10 queries, the benchmarks are statistical noise. Minimum 50, ideally 100+.
Trap 2: measuring precision only
Rerankers can improve precision at the expense of recall (or the other way around). Report both AND latency.
Trap 3: not measuring variability
A single benchmark run can give misleading results by chance. Run it 3 times and report the median of medians.
Trap 4: poor ground truth
If the ground truth was annotated by a single person in 10 minutes, the benchmarks reflect that bias. Use multiple annotators whenever you can.
Trap 5: comparing against an unfair baseline
If the baseline is "no re-rank but with optimized queries" vs "with re-rank but no optimization", the comparison is invalid.
Trap 6: deploying the winner without an A/B in production
Even if the benchmark says X wins, run an A/B test in production for 1-2 weeks before a 100% rollout.
Project summary
You built:
- ✅ A modular re-ranking system with a unified interface
- ✅ Three implementations tested side by side
- ✅ A robust eval set with ground truth
- ✅ A reproducible benchmark
- ✅ A report that justifies the decision with data
- ✅ An updated RAG pipeline, ready to deploy
Typical improvement expected with re-ranking integrated:
Metric No rerank With rerank Gain
Precision@5 71% 89% +18 pts
Recall@5 63% 78% +15 pts
MRR 0.587 0.812 +0.225
That closes the re-ranking module and sets you up for M05 (Hybrid Search), where you combine the re-ranking gain with BM25 for queries with exact keywords.
Next steps
Right away:
- Run the full benchmark over your eval set.
- Decide the final technique based on the results.
- Integrate it into the RAG pipeline and deploy to staging.
- A/B test in production for 1-2 weeks.
Module 5 (Hybrid Search):
Re-ranking refines the candidates, but it assumes the initial retrieval actually recovered them. If your corpus has technical content with exact identifiers (function names, error codes), the initial retrieval with cosine can miss relevant docs. M05 covers how to add BM25 to the initial retrieval so you don't lose those cases.
Resources
- Sentence Transformers — Cross-Encoders
- Cohere Rerank Documentation
- OpenAI Structured Outputs
- BEIR Benchmark
- LangChain RAG Evaluation
- Anthropic — Contextual Retrieval
Estimated time: 2-3 hours + 1 hour of analysis Next module: Module 5 — Hybrid Search