File size: 1,026 Bytes
1d1ee9b ee7e79e 1d1ee9b | 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 | """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]
|