| | import faiss |
| | import numpy as np |
| |
|
| | def generate_faiss_index(embeddings): |
| | |
| | embeddings = np.array(embeddings, dtype=np.float32) |
| | index = faiss.IndexFlatL2(768) |
| | index.add(embeddings) |
| | return index |
| |
|
| | def load_faiss_index_to_gpu(index): |
| | |
| | res = faiss.StandardGpuResources() |
| | gpu_index = faiss.index_cpu_to_gpu(res, 0, index) |
| | return gpu_index |
| |
|
| | def query_faiss_index(query_embedding, gpu_index): |
| | |
| | query_embedding = np.array(query_embedding, dtype=np.float32) |
| | distances, indices = gpu_index.search(query_embedding.reshape(1, -1), 1) |
| | return indices, distances |
| |
|