scripts/fetch_corpus.py

"""Build a fixed, reproducible corpus_sample.jsonl of arXiv abstracts.

Source: the `CShorten/ML-ArXiv-Papers` Hugging Face dataset (arXiv titles +
abstracts). We use a stable HF mirror rather than the live arXiv API because the
API rate-limits (HTTP 429) hard on bursty fetches. The output is checked into git
so runs are deterministic and the demo is never empty (spec, non-negotiable #3).

Requires: pip install datasets   (build-time only; not a runtime dependency)

Usage:
    python scripts/fetch_corpus.py --target 3000
"""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path

HF_DATASET = "CShorten/ML-ArXiv-Papers"
_WS = re.compile(r"\s+")


def clean(text: str) -> str:
    return _WS.sub(" ", (text or "")).strip()


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--target", type=int, default=3000, help="number of documents")
    ap.add_argument("--out", type=Path, default=Path("data/corpus_sample.jsonl"))
    args = ap.parse_args()

    from datasets import load_dataset

    print(f"Loading {HF_DATASET} ...")
    ds = load_dataset(HF_DATASET, split="train")

    args.out.parent.mkdir(parents=True, exist_ok=True)
    written = 0
    seen: set[str] = set()
    with args.out.open("w", encoding="utf-8") as f:
        for i, row in enumerate(ds):
            title = clean(row.get("title"))
            abstract = clean(row.get("abstract"))
            if not abstract or title in seen:
                continue
            seen.add(title)
            rec = {
                "id": f"arxiv:{i}",
                "title": title,
                "text": abstract,
                "source": "arxiv/ml",
            }
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            written += 1
            if written >= args.target:
                break

    print(f"Wrote {written} docs to {args.out}")


if __name__ == "__main__":
    main()