Spaces:
Running
Running
File size: 678 Bytes
5a38dc7 | 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 | import faiss
import numpy as np
class VectorStorage:
def __init__(self, dimension=384):
self.index = faiss.IndexFlatIP(dimension)
self.chunks = []
def add(self, vectors, chunks):
v_array = np.array(vectors).astype('float32')
faiss.normalize_L2(v_array)
self.index.add(v_array)
self.chunks.extend(chunks)
def search(self, query_vector, k=5):
q_array = np.array([query_vector]).astype('float32')
faiss.normalize_L2(q_array)
scores, indices = self.index.search(q_array, k)
return scores[0], indices[0] |