RobustDialogueDemo / semantic.py
DialogueRobust
First commit
e490ab5
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