File size: 2,292 Bytes
35676b4 2401ce5 35676b4 | 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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """storage/reranker.py — cross-encoder reranking for RAG chunks.
Wraps BAAI/bge-reranker-base via sentence-transformers.CrossEncoder.
Loaded lazily on first call so app startup stays fast (~100 MB model
downloaded once into ~/.cache/huggingface).
Used by storage/vector_store.search to re-rank Chroma's top-k candidates.
"""
from __future__ import annotations
from typing import Optional
_MODEL_NAME = "BAAI/bge-reranker-base"
_model = None # type: ignore[var-annotated]
def _get_model():
"""Lazy-load the cross-encoder. Returns None on import or load failure."""
global _model
if _model is None:
try:
from sentence_transformers import CrossEncoder
_model = CrossEncoder(_MODEL_NAME)
except Exception as exc: # pragma: no cover — defensive, not testable
import sys
print(f"[reranker] failed to load {_MODEL_NAME}: {exc}", file=sys.stderr)
_model = False # sentinel: tried and failed
return _model if _model is not False else None
def rerank(query: str, candidates: list[dict], top_k: int = 3) -> list[dict]:
"""Re-score (query, candidate.text) pairs with the cross-encoder, return top_k.
Each candidate is a dict with at least a 'text' field (matches the format
returned by storage.vector_store.search).
Falls back to the original ordering (truncated to top_k) if:
- candidates already <= top_k (no work to do)
- the cross-encoder model failed to load
- the predict call raises
"""
if not candidates or len(candidates) <= top_k:
return candidates[:top_k]
model = _get_model()
if model is None:
return candidates[:top_k]
try:
pairs = [(query, c.get("text", "")) for c in candidates]
scores = model.predict(pairs)
ranked = sorted(zip(candidates, scores), key=lambda x: float(x[1]), reverse=True)
return [c for c, _ in ranked[:top_k]]
except Exception:
return candidates[:top_k]
def warmup() -> None:
"""Eagerly load the cross-encoder model; call at app startup to avoid cold-start penalty."""
_get_model()
def reset_for_test() -> None:
"""Reset the cached model — used by unit tests that monkeypatch the loader."""
global _model
_model = None
|