234 lines
8.0 KiB
Python
234 lines
8.0 KiB
Python
"""
|
|
FAISS-backed ANN index over paper embeddings.
|
|
|
|
The index persists three things to disk so a query process never has to
|
|
re-embed the full corpus:
|
|
|
|
* ``<name>.index`` -- the FAISS flat L2 index (float32 vectors)
|
|
* ``<name>.meta.pkl`` -- parallel arrays: paper_id, year, category, url,
|
|
title, abstract (aligned to the FAISS rows)
|
|
* ``<name>.meta.json`` -- provenance (embedder name, dim, source, counts)
|
|
|
|
The index is *incremental*: ``build_index`` can be pointed at a fresh scan of
|
|
the metadata and will add any papers not already present (keyed by paper_id),
|
|
rather than rebuilding from scratch every time.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import pickle
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
|
|
try:
|
|
import faiss
|
|
except ImportError: # pragma: no cover - faiss is a hard dependency
|
|
faiss = None # type: ignore[assignment]
|
|
|
|
from data import DataStore, Paper
|
|
|
|
|
|
def _require_faiss() -> None:
|
|
if faiss is None:
|
|
raise RuntimeError("faiss-cpu is required for ANN indexing. Install it (see requirements.txt).")
|
|
|
|
|
|
@dataclass
|
|
class IndexMeta:
|
|
embedder: str
|
|
dim: int
|
|
source: str
|
|
count: int
|
|
built_at: str = ""
|
|
|
|
|
|
class AnnIndex:
|
|
"""
|
|
Read-only view over a persisted FAISS index, exposing the parallel metadata
|
|
arrays needed to attach titles / years / categories / URLs to ANN hits.
|
|
"""
|
|
|
|
def __init__(self, builder: IndexBuilder, meta: IndexMeta):
|
|
self.index = builder.index
|
|
self.meta = meta
|
|
self._ids = builder._ids
|
|
self._years = builder._years
|
|
self._cats = builder._cats
|
|
self._urls = builder._urls
|
|
self._titles = builder._titles
|
|
self._abstracts = builder._abstracts
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._ids)
|
|
|
|
|
|
class IndexBuilder:
|
|
"""
|
|
Build (and incrementally extend) a FAISS index from a DataStore scan.
|
|
|
|
Example
|
|
-------
|
|
>>> builder = IndexBuilder("arxiv", embedder=HashingEmbedder())
|
|
>>> store = DataStore()
|
|
>>> builder.build(store.scan_filtered(), store) # doctest: +SKIP
|
|
>>> builder.save("arxiv")
|
|
"""
|
|
|
|
def __init__(self, name: str, embedder, chunk_size: int = 4096):
|
|
_require_faiss()
|
|
self.name = name
|
|
self.embedder = embedder
|
|
self.dim = int(embedder.DIM)
|
|
self.chunk_size = chunk_size
|
|
self._existing: set[str] = set()
|
|
self._ids: list[str] = []
|
|
self._years: list[int] = []
|
|
self._cats: list[str] = []
|
|
self._urls: list[str] = []
|
|
self._titles: list[str] = []
|
|
self._abstracts: list[str] = []
|
|
self._loaded = False
|
|
|
|
# -- loading an existing index ------------------------------------------
|
|
def load(self, path: str | None = None) -> None:
|
|
path = path or self.default_path()
|
|
if not os.path.exists(path):
|
|
return
|
|
self.index = faiss.read_index(path)
|
|
with open(path.replace(".index", ".meta.pkl"), "rb") as fh:
|
|
blob = pickle.load(fh)
|
|
(
|
|
self._ids,
|
|
self._years,
|
|
self._cats,
|
|
self._urls,
|
|
self._titles,
|
|
self._abstracts,
|
|
) = blob
|
|
self._existing = set(self._ids)
|
|
self._loaded = True
|
|
|
|
def default_path(self, root: str | None = None) -> str:
|
|
root = root or os.path.join(os.path.dirname(__file__), "index")
|
|
return os.path.join(root, f"{self.name}.index")
|
|
|
|
# -- building -----------------------------------------------------------
|
|
def build(self, scan, store: DataStore) -> int:
|
|
"""
|
|
Embed every row of ``scan`` (a polars LazyFrame) and add new papers.
|
|
|
|
``scan`` must project the DISPLAY_COLUMNS; ``store`` is used only for its
|
|
provenance metadata. Returns the number of papers newly added.
|
|
"""
|
|
import datetime as _dt
|
|
|
|
if not self._loaded:
|
|
self.index = faiss.IndexFlatL2(self.dim)
|
|
added = 0
|
|
batch: list[str] = []
|
|
batch_ids: list[str] = []
|
|
batch_years: list[int] = []
|
|
batch_cats: list[str] = []
|
|
batch_urls: list[str] = []
|
|
batch_titles: list[str] = []
|
|
batch_abstracts: list[str] = []
|
|
|
|
def flush() -> None:
|
|
nonlocal added
|
|
if not batch:
|
|
return
|
|
# Only embed/add papers not already in the index.
|
|
keep = [i for i, pid in enumerate(batch_ids) if pid not in self._existing]
|
|
if not keep:
|
|
batch.clear()
|
|
batch_ids.clear()
|
|
batch_years.clear()
|
|
batch_cats.clear()
|
|
batch_urls.clear()
|
|
batch_titles.clear()
|
|
batch_abstracts.clear()
|
|
return
|
|
keep_ids = [batch_ids[i] for i in keep]
|
|
keep_texts = [batch[i] for i in keep]
|
|
vecs = self.embedder.embed(keep_texts)
|
|
vecs = np.ascontiguousarray(vecs, dtype=np.float32)
|
|
self.index.add(vecs)
|
|
for i in keep:
|
|
self._ids.append(batch_ids[i])
|
|
self._years.append(batch_years[i])
|
|
self._cats.append(batch_cats[i])
|
|
self._urls.append(batch_urls[i])
|
|
self._titles.append(batch_titles[i])
|
|
self._abstracts.append(batch_abstracts[i])
|
|
self._existing.add(batch_ids[i])
|
|
added += len(keep_ids)
|
|
batch.clear()
|
|
batch_ids.clear()
|
|
batch_years.clear()
|
|
batch_cats.clear()
|
|
batch_urls.clear()
|
|
batch_titles.clear()
|
|
batch_abstracts.clear()
|
|
|
|
# Collect all matching rows, then embed in chunks.
|
|
rows = scan.select(
|
|
["paper_id", "title", "primary_category", "first_version_date", "arxiv_abs_url", "abstract"]
|
|
).collect().to_dicts()
|
|
|
|
for r in rows:
|
|
rec = dict(r)
|
|
fvd = rec["first_version_date"]
|
|
year = fvd.year if isinstance(fvd, _dt.datetime) else (int(fvd) if fvd is not None else 0)
|
|
pid = str(rec["paper_id"])
|
|
if pid in self._existing:
|
|
continue
|
|
text = f"{rec['title']} {rec['abstract']}"
|
|
batch.append(text)
|
|
batch_ids.append(pid)
|
|
batch_years.append(year)
|
|
batch_cats.append(str(rec["primary_category"] or ""))
|
|
batch_urls.append(str(rec["arxiv_abs_url"] or ""))
|
|
batch_titles.append(str(rec["title"] or ""))
|
|
batch_abstracts.append(str(rec["abstract"] or ""))
|
|
if len(batch) >= self.chunk_size:
|
|
flush()
|
|
flush()
|
|
return added
|
|
|
|
# -- persistence --------------------------------------------------------
|
|
def save(self, path: str | None = None) -> IndexMeta:
|
|
path = path or self.default_path()
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
faiss.write_index(self.index, path)
|
|
meta_pkl = path.replace(".index", ".meta.pkl")
|
|
with open(meta_pkl, "wb") as fh:
|
|
pickle.dump(
|
|
(self._ids, self._years, self._cats, self._urls, self._titles, self._abstracts),
|
|
fh,
|
|
)
|
|
meta = IndexMeta(
|
|
embedder=type(self.embedder).__name__,
|
|
dim=self.dim,
|
|
source=os.environ.get("ARXIV_METADATA_PARQUET", "arxiv-complete/metadata"),
|
|
count=len(self._ids),
|
|
)
|
|
with open(path.replace(".index", ".meta.json"), "w") as fh:
|
|
json.dump(meta.__dict__, fh, indent=2)
|
|
return meta
|
|
|
|
|
|
def load_index(name: str, embedder, root: str | None = None) -> tuple["AnnIndex", IndexMeta]:
|
|
"""Load a previously-built index plus its embedder for querying."""
|
|
_require_faiss()
|
|
root = root or os.path.join(os.path.dirname(__file__), "index")
|
|
path = os.path.join(root, f"{name}.index")
|
|
builder = IndexBuilder(name, embedder)
|
|
builder.load(path)
|
|
meta_path = path.replace(".index", ".meta.json")
|
|
with open(meta_path) as fh:
|
|
meta = IndexMeta(**json.load(fh))
|
|
return AnnIndex(builder, meta), meta
|