merge: chore/cleanup-remove-bloat-and-secrets into main
This commit is contained in:
commit
bde2f3a97d
1173 changed files with 437609 additions and 0 deletions
417
app/cross_encoder_reranker.py
Normal file
417
app/cross_encoder_reranker.py
Normal file
|
|
@ -0,0 +1,417 @@
|
|||
"""
|
||||
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
|
||||
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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue