Create semantic/faiss_index.py

semantic/faiss_index.py

@@ -1,26 +1,21 @@
 import faiss
-import numpy as np
 from sentence_transformers import SentenceTransformer
 
-MODEL = SentenceTransformer("all-MiniLM-L6-v2")
-
 class SemanticIndex:
     def __init__(self):
+        self.model = SentenceTransformer("all-MiniLM-L6-v2")
         self.index = None
         self.texts = []
 
     def build(self, texts):
         self.texts = texts
-        embeddings = MODEL.encode(texts, convert_to_numpy=True)
-        dim = embeddings.shape[1]
-        self.index = faiss.IndexFlatL2(dim)
-        self.index.add(embeddings)
+        emb = self.model.encode(texts)
+        self.index = faiss.IndexFlatL2(len(emb[0]))
+        self.index.add(emb)
 
     def search(self, query, k=10):
         if not self.index:
             return []
-
-        q_emb = MODEL.encode([query], convert_to_numpy=True)
-        distances, indices = self.index.search(q_emb, k)
-
-        return [self.texts[i] for i in indices[0] if i < len(self.texts)]
+        q = self.model.encode([query])
+        _, idx = self.index.search(q, k)
+        return [self.texts[i] for i in idx[0]]