Spaces:
Running
Running
File size: 880 Bytes
8e72e1f | 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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | import faiss
import numpy as np
class VectorStore:
def __init__(self):
self.index = None
self.documents = []
def build(self, embeddings, documents):
dimension = embeddings.shape[1]
self.index = faiss.IndexFlatIP(
dimension
)
self.index.add(
np.array(embeddings)
)
self.documents = documents
def search(self, query_embedding, k=3):
scores, indexes = self.index.search(
np.array([query_embedding]),
k
)
results=[]
for score, idx in zip(
scores[0],
indexes[0]
):
results.append(
{
"score":float(score),
"document":
self.documents[idx]
}
)
return results |