File size: 552 Bytes
e490ab5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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