src/train.py

"""Train one nano-GPT per (corpus, scheme) cell.

    python src/train.py                 # train everything that's missing
    python src/train.py --force         # retrain all
    python src/train.py --quick         # fast smoke run (few iters)
    python src/train.py --only-corpus shakespeare --only-scheme bpe512

For each cell it: builds + saves the tokenizer, encodes the (capped) corpus,
trains a GPT, prints a sample, and saves model.pt + config.json into
artifacts/<corpus>/<scheme>/.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import time

import torch

import config
from model import GPT, GPTConfig
from tokenizer import build_tokenizer, load_tokenizer


# --------------------------------------------------------------------------- #
def get_batch(data, block_size, batch_size, device):
    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + 1 + block_size] for i in ix])
    return x.to(device, non_blocking=True), y.to(device, non_blocking=True)


@torch.no_grad()
def estimate_loss(model, splits, args, device):
    out = {}
    model.eval()
    for name, data in splits.items():
        losses = torch.zeros(args.eval_iters)
        for k in range(args.eval_iters):
            x, y = get_batch(data, args.block_size, args.batch_size, device)
            with _autocast(device):
                _, loss = model(x, y)
            losses[k] = loss.item()
        out[name] = losses.mean().item()
    model.train()
    return out


def _autocast(device):
    if device.type == "cuda":
        return torch.autocast(device_type="cuda", dtype=torch.bfloat16)
    return torch.autocast(device_type="cpu", enabled=False)


def lr_at(it, args):
    if it < args.warmup:
        return args.lr * (it + 1) / args.warmup
    if it > args.iters:
        return args.min_lr
    ratio = (it - args.warmup) / max(1, args.iters - args.warmup)
    coeff = 0.5 * (1.0 + math.cos(math.pi * ratio))
    return args.min_lr + coeff * (args.lr - args.min_lr)


def cell_is_done(adir: str) -> bool:
    return all(os.path.exists(os.path.join(adir, f))
               for f in ("tokenizer.json", "model.pt", "config.json"))


# --------------------------------------------------------------------------- #
def train_cell(corpus: dict, scheme: dict, args, device) -> None:
    adir = config.artifact_dir(corpus["name"], scheme["name"])
    tag = f"{corpus['name']}/{scheme['name']}"
    if cell_is_done(adir) and not args.force:
        print(f"[{tag}] already trained — skip (use --force to retrain)")
        return

    text_path = config.corpus_text_path(corpus["name"])
    if not os.path.exists(text_path):
        print(f"[{tag}] corpus text missing ({text_path}); run prepare_data.py first")
        return
    text = open(text_path, encoding="utf-8", errors="replace").read()
    if args.max_chars and len(text) > args.max_chars:
        text = text[:args.max_chars]

    print(f"\n[{tag}] building tokenizer over {len(text)/1e6:.2f}M chars ...")
    t0 = time.time()
    tok = build_tokenizer(scheme, text, verbose=args.verbose)
    os.makedirs(adir, exist_ok=True)
    tok.save(adir)
    print(f"[{tag}] tokenizer: vocab={tok.vocab_size}  "
          f"merges={len(getattr(tok, 'merges', []))}  ({time.time()-t0:.1f}s)")

    # encode corpus -> ids
    t0 = time.time()
    ids = torch.tensor(tok.encode(text), dtype=torch.long)
    n = len(ids)
    tok_per_char = n / len(text)
    split = int(0.9 * n)
    splits = {"train": ids[:split], "val": ids[split:]}
    print(f"[{tag}] encoded {n} tokens "
          f"({tok_per_char:.2f} tok/char)  ({time.time()-t0:.1f}s)")

    # model
    block_size = min(args.block_size, len(splits["val"]) - 1)
    cfg = GPTConfig(vocab_size=tok.vocab_size, block_size=block_size,
                    n_layer=args.n_layer, n_head=args.n_head,
                    n_embd=args.n_embd, dropout=args.dropout)
    model = GPT(cfg).to(device)
    print(f"[{tag}] model params: {model.num_params()/1e6:.2f}M  block={block_size}")

    opt = torch.optim.AdamW(model.parameters(), lr=args.lr,
                            betas=(0.9, 0.99), weight_decay=0.1)

    t0 = time.time()
    model.train()
    for it in range(args.iters + 1):
        for g in opt.param_groups:
            g["lr"] = lr_at(it, args)
        if it % args.eval_interval == 0 or it == args.iters:
            losses = estimate_loss(model, splits, args, device)
            print(f"[{tag}]   iter {it:5d}/{args.iters}  "
                  f"train {losses['train']:.3f}  val {losses['val']:.3f}  "
                  f"lr {opt.param_groups[0]['lr']:.2e}  ({time.time()-t0:.0f}s)")
        if it == args.iters:
            break
        x, y = get_batch(splits["train"], block_size, args.batch_size, device)
        with _autocast(device):
            _, loss = model(x, y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()

    # sample
    sample = generate_sample(model, tok, corpus.get("sample_text", "\n"), device)
    print(f"[{tag}] sample: {sample!r}")

    # bits-per-character makes loss comparable across tokenizers (per-token loss
    # is not — a bigger vocab inflates it). lower is better, and BPE usually wins.
    bits_per_char = losses["val"] * tok_per_char / math.log(2)

    # save
    torch.save(model.state_dict(), os.path.join(adir, "model.pt"))
    with open(os.path.join(adir, "config.json"), "w", encoding="utf-8") as f:
        json.dump({
            "corpus": corpus["name"],
            "scheme": scheme["name"],
            "model_config": cfg.to_dict(),
            "tokenizer": tok.to_meta(),
            "final_loss": losses,
            "tok_per_char": tok_per_char,
            "bits_per_char": bits_per_char,
            "sample": sample,
        }, f, ensure_ascii=False, indent=2)
    print(f"[{tag}] saved -> {adir}  (val {losses['val']:.3f} | {bits_per_char:.2f} bits/char)")


@torch.no_grad()
def generate_sample(model, tok, seed_text, device, n_new=200):
    model.eval()
    ids = tok.encode(seed_text) or tok.encode("\n") or [1]
    idx = torch.tensor([ids], dtype=torch.long, device=device)
    out = model.generate(idx, max_new_tokens=n_new, temperature=0.8, top_k=40)
    model.train()
    return tok.decode(out[0].tolist())


# --------------------------------------------------------------------------- #
def build_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--force", action="store_true")
    ap.add_argument("--quick", action="store_true", help="few iters, for smoke testing")
    ap.add_argument("--verbose", action="store_true")
    ap.add_argument("--only-corpus", default=None)
    ap.add_argument("--only-scheme", default=None)
    ap.add_argument("--device", default=None)
    ap.add_argument("--max-chars", type=int, default=1_500_000)
    ap.add_argument("--iters", type=int, default=3000)
    ap.add_argument("--batch-size", type=int, default=64)
    ap.add_argument("--block-size", type=int, default=128)
    ap.add_argument("--n-layer", type=int, default=4)
    ap.add_argument("--n-head", type=int, default=4)
    ap.add_argument("--n-embd", type=int, default=128)
    ap.add_argument("--dropout", type=float, default=0.1)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--min-lr", type=float, default=3e-5)
    ap.add_argument("--warmup", type=int, default=100)
    ap.add_argument("--eval-interval", type=int, default=250)
    ap.add_argument("--eval-iters", type=int, default=50)
    args = ap.parse_args()
    if args.quick:
        args.iters = 200
        args.eval_interval = 100
        args.warmup = 20
        args.max_chars = min(args.max_chars, 400_000)
    return args


def main():
    args = build_args()
    device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
    torch.manual_seed(1337)
    print(f"device: {device}  |  iters: {args.iters}  |  max_chars: {args.max_chars}")

    corpora = [c for c in config.corpora()
               if not args.only_corpus or c["name"] == args.only_corpus]
    schemes = [s for s in config.schemes()
               if not args.only_scheme or s["name"] == args.only_scheme]

    t0 = time.time()
    for corpus in corpora:
        for scheme in schemes:
            train_cell(corpus, scheme, args, device)
    print(f"\nAll cells done in {time.time()-t0:.0f}s.")


if __name__ == "__main__":
    main()