30 lines
1,021 B
Python
30 lines
1,021 B
Python
"""Cross-encoder re-ranking with bge-reranker-v2-m3.
|
|
|
|
The re-ranker reads each (profile description, paper) pair jointly and returns
|
|
a score in 0-1. It replaced ms-marco-MiniLM in ssearch because that model
|
|
scored journal-style text far outside its calibrated range. Highlighting
|
|
thresholds apply to this score.
|
|
"""
|
|
|
|
from art import env # noqa: F401 (sets HF offline variables first)
|
|
|
|
from functools import cache
|
|
|
|
|
|
@cache
|
|
def model():
|
|
from sentence_transformers import CrossEncoder
|
|
cfg = env.config()["screen"]
|
|
return CrossEncoder(cfg["rerank_model"], device=env.device(),
|
|
max_length=cfg["rerank_max_length"])
|
|
|
|
|
|
def model_name() -> str:
|
|
return env.config()["screen"]["rerank_model"]
|
|
|
|
|
|
def score(query: str, passages: list[str], batch_size: int = 16) -> list[float]:
|
|
if not passages:
|
|
return []
|
|
return [float(s) for s in model().predict([(query, p) for p in passages],
|
|
batch_size=batch_size, show_progress_bar=False)]
|