#!/usr/bin/env python3
"""Fetch WikiText-103-raw and write it as a flat character corpus.

WHY A STANDARD CORPUS IS THE POINT

Every number measured so far is on 683,065 characters of Cory's own logs. That corpus
cannot answer the scaling question for two separate reasons:

  1. At 10M+ parameters a 683K-character corpus is memorised, so every rung converges to
     the same overfit floor and the comparison silently becomes about regularisation
     rather than architecture.
  2. Results on a private corpus are not checkable by anyone else. WikiText-103 is the
     benchmark the field already uses, so a dyn12 advantage measured here is directly
     comparable to published work instead of being a claim about one person's log files.

One train shard is ~157 MB of parquet, which yields roughly a quarter of a billion
characters -- enough that a 30M-parameter model is data-limited rather than
memorisation-limited.

  python tools/fetch_wikitext.py [--shards 1]
"""
from __future__ import annotations

import sys
from pathlib import Path

sys.stdout.reconfigure(encoding="utf-8", errors="replace")

OUT = Path("01_HER_SOUL/corpus_snapshots/wikitext103_train.txt")
REPO = "Salesforce/wikitext"
SHARDS = ["wikitext-103-raw-v1/train-00000-of-00002.parquet",
          "wikitext-103-raw-v1/train-00001-of-00002.parquet"]


def main() -> int:
    n = 1
    if "--shards" in sys.argv:
        n = int(sys.argv[sys.argv.index("--shards") + 1])

    import pyarrow.parquet as pq
    from huggingface_hub import hf_hub_download

    OUT.parent.mkdir(parents=True, exist_ok=True)
    total = 0
    with OUT.open("w", encoding="utf-8", newline="\n") as f:
        for s in SHARDS[:n]:
            print(f"  downloading {s} ...", flush=True)
            local = hf_hub_download(REPO, s, repo_type="dataset")
            t = pq.read_table(local)
            col = t.column("text").to_pylist()
            # WikiText ships one row per line, blanks and " = Heading = " markers included.
            # Both are kept: they carry document structure a character model can learn.
            for line in col:
                if line:
                    f.write(line)
                    total += len(line)
            print(f"    {len(col):,} rows, running total {total/1e6:.1f}M chars", flush=True)

    size = OUT.stat().st_size
    print(f"\n  wrote {OUT}")
    print(f"  {total:,} characters, {size/1e6:.1f} MB on disk")

    import hashlib
    h = hashlib.sha256()
    with OUT.open("rb") as fh:
        for c in iter(lambda: fh.read(1 << 20), b""):
            h.update(c)
    print(f"  sha256 {h.hexdigest()[:16]}  <- freeze this in any result table")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())