akazmi commited on
Commit
3c058b0
·
verified ·
1 Parent(s): ccaec45

Create utils.py

Browse files
Files changed (1) hide show
  1. utils.py +21 -0
utils.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import faiss
2
+ import numpy as np
3
+
4
+ def generate_faiss_index(embeddings):
5
+ # Ensure that the embeddings are converted to np.float32 (FAISS expects float32)
6
+ embeddings = np.array(embeddings, dtype=np.float32)
7
+ index = faiss.IndexFlatL2(768) # Assuming 768-dimensional embeddings for a model like MiniLM
8
+ index.add(embeddings)
9
+ return index
10
+
11
+ def load_faiss_index_to_gpu(index):
12
+ # If you're using GPU, ensure the index is moved to the GPU
13
+ res = faiss.StandardGpuResources() # Create resources for the GPU
14
+ gpu_index = faiss.index_cpu_to_gpu(res, 0, index) # Load into GPU (assuming GPU 0 is available)
15
+ return gpu_index
16
+
17
+ def query_faiss_index(query_embedding, gpu_index):
18
+ # Query the FAISS index with the query embedding
19
+ query_embedding = np.array(query_embedding, dtype=np.float32) # Ensure the query is a np.array with the right type
20
+ distances, indices = gpu_index.search(query_embedding.reshape(1, -1), 1) # Reshaping as FAISS expects 2D array
21
+ return indices, distances