File size: 976 Bytes
7770c5f bdaeeeb 7770c5f bdaeeeb 7770c5f bdaeeeb 7770c5f bdaeeeb 7770c5f f42877a bdaeeeb | 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 | import logging
from FlagEmbedding import FlagReranker
from ..config import get_settings
logger = logging.getLogger(__name__)
_reranker: FlagReranker | None = None
def _get_reranker() -> FlagReranker:
global _reranker
if _reranker is None:
settings = get_settings()
logger.info(f"[Reranker] Loading model: {settings.reranker_model}")
_reranker = FlagReranker(settings.reranker_model, use_fp16=False)
logger.info(f"[Reranker] Model loaded successfully: {settings.reranker_model}")
return _reranker
def rerank(query: str, passages: list[str]) -> list[float]:
if not passages:
return []
reranker = _get_reranker()
pairs = [[query, p] for p in passages]
# batch_size=8 limits peak RAM usage — critical for cloud CPU-only environments
scores = reranker.compute_score(pairs, normalize=True, batch_size=8)
if isinstance(scores, float):
scores = [scores]
return [float(s) for s in scores]
|