Spaces:
Running
Running
| """Cross-encoder reranking using sentence-transformers. | |
| Provides high-precision re-scoring of (query, document) pairs using a | |
| cross-encoder model. This is the final stage of the retrieval pipeline, | |
| applied after fusion to re-sort candidates by cross-encoder relevance. | |
| Model: cross-encoder/ms-marco-MiniLM-L6-v2 (22M parameters, fast inference). | |
| """ | |
| from typing import Any | |
| import numpy as np | |
| class CrossEncoderReranker: | |
| """Lazy-loaded cross-encoder for reranking retrieval results.""" | |
| MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L6-v2" | |
| def __init__(self) -> None: | |
| self._model = None | |
| def model(self): | |
| """Lazy-load the cross-encoder model on first use.""" | |
| if self._model is None: | |
| from sentence_transformers import CrossEncoder | |
| self._model = CrossEncoder(self.MODEL_NAME) | |
| return self._model | |
| def rerank( | |
| self, query: str, candidates: list[dict[str, Any]], top_k: int | None = None | |
| ) -> list[dict[str, Any]]: | |
| """Re-score and re-sort candidates using cross-encoder. | |
| Args: | |
| query: The search query. | |
| candidates: List of result dicts with at least a 'text' field. | |
| top_k: If set, return only the top-k re-ranked results. | |
| Returns: | |
| Re-sorted candidates with 'cross_encoder_score' added to each. | |
| """ | |
| if not candidates: | |
| return [] | |
| # Build (query, document) pairs | |
| pairs = [(query, c["text"]) for c in candidates] | |
| # Score all pairs | |
| scores = self.model.predict(pairs) | |
| if isinstance(scores, np.ndarray): | |
| scores = scores.tolist() | |
| # Attach scores and sort descending | |
| for candidate, score in zip(candidates, scores, strict=True): | |
| candidate["cross_encoder_score"] = float(score) | |
| reranked = sorted(candidates, key=lambda x: x["cross_encoder_score"], reverse=True) | |
| if top_k is not None: | |
| reranked = reranked[:top_k] | |
| return reranked | |