common/vector_index.py

"""FAISS IndexFlatIP helpers for cosine-similarity search over paper embeddings."""
import os
import faiss
import numpy as np


def create_index(dim: int) -> faiss.IndexFlatIP:
    return faiss.IndexFlatIP(dim)


def load_index(path: str) -> faiss.IndexFlatIP:
    if not os.path.exists(path):
        raise FileNotFoundError(f"FAISS index not found at {path}")
    return faiss.read_index(path)


def save_index(index: faiss.IndexFlatIP, path: str) -> None:
    faiss.write_index(index, path)


def _normalize(vector: np.ndarray) -> np.ndarray:
    vector = vector.astype(np.float32).reshape(1, -1)
    faiss.normalize_L2(vector)
    return vector


def add_vector(index: faiss.IndexFlatIP, vector: np.ndarray) -> int:
    index.add(_normalize(vector))
    return index.ntotal - 1


def search(index: faiss.IndexFlatIP, query_vector: np.ndarray, k: int = 10) -> list[tuple[int, float]]:
    scores, ids = index.search(_normalize(query_vector), k)
    return [(int(i), float(s)) for i, s in zip(ids[0], scores[0]) if i != -1]