| from sentence_transformers import CrossEncoder | |
| from server.utils import setup_logger | |
| logger = setup_logger(__name__) | |
| MODEL_NAME = "cross-encoder/ms-marco-TinyBERT-L-2-v2" | |
| _model: CrossEncoder | None = None | |
| def load_reranker() -> CrossEncoder: | |
| """Load cross-encoder model (call at startup to avoid cold-start latency).""" | |
| global _model | |
| if _model is None: | |
| logger.info(f"Loading reranker: {MODEL_NAME}") | |
| _model = CrossEncoder(MODEL_NAME) | |
| logger.info("Reranker loaded") | |
| return _model | |
| def rerank(query: str, docs: list[dict], top_k: int = 5) -> list[dict]: | |
| """Score query-doc pairs; return top_k sorted by rerank_score desc.""" | |
| if not docs: | |
| return [] | |
| model = load_reranker() | |
| pairs = [(query, d["content"]) for d in docs] | |
| scores = model.predict(pairs, batch_size=4) | |
| ranked = sorted(zip(scores, docs), key=lambda x: x[0], reverse=True) | |
| result = [] | |
| for score, doc in ranked[:top_k]: | |
| doc = dict(doc) | |
| doc["rerank_score"] = round(float(score), 4) | |
| result.append(doc) | |
| return result | |