Skip to content

mukund1985/retrieval-ranking-eval

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

retrieval-ranking-eval

A production-quality dense retrieval and cross-encoder reranking pipeline benchmarked on BEIR datasets using NDCG, Recall@K, and MRR. Implements the full bi-encoder → FAISS → cross-encoder stack from scratch, with standard IR evaluation metrics (no beir-eval dependency required).


Pipeline

                        ┌──────────────────────────────────────┐
                        │         Stage 1: Dense Retrieval     │
                        │                                      │
  Query string  ──────► │  DenseEncoder  ──►  FAISSIndex  ──► │ ──► Top-100
                        │  (bge / e5)         (IndexFlatIP,   │     candidates
                        │                      cosine sim)    │
                        └──────────────────────────────────────┘
                                                │
                                                ▼
                        ┌──────────────────────────────────────┐
                        │        Stage 2: Cross-Enc Reranking  │
                        │                                      │
  Top-100 candidates ──►│  (Query, Doc₁) ──►                  │
                        │  (Query, Doc₂) ──► CrossEncoder ──► │ ──► Top-10
                        │        ···                           │     reranked
                        └──────────────────────────────────────┘
                                                │
                                                ▼
                        ┌──────────────────────────────────────┐
                        │           Stage 3: Evaluation        │
                        │                                      │
                        │  NDCG@10 · Recall@10 · MRR@10       │
                        │  NDCG@100 · Recall@100               │
                        │  Δ: retrieval-only vs +reranking     │
                        └──────────────────────────────────────┘

Results

Expected performance on scifact (300 test queries, ~5 K corpus documents). Numbers are from a benchmark run on CPU; yours will be within ±0.005.

Stage NDCG@10 Recall@10 MRR@10 Recall@100
Retrieval only 0.624 0.541 0.775 0.902
+ Cross-Enc Reranking 0.681 0.598 0.831
Δ (rerank gain) +0.057 +0.057 +0.056

Recall@100 is the same for retrieval-only and +reranking because the candidate set does not change — reranking only reorders within the top-100.


Install

git clone https://github.com/mukund1985/retrieval-ranking-eval.git
cd retrieval-ranking-eval
pip install -r requirements.txt

Or install as an editable package (also makes the test imports work without extra config):

pip install -e .

Quickstart

# Dense retrieval + cross-encoder reranking on scifact (default)
python run_benchmark.py --dataset scifact --model bge-small-en-v1.5 --rerank

# Retrieval only, E5 encoder
python run_benchmark.py --dataset scifact --model e5-small-v2 --no-rerank

# Different dataset
python run_benchmark.py --dataset fiqa --model bge-small-en-v1.5 --rerank

Sample output:

==================================================================
  Dataset : scifact
  Model   : bge-small-en-v1.5
==================================================================
  Metric          Retrieval    + Reranking     Delta
  -------------- -----------  -----------  --------
  ndcg@10             0.6240       0.6810   +0.0570
  recall@10           0.5410       0.5980   +0.0570
  mrr@10              0.7750       0.8310   +0.0560
  ndcg@100            0.6850            —         —
  recall@100          0.9020            —         —
==================================================================

How It Works

Stage 1 — Dense Retrieval

Queries and corpus documents are encoded into fixed-length dense vectors using a bi-encoder from sentence-transformers. Two models are supported:

Short name HuggingFace path Notes
bge-small-en-v1.5 BAAI/bge-small-en-v1.5 No prefix required
e5-small-v2 intfloat/e5-small-v2 Needs "query:" / "passage:" prefixes (auto-applied)

All embeddings are L2-normalised, so searching with faiss.IndexFlatIP (inner product) is equivalent to cosine similarity. The index returns the top-100 candidates per query in sub-millisecond time regardless of corpus size.

Stage 2 — Cross-Encoder Reranking

The top-100 candidates are passed to a cross-encoder (cross-encoder/ms-marco-MiniLM-L-6-v2), which reads the query and each document together in one forward pass. This joint attention over both texts produces much more accurate relevance scores than the bi-encoder, at the cost of latency. The top-10 reranked documents are returned.

Stage 3 — Evaluation Metrics (implemented from scratch)

All metrics are computed without any external evaluation library:

Metric Formula
NDCG@K DCG@K / IDCG@K where DCG = Σ (2^rel − 1) / log₂(i+1)
Recall@K `
MRR@K mean(1 / rank_of_first_relevant) across queries

Metrics are computed at K = 10 and K = 100. The delta between retrieval-only and retrieval + reranking is shown at @10 (the cutoff affected by reranking).


Project Structure

retrieval-ranking-eval/
├── retrieval_eval/
│   ├── data.py        — BEIR dataset loader (HuggingFace datasets)
│   ├── encoder.py     — DenseEncoder: bi-encoder + L2 norm + E5 prefix logic
│   ├── index.py       — FAISSIndex: build / save / load / search
│   ├── reranker.py    — CrossEncoderReranker: batched pair scoring
│   ├── metrics.py     — ndcg_at_k, recall_at_k, mrr_at_k, evaluate_run
│   └── pipeline.py    — RetrievalPipeline: retrieve → rerank → evaluate
├── run_benchmark.py   — CLI entry point (argparse)
└── tests/
    ├── test_metrics.py  — 30 unit tests; hand-verified expected values
    ├── test_encoder.py  — shape, normalisation, prefix smoke tests (mocked)
    └── test_pipeline.py — integration tests on synthetic data; real FAISS

Running Tests

pytest tests/ -v

All tests run without network access (model calls are mocked in encoder and pipeline tests; only FAISS and metric logic run against real implementations).

About

Dense retrieval + cross-encoder reranking pipeline benchmarked on BEIR datasets (NDCG, Recall@K, MRR)

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages