"""
Minimal GPT training with Trackio alerts for autonomous LR search.
Trains a tiny transformer on Shakespeare - runs in ~30s per run on CPU.

Usage:
uv run --with trackio --with torch python train.py --search
uv run --with trackio --with torch python train.py --lr 0.01 --steps 500
"""

import trackio
import torch
import torch.nn as nn
from torch.nn import functional as F
import argparse
import urllib.request
import os

# ---- Model ----
class TinyGPT(nn.Module):
def __init__(self, vocab_size, n_embd=64, n_head=4, n_layer=2, block_size=64):
super().__init__()
self.block_size = block_size
self.tok_emb = nn.Embedding(vocab_size, n_embd)
self.pos_emb = nn.Embedding(block_size, n_embd)
layer = nn.TransformerEncoderLayer(
d_model=n_embd, nhead=n_head, dim_feedforward=n_embd * 4,
dropout=0.0, batch_first=True,
)
self.transformer = nn.TransformerEncoder(layer, num_layers=n_layer)
self.ln_f = nn.LayerNorm(n_embd)
self.head = nn.Linear(n_embd, vocab_size, bias=False)

def forward(self, idx):
B, T = idx.shape
tok = self.tok_emb(idx)
pos = self.pos_emb(torch.arange(T, device=idx.device))
x = tok + pos
mask = torch.triu(torch.ones(T, T, device=idx.device), diagonal=1).bool()
x = self.transformer(x, mask=mask)
x = self.ln_f(x)
return self.head(x)

# ---- Data ----
def load_data(block_size=64):
data_path = os.path.join(os.path.dirname(__file__), "shakespeare.txt")
if not os.path.exists(data_path):
url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
print("Downloading Shakespeare dataset...")
urllib.request.urlretrieve(url, data_path)

with open(data_path, "r") as f:
text = f.read()

chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
return data[:n], data[n:], len(chars), stoi

def get_batch(data, block_size=64, batch_size=32):
ix = torch.randint(len(data) - block_size, (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, y

# ---- Training with Trackio ----
def train(learning_rate: float, max_steps: int = 300, val_interval: int = 25):
device = "cuda" if torch.cuda.is_available() else "cpu"
block_size, batch_size = 64, 32

train_data, val_data, vocab_size, stoi = load_data(block_size)
model = TinyGPT(vocab_size).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

trackio.init(
project="nanogpt-lr-search",
name=f"lr-{learning_rate}",
config={
"learning_rate": learning_rate,
"max_steps": max_steps,
"model": "TinyGPT-64d-2L",
"dataset": "tinyshakespeare",
"block_size": block_size,
"batch_size": batch_size,
"device": device,
},
)

prev_val_loss = None
val_loss_increases = 0 # Note: cumulative, not consecutive

try:
for step in range(max_steps):
model.train()
xb, yb = get_batch(train_data, block_size, batch_size)
xb, yb = xb.to(device), yb.to(device)
logits = model(xb)
loss = F.cross_entropy(logits.view(-1, vocab_size), yb.view(-1))
optimizer.zero_grad()
loss.backward()
optimizer.step()

trackio.log({"train/loss": loss.item(), "train/lr": learning_rate})

if step % 50 == 0:
print(f" step {step:4d} | train_loss: {loss.item():.4f}")

if step % val_interval == 0 and step \> 0:
model.eval()
val_losses = []
with torch.no_grad():
for _ in range(5):
xv, yv = get_batch(val_data, block_size, batch_size)
xv, yv = xv.to(device), yv.to(device)
vloss = F.cross_entropy(model(xv).view(-1, vocab_size), yv.view(-1))
val_losses.append(vloss.item())
val_loss = sum(val_losses) / len(val_losses)

trackio.log({"val/loss": val_loss})
print(f" step {step:4d} | val_loss: {val_loss:.4f}")

if prev_val_loss is not None and val_loss \> prev_val_loss:
val_loss_increases += 1
trackio.alert(
title="Validation loss increased",
text=f"Val loss: {prev_val_loss:.4f} -\> {val_loss:.4f} at step {step} (lr={learning_rate})",
level=trackio.AlertLevel.WARN,
)
if val_loss_increases \>= 3:
trackio.alert(
title="Training unstable - terminating",
text=f"Val loss increased {val_loss_increases}x. LR {learning_rate} too high.",
level=trackio.AlertLevel.ERROR,
)
print(f" [TERMINATE] LR {learning_rate} unstable after {val_loss_increases} increases")
return False

prev_val_loss = val_loss

print(f" [DONE] LR {learning_rate} completed {max_steps} steps")
return True
finally:
trackio.finish()

# ---- Autonomous search ----
def lr_search():
learning_rates = [1.0, 0.1, 0.01, 0.001]
print("=" * 50)
print(" Autonomous Learning Rate Search")
print("=" * 50)

for lr in learning_rates:
print(f"\n\>\>\> Testing LR = {lr}")
stable = train(learning_rate=lr, max_steps=300, val_interval=25)
if stable:
print(f"\n✓ First stable LR found: {lr}")
return lr

print("\nNo stable LR found.")
return None

if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Autonomous LR search with Trackio alerts")
parser.add_argument("--lr", type=float, default=None, help="Single LR to test")
parser.add_argument("--steps", type=int, default=300, help="Max training steps")
parser.add_argument("--search", action="store_true", help="Run autonomous LR search")
args = parser.parse_args()

if args.search:
lr_search()
elif args.lr is not None:
train(learning_rate=args.lr, max_steps=args.steps)
else:
parser.print_help()