Spaces:
Sleeping
Sleeping
| 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)] |