Spaces:
Sleeping
Sleeping
File size: 769 Bytes
5830944 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
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)] |