feat: reranker

src/scientific_rag/application/reranking/cross_encoder.py

@@ -0,0 +1,35 @@
+from loguru import logger
+from sentence_transformers import CrossEncoder
+
+from scientific_rag.domain.documents import PaperChunk
+from scientific_rag.settings import settings
+
+
+class CrossEncoderReranker:
+    def __init__(self):
+        self.model_name = settings.reranker_model_name
+        self.device = settings.embedding_device
+
+        logger.info(f"Loading Reranker model: {self.model_name} on {self.device}")
+        self.model = CrossEncoder(self.model_name, device=self.device)
+        logger.info("Reranker model loaded")
+
+    def rerank(self, query: str, chunks: list[PaperChunk], top_k: int = settings.rerank_top_k) -> list[PaperChunk]:
+        if not chunks:
+            return []
+
+        pairs = [[query, chunk.text] for chunk in chunks]
+        scores = self.model.predict(pairs)
+
+        reranked_chunks = []
+        for i, chunk in enumerate(chunks):
+            chunk_copy = chunk.model_copy()
+            chunk_copy.score = float(scores[i])
+            reranked_chunks.append(chunk_copy)
+
+        reranked_chunks.sort(key=lambda x: x.score, reverse=True)
+
+        return reranked_chunks[:top_k]
+
+
+reranker = CrossEncoderReranker()