Update src/retrieval.py

src/retrieval.py

@@ -1,19 +1,37 @@
-from langchain_community.vectorstores import FAISS
-from langchain.docstore.document import Document
-from langchain_huggingface import HuggingFaceEmbeddings
 import json
+import numpy as np
+from langchain_huggingface import HuggingFaceEmbeddings
 
 class RetrievalSystem:
     def __init__(self, document_path, embedder_model):
         self.embedder = HuggingFaceEmbeddings(model_name=embedder_model)
-        self.vectorstore = self._build_vectorstore(document_path)
+        self.documents = []
+        self.embeddings = None
+        self._load_documents(document_path)
+        self._build_index()
 
-    def _build_vectorstore(self, document_path):
+    def _load_documents(self, document_path):
         with open(document_path, "r") as f:
             docs_data = json.load(f)
-        documents = [Document(page_content=doc["content"], metadata={"category": doc["category"], "subkeyword": doc["subkeyword"]}) for doc in docs_data]
-        return FAISS.from_documents(documents, embedding=self.embedder)
+        self.documents = [(doc["content"], doc["metadata"]) for doc in docs_data]
+
+    def _build_index(self):
+        texts = [doc[0] for doc in self.documents]
+        self.embeddings = self.embedder.embed_documents(texts)
 
     def get_context(self, query, k=2):
-        docs = self.vectorstore.similarity_search(query, k=k)
-        return " ".join([doc.page_content for doc in docs])
+        # Embed the query
+        query_embedding = self.embedder.embed_query(query)
+
+        # Compute cosine similarity
+        embeddings = np.array(self.embeddings)
+        query_embedding = np.array(query_embedding)
+        similarities = np.dot(embeddings, query_embedding) / (
+            np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query_embedding)
+        )
+
+        # Get top k documents
+        top_k_indices = np.argsort(similarities)[-k:][::-1]
+        top_k_docs = [self.documents[i][0] for i in top_k_indices]
+
+        return " ".join(top_k_docs)