75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
"""Cross-encoder reranking of retrieved source units.
|
|
|
|
A cross-encoder reads the question and one candidate text together and scores
|
|
their relevance; it orders candidates that BM25 and E5 already retrieved. It
|
|
never adds a candidate, removes evidence of a named scope or writes text.
|
|
|
|
Model choice (measured on CPU, 60 question-text pairs): the multilingual
|
|
cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 needs about 1.2 s and ranked the
|
|
Slovak probes as well as BAAI/bge-reranker-v2-m3, which needs about 12 s.
|
|
The model is prepared into the image (runtime assets); RETRIEVAL_RERANKER=off
|
|
disables it for the documented A/B comparison.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
|
|
DEFAULT_MODEL = "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
|
|
MAX_PAIRS = 160
|
|
_lock = threading.Lock()
|
|
|
|
|
|
def model_name() -> str:
|
|
return os.getenv("RETRIEVAL_RERANKER_MODEL", DEFAULT_MODEL)
|
|
|
|
|
|
def enabled() -> bool:
|
|
return os.getenv("RETRIEVAL_RERANKER", "on").lower() != "off"
|
|
|
|
|
|
def model_path() -> str:
|
|
"""The prepared local copy in the image, else the Hugging Face cache."""
|
|
root = os.getenv("ZP_RUNTIME_ASSETS")
|
|
if root and (Path(root) / "reranker" / "config.json").is_file():
|
|
return str(Path(root) / "reranker")
|
|
return model_name()
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _model():
|
|
from sentence_transformers import CrossEncoder
|
|
local = os.getenv("HF_HUB_OFFLINE") == "1" or model_path() != model_name()
|
|
return CrossEncoder(model_path(), max_length=256, device="cpu", local_files_only=local)
|
|
|
|
|
|
def available() -> bool:
|
|
if not enabled():
|
|
return False
|
|
try:
|
|
_model()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def scores(question: str, texts: list[str]) -> list[float] | None:
|
|
"""Relevance score per text, or None when reranking is off or unavailable."""
|
|
if not texts or not available():
|
|
return None
|
|
texts = texts[:MAX_PAIRS]
|
|
with _lock:
|
|
values = _model().predict([(question, t[:1000]) for t in texts], batch_size=32, show_progress_bar=False)
|
|
return [float(v) for v in values]
|
|
|
|
|
|
def order(question: str, rows: list[dict], *, key=lambda r: r["text"]) -> list[dict]:
|
|
"""Rows sorted by relevance (stable for ties); unchanged when unavailable."""
|
|
values = scores(question, [key(r) for r in rows])
|
|
if values is None:
|
|
return rows
|
|
ranked = sorted(range(len(values)), key=lambda i: (-values[i], i))
|
|
return [rows[i] for i in ranked] + rows[len(values):]
|