Skip to content

Latest commit

 

History

History
171 lines (124 loc) · 5.17 KB

README.md

File metadata and controls

171 lines (124 loc) · 5.17 KB

Hybrid Search (BM25 + Semantic Reranking)

This module implements a two-stage retrieval system that combines the strengths of BM25 and dense vector retrieval.

Overview

The hybrid search approach:

  1. Stage 1 (Retrieval): Use BM25 to retrieve top-K candidates (e.g., top-1000)
  2. Stage 2 (Reranking): Rerank the candidates using semantic similarity (dot product of embeddings)

This combines:

  • BM25's strength: Fast retrieval and exact keyword matching
  • Vector search's strength: Semantic similarity and understanding

Usage

Generate Hybrid Search Results

# For dev queries
python -m python.reranking.hybrid_search \
  --bm25-results results/bm25_dev.txt \
  --queries data/queries_dev_eval_embeddings.h5 \
  --output results/hybrid_dev.txt

# For TREC 2019 eval
python -m python.reranking.hybrid_search \
  --bm25-results results/bm25_eval_one.txt \
  --queries data/queries_dev_eval_embeddings.h5 \
  --output results/hybrid_eval_one.txt

# For TREC 2020 eval
python -m python.reranking.hybrid_search \
  --bm25-results results/bm25_eval_two.txt \
  --queries data/queries_dev_eval_embeddings.h5 \
  --output results/hybrid_eval_two.txt

Parameters

  • --bm25-results: Path to BM25 results file
  • --queries: Path to query embeddings h5 file
  • --output: Output file path for reranked results
  • --doc-embeddings: Path to document embeddings (default: ./data/collection_embeddings_subset.h5)
  • --top-k: Number of results to return after reranking (default: 100)
  • --run-name: Run identifier in output (default: HYBRID)

How It Works

1. Load BM25 Candidates

The script reads BM25 results (typically top-100 or top-1000 per query):

query_id\tpassage_id\trank\tscore\tBM25

2. Load Embeddings

  • Query embeddings from h5 file
  • Document embeddings from h5 file
  • Both are normalized to unit length for dot product similarity

3. Rerank Candidates

For each query:

  1. Get the query embedding
  2. For each BM25 candidate:
    • Get the document embedding
    • Compute similarity: score = query_emb · doc_emb
  3. Sort by semantic similarity (descending)
  4. Return top-K results

4. Write Results

Output format matches BM25 and vector search:

query_id\tpassage_id\trank\tscore\tHYBRID

Performance

Typical performance on 1M document collection:

  • Reranking speed: ~300-700 queries/sec
  • Memory: ~4GB (for 1M documents with 384-dim embeddings)
  • Time:
    • Dev (1000 queries): ~1.5s
    • Eval (43-54 queries): ~0.15s

Evaluation Results

Based on MS-MARCO passage ranking:

Dev Set (1000 queries, binary relevance)

System MRR@10 Recall@100 MAP
BM25 0.3946 0.7860 0.3963
Vector 0.5709 0.9091 0.5696
Hybrid 0.5766 0.8763 0.5736

TREC 2019 (43 queries, graded relevance)

System MRR@10 Recall@100 NDCG@10 NDCG@100
BM25 0.7863 0.5172 0.4578 0.5358
Vector 0.9535 0.5582 0.6973 0.6544
Hybrid 0.9651 0.5723 0.7019 0.6619

TREC 2020 (54 queries, graded relevance)

System MRR@10 Recall@100 NDCG@10 NDCG@100
BM25 0.8308 0.5248 0.4969 0.5209
Vector 0.9230 0.6167 0.6632 0.6397
Hybrid 0.9259 0.6070 0.6592 0.6332

Key Observations

  1. MRR@10: Hybrid performs best, showing excellent top-result quality

    • Hybrid improves MRR over BM25 by ~46% (dev), ~23% (TREC 2019), ~11% (TREC 2020)
    • Slight improvements over pure vector search
  2. Recall@100: Vector search has the best recall

    • Vector: Best at finding all relevant documents
    • Hybrid: Good balance (better than BM25, slightly lower than Vector)
    • BM25: Lower recall due to vocabulary mismatch issues
  3. NDCG: Hybrid shows strong ranking quality

    • Consistently outperforms BM25
    • Competitive with or slightly better than pure vector search
  4. Trade-offs:

    • BM25: Fast, interpretable, but lower accuracy
    • Vector: Best recall, but expensive to search 1M docs
    • Hybrid: Best overall accuracy, efficient (only reranks BM25 candidates)

Implementation Details

Similarity Computation

Uses normalized dot product (equivalent to cosine similarity):

similarity = query_emb · doc_emb  # Both normalized to unit length

Memory Optimization

  • Loads document embeddings once at startup
  • Builds passage ID → index mapping for fast lookup
  • Processes queries in batches with progress tracking

Efficiency

The hybrid approach is more efficient than pure vector search because:

  1. BM25 quickly narrows down candidates (e.g., 1000 out of 1M)
  2. Only computes semantic similarity for candidates
  3. Avoids expensive nearest neighbor search over all 1M documents

Dependencies

  • h5py - Reading embeddings
  • numpy - Numerical operations
  • tyro - CLI argument parsing
  • tqdm - Progress bars

Related Scripts

  • python/semantic_search/search_index.py - Pure vector search
  • python/evaluate_batch.py - Batch evaluation of all systems
  • python/filter_results.py - Filter results by query IDs