reranker.py

from typing import List, Dict, Any
from sentence_transformers import CrossEncoder


class Reranker:
    def __init__(
        self,
        model_name: str = "cross-encoder/ms-marco-MiniLM-L-4-v2",
        device: str = "cpu",
    ):
        self.model = CrossEncoder(model_name, device=device)

    def rerank(
        self, query: str, documents: List[Dict[str, Any]], top_k: int = 10
    ) -> List[Dict[str, Any]]:
        if not documents:
            return []

        pairs = [[query, doc["text"]] for doc in documents]
        scores = self.model.predict(pairs, batch_size=32, show_progress_bar=False)

        for doc, score in zip(documents, scores):
            doc["rerank_score"] = float(score)

        ranked = sorted(documents, key=lambda x: x["rerank_score"], reverse=True)
        return ranked[:top_k]