"""Build per-fold memmap shards from a 10-fold_split fold's train.jsonl.zst. Same output format/planes as build_shards.py / build_bronze.py so the training loader works unchanged when STOICHEIA_DATA points at the fold root: /v1_punct/ tier in {pristine, repaired} (pristine rows first, then repaired — mirrors the flagship's canonical order so the idx%HOLDOUT_MOD holdout and eval/intrinsic sampling behave identically) /bronze_punct/ tier == bronze Records with tier == inscriptions (or anything else) are skipped and counted — the flagship recipe never trains on inscriptions. Record ids (including the #segN suffixes of excised-and-restitched train segments) are kept verbatim. Usage: python data/build_fold_shards.py --jsonl .../fold_0/train.jsonl.zst \ --out $STOICHEIA_DATA/folds/fold_0/shards --workers 16 """ from __future__ import annotations import argparse, json, shutil, subprocess, sys from collections import Counter from concurrent.futures import ProcessPoolExecutor from dataclasses import asdict from pathlib import Path import numpy as np import pyarrow as pa import pyarrow.parquet as pq sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from data.normalize import Stats, normalize_record PLANES = ("chars", "boundary", "dia", "cap", "punct") TIER2TARGET = {"pristine": "pri", "repaired": "rep", "bronze": "brz"} CLEAN = {"pristine": 1.0, "repaired": 0.75, "bronze": 0.5} INDEX_SCHEMA = pa.schema([("offset", pa.int64()), ("length", pa.int64()), ("tier", pa.string()), ("clean", pa.float64()), ("source", pa.string()), ("id", pa.string())]) def _dump_stats(stats: Stats, path: Path): st = asdict(stats) st["stripped"] = dict(stats.stripped) st["archaic"] = dict(stats.archaic) st["other_marks"] = dict(stats.other_marks) path.write_text(json.dumps(st)) def _load_stats(path: Path) -> Stats: s = json.loads(path.read_text()) st = Stats(**{k: s[k] for k in ("records_in", "records_kept", "records_dropped_nongreek", "records_dropped_empty", "letters", "words", "sentences", "mark_conflicts", "orphan_marks")}) st.stripped = Counter({int(k): v for k, v in s["stripped"].items()}) st.archaic = Counter(s["archaic"]) st.other_marks = Counter(s["other_marks"]) return st def process_chunk(args): chunk_file, tmpdir, cid = args tmpdir = Path(tmpdir) stats = {t: Stats() for t in ("pri", "rep", "brz")} bufs = {t: {p: [] for p in PLANES} for t in ("pri", "rep", "brz")} rows = {t: {"offset": [], "length": [], "tier": [], "clean": [], "source": [], "id": []} for t in ("pri", "rep", "brz")} off = {t: 0 for t in ("pri", "rep", "brz")} skipped = Counter() with open(chunk_file) as f: for line in f: try: rec = json.loads(line) except Exception: skipped["unparseable"] += 1 continue tier = rec.get("tier", "") t = TIER2TARGET.get(tier) if t is None: skipped[tier or "missing_tier"] += 1 continue r = normalize_record(rec.get("text", ""), stats[t], with_punct=True) if r is None: continue for p, a in zip(PLANES, r): bufs[t][p].append(a) rows[t]["offset"].append(off[t]) rows[t]["length"].append(len(r[0])) rows[t]["tier"].append(tier) rows[t]["clean"].append(CLEAN[tier]) rows[t]["source"].append(rec.get("source", "")) rows[t]["id"].append(rec.get("id", "")) off[t] += len(r[0]) for t in ("pri", "rep", "brz"): d = tmpdir / t d.mkdir(parents=True, exist_ok=True) for p in PLANES: if bufs[t][p]: np.concatenate(bufs[t][p]).tofile(d / f"{p}.bin") else: np.array([], np.uint8).tofile(d / f"{p}.bin") pq.write_table(pa.table(rows[t], schema=INDEX_SCHEMA), d / "index.parquet") _dump_stats(stats[t], d / "stats.json") (tmpdir / "skipped.json").write_text(json.dumps(dict(skipped))) return cid def assemble(out_dir: Path, part_dirs: list[Path]): """Concatenate per-chunk target dirs (in the given order) into one shard dir.""" out_dir.mkdir(parents=True, exist_ok=True) total = Stats() tables, cum = [], 0 for d in part_dirs: t = pq.read_table(d / "index.parquet") if t.num_rows: t = t.set_column(0, "offset", pa.array(t.column("offset").to_numpy() + cum, type=pa.int64())) tables.append(t) cum += int(np.sum(t.column("length").to_numpy())) total.merge(_load_stats(d / "stats.json")) for p in PLANES: with open(out_dir / f"{p}.bin", "wb") as fo: for d in part_dirs: fo.write((d / f"{p}.bin").read_bytes()) pq.write_table(pa.concat_tables(tables) if tables else INDEX_SCHEMA.empty_table(), out_dir / "index.parquet") st = asdict(total) st["stripped"] = {str(k): v for k, v in total.stripped.items()} st["archaic"] = dict(total.archaic) st["other_marks"] = dict(total.other_marks) (out_dir / "stats.json").write_text(json.dumps(st, indent=2, ensure_ascii=False)) return cum, total def main(): ap = argparse.ArgumentParser() ap.add_argument("--jsonl", required=True, help="fold_k/train.jsonl.zst") ap.add_argument("--out", required=True, help="fold root shards dir (gets v1_punct/, bronze_punct/)") ap.add_argument("--workers", type=int, default=16) ap.add_argument("--chunk", type=int, default=40000) a = ap.parse_args() out = Path(a.out) out.mkdir(parents=True, exist_ok=True) tmp = out / "tmp" if tmp.exists(): shutil.rmtree(tmp) (tmp / "chunks").mkdir(parents=True) # stream-decompress and split into line-chunk files (keeps memory flat) proc = subprocess.Popen(["zstdcat", a.jsonl], stdout=subprocess.PIPE, text=True) chunks, buf, cid = [], [], 0 for line in proc.stdout: buf.append(line) if len(buf) >= a.chunk: cf = tmp / "chunks" / f"c{cid:04d}.jsonl" cf.write_text("".join(buf)) chunks.append((str(cf), str(tmp / f"c{cid:04d}"), cid)) buf, cid = [], cid + 1 if buf: cf = tmp / "chunks" / f"c{cid:04d}.jsonl" cf.write_text("".join(buf)) chunks.append((str(cf), str(tmp / f"c{cid:04d}"), cid)) if proc.wait() != 0: raise RuntimeError(f"zstdcat failed on {a.jsonl}") print(f"{len(chunks)} chunks", flush=True) with ProcessPoolExecutor(max_workers=a.workers) as ex: list(ex.map(process_chunk, chunks)) order = [Path(c[1]) for c in chunks] # canonical order: all pristine parts, then all repaired parts (mirrors build_shards.py) lit_letters, lit_stats = assemble(out / "v1_punct", [d / "pri" for d in order] + [d / "rep" for d in order]) brz_letters, brz_stats = assemble(out / "bronze_punct", [d / "brz" for d in order]) skipped = Counter() for d in order: skipped.update(json.loads((d / "skipped.json").read_text())) prov = { "source_jsonl": str(Path(a.jsonl).resolve()), "source_mtime": Path(a.jsonl).stat().st_mtime, "v1_punct": {"letters": lit_letters, "records_kept": lit_stats.records_kept, "records_in": lit_stats.records_in}, "bronze_punct": {"letters": brz_letters, "records_kept": brz_stats.records_kept, "records_in": brz_stats.records_in}, "skipped": dict(skipped), } (out.parent / "provenance.json").write_text(json.dumps(prov, indent=2)) shutil.rmtree(tmp) print(f"v1_punct letters: {lit_letters/1e9:.3f}B kept: {lit_stats.records_kept}/{lit_stats.records_in}") print(f"bronze_punct letters: {brz_letters/1e9:.3f}B kept: {brz_stats.records_kept}/{brz_stats.records_in}") print(f"skipped (non-training tiers): {dict(skipped)}") if __name__ == "__main__": main()