| 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]] |