""" Cross-Encoder Reranker Module Local cross-encoder reranker using BAAI/bge-reranker-v2-m3 for CPU inference. Part of a 3-stage RAG pipeline: vector search (stage 1) -> cross-encoder (stage 2) -> LLM reranking (stage 3). Uses sentence-transformers CrossEncoder for batched scoring and combines cross-encoder relevance with original vector similarity. """ import asyncio import logging import time from typing import Optional from sentence_transformers import CrossEncoder logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- _MODEL_NAME = "BAAI/bge-reranker-v2-m3" _INFERENCE_TIMEOUT_SECS = 10 _CROSS_ENCODER_WEIGHT = 0.6 _VECTOR_SIM_WEIGHT = 0.4 # --------------------------------------------------------------------------- # Normalization helpers # --------------------------------------------------------------------------- def _min_max_normalize(scores: list[float]) -> list[float]: """Min-max normalize a list of floats to [0, 1]. If all scores are identical (range == 0), return 0.5 for every element so that the combined formula still degrades gracefully. """ if not scores: return [] min_s = min(scores) max_s = max(scores) rng = max_s - min_s if rng == 0: return [0.5] * len(scores) return [(s - min_s) / rng for s in scores] # --------------------------------------------------------------------------- # CrossEncoderReranker - singleton # --------------------------------------------------------------------------- class CrossEncoderReranker: """Cross-encoder reranker powered by BAAI/bge-reranker-v2-m3. Lazy-loads the model on first use (or via ``warm_up()``). Use ``get_reranker()`` to obtain the singleton instance. """ _instance: Optional["CrossEncoderReranker"] = None _lock = asyncio.Lock() def __init__(self) -> None: # Do NOT call _load_model here - lazy-load on first use. self._model: CrossEncoder | None = None self._model_loaded: bool = False self._model_type: str = "none" # "onnx", "fp32", or "none" self._load_time_secs: float | None = None self._memory_estimate_mb: float | None = None self._tokenizer = None # for ONNX path self._onnx_session = None # for ONNX path # ----------------------------------------------------------------------- # Singleton accessor # ----------------------------------------------------------------------- @classmethod async def get_reranker(cls) -> "CrossEncoderReranker": """Return the singleton ``CrossEncoderReranker``, creating it if needed.""" if cls._instance is None: async with cls._lock: # Double-check after acquiring lock if cls._instance is None: logger.info("Creating CrossEncoderReranker singleton") cls._instance = cls() return cls._instance # ----------------------------------------------------------------------- # Model loading # ----------------------------------------------------------------------- def _load_model(self) -> None: """Synchronously load the cross-encoder model into memory. Prefers ONNX INT8 quantized model if available (~500MB, 2-3x faster). Falls back to fp32 CrossEncoder from sentence-transformers. """ if self._model_loaded: return # Check for quantized ONNX model first import os as _os _onnx_path = "/app/data/models/bge-reranker-v2-m3-onnx" _onnx_model = _os.path.join(_onnx_path, "model.onnx") if _os.path.exists(_onnx_model): logger.info("Loading quantized ONNX reranker from %s ...", _onnx_path) start = time.perf_counter() try: import onnxruntime as ort from optimum.onnxruntime import ORTModelForSequenceClassification # noqa: F401 from transformers import AutoTokenizer # Use CPU execution provider self._tokenizer = AutoTokenizer.from_pretrained(_onnx_path) self._onnx_session = ort.InferenceSession( _onnx_model, providers=["CPUExecutionProvider"], ) self._model_type = "onnx" except ImportError: logger.warning("ONNX model found but optimum not installed. Falling back to fp32.") self._load_fp32_model() return elapsed = time.perf_counter() - start self._load_time_secs = round(elapsed, 2) self._model_loaded = True self._memory_estimate_mb = 500.0 # approximate for INT8 quantized logger.info( "ONNX reranker loaded in %.2fs (~500 MB INT8)", elapsed, ) return # Fallback: load full fp32 model self._load_fp32_model() def _load_fp32_model(self) -> None: """Load the full fp32 cross-encoder (2.1 GB, ~45s).""" logger.info("Loading cross-encoder model '%s' (CPU, fp32) ...", _MODEL_NAME) start = time.perf_counter() self._model = CrossEncoder(_MODEL_NAME, device="cpu") elapsed = time.perf_counter() - start self._load_time_secs = round(elapsed, 2) self._model_loaded = True self._model_type = "fp32" # Rough memory estimate: parameter count * 4 bytes (fp32) try: total_params = sum(p.numel() for p in self._model.model.parameters()) self._memory_estimate_mb = round(total_params * 4 / (1024 * 1024), 1) except Exception: self._memory_estimate_mb = 420.0 logger.info( "Cross-encoder model loaded in %.2fs (~%.0f MB parameters in fp32)", elapsed, self._memory_estimate_mb or 0, ) # ----------------------------------------------------------------------- # Public API - warm_up / health_check # ----------------------------------------------------------------------- async def warm_up(self) -> None: """Pre-load the model so first inference has no cold-start penalty. Call this at application startup. """ if self._model_loaded: logger.debug("warm_up() called but model already loaded - skipping") return # Model loading is CPU-bound; run in executor to avoid blocking the # async event loop. loop = asyncio.get_running_loop() await loop.run_in_executor(None, self._load_model) def health_check(self) -> dict: """Return model status, load time, and memory estimate.""" return { "model_name": _MODEL_NAME, "loaded": self._model_loaded, "load_time_secs": self._load_time_secs, "memory_estimate_mb": self._memory_estimate_mb, "device": "cpu", "inference_timeout_secs": _INFERENCE_TIMEOUT_SECS, "weights": { "cross_encoder": _CROSS_ENCODER_WEIGHT, "vector_similarity": _VECTOR_SIM_WEIGHT, }, } # ----------------------------------------------------------------------- # Core inference (internal) # ----------------------------------------------------------------------- def _score_pairs(self, query: str, texts: list[str]) -> list[float]: """Score (query, text) pairs via the cross-encoder in one forward pass. Uses ONNX Runtime if quantized model is loaded, else sentence-transformers. Returns raw cross-encoder scores (one per text). """ if self._model_type == "onnx" and self._onnx_session: scores = [] for text in texts: inputs = self._tokenizer( query, text, return_tensors="np", truncation=True, max_length=512, padding=True, ) ort_inputs = { "input_ids": inputs["input_ids"], "attention_mask": inputs["attention_mask"], } # Some ONNX models also expect token_type_ids if "token_type_ids" in inputs: ort_inputs["token_type_ids"] = inputs["token_type_ids"] outputs = self._onnx_session.run(None, ort_inputs) logits = outputs[0] # Cross-encoder output: logit for positive class score = float(logits[0][0]) if logits.ndim > 1 else float(logits[0]) scores.append(score) return scores # Fallback: sentence-transformers pairs = [(query, t) for t in texts] scores: list[float] = self._model.predict(pairs, batch_size=len(pairs)).tolist() return scores # ----------------------------------------------------------------------- # Public API - rerank # ----------------------------------------------------------------------- async def rerank( self, query: str, documents: list[dict], top_k: int = 10, ) -> list[dict]: """Re-rank *documents* against *query* combining cross-encoder + vector similarity. Parameters ---------- query: The user query string. documents: List of document dicts, each containing at least ``"content"`` and ``"similarity"`` keys. top_k: Number of top results to return. Returns ------- List[Dict] Documents sorted by combined score (descending), each augmented with ``"cross_score"``, ``"cross_score_normalized"``, and ``"combined_score"`` keys. A ``"reasoning"`` key describes why the document was ranked where it is. """ if not documents: return [] top_k = min(top_k, len(documents)) # Ensure model is loaded if not self._model_loaded: await self.warm_up() texts = [doc.get("content", "") for doc in documents] vector_sims = [float(doc.get("similarity", 0.0)) for doc in documents] # Attempt cross-encoder inference with timeout guard cross_scores: list[float] | None = None try: loop = asyncio.get_running_loop() cross_scores = await asyncio.wait_for( loop.run_in_executor(None, self._score_pairs, query, texts), timeout=_INFERENCE_TIMEOUT_SECS, ) logger.debug( "Cross-encoder scored %d documents for query %.80s", len(documents), query, ) except TimeoutError: logger.warning( "Cross-encoder inference timed out after %ds - falling back to vector-only ranking for query: %.80s", _INFERENCE_TIMEOUT_SECS, query, ) except Exception: logger.exception( "Cross-encoder inference failed - falling back to vector-only ranking for query: %.80s", query, ) # Build result list results: list[dict] = [] if cross_scores is not None: # Normalize cross-encoder scores to [0, 1] cross_normed = _min_max_normalize(cross_scores) for i, doc in enumerate(documents): ce_norm = cross_normed[i] vsim = vector_sims[i] combined = _CROSS_ENCODER_WEIGHT * ce_norm + _VECTOR_SIM_WEIGHT * vsim # Build a reasoning string if ce_norm > 0.8 and vsim > 0.8: reason = "Strong cross-encoder and vector agreement" elif ce_norm > 0.5 and vsim > 0.5: reason = "Moderate agreement from both signals" elif ce_norm > vsim + 0.3: reason = "Cross-encoder significantly boosted relevance over vector similarity" elif vsim > ce_norm + 0.3: reason = "Vector similarity significantly higher than cross-encoder score" else: reason = "Mixed signals with comparable strength" enriched = dict(doc) # shallow copy enriched["cross_score"] = round(cross_scores[i], 4) enriched["cross_score_normalized"] = round(ce_norm, 4) enriched["combined_score"] = round(combined, 4) enriched["reasoning"] = reason results.append(enriched) else: # Fallback: vector-similarity-only ranking for i, doc in enumerate(documents): enriched = dict(doc) enriched["cross_score"] = None enriched["cross_score_normalized"] = None enriched["combined_score"] = round(vector_sims[i], 4) enriched["reasoning"] = "Vector-only ranking (cross-encoder unavailable)" results.append(enriched) # Sort descending by combined_score results.sort(key=lambda d: d["combined_score"], reverse=True) return results[:top_k] # ----------------------------------------------------------------------- # Public API - rerank_only (simpler) # ----------------------------------------------------------------------- async def rerank_only( self, query: str, texts: list[str], top_k: int = 10, ) -> list[str]: """Score and sort *texts* by cross-encoder relevance to *query*. Simplified API that returns just the sorted text strings. Parameters ---------- query: The user query string. texts: List of text passages to rank. top_k: Number of top texts to return. Returns ------- List[str] Texts sorted by cross-encoder score (descending), truncated to *top_k*. """ if not texts: return [] top_k = min(top_k, len(texts)) if not self._model_loaded: await self.warm_up() try: loop = asyncio.get_running_loop() cross_scores = await asyncio.wait_for( loop.run_in_executor(None, self._score_pairs, query, texts), timeout=_INFERENCE_TIMEOUT_SECS, ) except TimeoutError: logger.warning( "rerank_only: cross-encoder timed out after %ds - returning texts in original order", _INFERENCE_TIMEOUT_SECS, ) return texts[:top_k] except Exception: logger.exception( "rerank_only: cross-encoder failed - returning texts in original order", ) return texts[:top_k] # Pair texts with scores, sort descending paired = list(zip(texts, cross_scores, strict=False)) paired.sort(key=lambda t: t[1], reverse=True) return [t for t, _ in paired[:top_k]] # --------------------------------------------------------------------------- # Module-level convenience (optional thin wrappers) # --------------------------------------------------------------------------- async def get_reranker() -> CrossEncoderReranker: """Async convenience function to obtain the singleton reranker.""" return await CrossEncoderReranker.get_reranker()