Engineer786 commited on
Commit
58ee1b3
·
verified ·
1 Parent(s): 7beb60d

Create utils.py

Browse files
Files changed (1) hide show
  1. utils.py +16 -0
utils.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import faiss
2
+ import numpy as np
3
+
4
+ def generate_faiss_index(embeddings):
5
+ index = faiss.IndexFlatL2(768) # Assuming 768-dimensional embeddings for a model like MiniLM
6
+ index.add(np.array(embeddings))
7
+ return index
8
+
9
+ def load_faiss_index_to_gpu(index):
10
+ res = faiss.StandardGpuResources()
11
+ gpu_index = faiss.index_cpu_to_gpu(res, 0, index) # Load into GPU (assuming GPU 0 is available)
12
+ return gpu_index
13
+
14
+ def query_faiss_index(query_embedding, gpu_index):
15
+ distances, indices = gpu_index.search(np.array([query_embedding]), 1)
16
+ return indices, distances