"""Export, document and publish GPT-U-20M to the Hugging Face Hub.

    python scripts/publish.py export    # checkpoints -> out/hf (step 19,840) and out/hf-stable-15872 (weights only)
    python scripts/publish.py card      # out/hf/README.md from eval results, data index and training logs
    python scripts/publish.py upload    # needs $HF_TOKEN: main branch + stable-15872 branch
    python scripts/publish.py verify    # fresh download into an empty cache, load with AutoModelForCausalLM, generate
"""
import argparse
import json
import os
import shutil
import sys
import tempfile
from datetime import datetime
from pathlib import Path

import torch
from huggingface_hub import HfApi
from tokenizers import Tokenizer, processors
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, PreTrainedTokenizerFast

from common import CKPT_DIR, DATA_DIR, LOG_DIR, OUT_DIR, ROOT, TOKENIZER_DIR, build_model

REPO_ID = "DedeProGames/GPT-U-20M"
FINAL_STEP, STABLE_STEP = 19_840, 15_872
STABLE_BRANCH = f"stable-{STABLE_STEP}"
HF_DIR = OUT_DIR / "hf"
STABLE_DIR = OUT_DIR / f"hf-{STABLE_BRANCH}"
STAGE_DIR = OUT_DIR / "publish_main"
EOT = "<|endoftext|>"
CHANCE = {"arc_easy": 25.0, "piqa": 50.0, "hellaswag": 25.0, "lambada_openai": 0.0, "winogrande": 50.0}


# ---------------------------------------------------------------- export
def export_model(step: int, out_dir: Path) -> None:
    ckpt = torch.load(CKPT_DIR / f"step_{step:05d}.pt", map_location="cpu", weights_only=False)
    model = build_model(init=False)
    model.load_state_dict(ckpt["model"])
    out_dir.mkdir(parents=True, exist_ok=True)
    model.to(torch.bfloat16).save_pretrained(out_dir)
    GenerationConfig(bos_token_id=0, eos_token_id=0, pad_token_id=0, do_sample=True, temperature=0.8, top_p=0.95,
                     max_new_tokens=256).save_pretrained(out_dir)


def export_tokenizer(out_dir: Path) -> None:
    """Same vocabulary/merges as training; the published copy prepends <|endoftext|> (the training BOS) on encode."""
    tok = Tokenizer.from_file(str(TOKENIZER_DIR / "tokenizer.json"))
    tok.post_processor = processors.TemplateProcessing(single=f"{EOT} $A", pair=f"{EOT} $A {EOT} $B",
                                                       special_tokens=[(EOT, 0)])
    PreTrainedTokenizerFast(tokenizer_object=tok, bos_token=EOT, eos_token=EOT, pad_token=EOT, model_max_length=1024,
                            clean_up_tokenization_spaces=False).save_pretrained(out_dir)


def cmd_export(_args) -> None:
    export_model(FINAL_STEP, HF_DIR)
    export_tokenizer(HF_DIR)
    export_model(STABLE_STEP, STABLE_DIR)
    for d in (HF_DIR, STABLE_DIR):
        print(d, sorted(p.name for p in d.iterdir()))


# ---------------------------------------------------------------- model card
def read_jsonl(path: Path) -> list[dict]:
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]


def lm_eval_rows() -> list[str]:
    files = sorted((OUT_DIR / "lm_eval").rglob("results_*.json"))
    if not files:
        return ["Benchmarks were not run."]
    results = json.loads(files[-1].read_text())["results"]
    rows = ["| Task | Metric | GPT-U-20M | Chance |", "|---|---|---:|---:|"]
    for task, chance in CHANCE.items():
        r = results.get(task, {})
        for metric in ("acc,none", "acc_norm,none", "perplexity,none"):
            if metric in r:
                name = metric.split(",")[0]
                value = f"{r[metric]:.2f}" if name == "perplexity" else f"{100 * r[metric]:.2f}"
                rows.append(f"| {task} | {name} | {value} | {'—' if name == 'perplexity' else f'{chance:.0f}'} |")
    return rows


def training_stats() -> dict:
    rows = read_jsonl(LOG_DIR / "train.jsonl")
    steps = [r for r in rows if "event" not in r]
    events = [r for r in rows if "event" in r]
    # Active time only: the run was paused for days while the machine was off, then resumed from checkpoints.
    times = sorted(datetime.fromisoformat(r["wall_clock"]) for r in steps)
    active = sum((b - a).total_seconds() for a, b in zip(times, times[1:]) if (b - a).total_seconds() < 300)
    return {
        "hours": active / 3600,
        "tok_s": sum(r["tok_s"] for r in steps) / len(steps),
        "mfu": sum(r["mfu"] for r in steps) / len(steps),
        "resumes": sum(r["event"] == "resume" for r in events),
        "loss_guard": sum(r["event"] == "loss_guard" for r in events),
        "compile": next((r["mode"] for r in reversed(events) if r["event"] == "compile"), "?"),
        "max_temp": max(max(t for t in (r.get("gpu0_temp"), r.get("gpu1_temp")) if t is not None) for r in steps),
    }


def build_card() -> str:
    data = json.loads((DATA_DIR / "index.json").read_text())
    ev = json.loads((OUT_DIR / "eval.json").read_text())
    tokrep = json.loads((LOG_DIR / "tokenizer_report.json").read_text())
    ts = training_stats()
    seen = FINAL_STEP * 131_072
    share = data["total"]["share_by_domain"]
    fert = tokrep["fertility"]
    val = data["val"]
    names = {"dclm": "DCLM-baseline", "edu": "FineWeb-Edu", "code": "The Stack v3 (code)"}

    lines = f"""---
license: apache-2.0
datasets:
- mlfoundations/dclm-baseline-1.0
- HuggingFaceFW/fineweb-edu
- HuggingFaceCode/stack-v3-train
language: en
library_name: transformers
pipeline_tag: text-generation
tags:
- gpt-u
- tiny-lm
- pretrained-from-scratch
---

# GPT-U-20M

A 20.5M-parameter Llama-architecture language model pretrained from scratch on 2.6B tokens of web text, educational text and code.

## Architecture

| | |
|---|---|
| Type | decoder-only transformer, `LlamaForCausalLM`: RoPE, SwiGLU, RMSNorm, pre-norm, no biases |
| Parameters | 20,453,760 (6,291,456 tied embedding + 8 × 1,770,240 per block + 384 final norm) |
| Layers | 8 |
| Hidden size | 384 |
| Attention | 6 heads × 64 dims, full multi-head (no GQA), SDPA |
| MLP | SwiGLU, intermediate size 1,024 |
| Vocabulary | 16,384 (byte-level BPE, trained on the same mix) |
| Context length | 1,024 tokens |
| RoPE θ | 10,000 |
| Input/output embeddings | tied |

The 16,384-token vocabulary keeps the embedding at 31% of the parameters (a 50k vocabulary would need 19.3M parameters at width 384) and lets tokens be stored as `uint16`.

**Tokenizer.** Byte-level BPE with a cl100k-style split that keeps whitespace runs whole (Python indentation), one token per digit, and `<|endoftext|>` (id 0) as BOS/EOS/padding. Round-trip `decode(encode(x)) == x` holds on 1,000 held-out samples (tabs, CRLF, emoji, CJK, RTL, control characters).

| Held-out domain | Tokens / word | Bytes / token |
|---|---:|---:|
| DCLM | {fert['dclm']['tokens_per_word']:.2f} | {fert['dclm']['bytes_per_token']:.2f} |
| FineWeb-Edu | {fert['edu']['tokens_per_word']:.2f} | {fert['edu']['bytes_per_token']:.2f} |
| Code (27% Python) | {fert['code']['tokens_per_word']:.2f} | {fert['code']['bytes_per_token']:.2f} |
| Python only | {fert['code_python_only']['tokens_per_word']:.2f} | {fert['code_python_only']['bytes_per_token']:.2f} |

## Training data

One epoch over 2.6B tokens, no repetition:

| Dataset | Target share | Realized share | Tokens seen (≈) |
|---|---:|---:|---:|
| [mlfoundations/dclm-baseline-1.0](https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0) | 45% | {100 * share['dclm']:.2f}% | {share['dclm'] * seen / 1e9:.2f}B |
| [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) (`sample-10BT`) | 35% | {100 * share['edu']:.2f}% | {share['edu'] * seen / 1e9:.2f}B |
| [HuggingFaceCode/stack-v3-train](https://huggingface.co/datasets/HuggingFaceCode/stack-v3-train) | 20% | {100 * share['code']:.2f}% | {share['code'] * seen / 1e9:.2f}B |

- Each source is streamed from the Hub and tokenized up to its token budget; the documents are then interleaved **per document** in a seeded random order, i.e. with fixed proportions for the whole run (no domain curriculum), so the realized shares match the targets exactly. (Holding the four streams open at once was avoided: it needed 7–8 GB of RAM and destabilized the training machine.)
- The Stack v3 is grouped by repository; repositories were split into individual files, vendored files dropped, and Python held at {100 * data['total']['python_share_of_code']:.1f}% of the code tokens (it is ~4% of Stack v3 naturally).
- Documents are separated by `<|endoftext|>` and packed into contiguous 1,025-token blocks (1,024 inputs + 1 shifted target), no padding. Batches are sampled from a seeded permutation of all blocks.
- Validation holdouts (2M tokens per domain) come from source files that never fed training or the tokenizer.

## Training

| Hyperparameter | Value |
|---|---|
| Tokens per step | 131,072 (2 GPUs × 32 sequences × 1,024 tokens × 2 accumulation steps) |
| Steps | 19,840 |
| Optimizer | AdamW, β = (0.9, 0.95), ε = 1e-8 |
| Peak learning rate | 1.8e-3 |
| Schedule | WSD: 500 warmup steps → constant to step 15,872 → linear decay to 1.8e-4 at step 19,840 |
| Weight decay | 0.1 on 2D matrices only (none on norms and embeddings) |
| Gradient clipping | 1.0 |
| z-loss | 1e-4 |
| Dropout | 0 |
| Initialization | N(0, 0.02); `o_proj` and `down_proj` scaled by 1/√(2·8) |
| Precision | bf16 autocast, fp32 master weights |

Trained on 2× RTX 3060 12GB (no NVLink) under WSL2 with PyTorch DDP and `torch.compile` (mode: {ts['compile']}): {ts['hours']:.1f} h of active training time, {ts['tok_s']:,.0f} tokens/s on average (MFU ≈ {100 * ts['mfu']:.0f}% of the GPUs' measured bf16 matmul throughput), max GPU temperature {ts['max_temp']}°C, {ts['resumes']} automatic resume(s) from checkpoints after the machine was shut down, {ts['loss_guard']} loss-spike rollback(s).

## Results

Validation loss on the 2M-token holdouts (1,600 fixed 1,024-token windows per domain) and bits per byte, `bpb = loss_nats × tokens / (bytes × ln 2)`, which is comparable across tokenizers:

| Domain | Val loss (nats/token) | Bits per byte |
|---|---:|---:|
""".splitlines()
    for d in ("dclm", "edu", "code"):
        lines.append(f"| {names[d]} | {ev[d]['loss']:.4f} | {ev[d]['bpb']:.4f} |")
    lines.append(f"| **Aggregate** | **{ev['aggregate']['loss']:.4f}** | **{ev['aggregate']['bpb']:.4f}** |")
    lines += ["", "The aggregate is the equal-weight mean of the three holdouts (same number of tokens each). Code tokens "
              "are far easier to predict than prose with this tokenizer, which pulls the aggregate down; the prose domains "
              f"({ev['edu']['loss']:.2f}–{ev['dclm']['loss']:.2f} nats/token) are the fairer comparison point for "
              "web-text models, and bits per byte is the fairer one across tokenizers."]
    lines += ["", "Zero-shot benchmarks ([lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness)). "
              "`<|endoftext|>` is prepended as BOS. hellaswag and winogrande sit at chance, as expected at 20M parameters; "
              "arc_easy and piqa land clearly above it, which reflects surface lexical knowledge more than reasoning. "
              "They are reported for transparency, not as a quality signal.",
              ""]
    lines += lm_eval_rows()
    lines += f"""
Sample generations (temperature 0.8, top-p 0.95) are in [`out/samples.md`](out/samples.md).

## Loss curve

![Training and validation loss](out/loss.png)

## Usage

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("{REPO_ID}")
model = AutoModelForCausalLM.from_pretrained("{REPO_ID}")

inputs = tokenizer("The history of the Roman Empire", return_tensors="pt")  # prepends <|endoftext|> as BOS
output = model.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=0.8, top_p=0.95)
print(tokenizer.decode(output[0], skip_special_tokens=True))
```

## Checkpoints

- `main`: final weights (step 19,840, end of the decay phase).
- [`{STABLE_BRANCH}`](https://huggingface.co/{REPO_ID}/tree/{STABLE_BRANCH}): weights at step 15,872, the end of the constant-learning-rate phase, meant for continued pretraining (e.g. extending to 4–5B tokens before decaying again). Load with `revision="{STABLE_BRANCH}"`.

## Limitations

- **It is tiny.** 20M parameters is a research and teaching scale; the model has shallow world knowledge and weak reasoning.
- **Heavily over-trained for its size by design:** 127 tokens per parameter, about 6× the Chinchilla-optimal ratio (~20).
- **Not for production use.** It hallucinates freely and produces fluent but often false statements.
- **Generated code is not functional.** It imitates the surface form of code; do not run it.
- English only (web and educational text); no instruction tuning, no safety tuning. Web data can carry biases and offensive content.
- Context is limited to 1,024 tokens.

## Data provenance and licenses

| Source | License | Notes |
|---|---|---|
| DCLM-baseline 1.0 | CC-BY-4.0 | Common Crawl web text filtered by DataComp-LM; see the dataset card. |
| FineWeb-Edu (`sample-10BT`) | ODC-By 1.0 | Common Crawl pages filtered for educational value; use is also subject to Common Crawl's terms of use. |
| The Stack v3 (train) | ODC-By | GitHub code as of August 2025; files are permissively licensed or carry no detected license (non-permissive files excluded upstream), PII-redacted, with an opt-out process. Individual files keep their original licenses. |

The model weights are released under **Apache-2.0**. Tokenized data is not redistributed.

## Reproducibility

Everything used to build this model is in the repository: `scripts/` (tokenizer training, data preparation, training, evaluation, publishing), `configs/model.json`, `logs/train.jsonl` (every 10 steps: loss, lr, grad norm, throughput, MFU, GPU temperatures, memory), `logs/tokenizer_report.json` and `out/data_index.json` (token counts per shard and domain, holdout bytes). Raw evaluation outputs: `out/eval.json` and `out/lm_eval_results.json`.
""".splitlines()
    return "\n".join(lines) + "\n"


def cmd_card(_args) -> None:
    card = build_card()
    (HF_DIR / "README.md").write_text(card, encoding="utf-8")
    print(card)


# ---------------------------------------------------------------- upload
def stage_main() -> Path:
    if STAGE_DIR.exists():
        shutil.rmtree(STAGE_DIR)
    shutil.copytree(HF_DIR, STAGE_DIR)
    shutil.copytree(ROOT / "scripts", STAGE_DIR / "scripts", ignore=shutil.ignore_patterns("__pycache__"))
    (STAGE_DIR / "configs").mkdir()
    shutil.copy2(ROOT / "configs" / "model.json", STAGE_DIR / "configs" / "model.json")
    (STAGE_DIR / "logs").mkdir()
    for name in ("train.jsonl", "tokenizer_report.json"):
        shutil.copy2(LOG_DIR / name, STAGE_DIR / "logs" / name)
    (STAGE_DIR / "out").mkdir()
    for name in ("loss.png", "samples.md", "eval.json"):
        shutil.copy2(OUT_DIR / name, STAGE_DIR / "out" / name)
    shutil.copy2(DATA_DIR / "index.json", STAGE_DIR / "out" / "data_index.json")
    results = sorted((OUT_DIR / "lm_eval").rglob("results_*.json"))
    if results:
        shutil.copy2(results[-1], STAGE_DIR / "out" / "lm_eval_results.json")
    return STAGE_DIR


def cmd_upload(_args) -> None:
    token = os.environ.get("HF_TOKEN")
    if not token:
        sys.exit("HF_TOKEN is not set: nothing uploaded (out/hf and out/hf-stable-15872 are ready).")
    api = HfApi(token=token)
    api.create_repo(REPO_ID, repo_type="model", private=False, exist_ok=True)
    if STABLE_BRANCH not in {b.name for b in api.list_repo_refs(REPO_ID).branches}:
        first_commit = api.list_repo_commits(REPO_ID)[-1].commit_id  # branch from the empty initial commit
        api.create_branch(REPO_ID, branch=STABLE_BRANCH, revision=first_commit)
    api.upload_folder(repo_id=REPO_ID, folder_path=STABLE_DIR, revision=STABLE_BRANCH,
                      commit_message=f"Weights at step {STABLE_STEP:,} (end of the stable WSD phase)")
    api.upload_folder(repo_id=REPO_ID, folder_path=stage_main(),
                      commit_message="GPT-U-20M: final weights, tokenizer, model card, scripts and logs")
    print(f"uploaded: https://huggingface.co/{REPO_ID}")


# ---------------------------------------------------------------- verification
def cmd_verify(_args) -> None:
    with tempfile.TemporaryDirectory() as cache:
        tokenizer = AutoTokenizer.from_pretrained(REPO_ID, cache_dir=cache)
        model = AutoModelForCausalLM.from_pretrained(REPO_ID, cache_dir=cache)
        n_params = sum(p.numel() for p in model.parameters())
        inputs = tokenizer("The history of the Roman Empire", return_tensors="pt")
        assert inputs.input_ids[0, 0].item() == 0, "tokenizer should prepend <|endoftext|>"
        torch.manual_seed(0)
        out = model.generate(**inputs, max_new_tokens=60, do_sample=True, temperature=0.8, top_p=0.95)
        print(f"main: {n_params:,} params | {tokenizer.decode(out[0], skip_special_tokens=True)!r}")
        stable = AutoModelForCausalLM.from_pretrained(REPO_ID, revision=STABLE_BRANCH, cache_dir=cache)
        print(f"{STABLE_BRANCH}: {sum(p.numel() for p in stable.parameters()):,} params loaded")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("command", choices=["export", "card", "upload", "verify"])
    args = ap.parse_args()
    {"export": cmd_export, "card": cmd_card, "upload": cmd_upload, "verify": cmd_verify}[args.command](args)


if __name__ == "__main__":
    main()