Query Expansion & Reranking: Two-Stage Retrieval

Descripción

En esta cápsula implementarás Query Expansion y Cross-Encoder Reranking para crear un sistema de retrieval de dos etapas que maximiza tanto recall como precision.

Two-Stage Retrieval:

1. Stage 1 (FAISS): Recupera muchos candidatos (k=20-50) rápidamente
2. Stage 2 (Reranker): Reordena candidatos con modelo más preciso, retorna top-K final

Este approach es estándar en producción: FAISS da recall alto (~95%), reranker mejora precision final (~10-20% improvement en nDCG).

Duración estimada: 40-50 minutos

---

Objetivos

Al completar esta cápsula, serás capaz de:

• ✅ Implementar query expansion con embedding arithmetic
• ✅ Usar cross-encoder para reranking
• ✅ Implementar two-stage retrieval pipeline
• ✅ Comparar single-stage vs two-stage
• ✅ Optimizar latencia del reranker

---

Parte 1: Query Expansion

¿Por qué Query Expansion?

Problema: Query corto puede no capturar toda la intención

query = "Python install"
# Embedding de "Python install" solo captura esos términos

Solución: Expandir query con términos relacionados

query = "Python install"
expansions = ["Python installation", "setup Python", "install Python macOS"]
# Combinar embeddings para query más rico

---

1.1: Implementación con embedding arithmetic

Crear src/retrieval/query_expander.py:

"""
Query Expander - RAG System
Expande queries para mejor recall
"""

import numpy as np
from typing import List
import logging


class QueryExpander:
    """
    Expande queries usando embedding arithmetic
    
    Strategies:
    - Synonym expansion: Agregar sinónimos
    - Concept expansion: Agregar términos relacionados
    - Weighted blend: Combinar query original con expansiones
    
    Example:
        expander = QueryExpander(embedder)
        
        expanded_emb = expander.expand(
            query="Python install",
            expansions=["setup Python", "Python installation"]
        )
    """
    
    def __init__(self, embedder, alpha: float = 0.7):
        """
        Inicializar QueryExpander
        
        Args:
            embedder: EmbeddingClient instance
            alpha: Peso del query original (0-1)
                - alpha=1.0: Solo query original
                - alpha=0.5: 50% query, 50% expansions
                - alpha=0.7: Recomendado (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:
        """
        Expandir query con términos adicionales
        
        Args:
            query: Query original
            expansions: Lista de términos de expansión
            alpha: Override de alpha (opcional)
        
        Returns:
            Embedding expandido
        """
        if alpha is None:
            alpha = self.alpha
        
        # Embedding del query original
        query_emb = self.embedder.embed(query)
        
        if not expansions:
            return query_emb
        
        # Embeddings de expansiones
        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
        
        # Normalizar (importante para 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-expansión usando top-K results iniciales
        
        Args:
            query: Query original
            index: FAISS index
            k_initial: Cuántos docs usar para expansión
            n_expansions: Cuántos términos extraer
        
        Returns:
            Embedding expandido
        """
        # Búsqueda inicial
        query_emb = self.embedder.embed(query)
        results = index.search(query_emb, k=k_initial)
        
        # Extraer términos de top results (simplificado)
        # En producción: usar TF-IDF, entity extraction, etc.
        expansions = []
        for r in results[:n_expansions]:
            text = r['chunk'].text
            # Tomar primeras palabras como 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)
    
    # Expandir 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}")

---

Parte 2: Cross-Encoder Reranking

¿Por qué Reranking?

Problema: Bi-encoders (FAISS) encode query y docs separadamente

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

Ventaja: Muy rápido (pre-computed doc embeddings) Desventaja: No ve interacción query-document

---

Solución: Cross-encoder ve query + document juntos

# Cross-encoder (reranker):
score = model([query, document])  # Together!
# Ve interacción directa → más preciso

Trade-off:

• Bi-encoder: 1000x más rápido, ~95% precision
• Cross-encoder: 1000x más lento, ~100% precision

---

2.1: Implementación del Reranker

Crear src/retrieval/reranker.py:

"""
Reranker - RAG System
Cross-encoder reranking para 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'
    ):
        """
        Inicializar Reranker
        
        Args:
            model_name: Nombre del modelo cross-encoder
        """
        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 resultados con cross-encoder
        
        Args:
            query: Query del usuario
            results: Lista de results de FAISS
            top_k: Cantidad de resultados finales
        
        Returns:
            Lista de results rerankeados
        """
        if not results:
            return []
        
        # Preparar pares (query, document)
        pairs = []
        for r in results:
            doc_text = r['chunk'].text
            pairs.append([query, doc_text])
        
        # Score con cross-encoder
        scores = self.model.predict(pairs)
        
        # Combinar results con nuevos 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)
        
        # Re-ordenar por 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 múltiples queries en batch
        
        Args:
            queries: Lista de queries
            results_list: Lista de results por query
            top_k: Cantidad de resultados finales
        
        Returns:
            Lista de results rerankeados por 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()
    
    # Results de ejemplo
    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}")

---

Parte 3: Two-Stage Retrieval Pipeline

3.1: Pipeline completo

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: Query del usuario
            k_candidates: Candidatos en Stage 1
            k_final: Resultados finales en Stage 2
            expansions: Términos de expansión (opcional)
        
        Returns:
            Top-K results finales
        """
        # Query expansion (opcional)
        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 (opcional)
        if self.use_reranking:
            final_results = self.reranker.rerank(query, candidates, top_k=k_final)
        else:
            final_results = candidates[:k_final]
        
        return final_results


# Uso
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]}...")

---

Comparación: Single-Stage vs Two-Stage

Benchmark

def benchmark_retrieval_strategies():
    """Compare single-stage vs two-stage"""
    
    # Single-stage (solo FAISS)
    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)
    
    # Evaluar con 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")

Resultados típicos:

• nDCG improvement: +10-20%
• Latency increase: 5ms → 50ms (10x)
• Trade-off: Vale la pena si precision es crítica

---

Troubleshooting

Problema 1: Reranker muy lento

Causa: Cross-encoder procesa cada par individualmente

Solución:

# Reducir candidatos
results = pipeline.retrieve(
    query=query,
    k_candidates=10,  # En vez de 20
    k_final=5
)

# O usar modelo más rápido
reranker = Reranker(model_name='cross-encoder/ms-marco-TinyBERT-L-2-v2')

---

Problema 2: Query expansion empeora resultados

Causa: Expansiones no relevantes diluyen query original

Solución:

# Aumentar alpha (más peso al query original)
expander = QueryExpander(embedder, alpha=0.9)  # 90% query, 10% expansions

---

Resumen

En esta cápsula implementaste:

• ✅ Query expansion con embedding arithmetic
• ✅ Cross-encoder reranking
• ✅ Two-stage retrieval pipeline
• ✅ Benchmark single-stage vs two-stage
• ✅ Optimización de latencia

Próxima cápsula: Production Deployment - Docker, FastAPI, Kubernetes.

---

Recursos Adicionales

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

---

Módulo 8 - Cápsula 06