Spaces:
Sleeping
Sleeping
| import faiss | |
| import numpy as np | |
| def generate_faiss_index(embeddings): | |
| # Ensure that the embeddings are converted to np.float32 (FAISS expects float32) | |
| embeddings = np.array(embeddings, dtype=np.float32) | |
| index = faiss.IndexFlatL2(768) # Assuming 768-dimensional embeddings for a model like MiniLM | |
| index.add(embeddings) | |
| return index | |
| def load_faiss_index_to_gpu(index): | |
| # If you're using GPU, ensure the index is moved to the GPU | |
| res = faiss.StandardGpuResources() # Create resources for the GPU | |
| gpu_index = faiss.index_cpu_to_gpu(res, 0, index) # Load into GPU (assuming GPU 0 is available) | |
| return gpu_index | |
| def query_faiss_index(query_embedding, gpu_index): | |
| # Query the FAISS index with the query embedding | |
| query_embedding = np.array(query_embedding, dtype=np.float32) # Ensure the query is a np.array with the right type | |
| distances, indices = gpu_index.search(query_embedding.reshape(1, -1), 1) # Reshaping as FAISS expects 2D array | |
| return indices, distances |