Update src/retrieval.py

src/retrieval.py

@@ -1,29 +1,23 @@
-import numpy as np
 from sentence_transformers import SentenceTransformer
+import numpy as np
 import json
-import os
-from typing import List, Dict
 
 class DocumentRetriever:
     def __init__(self, model_name='all-MiniLM-L6-v2'):
         self.model = SentenceTransformer(model_name)
         self.documents = self._load_documents()
         self.doc_embeddings = self._embed_documents()
-
-    def _load_documents(self) -> List[Dict]:
+    
+    def _load_documents(self):
         with open('data/rupeia_document.json', 'r') as f:
             return json.load(f)
-
-    def _embed_documents(self) -> np.ndarray:
+    
+    def _embed_documents(self):
         texts = [doc['content'] for doc in self.documents]
         return self.model.encode(texts)
-
-    def retrieve(self, query: str, top_k: int = 3) -> List[Dict]:
+    
+    def retrieve(self, query: str, top_k: int = 3):
         query_embedding = self.model.encode(query)
         scores = np.dot(self.doc_embeddings, query_embedding)
         top_indices = np.argsort(scores)[-top_k:][::-1]
-        return [self.documents[i] for i in top_indices]
-
-def retrieve_relevant_documents(query: str) -> List[Dict]:
-    retriever = DocumentRetriever()
-    return retriever.retrieve(query)
+        return [self.documents[i] for i in top_indices]