import logging import os import re from pathlib import Path from typing import Any, Dict, List import numpy as np import torch from rank_bm25 import BM25Okapi from sentence_transformers.cross_encoder import CrossEncoder import config logger = logging.getLogger(__name__) _CROSS_ENCODER_INSTANCE = None class ONNXCrossEncoderRanker: """ ONNX Runtime accelerated CrossEncoder for multilingual cross-encoders. Supports dynamic INT8 quantization and context bounding (<35ms on CPU). """ def __init__(self, model_name: str = config.CROSS_ENCODER_MODEL_NAME): self.model_name = model_name self.onnx_dir = Path(getattr(config, "ONNX_MODELS_DIR", config.DATA_DIR / "onnx_models")) self.onnx_dir.mkdir(parents=True, exist_ok=True) sanitized_name = re.sub(r'[^a-zA-Z0-9_]', '_', model_name) self.onnx_fp32_path = self.onnx_dir / f"ce_{sanitized_name}.onnx" self.onnx_int8_path = self.onnx_dir / f"ce_{sanitized_name}_int8.onnx" from transformers import AutoTokenizer self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.has_token_type_ids = "token_type_ids" in (self.tokenizer.model_input_names or []) # Ensure ONNX model exists and is INT8 quantized self._ensure_onnx_model() import onnxruntime as ort opts = ort.SessionOptions() num_threads = getattr(config, "ONNX_NUM_THREADS", 2) opts.intra_op_num_threads = num_threads opts.inter_op_num_threads = 1 opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL load_path = self.onnx_int8_path if self.onnx_int8_path.exists() else self.onnx_fp32_path logger.info(f"Loading ONNX CrossEncoder from: {load_path} (threads={num_threads})") self.session = ort.InferenceSession(str(load_path), opts, providers=["CPUExecutionProvider"]) # Warmup ONNX inference graph to avoid cold-start JIT latency try: self.score_pairs("warmup query", ["warmup passage"], max_length=64) except Exception: pass logger.info("ONNX CrossEncoder session initialized and warmed up successfully.") def _ensure_onnx_model(self): """Auto-export CrossEncoder to ONNX and apply INT8 dynamic quantization.""" if self.onnx_int8_path.exists() or self.onnx_fp32_path.exists(): if not self.onnx_int8_path.exists() and self.onnx_fp32_path.exists(): try: from onnxruntime.quantization import quantize_dynamic, QuantType logger.info(f"Quantizing ONNX CrossEncoder to INT8 format: {self.onnx_int8_path}...") quantize_dynamic( str(self.onnx_fp32_path), str(self.onnx_int8_path), weight_type=QuantType.QInt8, ) except Exception as q_err: logger.warning(f"INT8 CrossEncoder quantization skipped: {q_err}") return logger.info(f"Exporting CrossEncoder '{self.model_name}' to ONNX format at {self.onnx_fp32_path}...") import torch.nn as nn from transformers import AutoModelForSequenceClassification from torch.export import Dim has_tt = self.has_token_type_ids class CEWrapper(nn.Module): def __init__(self, m, use_tt): super().__init__() self.m = m self.use_tt = use_tt def forward(self, input_ids, attention_mask, token_type_ids=None): if self.use_tt and token_type_ids is not None: return self.m(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, return_dict=False)[0] return self.m(input_ids=input_ids, attention_mask=attention_mask, return_dict=False)[0] base_model = AutoModelForSequenceClassification.from_pretrained(self.model_name) base_model.eval() wrapper = CEWrapper(base_model, has_tt) wrapper.eval() b = Dim("batch") s = Dim("seq") dummy = self.tokenizer(["query 1", "query 2"], ["passage 1", "passage 2"], padding=True, return_tensors="pt") try: if has_tt and "token_type_ids" in dummy: torch.onnx.export( wrapper, (dummy["input_ids"], dummy["attention_mask"], dummy["token_type_ids"]), str(self.onnx_fp32_path), input_names=["input_ids", "attention_mask", "token_type_ids"], output_names=["logits"], dynamic_shapes={ "input_ids": {0: b, 1: s}, "attention_mask": {0: b, 1: s}, "token_type_ids": {0: b, 1: s}, }, opset_version=18, do_constant_folding=True, ) else: torch.onnx.export( wrapper, (dummy["input_ids"], dummy["attention_mask"]), str(self.onnx_fp32_path), input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_shapes={ "input_ids": {0: b, 1: s}, "attention_mask": {0: b, 1: s}, }, opset_version=18, do_constant_folding=True, ) logger.info("Exported FP32 ONNX CrossEncoder model.") # Perform INT8 dynamic quantization try: from onnxruntime.quantization import quantize_dynamic, QuantType logger.info(f"Quantizing ONNX CrossEncoder to INT8 format: {self.onnx_int8_path}...") quantize_dynamic( str(self.onnx_fp32_path), str(self.onnx_int8_path), weight_type=QuantType.QInt8, ) except Exception as q_err: logger.warning(f"INT8 CrossEncoder quantization failed: {q_err}") except Exception as e: logger.warning(f"ONNX CrossEncoder export failed: {e}. PyTorch fallback will be used.") if self.onnx_fp32_path.exists(): try: self.onnx_fp32_path.unlink() except Exception: pass raise e def score_pairs(self, query: str, passages: List[str], max_length: int = 64) -> np.ndarray: """ Scores (query, passage) pairs using ONNX Runtime with Context Bounding (64 tokens). """ if not passages: return np.array([], dtype=np.float32) bound_len = min(max_length, getattr(config, "CONTEXT_BOUNDING_MAX_TOKENS", 64)) pairs = [[query, p[:150]] for p in passages] inputs = self.tokenizer( pairs, padding=True, truncation=True, max_length=bound_len, return_tensors="np", ) ort_inputs = { "input_ids": inputs["input_ids"].astype(np.int64), "attention_mask": inputs["attention_mask"].astype(np.int64), } if self.has_token_type_ids and "token_type_ids" in inputs: ort_inputs["token_type_ids"] = inputs["token_type_ids"].astype(np.int64) logits = self.session.run(None, ort_inputs)[0] return np.asarray(logits.flatten(), dtype=np.float32) class PyTorchCrossEncoderRanker: """ PyTorch fallback wrapper for sentence-transformers CrossEncoder. """ def __init__(self, model_name: str = config.CROSS_ENCODER_MODEL_NAME): load_path = model_name local_cache = getattr(config, "CROSS_ENCODER_LOCAL_CACHE", None) if local_cache and Path(local_cache).exists(): load_path = str(local_cache) logger.info(f"Loading PyTorch CrossEncoder from local cache: {load_path}") else: logger.info(f"Loading PyTorch CrossEncoder from model name: {load_path}") try: self.model = CrossEncoder(load_path) logger.info("PyTorch CrossEncoder loaded successfully.") except Exception as e: logger.warning(f"Failed to load CrossEncoder from '{load_path}': {e}. Falling back to default '{model_name}'.") self.model = CrossEncoder(model_name) def score_pairs(self, query: str, passages: List[str], max_length: int = 64) -> np.ndarray: if not passages: return np.array([], dtype=np.float32) bound_len = min(max_length, getattr(config, "CONTEXT_BOUNDING_MAX_TOKENS", 64)) pairs = [[query, p[:150]] for p in passages] with torch.inference_mode(): scores = self.model.predict( pairs, show_progress_bar=False, batch_size=len(pairs), max_length=bound_len, ) return np.asarray(scores, dtype=np.float32) def get_cross_encoder(): """ Get or initialize the global singleton cross-encoder instance with ONNX-first policy. """ global _CROSS_ENCODER_INSTANCE if _CROSS_ENCODER_INSTANCE is None: if getattr(config, "ENABLE_ONNX_CROSS_ENCODER", True): try: _CROSS_ENCODER_INSTANCE = ONNXCrossEncoderRanker() except Exception as e: logger.warning(f"Failed to initialize ONNX CrossEncoder: {e}. Falling back to PyTorch.") _CROSS_ENCODER_INSTANCE = PyTorchCrossEncoderRanker() else: _CROSS_ENCODER_INSTANCE = PyTorchCrossEncoderRanker() return _CROSS_ENCODER_INSTANCE STOPWORDS = { "what", "is", "the", "of", "in", "and", "how", "do", "does", "are", "for", "to", "a", "an", "why", "who", "which", "can", "with", "from", "by", "on", "as", "between", "explain", "difference", "best", "make", "made", "rule", "rules", "basic", "basics", "way", "ways", "method", "methods", "tell", "give", "show", "know", "anyone", "someone", "recipe", "baking", "baker", "क्या", "है", "हैं", "के", "की", "का", "में", "और", "से", "होता", "होती", "होते", "कैसे", "क्यों", "किए", "किया", "जाता", "जाती", "गया", "गई", "को", "पर", "लिए", "एक", "या", "बारे", "नियम", "विधि", "तरीका", "बुनियादी", "आसान", "सबसे", "बनाने", "बताएं", "என்ன", "எவ்வாறு", "ஏன்", "மற்றும்", "ஒரு", "ஆகும்", "உள்ளது", "என்பது", "யாவை", "பற்றி", "செய்கிறது", "செய்யப்படுகிறது", "எப்படி", "செய்வது", "வழிமுறை", "வழிமுறைகள்", "அடிப்படை", "விதிகள்" } PUNCT_REGEX = re.compile(r'[\s!\"#$%&\'()*+,\-./:;<=>?@\[\\\]^_`{|}~।॥]+') def detect_script(text: str) -> str: """Fast Unicode script identifier for adaptive multilingual routing.""" for char in text: code = ord(char) if 0x0900 <= code <= 0x097F: return "Deva" if 0x0980 <= code <= 0x09FF: return "Beng" if 0x0A00 <= code <= 0x0A7F: return "Guru" if 0x0A80 <= code <= 0x0AFF: return "Gujr" if 0x0B00 <= code <= 0x0B7F: return "Orya" if 0x0B80 <= code <= 0x0BFF: return "Taml" if 0x0C00 <= code <= 0x0C7F: return "Telu" if 0x0C80 <= code <= 0x0CFF: return "Knda" if 0x0D00 <= code <= 0x0D7F: return "Mlym" if 0x0600 <= code <= 0x06FF or 0x0750 <= code <= 0x077F or 0xFB50 <= code <= 0xFDFF or 0xFE70 <= code <= 0xFEFF: return "Arab" if bool(re.search(r"[a-zA-Z]", text)): return "Latn" return "Latn" def tokenize_indic(text: str) -> List[str]: """Clean and split multilingual and Indic text without breaking ligatures.""" if not text: return [] clean = PUNCT_REGEX.sub(' ', text.lower()).strip() return [w for w in clean.split() if len(w) > 1] def tokenize_for_bm25(text: str) -> List[str]: """Multilingual tokenization for BM25 scoring.""" return tokenize_indic(text) def rerank_bm25_hybrid( query_text: str, candidates: List[Dict[str, Any]], bm25_weight: float = config.HYBRID_BM25_WEIGHT, top_k: int = config.RERANK_TOP_K, ) -> List[Dict[str, Any]]: """ Adaptive Script-Aware Hybrid re-ranking: - Same-Script (e.g. Hindi -> Hindi, English -> English): Blends BM25 lexical precision with dense vector score. - Cross-Script (e.g. English -> Hindi, Hindi -> English): Dynamically bypasses BM25 to prevent false lexical penalties. """ if not candidates: return [] if len(candidates) == 1: c = candidates[0].copy() c["final_score"] = float(c.get("score", c.get("dense_score", 1.0))) c["bm25_score"] = 1.0 c["confidence"] = float(c.get("dense_score", 0.9)) return [c] query_script = detect_script(query_text) query_tokens = tokenize_indic(query_text) if not query_tokens: query_tokens = query_text.lower().split() # Build BM25 corpus from candidate texts corpus_tokens = [tokenize_indic(c.get("text", "")) for c in candidates] bm25 = BM25Okapi(corpus_tokens) raw_bm25_scores = bm25.get_scores(query_tokens) # Normalize BM25 scores to [0, 1] max_bm25 = float(np.max(raw_bm25_scores)) if len(raw_bm25_scores) > 0 else 0.0 min_bm25 = float(np.min(raw_bm25_scores)) if len(raw_bm25_scores) > 0 else 0.0 bm25_range = max_bm25 - min_bm25 # Extract dense scores raw_dense_scores = [float(c.get("score", c.get("dense_score", 0.0))) for c in candidates] max_dense = max(raw_dense_scores) if raw_dense_scores else 1.0 min_dense = min(raw_dense_scores) if raw_dense_scores else 0.0 dense_range = max_dense - min_dense q_words = [w for w in query_tokens if w not in STOPWORDS and len(w) > 2] if not q_words: q_words = query_tokens reranked = [] for idx, cand in enumerate(candidates): cand_text = cand.get("text", "") cand_script = detect_script(cand_text) is_cross_script = (query_script != cand_script) # Normalized BM25 if bm25_range > 1e-6: norm_bm25 = (raw_bm25_scores[idx] - min_bm25) / bm25_range else: norm_bm25 = 1.0 if max_bm25 > 0 else 0.0 # Normalized Dense if dense_range > 1e-6: norm_dense = (raw_dense_scores[idx] - min_dense) / dense_range else: norm_dense = max(0.0, min(1.0, raw_dense_scores[idx])) # Adaptive Hybrid Combination: # For cross-script matches, avoid penalizing with a 0 BM25 score. if is_cross_script: final_score = norm_dense else: final_score = (1.0 - bm25_weight) * norm_dense + bm25_weight * norm_bm25 # Absolute confidence computation dense_val = float(raw_dense_scores[idx]) if is_cross_script: # Cross-script candidate relies directly on dense semantic alignment confidence = dense_val else: p_tokens = set(tokenize_indic(cand_text)) p_clean_text = " ".join(tokenize_indic(cand_text)) matched = 0 for qw in q_words: stem = qw[:5] if len(qw) > 5 else qw if qw in p_tokens or (len(qw) > 4 and stem in p_clean_text) or any(pt.startswith(stem) for pt in p_tokens if len(stem) > 3): matched += 1 overlap = matched / len(q_words) if q_words else 0.0 if matched > 0: confidence = (dense_val * 0.70) + (0.30 * min(1.0, overlap + 0.2)) else: confidence = dense_val * 0.75 item = cand.copy() item["dense_score"] = dense_val item["bm25_score"] = float(raw_bm25_scores[idx]) item["final_score"] = float(final_score) item["confidence"] = round(float(confidence), 4) item["is_cross_script"] = is_cross_script reranked.append(item) # Sort by calibrated confidence descending, breaking ties with final_score reranked = sorted(reranked, key=lambda x: (x["confidence"], x["final_score"]), reverse=True) return reranked[:top_k] def rerank_cross_encoder( query_text: str, candidates: List[Dict[str, Any]], top_k: int = config.CROSS_ENCODER_TOP_K, ) -> List[Dict[str, Any]]: """ Applies deep cross-attention re-ranking over candidate passages. Attaches `cross_encoder_score` and recalibrates ranking. """ if not candidates: return [] try: ranker = get_cross_encoder() passages = [c.get("text", "") for c in candidates[:top_k]] ce_scores = ranker.score_pairs(query_text, passages) scored_candidates = [] for idx, cand in enumerate(candidates[:top_k]): c = cand.copy() raw_ce = float(ce_scores[idx]) if idx < len(ce_scores) else -10.0 c["cross_encoder_score"] = round(raw_ce, 4) # Normalize cross-encoder output (handle both pre-calibrated [0, 1] probabilities and unbounded logits) if 0.0 <= raw_ce <= 1.0: sig_ce = raw_ce else: sig_ce = 1.0 / (1.0 + np.exp(-raw_ce)) c["ce_prob"] = round(float(sig_ce), 4) # Recalibrate composite confidence blending cross-encoder with dense score dense_val = float(c.get("dense_score", 0.5)) c["confidence"] = round(0.70 * sig_ce + 0.30 * dense_val, 4) c["final_score"] = round(raw_ce, 4) scored_candidates.append(c) # Append remaining candidates beyond top_k if any if len(candidates) > top_k: for cand in candidates[top_k:]: c = cand.copy() c["cross_encoder_score"] = -10.0 c["ce_prob"] = 0.0 scored_candidates.append(c) # Sort by cross-encoder score descending scored_candidates = sorted(scored_candidates, key=lambda x: x.get("cross_encoder_score", -10.0), reverse=True) return scored_candidates except Exception as e: logger.warning(f"Cross-encoder reranking failed: {e}. Falling back to BM25-hybrid ranking.") return candidates