Upload retriever.py

api/retriever.py

@@ -0,0 +1,26 @@
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+TOP_K = 5
+
+class ChunkRetriever:
+    """Stage 1 Bi-Encoder: quickly narrows down hundreds of chunks
+    to the few that are actually semantically relevant to the LLM output."""
+
+    def __init__(self):
+        self.model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+        print("Retriever (MiniLM-L6-v2) loaded.")
+
+    def get_top_chunks(self, llm_output: str, chunks: list[str], top_k: int = TOP_K) -> list[str]:
+        """Embeds everything, ranks by cosine similarity, returns the top_k chunks."""
+        if len(chunks) <= top_k:
+            return chunks
+
+        query_embedding = self.model.encode(llm_output, normalize_embeddings=True)
+        chunk_embeddings = self.model.encode(chunks, normalize_embeddings=True, batch_size=32)
+
+        # cosine sim is just dot product when vectors are already L2-normalized
+        similarities = np.dot(chunk_embeddings, query_embedding)
+
+        top_indices = np.argsort(similarities)[::-1][:top_k]
+        return [chunks[i] for i in top_indices]