| """ |
| 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 |
| |
| return {idx: 1.0 / (1.0 + math.exp(-float(s))) for (idx, _), s in zip(docs, raw)} |
|
|
