"""Lesson 7: train a tiny GPT on your own text (sample.txt or any text file). Saves tinygpt.pt.
Usage: python lesson-07-train-tiny-gpt.py [textfile] [steps]"""
import sys, time, torch
from tinygpt import CharTokenizer, TinyGPT, get_batch, generate
torch.manual_seed(42)
text = open(sys.argv[1] if len(sys.argv) > 1 else "sample.txt", encoding="utf-8").read()
tok = CharTokenizer(text); data = torch.tensor(tok.encode(text))
split = int(0.9 * len(data)); train, val = data[:split], data[split:]
steps = int(sys.argv[2]) if len(sys.argv) > 2 else 300
cfg = dict(d_model=96, n_heads=4, n_layers=3, block_size=64)
model = TinyGPT(tok.vocab_size, **cfg)
opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
print(f"vocab {tok.vocab_size}, params {sum(p.numel() for p in model.parameters()):,}")
t0 = time.time()
for step in range(steps + 1):
    model.train(); xb, yb = get_batch(train, 64, 32)
    _, loss = model(xb, yb); opt.zero_grad(); loss.backward(); opt.step()
    if step % 100 == 0:
        with torch.no_grad():
            model.eval(); _, vl = model(*get_batch(val, 64, 32))
        print(f"step {step:4d}  train {loss.item():.3f}  val {vl.item():.3f}")
print(f"trained in {time.time() - t0:.0f}s")
torch.save({"model": model.state_dict(), "chars": tok.chars, "cfg": cfg}, "tinygpt.pt")
print(tok.decode(generate(model, torch.tensor([tok.encode("Excel")]), 80, temperature=0.8)[0].tolist()))