Upload src/labdaps/retrieval/retriever.py with huggingface_hub

src/labdaps/retrieval/retriever.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass
+import numpy as np
+from src.labdaps.config import TOP_K_CANDIDATES, TOP_K_FINAL
+from src.labdaps.retrieval.vector_store import query_store
+from src.labdaps.ingestion.embedder import Embedder
+
+
+@dataclass
+class RetrievedChunk:
+    text: str
+    source_file: str
+    page_number: int
+    score: float
+
+
+def _cosine_similarity(a: list[float], b: list[float]) -> float:
+    a, b = np.array(a), np.array(b)
+    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))
+
+
+def _mmr(
+    query_embedding: list[float],
+    candidate_embeddings: list[list[float]],
+    candidates: list[dict],
+    top_k: int,
+    lambda_param: float = 0.6,
+) -> list[int]:
+    selected = []
+    remaining = list(range(len(candidates)))
+    query_sims = [_cosine_similarity(query_embedding, e) for e in candidate_embeddings]
+
+    while len(selected) < top_k and remaining:
+        if not selected:
+            best = max(remaining, key=lambda i: query_sims[i])
+        else:
+            selected_embeddings = [candidate_embeddings[i] for i in selected]
+            best = max(
+                remaining,
+                key=lambda i: lambda_param * query_sims[i]
+                - (1 - lambda_param) * max(
+                    _cosine_similarity(candidate_embeddings[i], se)
+                    for se in selected_embeddings
+                ),
+            )
+        selected.append(best)
+        remaining.remove(best)
+
+    return selected
+
+
+def retrieve(query: str, embedder: Embedder) -> list[RetrievedChunk]:
+    query_embedding = embedder.embed_query(query)
+    documents, metadatas, distances = query_store(query_embedding, TOP_K_CANDIDATES)
+
+    if not documents:
+        return []
+
+    candidate_embeddings = embedder.embed_passages(documents)
+    selected_indices = _mmr(query_embedding, candidate_embeddings, metadatas, TOP_K_FINAL)
+
+    return [
+        RetrievedChunk(
+            text=documents[idx],
+            source_file=metadatas[idx]["source_file"],
+            page_number=metadatas[idx]["page_number"],
+            score=1.0 - distances[idx],
+        )
+        for idx in selected_indices
+    ]