spb/wpai0002_uqo

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