Update src/retrieval.py

src/retrieval.py

@@ -1,22 +1,23 @@
-from sentence_transformers import SentenceTransformer
 import numpy as np
+from sentence_transformers import SentenceTransformer
 import json
+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):
+
+    def _load_documents(self) -> List[Dict]:
         with open('data/rupeia_document.json', 'r') as f:
             return json.load(f)
-    
-    def _embed_documents(self):
+
+    def _embed_documents(self) -> np.ndarray:
         texts = [doc['content'] for doc in self.documents]
         return self.model.encode(texts)
-    
-    def retrieve(self, query: str, top_k: int = 3):
+
+    def retrieve(self, query: str, top_k: int = 3) -> List[Dict]:
         query_embedding = self.model.encode(query)
         scores = np.dot(self.doc_embeddings, query_embedding)
         top_indices = np.argsort(scores)[-top_k:][::-1]