File size: 5,443 Bytes
45383d0 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """Load a FOTO caption index: vectors (FAISS or npy) plus row-aligned metadata,
joining full captions/titles from astro_captions.parquet when missing."""
import json
from pathlib import Path
import numpy as np
import pandas as pd
FIG_COL_ALIASES = ["fig_index", "figure_index", "fig_idx", "figure_idx",
"fig_num", "caption_index"]
CAPTION_ALIASES = ["caption", "caption_text", "text", "caption_prefix"]
class LoadedIndex:
"""Uniform search interface over a FAISS index or a raw numpy matrix."""
def __init__(self, faiss_index=None, matrix=None):
assert (faiss_index is None) != (matrix is None)
self.faiss = faiss_index
self.matrix = matrix
self.n = faiss_index.ntotal if faiss_index is not None else len(matrix)
def search(self, query_vecs, k):
q = np.ascontiguousarray(query_vecs, dtype=np.float32)
if q.ndim == 1:
q = q[None, :]
if self.faiss is not None:
return self.faiss.search(q, k)
sims = q @ self.matrix.T
ids = np.argsort(-sims, axis=1)[:, :k]
return np.take_along_axis(sims, ids, axis=1), ids
def vectors(self, row_ids):
if self.faiss is not None:
import faiss
ids = np.asarray(row_ids, dtype=np.int64)
if hasattr(self.faiss, "reconstruct_batch"):
return self.faiss.reconstruct_batch(ids)
return np.stack([self.faiss.reconstruct(int(i)) for i in ids])
return np.asarray(self.matrix[np.asarray(row_ids)])
def _find_one(index_dir, patterns, what):
hits = [p for pat in patterns for p in sorted(index_dir.glob(pat))]
if not hits:
raise FileNotFoundError(
f"no {what} found in {index_dir} (looked for {patterns}); "
f"contents: {[p.name for p in index_dir.iterdir()]}")
if len(hits) > 1:
print(f" multiple {what} candidates {[h.name for h in hits]}, "
f"using {hits[0].name}")
return hits[0]
def load_caption_index(index_dir, captions_parquet=None):
"""Returns (LoadedIndex, meta DataFrame with columns
arxiv_id, fig_index, title, caption -- row-aligned to the index)."""
index_dir = Path(index_dir)
info_path = index_dir / "info.json"
if info_path.exists():
print(f"info.json: {json.load(open(info_path))}")
try:
vec_path = _find_one(index_dir, ["*.faiss", "*.index"], "FAISS index")
import faiss
loaded = LoadedIndex(faiss_index=faiss.read_index(str(vec_path)))
print(f"vectors: {vec_path.name} (FAISS, {loaded.n} rows)")
except FileNotFoundError:
vec_path = _find_one(index_dir, ["*.npy"], "embedding matrix")
loaded = LoadedIndex(matrix=np.load(vec_path, mmap_mode="r"))
print(f"vectors: {vec_path.name} (numpy, {loaded.n} rows)")
meta_path = _find_one(index_dir, ["*.parquet", "*.jsonl"], "metadata")
meta = pd.read_parquet(meta_path) if meta_path.suffix == ".parquet" \
else pd.read_json(meta_path, lines=True)
print(f"metadata: {meta_path.name}, columns {list(meta.columns)}")
if len(meta) != loaded.n:
raise ValueError(f"metadata rows ({len(meta)}) != index rows "
f"({loaded.n}); wrong file pair?")
if "arxiv_id" not in meta.columns:
raise ValueError(f"no arxiv_id column in {meta_path.name}")
fig_col = next((c for c in FIG_COL_ALIASES if c in meta.columns), None)
if fig_col is None:
raise ValueError(f"no figure-index column in {meta_path.name} "
f"(tried {FIG_COL_ALIASES})")
meta = meta.rename(columns={fig_col: "fig_index"})
cap_col = next((c for c in CAPTION_ALIASES if c in meta.columns), None)
if cap_col and cap_col != "caption":
meta = meta.rename(columns={cap_col: "caption"})
need_caption = "caption" not in meta.columns or cap_col == "caption_prefix"
need_title = "title" not in meta.columns
if need_caption or need_title:
if captions_parquet is None:
raise ValueError(
"index metadata lacks full caption/title columns; pass "
"captions_parquet to join them")
src = pd.read_parquet(captions_parquet)
if "captions" not in src.columns:
raise ValueError(f"{captions_parquet} has no 'captions' column; "
f"columns: {list(src.columns)}")
long = src[["arxiv_id", "title", "captions"]].explode("captions")
# the index uses 1-based figure positions (generate_subscripts)
long["fig_index"] = long.groupby("arxiv_id").cumcount() + 1
long = long.rename(columns={"captions": "caption_full"})
meta = meta.merge(long, on=["arxiv_id", "fig_index"], how="left",
suffixes=("", "_joined"))
if need_caption:
missing = int(meta.caption_full.isna().sum())
if missing:
print(f" WARNING: {missing} rows had no caption after join")
meta["caption"] = meta.caption_full.fillna("")
if need_title:
tcol = "title_joined" if "title_joined" in meta.columns else "title"
meta["title"] = meta[tcol].fillna("")
print(f"joined captions/titles from {captions_parquet}")
meta["fig_index"] = meta.fig_index.astype(int)
meta = meta[["arxiv_id", "fig_index", "title", "caption"]]
return loaded, meta.reset_index(drop=True)
|