Spaces:
Running
Running
File size: 2,042 Bytes
ed65693 | 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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """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
@property
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
|