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]