FOIA_Doc_Search / core /faiss_vector.py
GodsDevProject's picture
Upload 20 files
5830944 verified
raw
history blame
769 Bytes
try:
import faiss
from sentence_transformers import SentenceTransformer
except ImportError:
faiss = None
class FaissIndex:
def __init__(self, model_name="all-MiniLM-L6-v2"):
if faiss is None:
raise RuntimeError("FAISS not installed")
self.model = SentenceTransformer(model_name)
self.index = None
self.docs = []
def add(self, texts):
emb = self.model.encode(texts)
if self.index is None:
self.index = faiss.IndexFlatL2(emb.shape[1])
self.index.add(emb)
self.docs.extend(texts)
def search(self, query, k=5):
emb = self.model.encode([query])
D, I = self.index.search(emb, k)
return [self.docs[i] for i in I[0] if i < len(self.docs)]