File size: 1,603 Bytes
22b7d63 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """
rerank.py — cross-encoder reranking of the top fused candidates.
A bi-encoder (bge) retrieves fast but scores query and document independently;
a cross-encoder reads the (query, document) pair jointly and is far more precise
on the final shortlist. We rerank only `rerank_pool` candidates to keep latency
low. Loads lazily and degrades gracefully (returns None) if unavailable.
"""
from __future__ import annotations
import math
class Reranker:
def __init__(self, model_name: str):
self.model_name = model_name
self._model = None
self.available = True
@property
def model(self):
if self._model is None:
from sentence_transformers import CrossEncoder
self._model = CrossEncoder(self.model_name, max_length=384)
return self._model
def warmup(self) -> None:
try:
self.model.predict([("a", "b")])
except Exception:
self.available = False
def rerank(self, query: str, docs: list[tuple[int, str]]) -> dict[int, float] | None:
"""Return {idx: score in 0..1} or None if the reranker is unavailable."""
if not docs:
return {}
try:
pairs = [(query, text) for _, text in docs]
raw = self.model.predict(pairs, convert_to_numpy=True, show_progress_bar=False)
except Exception:
self.available = False
return None
# ms-marco cross-encoders emit an unbounded relevance logit; squash to 0..1
return {idx: 1.0 / (1.0 + math.exp(-float(s))) for (idx, _), s in zip(docs, raw)}
|