code/ts300.py

#!/usr/bin/env python3
"""Train a ~10M-parameter GPT on tinystories-top300, then sample stories from it and from
roneneldan/TinyStories-33M for blind judging.

    export HF_TOKEN=hf_...          # read access to the private top-300 repos
    python ts300.py train  --out runs/ts300_10m
    python ts300.py sample --ckpt runs/ts300_10m/best --out runs/ts300_10m/samples
    python ts300.py score  --key runs/ts300_10m/samples/judge_key.json --judge judge_replies.txt

`train` uses the same architecture as your model.py (pre-RMSNorm, RoPE, GELU MLP, tied
embeddings, no biases) and saves model.safetensors + config.json in the same layout.
Defaults: d_model 320, 8 layers, 5 heads -> ~9.8M non-embedding, ~10.7M total params,
10 epochs over the 43.0M training tokens.

`sample` continues the same 20 story beginnings (every word inside the 300-lemma
lexicon) with both models and writes judge_packet_partN.md (blind: A/B order randomized
per item) plus judge_key.json (keep this out of the judging chat).

`score` reads the judge's replies (pasted into one text file), unblinds them with the key
and prints per-model averages and pairwise win rates.

Needs: torch, numpy, pyarrow, huggingface_hub, tokenizers, safetensors, transformers.
"""
import argparse
import csv
import io
import json
import math
import os
import random
import re
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

# Everything (tokenized blocks, tokenizer, lexicon) lives in one repo: huggingface.co/a-kemeklis/constraining-pretraining-data
DATA_REPO = "a-kemeklis/constraining-pretraining-data"
DATA_REV = "main"  # pin a commit hash here for exact reproducibility
DATA_DIR = "data/blocks512"
TOK_REPO = DATA_REPO
TOK_REV = DATA_REV
REF_MODEL = "roneneldan/TinyStories-33M"
EOS_ID = 0  # <|endoftext|> in the top-300 tokenizer; also written before every story
OURS = "ts300-10m"

# Story beginnings for sampling. Every word is in the 300-lemma lexicon (or its 50 names),
# so the top-300 model never sees out-of-vocabulary text; `sample` re-checks this.
PROMPTS = [
    "Once upon a time, there was a little girl named Lily. She had a red ball and a blue box. One day, she put the ball in the box and went to the park. When she came back home, she",
    "Tom and his dog Buddy liked to play in the garden. One day, Buddy found a big stick under a tree. Tom wanted the stick too, so",
    "Mia was very sad because she could not find her doll. She looked in her room and in the garden, but",
    "Ben saw a little bird on the ground. The bird could not fly. Ben wanted to help the bird, so he",
    "One day, Sue and her mom went to the forest. Sue saw a fluffy bunny near a big rock. She said, \"",
    "Jack had a new toy car. His friend Sam wanted to play with it, but Jack did not want to share. Sam was sad and",
    "Lucy and Ben were at the park. Lucy wanted to go on the swing, and Ben wanted to go on the slide. They",
    "Once upon a time, a little frog lived near a big tree. Every day, the frog liked to jump and sing. One day, a loud noise",
    "Anna made a cake for her mom. She wanted to give it to her mom, but the cake fell on the ground. Anna",
    "Tim found a shiny rock in the garden. He showed it to his dad and said, \"",
    "The sun was out and the sky was blue. Lily and her friend Jane wanted to",
    "A little mouse lived in a hole in a big house. The mouse was scared of the cat. One day, the cat went outside, so",
    "Max and Lily had a boat. They put the boat in the water and watched it go. Then the wind",
    "Sara wanted to learn how to climb a tree. Her dad said, \"Be careful!\" Sara started to climb, but",
    "Once upon a time, there was a kind man who had a big truck. Every day, he helped people in the town. One day,",
    "Ben and Mia were angry at each other because they both wanted the red ball. Mom said,",
    "Lily liked to sing and dance with her friends. One day, she could not find her friends. She looked in the park and in the forest, but",
    "Tom had two toys, a ball and a doll. He gave the doll to his friend Ann. Now Tom had",
    "There was a little dog named Buddy. Buddy liked to run fast. One day, Buddy saw a cat near the door. The cat",
    "Emma had a little green bug in a box. She showed the bug to her friend Leo. Leo said, \"Can I hold it?\" Emma",
]


# --------------------------------------------------------------------------- model
# Same architecture and parameter names as model.py in your HF repos.

@dataclass
class GPTConfig:
    vocab_size: int = 2700
    n_layer: int = 8
    n_head: int = 5
    d_model: int = 320
    seq_len: int = 512
    dropout: float = 0.0


class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        return self.w * x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype)


def rope_cache(T, hd, base=10000.0, device=None):
    inv = 1.0 / (base ** (torch.arange(0, hd, 2, device=device).float() / hd))
    f = torch.outer(torch.arange(T, device=device).float(), inv)
    return f.cos(), f.sin()


def apply_rope(x, cos, sin):  # x: B,H,T,hd
    x1, x2 = x[..., ::2], x[..., 1::2]
    T = x.shape[-2]
    c, s = cos[:T].to(x.dtype), sin[:T].to(x.dtype)
    return torch.stack((x1 * c - x2 * s, x1 * s + x2 * c), -1).flatten(-2)


class Block(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        d, self.h = cfg.d_model, cfg.n_head
        self.n1, self.n2 = RMSNorm(d), RMSNorm(d)
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.proj = nn.Linear(d, d, bias=False)
        self.fc = nn.Linear(d, 4 * d, bias=False)
        self.fc2 = nn.Linear(4 * d, d, bias=False)
        self.drop = cfg.dropout

    def forward(self, x, cos, sin):
        B, T, D = x.shape
        q, k, v = self.qkv(self.n1(x)).view(B, T, 3, self.h, D // self.h).permute(2, 0, 3, 1, 4)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                                           dropout_p=self.drop if self.training else 0.0)
        x = x + F.dropout(self.proj(a.transpose(1, 2).reshape(B, T, D)), self.drop, self.training)
        x = x + F.dropout(self.fc2(F.gelu(self.fc(self.n2(x)), approximate="tanh")), self.drop, self.training)
        return x


class GPT(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.cfg = cfg
        self.emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.norm = RMSNorm(cfg.d_model)
        cos, sin = rope_cache(cfg.seq_len, cfg.d_model // cfg.n_head)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)
        self.apply(self._init)
        for n, p in self.named_parameters():  # GPT-2 style scaled init on residual projections
            if n.endswith("proj.weight") or n.endswith("fc2.weight"):
                nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * cfg.n_layer))

    @staticmethod
    def _init(m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, 0.0, 0.02)

    def forward(self, idx, targets=None):
        x = F.dropout(self.emb(idx), self.cfg.dropout, self.training)
        for b in self.blocks:
            x = b(x, self.cos, self.sin)
        logits = F.linear(self.norm(x), self.emb.weight)  # tied head
        if targets is None:
            return logits
        return logits, F.cross_entropy(logits.float().reshape(-1, logits.size(-1)), targets.reshape(-1))

    def n_params(self, non_embedding=True):
        n = sum(p.numel() for p in self.parameters())
        return n - self.emb.weight.numel() if non_embedding else n


# --------------------------------------------------------------------------- helpers

def pick_device(name):
    if name != "auto":
        return torch.device(name)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def autocast_ctx(device):
    if device.type == "cuda" and torch.cuda.is_bf16_supported():
        return torch.autocast("cuda", dtype=torch.bfloat16)
    return torch.autocast(device.type, enabled=False)


def hf_token():
    return os.environ.get("HF_TOKEN") or None  # None -> use `huggingface-cli login` cache


def fetch_data(args, train_shards=0, include_train=True):
    """Download data/blocks512 (train/validation shards, meta, tokenizer) and return the local dir.
    train_shards=0 means all 16 shards; include_train=False fetches only meta, tokenizer, validation."""
    from huggingface_hub import snapshot_download
    patterns = [f"{DATA_DIR}/meta.json", f"{DATA_DIR}/tokenizer/*", f"{DATA_DIR}/validation-*"]
    if include_train and train_shards:
        patterns += [f"{DATA_DIR}/train-{i:05d}-of-00016.parquet" for i in range(train_shards)]
    elif include_train:
        patterns += [f"{DATA_DIR}/train-*"]
    root = snapshot_download(DATA_REPO, repo_type="model", revision=args.data_revision,
                             allow_patterns=patterns, token=hf_token())
    return Path(root) / DATA_DIR


def load_blocks(files):
    import pyarrow.parquet as pq
    arrays = []
    for f in sorted(files):
        col = pq.read_table(f, columns=["input_ids"]).column("input_ids").combine_chunks()
        width = col.type.list_size
        arrays.append(col.values.to_numpy(zero_copy_only=False).reshape(-1, width))
    return np.concatenate(arrays).astype(np.uint16)


def save_model(model, cfg, path: Path, extra: dict):
    from safetensors.torch import save_file
    path.mkdir(parents=True, exist_ok=True)
    state = {k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()}
    save_file(state, str(path / "model.safetensors"))
    (path / "config.json").write_text(json.dumps({**asdict(cfg), **extra}, indent=1))


def load_model(path: Path, device):
    from safetensors.torch import load_file
    raw = json.loads((path / "config.json").read_text())
    cfg = GPTConfig(**{k: v for k, v in raw.items() if k in GPTConfig.__dataclass_fields__})
    model = GPT(cfg)
    model.load_state_dict(load_file(str(path / "model.safetensors")))
    return model.to(device).eval(), cfg


# --------------------------------------------------------------------------- train

@torch.no_grad()
def evaluate(model, val, device, batch=32):
    model.eval()
    total, n = 0.0, 0
    for i in range(0, len(val), batch):
        b = torch.from_numpy(val[i:i + batch].astype(np.int64)).to(device)
        with autocast_ctx(device):
            _, loss = model(b[:, :-1], b[:, 1:])
        total += loss.item() * len(b)
        n += len(b)
    model.train()
    return total / n


def train(args):
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    device = pick_device(args.device)
    if device.type == "cuda":
        torch.set_float32_matmul_precision("high")
        print(f"GPU: {torch.cuda.get_device_name(device)} | torch {torch.__version__} | "
              f"bf16 {torch.cuda.is_bf16_supported()}")
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    t_start = time.time()

    data_dir = fetch_data(args, args.train_shards)
    meta = json.loads((data_dir / "meta.json").read_text())
    train_blocks = load_blocks(data_dir.glob("train-*.parquet"))
    val_blocks = load_blocks(data_dir.glob("validation-*.parquet"))
    seq_len = train_blocks.shape[1]
    print(f"train blocks {len(train_blocks):,} x {seq_len} = {train_blocks.size / 1e6:.1f}M tokens | "
          f"val blocks {len(val_blocks):,}")

    cfg = GPTConfig(vocab_size=meta["vocab_size"], n_layer=args.n_layer, n_head=args.n_head,
                    d_model=args.d_model, seq_len=seq_len, dropout=args.dropout)
    assert cfg.d_model % cfg.n_head == 0 and (cfg.d_model // cfg.n_head) % 2 == 0, "head dim must be even"
    model = GPT(cfg).to(device)
    n_ne, n_tot = model.n_params(True), model.n_params(False)
    print(f"model: d_model {cfg.d_model}, {cfg.n_layer} layers, {cfg.n_head} heads | "
          f"{n_ne / 1e6:.2f}M non-embedding, {n_tot / 1e6:.2f}M total params | device {device}")

    steps_per_epoch = len(train_blocks) // args.batch
    total_steps = int(args.epochs * steps_per_epoch)
    if args.max_steps:
        total_steps = min(total_steps, args.max_steps)
    warmup = max(1, int(args.warmup_frac * total_steps))
    accum = args.batch // args.micro
    assert accum * args.micro == args.batch, "--batch must be a multiple of --micro"
    tokens_total = total_steps * args.batch * (seq_len - 1)
    print(f"{total_steps:,} steps x {args.batch} blocks = {tokens_total / 1e6:.0f}M tokens "
          f"({tokens_total / (len(train_blocks) * (seq_len - 1)):.1f} epochs, "
          f"{tokens_total / n_tot:.0f} tokens/param)")

    decay = [p for p in model.parameters() if p.dim() >= 2]
    no_decay = [p for p in model.parameters() if p.dim() < 2]
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": args.wd},
                             {"params": no_decay, "weight_decay": 0.0}],
                            lr=args.lr, betas=(0.9, 0.95), fused=(device.type == "cuda"))

    def lr_at(step):  # linear warmup, cosine decay to min_lr_frac * lr
        if step < warmup:
            return args.lr * (step + 1) / warmup
        t = (step - warmup) / max(1, total_steps - warmup)
        return args.lr * (args.min_lr_frac + (1 - args.min_lr_frac) * 0.5 * (1 + math.cos(math.pi * t)))

    fwd = torch.compile(model) if (args.compile and device.type == "cuda") else model
    rng = np.random.default_rng(args.seed)
    order = rng.permutation(len(train_blocks))
    pos = 0
    best = float("inf")
    log = open(out / "log.jsonl", "a")
    t0, tok_since = time.time(), 0
    model.train()
    for step in range(total_steps):
        for g in opt.param_groups:
            g["lr"] = lr_at(step)
        for _ in range(accum):
            if pos + args.micro > len(order):
                order, pos = rng.permutation(len(train_blocks)), 0
            b = train_blocks[order[pos:pos + args.micro]].astype(np.int64)
            pos += args.micro
            b = torch.from_numpy(b).to(device, non_blocking=True)
            with autocast_ctx(device):
                _, loss = fwd(b[:, :-1], b[:, 1:])
            (loss / accum).backward()
            tok_since += b[:, 1:].numel()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        opt.zero_grad(set_to_none=True)

        last = step == total_steps - 1
        if step % args.log_every == 0 or last:
            dt = time.time() - t0
            print(f"step {step:6d}/{total_steps} | epoch {step / steps_per_epoch:5.2f} | "
                  f"loss {loss.item():.4f} | lr {lr_at(step):.2e} | {tok_since / max(dt, 1e-9):,.0f} tok/s")
            t0, tok_since = time.time(), 0
        if (step > 0 and step % args.eval_every == 0) or last:
            vl = evaluate(model, val_blocks, device)
            rec = {"step": step, "epoch": step / steps_per_epoch, "train_loss": loss.item(), "val_loss": vl}
            log.write(json.dumps(rec) + "\n")
            log.flush()
            print(f"  val loss {vl:.4f}" + ("  (best)" if vl < best else ""))
            extra = {"step": step, "val_loss": vl, "n_params_nonembedding": n_ne, "n_params_total": n_tot,
                     "data": f"{DATA_REPO}@{args.data_revision}", "train_args": vars(args)}
            if vl < best:
                best = vl
                save_model(model, cfg, out / "best", extra)
            if last:
                save_model(model, cfg, out / "final", extra)
    for name in ("best", "final"):
        for f in ("tokenizer.json", "tokenizer_config.json"):
            src = data_dir / "tokenizer" / f
            if src.exists():
                (out / name / f).write_bytes(src.read_bytes())
    best_cfg = json.loads((out / "best" / "config.json").read_text())
    minutes = (time.time() - t_start) / 60
    summary = {"n_params_nonembedding": n_ne, "n_params_total": n_tot, "config": asdict(cfg),
               "steps": total_steps, "batch_blocks": args.batch, "tokens_seen": tokens_total,
               "epochs": tokens_total / (len(train_blocks) * (seq_len - 1)),
               "best_val_loss": best, "best_step": best_cfg["step"], "final_val_loss": vl,
               "wall_minutes": round(minutes, 1), "device": str(device),
               "gpu": torch.cuda.get_device_name(device) if device.type == "cuda" else None,
               "torch": torch.__version__, "compiled": fwd is not model, "train_args": vars(args)}
    (out / "run_summary.json").write_text(json.dumps(summary, indent=1))
    print(f"done in {minutes:.1f} min. best val loss {best:.4f} at step {best_cfg['step']}, "
          f"final {vl:.4f}; checkpoints in {out}/best and {out}/final; summary in {out}/run_summary.json")


# --------------------------------------------------------------------------- sample

def placeholder_ids(tok):
    """Ids of <UNK_k>/<NAME_k>/<NUM_k> tokens, plus the bare "UNK"/"NAME"/"NUM" pieces that could spell one out."""
    return sorted(i for t, i in tok.get_vocab().items()
                  if re.search(r"<(UNK|NAME|NUM)_\d+>", t) or t.lstrip("\u0120 ") in ("UNK", "NAME", "NUM"))


def sample_logits(logits, temperature, top_p, gen):
    if temperature <= 0:
        return int(logits.argmax(-1))
    probs = torch.softmax(logits / temperature, -1)
    if top_p < 1.0:
        p_sorted, idx = probs.sort(descending=True)
        keep = p_sorted.cumsum(-1) - p_sorted < top_p
        probs = torch.zeros_like(probs).scatter(-1, idx[keep], p_sorted[keep])
        probs = probs / probs.sum()
    return int(torch.multinomial(probs, 1, generator=gen))


RESAMPLE_STATS = {"rollbacks": 0, "masked_sentences": 0}


@torch.no_grad()
def generate_ours(model, cfg, prompt_ids, args, seed, device, banned=None, sent_end=frozenset()):
    """Sample a continuation. With `banned`, placeholder tokens are masked; with --resample-placeholders N,
    a sampled placeholder instead rolls the current sentence back and resamples it, masking only after N tries."""
    gen = torch.Generator(device="cpu").manual_seed(seed)
    prefix = [EOS_ID] + prompt_ids  # stories in training follow an <|endoftext|>
    banned_set = set(banned.tolist()) if banned is not None else set()
    retries = getattr(args, "resample_placeholders", 0)
    out, sent_start, tries = [], 0, 0
    while len(out) < args.max_new_tokens:
        x = torch.tensor([(prefix + out)[-cfg.seq_len:]], device=device)
        with autocast_ctx(device):
            logits = model(x)[0, -1].float().cpu()
        if banned is not None and tries >= retries:
            if tries == retries and retries:
                RESAMPLE_STATS["masked_sentences"] += 1
                tries += 1
            logits[banned] = -float("inf")
        t = sample_logits(logits, args.temperature, args.top_p, gen)
        if t in banned_set:  # only reachable while resampling is still allowed
            RESAMPLE_STATS["rollbacks"] += 1
            tries += 1
            del out[sent_start:]
            continue
        if t == EOS_ID:
            return out, True
        out.append(t)
        if t in sent_end:
            sent_start, tries = len(out), 0
    return out, False


@torch.no_grad()
def generate_ref(model, tok, prompt, args, seed, device):
    torch.manual_seed(seed)
    enc = tok(prompt, return_tensors="pt").to(device)
    eos = tok.eos_token_id
    out = model.generate(**enc, do_sample=args.temperature > 0, temperature=max(args.temperature, 1e-6),
                         top_p=args.top_p, top_k=0, max_new_tokens=args.max_new_tokens,
                         eos_token_id=eos, pad_token_id=eos)
    new = out[0, enc["input_ids"].shape[1]:].tolist()
    finished = eos in new
    if finished:
        new = new[:new.index(eos)]
    return tok.decode(new, skip_special_tokens=True), finished


def trim_words(text, max_words):
    """Cut to at most max_words words, backing off to the last sentence end."""
    words = text.split()
    if len(words) <= max_words:
        return text.strip(), False
    cut = " ".join(words[:max_words])
    end = max(cut.rfind(c) for c in ".!?")
    if end > len(cut) // 2:
        cut = cut[:end + 1] + ('"' if cut[end + 1:end + 2] == '"' else "")
    return cut.strip(), True


def check_prompts(prompts):
    """Warn about prompt words outside the 300-lemma lexicon (needs lexicon.json from the tokenizer repo)."""
    try:
        from huggingface_hub import hf_hub_download
        lex = json.loads(Path(hf_hub_download(TOK_REPO, "data/lexicon.json", revision=TOK_REV, token=hf_token())).read_text())
    except Exception as e:  # noqa: BLE001 - optional check
        print(f"(skipping lexicon check: {e})")
        return
    forms, names = {f.lower() for f in lex["forms"]}, set(lex["names"])
    for i, p in enumerate(prompts, 1):
        oov = [w for w in re.findall(r"[A-Za-z]+(?:'[A-Za-z]+)?", p)
               if not ({w, w.split("'")[0]} & names) and not ({w.lower(), w.split("'")[0].lower()} & forms)]
        if oov:
            print(f"WARNING prompt {i} has words outside the lexicon: {oov}")


RUBRIC = """You are grading short story continuations written by two small language models trained on simple children's stories. Each item gives a story beginning (written by a person) and two continuations, A and B, from different models. The A/B order is random.

Some continuations contain placeholders such as <UNK_3>, <NAME_2> or <NUM_1>. Each stands for a word, name or number outside a restricted vocabulary. Don't count a placeholder as an error by itself. Do count it as a consistency error if the same placeholder is clearly used for two different things. "[...]" marks text cut off at a length limit; treat that continuation as not reaching an ending.

For each continuation, judged against its beginning:
- grammar_errors: number of grammatical errors (agreement, tense, missing or extra words, broken sentences)
- consistency_errors: number of contradictions with the beginning or with earlier text (wrong names, objects or pronouns, facts that change, actions that make no sense given what happened)
- ending: 1 if the story reaches a sensible ending, otherwise 0
- grammar, consistency, plot: scores from 1 to 10
Then decide which continuation of the item is better overall: A, B or tie.

For each item, first write one or two lines noting the main errors you see. End your reply with a CSV block using exactly this header, one row per continuation (two rows per item, with the same "preferred" value on both rows):

item,label,grammar_errors,consistency_errors,ending,grammar,consistency,plot,preferred
"""


def sample(args):
    from tokenizers import Tokenizer
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from transformers.utils import logging as hf_logging

    hf_logging.set_verbosity_error()  # TinyStories-33M has stale attention-bias buffers; harmless
    device = pick_device(args.device)
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    prompts = PROMPTS[:args.n_prompts] if args.n_prompts else PROMPTS
    check_prompts(prompts)

    ckpt = Path(args.ckpt)
    model, cfg = load_model(ckpt, device)
    tok_path = ckpt / "tokenizer.json"
    if not tok_path.exists():
        tok_path = fetch_data(args, include_train=False) / "tokenizer" / "tokenizer.json"
    tok = Tokenizer.from_file(str(tok_path))
    print(f"{OURS}: {model.n_params(True) / 1e6:.2f}M non-embedding params from {ckpt}")
    banned = torch.tensor(placeholder_ids(tok)) if (args.ban_placeholders or args.resample_placeholders) else None
    sent_end = frozenset(i for i in range(tok.get_vocab_size()) if re.search(r"[.!?\n]\W*$", tok.decode([i])))
    if banned is not None:
        print(f"banning {len(banned)} placeholder tokens while decoding {OURS}")

    ref_tok = AutoTokenizer.from_pretrained(args.ref)
    ref = AutoModelForCausalLM.from_pretrained(args.ref).to(device).eval()
    print(f"{args.ref}: {sum(p.numel() for p in ref.parameters()) / 1e6:.1f}M params (incl. embeddings)")

    records = []
    for i, p in enumerate(prompts, 1):
        p_ids = tok.encode(p).ids
        for s in range(args.samples):
            seed = args.seed + 1000 * i + s
            ids, fin = generate_ours(model, cfg, p_ids, args, seed, device, banned, sent_end)
            full = tok.decode(p_ids + ids, skip_special_tokens=False)
            prefix = tok.decode(p_ids, skip_special_tokens=False)
            ours = full[len(prefix):] if full.startswith(prefix) else tok.decode(ids, skip_special_tokens=False)
            ref_text, ref_fin = generate_ref(ref, ref_tok, p, args, seed, device)
            for name, text, done in ((OURS, ours, fin), (args.ref.split("/")[-1], ref_text, ref_fin)):
                text, cut = trim_words(text, args.max_words)
                records.append({"item": f"P{i:02d}" + (f"s{s + 1}" if args.samples > 1 else ""),
                                "prompt": p, "model": name, "continuation": text,
                                "finished": done and not cut, "n_words": len(text.split())})
        print(f"prompt {i}/{len(prompts)} done")

    with open(out / "samples.jsonl", "w") as f:
        for r in records:
            f.write(json.dumps(r) + "\n")

    items = {}
    for r in records:
        items.setdefault(r["item"], []).append(r)
    # balanced blind order: each model is "A" in half the items, in shuffled positions
    swaps = [i % 2 == 1 for i in range(len(items))]
    random.Random(args.seed).shuffle(swaps)
    key, blocks = {}, []
    for (item, rs), swap in zip(items.items(), swaps):
        if swap:
            rs = rs[::-1]
        key[item] = {"A": rs[0]["model"], "B": rs[1]["model"]}
        body = [f"### Item {item}", "", f"**Beginning:** {rs[0]['prompt']}", ""]
        for label, r in zip("AB", rs):
            text = r["continuation"] + ("" if r["finished"] else " [...]")
            body += [f"**Continuation {label}:**", text, ""]
        blocks.append("\n".join(body))
    (out / "judge_key.json").write_text(json.dumps(key, indent=1))

    n_parts = math.ceil(len(blocks) / args.items_per_part)
    for k in range(n_parts):
        part = blocks[k * args.items_per_part:(k + 1) * args.items_per_part]
        head = f"# Story judging packet, part {k + 1} of {n_parts}\n\n{RUBRIC}\n"
        (out / f"judge_packet_part{k + 1}.md").write_text(head + "\n" + "\n".join(part))

    fin_rate = {m: np.mean([r["finished"] for r in records if r["model"] == m]) for m in {r["model"] for r in records}}
    words = {m: np.mean([r["n_words"] for r in records if r["model"] == m]) for m in fin_rate}
    for m in fin_rate:
        print(f"{m}: ended naturally {fin_rate[m]:.0%}, mean {words[m]:.0f} words")
    (out / "sample_summary.json").write_text(json.dumps({
        "checkpoint": str(ckpt), "reference": args.ref, "prompts": len(prompts), "samples_per_prompt": args.samples,
        "items": len(items), "packet_parts": n_parts, "temperature": args.temperature, "top_p": args.top_p,
        "max_new_tokens": args.max_new_tokens, "max_words": args.max_words, "seed": args.seed,
        "ban_placeholders": args.ban_placeholders, "resample_placeholders": args.resample_placeholders,
        "resample_stats": RESAMPLE_STATS,
        "ended_naturally": {m: float(v) for m, v in fin_rate.items()},
        "mean_words": {m: float(v) for m, v in words.items()}}, indent=1))
    print(f"wrote {out}/samples.jsonl, {n_parts} judge packet part(s) and judge_key.json "
          f"(don't paste the key into the judging chat)")


# --------------------------------------------------------------------------- score

def score(args):
    key = json.loads(Path(args.key).read_text())
    text = Path(args.judge).read_text()
    rows, header = [], None
    for line in text.splitlines():
        line = line.strip().strip("`")
        if line.startswith("item,label,"):
            header = next(csv.reader([line]))
            continue
        if header and line and line.count(",") >= len(header) - 1:
            vals = next(csv.reader(io.StringIO(line)))
            if len(vals) == len(header) and vals[0] in key:
                rows.append(dict(zip(header, vals)))
    if not rows:
        raise SystemExit("no CSV rows found; the judge output needs the header line "
                         "'item,label,grammar_errors,...'")
    nums = ["grammar_errors", "consistency_errors", "ending", "grammar", "consistency", "plot"]
    by_model, prefs = {}, {}
    for r in rows:
        model = key[r["item"]][r["label"].strip().upper()]
        d = by_model.setdefault(model, {c: [] for c in nums})
        for c in nums:
            try:
                d[c].append(float(r[c]))
            except ValueError:
                pass
        pref = r["preferred"].strip().upper()
        prefs[r["item"]] = key[r["item"]].get(pref, "tie")
    models = sorted(by_model)
    print(f"{len(rows)} rows over {len(prefs)} items\n")
    print(f"{'metric':20s}" + "".join(f"{m:>22s}" for m in models))
    for c in nums:
        print(f"{c:20s}" + "".join(f"{np.mean(by_model[m][c]):22.2f}" for m in models))
    print(f"{'preferred (items)':20s}" + "".join(f"{sum(v == m for v in prefs.values()):22d}" for m in models)
          + f"   ties: {sum(v == 'tie' for v in prefs.values())}")


# --------------------------------------------------------------------------- cli

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    t = sub.add_parser("train", help="train the small model on tinystories-top300")
    t.add_argument("--out", default="runs/ts300_10m")
    t.add_argument("--d-model", type=int, default=320)
    t.add_argument("--n-layer", type=int, default=8)
    t.add_argument("--n-head", type=int, default=5)
    t.add_argument("--dropout", type=float, default=0.0)
    t.add_argument("--epochs", type=float, default=10)
    t.add_argument("--batch", type=int, default=64, help="blocks of 512 tokens per optimizer step")
    t.add_argument("--micro", type=int, default=64, help="blocks per forward pass (gradient accumulation)")
    t.add_argument("--lr", type=float, default=1e-3)
    t.add_argument("--min-lr-frac", type=float, default=0.1)
    t.add_argument("--warmup-frac", type=float, default=0.02)
    t.add_argument("--wd", type=float, default=0.1)
    t.add_argument("--eval-every", type=int, default=500)
    t.add_argument("--log-every", type=int, default=50)
    t.add_argument("--max-steps", type=int, default=0, help="cap steps (for quick tests)")
    t.add_argument("--train-shards", type=int, default=0, help="use only the first N of 16 shards (tests)")
    t.add_argument("--compile", action=argparse.BooleanOptionalAction, default=True)

    s = sub.add_parser("sample", help="generate stories from both models and write the judge packet")
    s.add_argument("--ckpt", default="runs/ts300_10m/best")
    s.add_argument("--out", default="runs/ts300_10m/samples")
    s.add_argument("--ref", default=REF_MODEL)
    s.add_argument("--samples", type=int, default=1, help="continuations per prompt per model")
    s.add_argument("--n-prompts", type=int, default=0, help="use only the first N prompts (tests)")
    s.add_argument("--temperature", type=float, default=1.0)
    s.add_argument("--top-p", type=float, default=1.0)
    s.add_argument("--max-new-tokens", type=int, default=400)
    s.add_argument("--max-words", type=int, default=250)
    s.add_argument("--items-per-part", type=int, default=10)
    s.add_argument("--ban-placeholders", action="store_true",
                   help="mask <UNK_k>/<NAME_k>/<NUM_k> tokens when decoding our model")
    s.add_argument("--resample-placeholders", type=int, default=0, metavar="N",
                   help="on a sampled placeholder, resample the current sentence up to N times before masking")

    c = sub.add_parser("score", help="unblind the judge's CSV and print per-model results")
    c.add_argument("--key", required=True)
    c.add_argument("--judge", required=True, help="text file with the judge's replies pasted in")

    for p in (t, s):
        p.add_argument("--device", default="auto")
        p.add_argument("--seed", type=int, default=0)
        p.add_argument("--data-revision", default=DATA_REV)

    args = ap.parse_args()
    {"train": train, "sample": sample, "score": score}[args.cmd](args)


if __name__ == "__main__":
    main()