File size: 1,689 Bytes
1200d57 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import os
import faiss
import numpy as np
from src.customlogger import log_time, logger
# Type alias to decouple from FAISS
Index = faiss.IndexFlat
INDEX_FILE = "proverbs.index"
# This were the index types tested
INDEX_TYPES = [faiss.IndexFlatL2, faiss.IndexFlatIP]
# This is the pooling method used in the final iteration
DEFAULT_INDEX_TYPE = faiss.IndexFlatL2
def index_exists(index_file: str = INDEX_FILE) -> bool:
"""Check if the index file exists."""
return os.path.exists(index_file)
@log_time
def create_index(embeddings: np.ndarray, index_type: type = None, index_file: str = INDEX_FILE) -> Index:
"""Create a FAISS index and store the given embeddings."""
if not index_type:
index_type = DEFAULT_INDEX_TYPE
dimension = embeddings.shape[1]
logger.debug(
f"Creating FAISS index with {len(embeddings)} {embeddings.shape[1]}-dimensional embeddings...")
index = index_type(dimension)
index.add(embeddings)
logger.debug(f"Saving FAISS index to '{index_file}'...")
faiss.write_index(index, index_file)
return index
@log_time
def load_index(index_file: str = INDEX_FILE) -> Index:
"""Load the FAISS index from the specified file."""
logger.debug(f"Loading FAISS index from '{index_file}'...")
index = faiss.read_index(index_file)
return index
@log_time
def find_closest(embeddings, index: Index, k=5) -> tuple[np.ndarray, np.ndarray]:
"""Find the closest k vectors in the index for the given embeddings."""
logger.debug(
f"Performing search for the top {k} matches of {len(embeddings)} embedding...")
distances, indices = index.search(embeddings, k)
return distances, indices
|