Spaces:
Sleeping
Sleeping
| import faiss | |
| import numpy as np | |
| import pickle | |
| INDEX_PATH = "data/index.faiss" | |
| DOCS_PATH = "data/docs.pkl" | |
| class SemanticSearch: | |
| def __init__(self): | |
| self.index = faiss.read_index(INDEX_PATH) | |
| with open(DOCS_PATH, "rb") as f: | |
| self.docs = pickle.load(f) | |
| def search(self, query_emb, k=3): | |
| D, I = self.index.search(np.array([query_emb]).astype("float32"), k) | |
| results = [] | |
| for idx in I[0]: | |
| if idx >= 0: | |
| results.append(self.docs[idx]['section']) | |
| return results | |