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)]