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_results(self, query: str, results: list[dict], top_n: int = 5) -> list[dict]:
        pairs = [(query, r["text"]) for r in results if r.get("text")]
        scores = self.model.predict(pairs)
        scored_results = sorted(zip(scores, results), key=lambda x: x[0], reverse=True)
        return [r for _, r in scored_results[:top_n]]