FOIA_Doc_Search / semantic /faiss_index.py
GodsDevProject's picture
Create semantic/faiss_index.py
8aeaa64 verified
import faiss
from sentence_transformers import SentenceTransformer
class SemanticIndex:
def __init__(self):
self.model = SentenceTransformer("all-MiniLM-L6-v2")
self.index = None
self.texts = []
def build(self, texts):
self.texts = texts
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 = self.model.encode([query])
_, idx = self.index.search(q, k)
return [self.texts[i] for i in idx[0]]