"""Phase 5: the final weights -> a public model repo with a card that only states measured facts.

Everything quoted in the card is read back off the Hub at run time — the mix manifest, `audit.json`, the
trainer's `run_summary.json` — rather than retyped from memory, because the card is the artifact a stranger
will trust and `final-report.md` will cite (§3.12: the claim and the evidence must be the same object).

Gate 5 is not this script: it is `--verify-only` running on a *fresh instance with an empty cache*, which is
the only way to prove nothing the model needs lives on a Kaggle disk. This script therefore ends by printing
the exact command for that second job.
"""

import argparse
import hashlib
import json
import os
import shutil
import struct
import sys
import urllib.request

MIX = "Cion-lab/ounce100m-mix-v1"
# A default, not a constant: this run's checkpoint repo, hard-coded as the only source, let the publisher
# fetch `latest.json` from a repo the operator never pointed at -- for a second model that is publishing one
# run's weights into another run's repo. Override with --ckpt-repo, or bypass it with --from-dir.
CKPT = "Cion-lab/ounce100m-ckpt"
CARD_FILES = ("config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json",
              "special_tokens_map.json", "vocab.json", "merges.txt", "cursor.json",
              "run_summary.json", "generation_config.json")


def sha(p):
    h = hashlib.sha256()
    with open(p, "rb") as fh:
        for c in iter(lambda: fh.read(1 << 20), b""):
            h.update(c)
    return h.hexdigest()


def safetensors_counts(p):
    """(n_tensors, n_params, dtypes) straight from the file header — no torch, so this check runs on a
    CPU box and on any image, and it proves the *published bytes* have the parameter count the card claims."""
    with open(p, "rb") as fh:
        (n,) = struct.unpack("<Q", fh.read(8))
        hdr = json.loads(fh.read(n))
    total, dtypes = 0, set()
    for k, v in hdr.items():
        if k == "__metadata__":
            continue
        cnt = 1
        for d in v["shape"]:
            cnt *= d
        total += cnt
        dtypes.add(v["dtype"])
    return len([k for k in hdr if k != "__metadata__"]), total, sorted(dtypes)


def get_json(url, token=None):
    req = urllib.request.Request(url)
    if token:
        req.add_header("Authorization", f"Bearer {token}")
    with urllib.request.urlopen(req, timeout=120) as r:
        return json.loads(r.read().decode())


def card(cfg, man, aud, run, repo, files, n_params):
    p = cfg.get("num_hidden_layers"), cfg.get("hidden_size"), cfg.get("vocab_size")
    # The card states the architecture; half of it was typed rather than read, so a variant export would
    # have published a description of a model it is not. Derive it and refuse on drift.
    frozen = {"num_hidden_layers": 22, "hidden_size": 576, "intermediate_size": 1536,
              "num_attention_heads": 9, "num_key_value_heads": 3, "vocab_size": 49152,
              "tie_word_embeddings": True, "rms_norm_eps": 1e-5}
    drift = {k: (cfg.get(k), v) for k, v in frozen.items() if cfg.get(k) != v}
    if drift:
        raise SystemExit("config.json does not match the frozen architecture, so the card would be a false "
                         "statement about this model: " + json.dumps(drift, default=str)[:600])
    rope = cfg.get("rope_theta") or (cfg.get("rope_parameters") or {}).get("rope_theta")
    cf = man.get('contamination_filter') or {}
    tok_consumed = int(run.get("tokens_consumed") or 0)
    loss = " -> ".join(str(x) for x in (run.get("prior_loss"), run.get("final_loss")) if x is not None)
    # Same rule as the architecture: anything the card states about *this* model comes from its own
    # `run_summary.json`, and a field that is not there stops the publish rather than rendering as None.
    # The Optimiser/Precision/Batch rows used to be typed prose, which is how a card ends up describing the
    # recipe somebody meant to run (review E-045/2).
    need = ("optim", "lr", "warmup_frac", "decay_frac", "world", "tokens_per_step", "seq_len",
            "micro_batch", "accum", "attn", "grad_ckpt", "tok_per_s", "seed", "data_seed",
            "tokens_consumed", "session_seconds")
    absent = [k for k in need if run.get(k) is None]
    if absent:
        raise SystemExit(f"run_summary.json carries no {absent}, so the card cannot state them without "
                         "guessing. Publish the summary the trainer wrote, or say in the card what is "
                         "unknown -- never both.")
    geom = run["micro_batch"] * run["accum"] * run["world"] * run["seq_len"]
    if geom != run["tokens_per_step"]:
        raise SystemExit(f"the batch row would be a false statement: {run['micro_batch']} x "
                         f"{run['accum']} x {run['world']} x {run['seq_len']} = {geom} != "
                         f"tokens_per_step {run['tokens_per_step']}")
    if run.get("val_ppl") is None and run.get("val_skipped") and run.get("val_loss_error") is None:
        val_row = ("not measured on this leg -- a session that stops at a checkpoint boundary skips the "
                   "validation pass by design; the run's PPL comes from the horizon export")
    elif run.get("val_ppl") is None:
        val_row = f"NOT MEASURED ({run.get('val_loss_error')})"
    else:
        val_row = (f"**{run['val_ppl']:.2f}** on {MIX}'s held-out `val/` shards "
                   f"({man['val_tokens']:,} tokens, never trained on)")
    rows = "\n".join(f"| `{k}` | {v['sha256'][:16]}… | {v['bytes']:,} |" for k, v in sorted(files.items()))
    return f"""---
license: cc-by-sa-4.0
base_model: {repo}
language: en
library_name: transformers
pipeline_tag: text-generation
tags: [ounce100m, from-scratch, base, decoder-only, {man['tokenizer']['id'].split('/')[0].lower()}]
---

# {repo.split('/')[-1]}

A **{n_params:,}-parameter decoder-only transformer trained from scratch on
{tok_consumed:,} tokens** of `{MIX}`, in {run.get('steps_planned', '?')} optimiser steps on
**2×Nvidia Tesla T4** inside Kaggle notebook sessions. It is the deliverable of a project whose whole point
was to find out what a ~100 M-parameter model reaches from ~1 B tokens of a curated public mix, and to
report it honestly, including where that is near chance.

## Measured facts

| | |
|---|---|
| Parameters | **{n_params:,}** (recomputed from `model.safetensors`'s own header, not asserted) |
| Architecture | {p[0]} layers, hidden {p[1]}, GQA {cfg["num_attention_heads"]} Q / {cfg["num_key_value_heads"]} KV heads, SwiGLU {cfg["intermediate_size"]}, RMSNorm {cfg["rms_norm_eps"]}, RoPE θ={rope}, **tied** embeddings, vocab {p[2]:,} |
| Sequence length | {run['seq_len']} — which is why ARC's 25-shot prompts are truncated from the left by the eval harness |
| Tokens consumed | **{tok_consumed:,}** ({100.0 * tok_consumed / 1e9:.2f} % of the 1.0 B target) |
| Batch | micro {run['micro_batch']} × accum {run['accum']} × {run['world']} ranks = {run['tokens_per_step']:,} tokens/step |
| Optimiser | `{run['optim']}`, LR {run['lr']}, warmup fraction {run['warmup_frac']}, trapezoid decay over the final {100.0 * (1.0 - (run['decay_frac'] or 0)):.1f} % of the schedule |
| Precision | fp16 autocast + fp32 master weights + GradScaler — asserted at the start of every session, and each session's log line `precision: fp16=True bf16=False scaler=GradScaler` is the evidence |
| Throughput | {run['tok_per_s']:,} tok/s in the session that produced this export |
| Attention | `{run.get('effective_attn') or run['attn']}` |
| Gradient checkpointing | `{run['grad_ckpt']}` |
| Randomness | model seed {run['seed']}, data seed {run['data_seed']} |
| Wall clock | {run.get('session_seconds', run.get('seconds', 0)) / 3600.0:.2f} h in the session that produced this export. The run is split across sessions at checkpoint boundaries, so this is **not** the total training time |
| Loss | {loss} |
| Validation PPL | {val_row} |

## Data

`{MIX}`: {man['total_tokens']:,} training tokens in {man['n_shards']} shards + {man['val_tokens']:,}
held-out validation tokens, cut with `{man['tokenizer']['id']}` ({man['tokenizer']['vocab_size']:,} vocab,
Apache-2.0), from {len(man.get('stage') or {})} public sources at ≥{man.get('distinct_sources_per_shard_min')}
per shard. Format and loading code are in that dataset's card.

**Contamination.** Overlap was measured mechanically against benchmark **train/validation/dev** material of
the eight evaluation tasks (13-token windows, counts only — no benchmark item text was ever read, and test
splits stayed untouched until evaluation). `audit.json` in the dataset repo records
`overlap_total = {aud.get('overlap_total')}` over {aud.get('mix_documents_audited'):,} documents and
{aud.get('val_documents_audited'):,} held-out documents, `tasks_covered` = {len(aud.get('tasks_covered') or [])}/8.
Building it required dropping {cf.get('docs_dropped', '?')} documents
({cf.get('tokens_dropped', 0):,} tokens)
that matched reference material; the exact exclusion masks ship in `filter/` so the mix is reproducible.

## Evaluate it

```bash
pip install "lm-eval==0.4.13"
lm_eval --model hf --model_args pretrained={repo},dtype=float16,trust_remote_code=False \\
  --tasks arc_challenge,arc_easy,hellaswag,mmlu,piqa,truthfulqa_mc1,truthfulqa_mc2,winogrande,gsm8k \\
  --batch_size 8 --seed 42 --log_samples
```

The run's own results, with shot counts and chance levels beside every number, are in `eval/` here. Four of
the eight tasks (Winogrande, HellaSwag, PIQA, TruthfulQA) are scored by the harness on *validation* — the
splits this project was allowed to read — and the other four on test.

## Files

| file | sha256 | bytes |
|---|---|---|
{rows}

`cursor.json` is the data position this checkpoint corresponds to: samples consumed, step, seed, the
corpus fingerprint and a hash of the visit-order permutation, so the resume can be checked rather than
trusted.

## Not in this model

No instruction tuning, no RLHF, no chat template — it is a base model, deliberately. It has ~1 B tokens of
training against the ~300 B of the comparable published reference models, so expect scores at or near
chance on several tasks; that is the experiment's result, not a defect to be engineered around.
"""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="Cion-lab/ounce100m-v1")
    ap.add_argument("--src", default="/kaggle/working/mixroot", help="mix root (for tokenizer.json)")
    ap.add_argument("--from-dir", default="", help="local dir holding the trainer's final/ export")
    ap.add_argument("--ckpt-repo", default=CKPT,
                    help="dataset repo whose latest.json names the export to publish (ignored with --from-dir)")
    ap.add_argument("--verify-only", action="store_true",
                    help="fresh-instance clean-room test: download the published repo and load it")
    ap.add_argument("--dry-run", action="store_true")
    a = ap.parse_args()

    # Order matters twice over. huggingface_hub resolves its cache locations at *import* time, and
    # `ounce100m_credentials` imports huggingface_hub -- so a clean room configured after that import is
    # reading Kaggle's warm disk (review finding A6, which the first fix for it did not actually fix:
    # E-046/1). Set the environment before anything is imported, and pass cache_dir explicitly as well,
    # so the claim does not depend on which of the two the library honours.
    if a.verify_only:
        os.environ["HF_HOME"] = "/tmp/cleanroom_hf"
        os.environ["HF_HUB_CACHE"] = "/tmp/cleanroom_hf/hub"
        os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "120"
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))  # E-028
    import ounce100m_credentials
    ounce100m_credentials.install(verify=True)
    tok = os.environ["HF_TOKEN"]
    from huggingface_hub import HfApi, create_repo, snapshot_download

    man = get_json(f"https://huggingface.co/datasets/{MIX}/resolve/main/manifest.json")
    aud = get_json(f"https://huggingface.co/datasets/{MIX}/resolve/main/audit.json")

    if a.verify_only:
        dest = "/tmp/cleanroom_" + os.path.basename(a.repo)
        shutil.rmtree(dest, ignore_errors=True)
        shutil.rmtree("/tmp/cleanroom_hf", ignore_errors=True)
        # Gate 5's whole claim is "a stranger can load this". `install()` above puts HF_TOKEN in the
        # environment and `huggingface_hub` picks it up from there, so without this the clean room would
        # be authenticated and a private or gated repo would pass (review E-045/1).
        for k in ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGING_FACE_TOKEN"):
            os.environ.pop(k, None)
        got = snapshot_download(a.repo, local_dir=dest, cache_dir="/tmp/cleanroom_hf/hub",
                                token=False)
        print("clean-room download (anonymous) ->", got, flush=True)
        import json as _j
        cfg = _j.load(open(os.path.join(dest, "config.json")))
        n_t, n_p, dt = safetensors_counts(os.path.join(dest, "model.safetensors"))
        print("VERIFY tensors", n_t, "params", f"{n_p:,}", "dtypes", dt, flush=True)
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        tkn = AutoTokenizer.from_pretrained(dest)
        try:                          # transformers 5 renamed torch_dtype; Gate 5 is this very call
            m = AutoModelForCausalLM.from_pretrained(dest, dtype=torch.float16)
        except TypeError:
            m = AutoModelForCausalLM.from_pretrained(dest, torch_dtype=torch.float16)
        assert m.num_parameters() == n_p, (m.num_parameters(), n_p)
        assert cfg["tie_word_embeddings"] and m.get_input_embeddings().weight.data_ptr() == \
            m.get_output_embeddings().weight.data_ptr(), "tied embeddings did not survive the round trip"
        g = torch.Generator(device="cpu").manual_seed(20260919)
        torch.manual_seed(20260919)
        ids = tkn("The city of Paris is", return_tensors="pt").input_ids
        out = m.generate(ids, max_new_tokens=24, do_sample=True, temperature=1.0, top_k=50,
                         top_p=0.9, pad_token_id=tkn.eos_token_id)
        text = tkn.decode(out[0], skip_special_tokens=True)
        print("GENERATED:", text, flush=True)
        print("GEN_SHA", hashlib.sha256(text.encode()).hexdigest()[:16], flush=True)
        print("VERDICT GATE5 load_ok=True generate_ok=%s params_ok=%s tied_ok=True vocab=%d"
              % (len(out[0]) > ids.shape[1], m.num_parameters() == n_p, cfg["vocab_size"]), flush=True)
        return

    src = a.from_dir
    if not src:
        ptr = get_json(f"https://huggingface.co/datasets/{a.ckpt_repo}/resolve/main/latest.json", tok)
        print("ckpt latest.json ->", {"repo": a.ckpt_repo,
                                      **{k: ptr.get(k) for k in ("step", "path_in_repo")}}, flush=True)
        src = snapshot_download(a.ckpt_repo, repo_type="dataset",
                                allow_patterns=[ptr["path_in_repo"] + "/*"],
                                local_dir="/kaggle/working/from_hub")
        src = os.path.join(src, ptr["path_in_repo"])
    have = {f for f in os.listdir(src) if os.path.isfile(os.path.join(src, f))}
    need = {"config.json", "model.safetensors"}
    if not need <= have:
        raise SystemExit(f"{src} is missing {sorted(need - have)} -- nothing to publish")
    if "tokenizer.json" not in have and os.path.exists(os.path.join(a.src, "tokenizer.json")):
        # No `import shutil` here: a module-level name re-imported inside this function makes it *local*
        # to main(), so the verify-only branch's earlier `shutil.rmtree` raised UnboundLocalError and
        # Gate 5 could never run (found by running it, 2026-09-21T13:40Z).
        shutil.copy(os.path.join(a.src, "tokenizer.json"), src)     # ship the tokenizer that cut the mix
        have.add("tokenizer.json")
    # CARD_FILES already promised tokenizer_config.json and special_tokens_map.json, but nothing ever put
    # them in the export, so the repo shipped a bare tokenizer.json. Without them AutoTokenizer falls back
    # to a literal unknown-token name that this BPE vocab does not contain, so decoding raised
    # "Unk token not found in the vocabulary" and Gate 5 could not load the tokenizer (2026-09-21T13:48Z).
    # They are fetched from the tokenizer's own hub repo, so the published tokenizer is byte-complete.
    TOK_SOURCE = "HuggingFaceTB/SmolLM2-135M"
    tk = man.get("tokenizer")                       # manifest may hold a string or a mapping
    if isinstance(tk, dict):
        TOK_SOURCE = tk.get("hub") or tk.get("repo") or TOK_SOURCE
    elif isinstance(tk, str) and "/" in tk:
        TOK_SOURCE = tk
    print("tokenizer source for config files:", TOK_SOURCE, flush=True)
    for extra in ("tokenizer_config.json", "special_tokens_map.json"):
        if extra in have:
            continue
        blob = urllib.request.urlopen(f"https://huggingface.co/{TOK_SOURCE}/resolve/main/{extra}",
                                       timeout=120).read()
        dec = json.loads(blob.decode("utf-8"))
        u = dec.get("unk_token")
        u = (u.get("content") if isinstance(u, dict) else u)
        if u and u not in json.load(open(os.path.join(src, "tokenizer.json"), encoding="utf-8")
                                    )["model"]["vocab"]:
            dec["unk_token"] = None        # never declare a token the vocab lacks
            blob = json.dumps(dec, indent=2, ensure_ascii=False).encode("utf-8")
        with open(os.path.join(src, extra), "wb") as fh:
            fh.write(blob)
        have.add(extra)
        print("tokenizer config", extra, len(blob), "bytes sha256",
              sha(os.path.join(src, extra))[:16], flush=True)
    for required in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"):
        if required not in have:
            raise SystemExit(f"tokenizer is incomplete: {required} is missing from the export")
    cfg = json.load(open(os.path.join(src, "config.json")))
    run = {}
    if os.path.exists(os.path.join(src, "run_summary.json")):
        run = json.load(open(os.path.join(src, "run_summary.json")))
    n_t, n_p, dt = safetensors_counts(os.path.join(src, "model.safetensors"))
    # `params` in run_summary.json is a dict {sum_numel, unique, trainable}; comparing that to an integer
    # is always unequal and would have failed every real export (review finding A1).
    claim = run.get("params_total") or run.get("params")
    if isinstance(claim, dict):
        claim = claim.get("sum_numel")
    if claim and int(claim) != n_p:
        raise SystemExit(f"run_summary claims {claim} params, the published safetensors file has {n_p}")
    files = {f: {"sha256": sha(os.path.join(src, f)), "bytes": os.path.getsize(os.path.join(src, f))}
             for f in sorted(have) if f in CARD_FILES or f.endswith(".json")}
    print("tensors", n_t, "params", f"{n_p:,}", "dtypes", dt, flush=True)
    if not a.dry_run and not run:
        raise SystemExit("no run_summary.json next to the export: the card would be full of blanks. "
                         "Pass --from-dir at the trainer's final/ export, or --dry-run to preview.")
    readme = card(cfg, man, aud, dict(run, params_total=n_p), a.repo, files, n_p)
    print("card:", f"{n_p:,} params from the safetensors header; {len(files)} files hashed", flush=True)
    if a.dry_run:
        # A preview that writes into the export is not a preview: it leaves a README behind that the next
        # real run then treats as one of the files to publish (review E-045, minor).
        print(readme.encode("utf-8", "replace").decode("utf-8", "replace"))
        return
    open(os.path.join(src, "README.md"), "w", encoding="utf-8").write(readme)

    api = HfApi(token=tok)
    create_repo(repo_id=a.repo, repo_type="model", private=False, exist_ok=True, token=tok)
    up = [(f, os.path.join(src, f)) for f in sorted(os.listdir(src))
          if os.path.isfile(os.path.join(src, f))]
    local_sha = {f: sha(p) for f, p in up}
    remote = {}
    try:
        for attempt in (1, 2, 3):
            try:
                for ent in api.model_info(a.repo, files_metadata=True).siblings or []:
                    lfs = getattr(ent, "lfs", None)
                    raw = (lfs.get("sha256") or lfs.get("oid")) if isinstance(lfs, dict) else None
                    remote[ent.rfilename] = raw if isinstance(raw, str) and len(raw) == 64 else None
                break
            except Exception as e:
                # E-034 again: "this repo has nothing in it yet" and "the Hub did not answer" are different
                # facts. The first means upload everything; the second, swallowed into `remote = {}` here,
                # means silently re-uploading 1.3 GB of bytes that are already on the Hub.
                if type(e).__name__ == "RepositoryNotFoundError":
                    print("(remote listing 404 -- this repo has nothing in it yet)", flush=True)
                    break
                if attempt == 3:
                    raise
                print(f"(listing failed {type(e).__name__}, retry {attempt}/3)", flush=True)
                import time as _t
                _t.sleep(10 * attempt)
    except Exception as e:
        raise SystemExit(f"cannot read the remote listing after 3 tries ({type(e).__name__}: "
                         f"{str(e)[:160]}) -- refusing to guess what is already uploaded")
    todo = [f for f, _ in up if remote.get(f) != local_sha[f]]
    print(f"uploading {len(todo)} of {len(up)} files to {a.repo}", flush=True)
    if todo:
        for i in range(0, len(todo), 120):        # one commit per 120 files: E-033's rate ceiling
            b = todo[i:i + 120]
            # `max_workers` is accepted by the Hub client on this workspace but not in the Kaggle image,
            # where it raised TypeError and killed the first publish attempt (E-007's class: an API that
            # exists locally and not in the job). Default concurrency is fine for a 7-file repo.
            api.upload_folder(repo_id=a.repo, folder_path=src, path_in_repo="",
                              allow_patterns=b, repo_type="model",
                              commit_message=f"publish weights: batch {i // 120 + 1} of {len(b)} file(s)")
    back = set(api.list_repo_files(repo_id=a.repo, repo_type="model"))
    missing = sorted({f for f, _ in up} - back)
    R = {"repo": a.repo, "params": n_p, "tensors": n_t, "files": len(back), "missing": missing}
    # Every file, not just the weights and the config: `list_repo_files` proves a name exists, and a
    # tokenizer whose bytes did not land loads as a *different* tokenizer, which is the failure that
    # silently invalidates every number in Phase 6 (E-031's rule, applied to the model repo).
    anon = {}
    for f, _ in up:
        try:
            h, nbytes = hashlib.sha256(), 0      # stream it: 1.7 GB does not belong in RAM
            with urllib.request.urlopen(f"https://huggingface.co/{a.repo}/resolve/main/{f}",
                                        timeout=900) as r:
                for chunk in iter(lambda: r.read(1 << 20), b""):
                    nbytes += len(chunk)
                    h.update(chunk)
            anon[f] = {"bytes": nbytes, "sha_ok": h.hexdigest() == local_sha[f]}
        except Exception as e:
            anon[f] = {"error": f"{type(e).__name__}: {str(e)[:120]}"}
    R["anon_readback"] = {k: (v.get("sha_ok") if "sha_ok" in v else v.get("error")) for k, v in anon.items()}
    bad = sorted(k for k, v in anon.items() if not v.get("sha_ok"))
    R["published_verified"] = bool(not missing and not bad)
    if bad:
        R["anon_failed"] = bad[:10]
    print("PUBLISH_JSON_BEGIN")
    print(json.dumps(R, indent=1, default=str))
    print("PUBLISH_JSON_END")
    print("NEXT (Gate 5, on a fresh instance with an empty cache): "
          f"python build/publish_model.py --repo {a.repo} --verify-only", flush=True)
    raise SystemExit(0 if R["published_verified"] else 11)


if __name__ == "__main__":
    main()