backend/scripts/build_corpus.py

"""Build the compact local corpus from the downloaded HF parquet shards.

The full kiyer/pathfinder_arxiv_data dataset is ~20 GB (float64 embeddings +
citation graphs). This machine has 16 GB RAM and limited disk, so we keep only
what retrieval needs:
  - text columns -> backend/data/corpus/data-NNNNN.parquet  (~1.3 GB total)
  - embeddings   -> a float16-quantized faiss index          (~3.8 GB)

Each source shard is DELETED from the HF cache after it is consumed (pass
--keep-sources to disable) so peak disk usage decreases as the build runs.
Shards are re-downloadable from the Hub if ever needed again.

Run from backend/:  uv run python scripts/build_corpus.py
"""

import argparse
import glob
import os
import sys
from pathlib import Path

import faiss
import numpy as np
import pyarrow.parquet as pq

TEXT_COLUMNS = ["title", "abstract", "authors", "date", "keywords", "bibcode", "arxiv_id"]
EMBED_DIM = 1536

DATA_DIR = Path(__file__).resolve().parent.parent / "data"
HUB_GLOB = os.path.expanduser(
    "~/.cache/huggingface/hub/datasets--kiyer--pathfinder_arxiv_data/snapshots/*/data/*.parquet"
)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--keep-sources", action="store_true",
                    help="do not delete consumed HF cache shards")
    args = ap.parse_args()

    shards = sorted(glob.glob(HUB_GLOB))
    if not shards:
        print(f"no source shards found at {HUB_GLOB}", file=sys.stderr)
        return 1

    corpus_dir = DATA_DIR / "corpus"
    corpus_dir.mkdir(parents=True, exist_ok=True)
    faiss_path = DATA_DIR / "astroparse_fp16.faiss"

    index = faiss.IndexScalarQuantizer(
        EMBED_DIM, faiss.ScalarQuantizer.QT_fp16, faiss.METRIC_L2
    )

    total_rows = 0
    for i, shard in enumerate(shards):
        table = pq.read_table(shard, columns=TEXT_COLUMNS + ["embed"])
        flat = table.column("embed").combine_chunks().flatten().to_numpy(zero_copy_only=False)
        embeds = np.ascontiguousarray(flat.reshape(-1, EMBED_DIM), dtype=np.float32)
        del flat
        if not index.is_trained:
            index.train(embeds)  # no-op statistics pass for QT_fp16
        index.add(embeds)
        del embeds

        pq.write_table(table.select(TEXT_COLUMNS), corpus_dir / f"data-{i:05d}.parquet")
        total_rows += table.num_rows
        del table

        # real (resolved) file behind the cache symlink, then the symlink itself
        if not args.keep_sources:
            real = os.path.realpath(shard)
            os.remove(shard)
            if real != shard and os.path.exists(real):
                os.remove(real)
        print(f"[{i + 1}/{len(shards)}] {Path(shard).name}: {total_rows} rows total, "
              f"index size {index.ntotal}", flush=True)

    faiss.write_index(index, str(faiss_path))
    print(f"DONE: {total_rows} rows, faiss index at {faiss_path} "
          f"({faiss_path.stat().st_size / 1e9:.2f} GB)", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())