reranker.py

# reranker.py
from sentence_transformers import CrossEncoder


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

    def rerank(self, query, candidates):
        pairs = [(query, c["text"]) for c in candidates]
        scores = self.model.predict(pairs)

        sorted_items = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
        top = [item[0] for item in sorted_items]
        return top