Query Expansion & Reranking: Two-Stage Retrieval

Description

In this capsule you'll implement Query Expansion and Cross-Encoder Reranking to build a two-stage retrieval system that maximizes both recall and precision.

Two-Stage Retrieval:

1. Stage 1 (FAISS): Retrieves many candidates (k=20-50) quickly
2. Stage 2 (Reranker): Reorders the candidates with a more precise model, returns the final top-K

This approach is standard in production: FAISS gives high recall (~95%), the reranker improves final precision (~10-20% improvement in nDCG).

Estimated duration: 40-50 minutes

---

Objectives

By completing this capsule, you'll be able to:

• ✅ Implement query expansion with embedding arithmetic
• ✅ Use a cross-encoder for reranking
• ✅ Implement a two-stage retrieval pipeline
• ✅ Compare single-stage vs two-stage
• ✅ Optimize the reranker latency

---

Part 1: Query Expansion

Why Query Expansion?

Problem: A short query may not capture the full intent

query = "Python install"
# The embedding of "Python install" only captures those terms

Solution: Expand the query with related terms

query = "Python install"
expansions = ["Python installation", "setup Python", "install Python macOS"]
# Combine embeddings for a richer query

---

1.1: Implementation with embedding arithmetic

Create src/retrieval/query_expander.py:

"""
Query Expander - RAG System
Expands queries for better recall
"""

import numpy as np
from typing import List
import logging


class QueryExpander:
    """
    Expands queries using embedding arithmetic
    
    Strategies:
    - Synonym expansion: Add synonyms
    - Concept expansion: Add related terms
    - Weighted blend: Combine the original query with expansions
    
    Example:
        expander = QueryExpander(embedder)
        
        expanded_emb = expander.expand(
            query="Python install",
            expansions=["setup Python", "Python installation"]
        )
    """
    
    def __init__(self, embedder, alpha: float = 0.7):
        """
        Initialize QueryExpander
        
        Args:
            embedder: EmbeddingClient instance
            alpha: Weight of the original query (0-1)
                - alpha=1.0: Only the original query
                - alpha=0.5: 50% query, 50% expansions
                - alpha=0.7: Recommended (70% query, 30% expansions)
        """
        self.embedder = embedder
        self.alpha = alpha
        self.logger = logging.getLogger(__name__)
    
    def expand(
        self,
        query: str,
        expansions: List[str],
        alpha: float = None
    ) -> np.ndarray:
        """
        Expand the query with additional terms
        
        Args:
            query: Original query
            expansions: List of expansion terms
            alpha: Alpha override (optional)
        
        Returns:
            Expanded embedding
        """
        if alpha is None:
            alpha = self.alpha
        
        # Embedding of the original query
        query_emb = self.embedder.embed(query)
        
        if not expansions:
            return query_emb
        
        # Embeddings of the expansions
        expansion_embs = [self.embedder.embed(term) for term in expansions]
        expansion_avg = np.mean(expansion_embs, axis=0)
        
        # Weighted blend
        expanded = alpha * query_emb + (1 - alpha) * expansion_avg
        
        # Normalize (important for FAISS dot product)
        expanded = expanded / np.linalg.norm(expanded)
        
        self.logger.info(
            f"Expanded query with {len(expansions)} terms (alpha={alpha})"
        )
        
        return expanded
    
    def auto_expand(
        self,
        query: str,
        index,
        k_initial: int = 10,
        n_expansions: int = 3
    ) -> np.ndarray:
        """
        Auto-expansion using the initial top-K results
        
        Args:
            query: Original query
            index: FAISS index
            k_initial: How many docs to use for expansion
            n_expansions: How many terms to extract
        
        Returns:
            Expanded embedding
        """
        # Initial search
        query_emb = self.embedder.embed(query)
        results = index.search(query_emb, k=k_initial)
        
        # Extract terms from the top results (simplified)
        # In production: use TF-IDF, entity extraction, etc.
        expansions = []
        for r in results[:n_expansions]:
            text = r['chunk'].text
            # Take the first words as a pseudo-expansion
            words = text.split()[:5]
            expansions.append(' '.join(words))
        
        return self.expand(query, expansions)


# Demo
if __name__ == "__main__":
    from src.embeddings.embedding_client import EmbeddingClient
    
    embedder = EmbeddingClient()
    expander = QueryExpander(embedder, alpha=0.7)
    
    # Expand the query
    query = "Python installation"
    expansions = [
        "install Python on macOS",
        "Python setup guide",
        "download Python"
    ]
    
    expanded_emb = expander.expand(query, expansions)
    print(f"✅ Expanded embedding: {expanded_emb.shape}")

---

Part 2: Cross-Encoder Reranking

Why Reranking?

Problem: Bi-encoders (FAISS) encode query and docs separately

# Bi-encoder (FAISS):
query_emb = embed(query)      # Separate
doc_emb = embed(document)     # Separate
score = cosine(query_emb, doc_emb)

Advantage: Very fast (pre-computed doc embeddings) Disadvantage: Doesn't see the query-document interaction

---

Solution: A cross-encoder sees the query + document together

# Cross-encoder (reranker):
score = model([query, document])  # Together!
# Sees the direct interaction → more precise

Trade-off:

• Bi-encoder: 1000x faster, ~95% precision
• Cross-encoder: 1000x slower, ~100% precision

---

2.1: Reranker implementation

Create src/retrieval/reranker.py:

"""
Reranker - RAG System
Cross-encoder reranking for two-stage retrieval
"""

from sentence_transformers import CrossEncoder
from typing import List, Dict
import logging


class Reranker:
    """
    Cross-encoder reranker
    
    Models:
    - ms-marco-MiniLM-L-6-v2: Fast (40ms/pair), good quality
    - ms-marco-electra-base: Slower (100ms/pair), best quality
    
    Example:
        reranker = Reranker()
        
        # Stage 1: FAISS retrieve 20 candidates
        candidates = index.search(query_emb, k=20)
        
        # Stage 2: Rerank to top-5
        final_results = reranker.rerank(query, candidates, top_k=5)
    """
    
    def __init__(
        self,
        model_name: str = 'cross-encoder/ms-marco-MiniLM-L-6-v2'
    ):
        """
        Initialize Reranker
        
        Args:
            model_name: Name of the cross-encoder model
        """
        self.model_name = model_name
        self.model = CrossEncoder(model_name)
        self.logger = logging.getLogger(__name__)
    
    def rerank(
        self,
        query: str,
        results: List[Dict],
        top_k: int = 5
    ) -> List[Dict]:
        """
        Rerank results with a cross-encoder
        
        Args:
            query: User query
            results: List of results from FAISS
            top_k: Number of final results
        
        Returns:
            List of reranked results
        """
        if not results:
            return []
        
        # Prepare (query, document) pairs
        pairs = []
        for r in results:
            doc_text = r['chunk'].text
            pairs.append([query, doc_text])
        
        # Score with the cross-encoder
        scores = self.model.predict(pairs)
        
        # Combine results with the new scores
        reranked = []
        for r, score in zip(results, scores):
            r_copy = r.copy()
            r_copy['cross_encoder_score'] = float(score)
            r_copy['original_score'] = r['score']
            reranked.append(r_copy)
        
        # Reorder by cross-encoder score
        reranked.sort(key=lambda x: x['cross_encoder_score'], reverse=True)
        
        self.logger.info(
            f"Reranked {len(results)} candidates, returning top-{top_k}"
        )
        
        return reranked[:top_k]
    
    def batch_rerank(
        self,
        queries: List[str],
        results_list: List[List[Dict]],
        top_k: int = 5
    ) -> List[List[Dict]]:
        """
        Rerank multiple queries in a batch
        
        Args:
            queries: List of queries
            results_list: List of results per query
            top_k: Number of final results
        
        Returns:
            List of reranked results per query
        """
        reranked_all = []
        
        for query, results in zip(queries, results_list):
            reranked = self.rerank(query, results, top_k)
            reranked_all.append(reranked)
        
        return reranked_all


# Demo
if __name__ == "__main__":
    reranker = Reranker()
    
    # Example results
    results = [
        {
            'chunk': type('obj', (object,), {'text': 'Python is a programming language'})(),
            'score': 0.85
        },
        {
            'chunk': type('obj', (object,), {'text': 'JavaScript for web development'})(),
            'score': 0.82
        }
    ]
    
    query = "What is Python?"
    reranked = reranker.rerank(query, results, top_k=2)
    
    for i, r in enumerate(reranked, 1):
        print(f"{i}. Cross-encoder score: {r['cross_encoder_score']:.3f}")
        print(f"   Original score: {r['original_score']:.3f}")

---

Part 3: Two-Stage Retrieval Pipeline

3.1: Complete pipeline

from src.retrieval.query_expander import QueryExpander
from src.retrieval.reranker import Reranker

class TwoStageRetrieval:
    """
    Two-stage retrieval pipeline
    
    Stage 1: FAISS (fast, high recall)
    Stage 2: Cross-encoder (precise, high precision)
    """
    
    def __init__(
        self,
        embedder,
        index,
        use_expansion: bool = True,
        use_reranking: bool = True
    ):
        self.embedder = embedder
        self.index = index
        self.use_expansion = use_expansion
        self.use_reranking = use_reranking
        
        if use_expansion:
            self.expander = QueryExpander(embedder)
        
        if use_reranking:
            self.reranker = Reranker()
    
    def retrieve(
        self,
        query: str,
        k_candidates: int = 20,
        k_final: int = 5,
        expansions: List[str] = None
    ) -> List[Dict]:
        """
        Two-stage retrieval
        
        Args:
            query: User query
            k_candidates: Candidates in Stage 1
            k_final: Final results in Stage 2
            expansions: Expansion terms (optional)
        
        Returns:
            Final top-K results
        """
        # Query expansion (optional)
        if self.use_expansion and expansions:
            query_emb = self.expander.expand(query, expansions)
        else:
            query_emb = self.embedder.embed(query)
        
        # Stage 1: FAISS retrieval
        candidates = self.index.search(query_emb, k=k_candidates)
        
        # Stage 2: Reranking (optional)
        if self.use_reranking:
            final_results = self.reranker.rerank(query, candidates, top_k=k_final)
        else:
            final_results = candidates[:k_final]
        
        return final_results


# Usage
pipeline = TwoStageRetrieval(
    embedder=embedder,
    index=faiss_index,
    use_expansion=True,
    use_reranking=True
)

query = "How to install Python on macOS?"
expansions = ["Python installation guide", "setup Python"]

results = pipeline.retrieve(
    query=query,
    k_candidates=20,
    k_final=5,
    expansions=expansions
)

for i, r in enumerate(results, 1):
    print(f"\n{i}. Score: {r['cross_encoder_score']:.3f}")
    print(f"   {r['chunk'].text[:200]}...")

---

Comparison: Single-Stage vs Two-Stage

Benchmark

def benchmark_retrieval_strategies():
    """Compare single-stage vs two-stage"""
    
    # Single-stage (FAISS only)
    results_single = index.search(query_emb, k=5)
    
    # Two-stage (FAISS + reranker)
    candidates = index.search(query_emb, k=20)
    results_two_stage = reranker.rerank(query, candidates, top_k=5)
    
    # Evaluate with metrics
    evaluator = RetrievalEvaluator()
    
    # Single-stage
    retrieved_single = [r['chunk'].id for r in results_single]
    metrics_single = evaluator.evaluate_query(retrieved_single, relevant, k=5)
    
    # Two-stage
    retrieved_two = [r['chunk'].id for r in results_two_stage]
    metrics_two = evaluator.evaluate_query(retrieved_two, relevant, k=5)
    
    print("📊 Single-Stage vs Two-Stage:")
    print(f"\nSingle-Stage (FAISS only):")
    print(f"  nDCG@5: {metrics_single['ndcg@k']:.3f}")
    print(f"  Latency: ~5ms")
    
    print(f"\nTwo-Stage (FAISS + Reranker):")
    print(f"  nDCG@5: {metrics_two['ndcg@k']:.3f}")
    print(f"  Improvement: {((metrics_two['ndcg@k'] - metrics_single['ndcg@k']) / metrics_single['ndcg@k'] * 100):.1f}%")
    print(f"  Latency: ~50ms")

Typical results:

• nDCG improvement: +10-20%
• Latency increase: 5ms → 50ms (10x)
• Trade-off: Worth it if precision is critical

---

Troubleshooting

Problem 1: Reranker too slow

Cause: The cross-encoder processes each pair individually

Solution:

# Reduce candidates
results = pipeline.retrieve(
    query=query,
    k_candidates=10,  # Instead of 20
    k_final=5
)

# Or use a faster model
reranker = Reranker(model_name='cross-encoder/ms-marco-TinyBERT-L-2-v2')

---

Problem 2: Query expansion worsens the results

Cause: Irrelevant expansions dilute the original query

Solution:

# Increase alpha (more weight on the original query)
expander = QueryExpander(embedder, alpha=0.9)  # 90% query, 10% expansions

---

Summary

In this capsule you implemented:

• ✅ Query expansion with embedding arithmetic
• ✅ Cross-encoder reranking
• ✅ Two-stage retrieval pipeline
• ✅ Single-stage vs two-stage benchmark
• ✅ Latency optimization

Next capsule: Production Deployment - Docker, FastAPI, Kubernetes.

---

Additional Resources

1. Cross-Encoders for Reranking - SBERT docs
2. Query Expansion Techniques - Wikipedia
3. Two-Stage Retrieval - DPR paper
4. MS MARCO Dataset - Reranking benchmark
5. Sentence-BERT Cross-Encoders - Pre-trained models
6. RAG Reranking Guide - Pinecone tutorial
7. Hybrid Search Strategies - DeepLearning.AI

---

Module 8 - Capsule 06