Repository Public Research preview
1#!/usr/bin/env python32"""Leakage-safe document splits (cluster + author-group aware; held-out discipline / period)."""3from __future__ import annotations45import argparse6import sys7from pathlib import Path89ROOT = Path(__file__).resolve().parents[1]10sys.path.insert(0, str(ROOT / "src"))11from htr.data.splits import SplitConfig, assign_splits, check_no_group_leak, split_summary # noqa: E40212from htr.utils.io import read_json, read_jsonl, write_json, write_jsonl # noqa: E40213from htr.utils.provenance import run_header # noqa: E402141516def main() -> None:17 ap = argparse.ArgumentParser()18 ap.add_argument("--corpus", default=str(ROOT / "data/parsed/pilot"))19 ap.add_argument("--out", default=None)20 ap.add_argument("--seed", type=int, default=20261109)21 ap.add_argument("--heldout-disciplines", default="")22 ap.add_argument("--heldout-before-year", type=int, default=None)23 ap.add_argument("--no-author-safe", action="store_true")24 ap.add_argument("--venue-safe", action="store_true", help="also keep every venue (journal) within one split")25 ap.add_argument("--fixed", default=None, help="previous splits.jsonl whose assignments must be preserved")26 a = ap.parse_args()27 corpus = Path(a.corpus)28 docs = list(read_jsonl(corpus / "documents.jsonl"))29 clusters = read_json(corpus / "clusters.json") if (corpus / "clusters.json").exists() else {}30 cfg = SplitConfig(seed=a.seed, heldout_disciplines=[d for d in a.heldout_disciplines.split(",") if d], heldout_before_year=a.heldout_before_year, author_safe=not a.no_author_safe, venue_safe=a.venue_safe)31 fixed = {a["doc_id"]: a["split"] for a in read_jsonl(a.fixed)} if a.fixed else None32 asg = assign_splits(docs, clusters, cfg, fixed)33 leaks = check_no_group_leak(asg)34 out = Path(a.out) if a.out else ROOT / "data/splits" / corpus.name35 write_jsonl(out / "splits.jsonl", asg)36 summ = {**run_header({"config": cfg.__dict__}), "summary": split_summary(asg), "n_docs": len(docs), "group_leaks": leaks, "n_fixed": len(fixed or {})}37 write_json(out / "splits_summary.json", summ)38 print(summ["summary"], "leaks:", leaks)394041if __name__ == "__main__":42 main()43