HF-LLM-Intent-Detection / src /E_OLD_Faiss_utils.py
georgeek's picture
Transfer
5ecde30
import faiss
import numpy as np
def save_faiss_embeddings_index(embeddings, file_name):
# Ensure embeddings are in float32 format
if not isinstance(embeddings, np.ndarray):
embeddings = embeddings.numpy()
embeddings = embeddings.astype('float32')
# Create a FAISS index
index = faiss.IndexFlatL2(embeddings.shape[1]) # L2 distance
index.add(embeddings)
# Save the FAISS index
faiss.write_index(index, file_name)
def load_faiss_index(index_path):
index = faiss.read_index(index_path)
return index
def normalize_embeddings(embeddings):
# Normalize embeddings
embeddings = embeddings / np.linalg.norm(embeddings, axis=1)[:, None]
return embeddings
def search_faiss_index(index, query_embedding, k=5):
# Perform similarity search
D, I = index.search(query_embedding, k) # D: distances, I: indices
return D, I
def Z_load_embeddings_and_index(file_name):
# Load embeddings from .npy file
embeddings = np.load(f"{file_name}_embeddings.npy")
# Load FAISS index from .index file
index = faiss.read_index(file_name)
return embeddings, index