Cross-Encoder Reranking¶

class CrossEncoderReranker:
    """Cross-encoder reranker using sentence-transformers models.

    Provides both synchronous and asynchronous reranking with:
    - Automatic model loading and caching
    - Error handling and validation
    - Performance metrics logging
    - Type-safe QueryResult integration
    """

    def __init__(self, model_name: str | None = None):
        """Initialize with optional model override."""

    def rerank(
        self,
        query: str,
        results: list[QueryResult],
        top_k: int | None = None,
    ) -> list[QueryResult]:
        """Synchronous reranking (CPU-bound operation)."""

    async def rerank_async(
        self,
        query: str,
        results: list[QueryResult],
        top_k: int | None = None,
    ) -> list[QueryResult]:
        """Async reranking (runs in thread pool executor)."""

from rag_solution.retrieval.reranker import CrossEncoderReranker
from vectordbs.data_types import QueryResult

# Initialize reranker (model loads once, ~7s first time, <1s cached)
reranker = CrossEncoderReranker()

# Rerank search results
query = "What is machine learning?"
reranked_results = reranker.rerank(
    query=query,
    results=search_results,  # list[QueryResult] from vector search
    top_k=10  # Return top 10 reranked results
)

# Results are sorted by cross-encoder scores
for result in reranked_results:
    print(f"Score: {result.score:.4f} - {result.chunk.text[:100]}")

import asyncio
from rag_solution.retrieval.reranker import CrossEncoderReranker

async def search_with_reranking(query: str, collection_id: str):
    reranker = CrossEncoderReranker()

    # Get initial results from vector store
    vector_results = await vector_store.search(query, collection_id, top_k=100)

    # Rerank asynchronously (non-blocking)
    reranked_results = await reranker.rerank_async(
        query=query,
        results=vector_results,
        top_k=10
    )

    return reranked_results

# Run async
results = asyncio.run(search_with_reranking("ML basics", "col_123"))

# Use fastest model for low-latency production
fast_reranker = CrossEncoderReranker(
    model_name="cross-encoder/ms-marco-TinyBERT-L-2-v2"
)

# Use most accurate model for critical queries
accurate_reranker = CrossEncoderReranker(
    model_name="cross-encoder/ms-marco-MiniLM-L-12-v2"
)

async def search(self, search_input: SearchInput) -> SearchResponse:
    # 1. Vector search (top 100 candidates)
    vector_results = await self.vector_store.search(
        query=query,
        collection_id=collection_id,
        top_k=100  # Over-fetch for reranking
    )

    # 2. Cross-encoder reranking (if enabled)
    if pipeline_config.reranker_enabled:
        reranked_results = await self.reranker.rerank_async(
            query=query,
            results=vector_results,
            top_k=10  # Final top-k
        )

    # 3. Return final results
    return SearchResponse(results=reranked_results)

# In PipelineService or configuration
pipeline_config = {
    "reranker_enabled": True,  # Enable/disable reranking
    "reranker_model": "cross-encoder/ms-marco-MiniLM-L-6-v2",  # Model name
    "reranker_top_k": 10,  # Final number of results
    "vector_top_k": 100,  # Initial candidates to rerank
}

try:
    reranked = reranker.rerank(query, results)
except ValueError as e:
    # Wrapped exception with context
    # Example: "Reranking failed for model ms-marco-MiniLM-L-6-v2: OOM error"
    logger.error(f"Reranking failed: {e}")
    # Fall back to original vector search results
    return results

import logging
logging.basicConfig(level=logging.DEBUG)

# Test with real data
from rag_solution.retrieval.reranker import CrossEncoderReranker
from vectordbs.data_types import *

# Create test data
metadata = DocumentChunkMetadata(source=Source.PDF, document_id="test_doc")
chunk = DocumentChunkWithScore(
    chunk_id="chunk_1",
    text="Machine learning is a subset of artificial intelligence.",
    metadata=metadata
)
result = QueryResult(chunk=chunk, score=0.85)

# Test reranking
reranker = CrossEncoderReranker()
reranked = reranker.rerank("What is ML?", [result])
print(f"Original score: 0.85, Reranked score: {reranked[0].score:.4f}")

# Retrieve more candidates than needed
vector_results = vector_search(query, top_k=100)
# Rerank to final top-k
final_results = reranker.rerank(query, vector_results, top_k=10)

# Non-blocking reranking in FastAPI
@app.post("/search")
async def search_endpoint(query: str):
    results = await vector_search(query)
    reranked = await reranker.rerank_async(query, results)
    return reranked

# Rerank multiple queries concurrently
tasks = [
    reranker.rerank_async(q, results)
    for q in queries
]
all_results = await asyncio.gather(*tasks)

# Old: Slow LLM-based reranking (~20-30s)
reranked = await llm_reranker.rerank(query, results, top_k=10)

# New: Fast cross-encoder reranking (~80ms)
reranker = CrossEncoderReranker()
reranked = await reranker.rerank_async(query, results, top_k=10)

# In main.py or startup event
@app.on_event("startup")
async def warmup_reranker():
    reranker = CrossEncoderReranker()
    logger.info("Cross-encoder model pre-loaded")