train.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
chat_sprint_standalone.py
One-file pipeline: collect datasets -> reformat as You:/Bot: -> train tiny GPT (CUDA) -> sample & save model

Requirements (install once):
  pip install torch datasets sentencepiece tqdm numpy

Run:
  python chat_sprint_standalone.py
"""

import os, re, time, math, random, json
from pathlib import Path
from typing import List, Optional, Tuple
from itertools import islice
from contextlib import nullcontext

import numpy as np
from tqdm import tqdm
from datasets import load_dataset, get_dataset_config_names
import sentencepiece as spm

import torch
import torch.nn as nn
import torch.nn.functional as F

# --------------------------
# Global config (tweak here)
# --------------------------
SEED = 1337
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if DEVICE.type == "cuda":
    torch.set_float32_matmul_precision("high")
    torch.backends.cuda.matmul.allow_tf32 = True

SAVE_DIR = Path("./chat_sprint_artifacts")
SAVE_DIR.mkdir(parents=True, exist_ok=True)

# How many formatted You:/Bot: pairs to KEEP from each dataset (first run)
CAPS = {
    "shakespeare": 15000,
    "jokes":       20000,
    "dadjokes":     8000,
    "rsarcasm":     8000,   # Thewillonline/reddit-sarcasm
    "figlang":     12000,
    "shower":       4000,   # HuggingFaceGECLM/REDDIT_submissions split "Showerthoughts"
    "personas":     2000,
    "tweeteval":    4000,
    "fourchan":      500,   # dataset is tiny (~195)
    "elonvtrump":   3000,
}

# Upper bound on rows to SCAN in each streaming dataset
SCAN = {
    "jokes":       120_000,
    "dadjokes":     60_000,
    "rsarcasm":    120_000,
    "figlang":     150_000,
    "shower":      250_000,
    "personas":     30_000,
    "tweeteval":    60_000,
    "fourchan":      2_000,
    "elonvtrump":   60_000,
}

MAX_TOTAL_PAIRS = 60_000
MAX_LEN = 120

# Tokenizer
VOCAB_SIZE = 1500
TOKENIZER_PREFIX = SAVE_DIR / "spm_chat"
USER_SYMBOLS = ["You:", "Bot:", "[STYLE=Snark]", "[FORM=TWEET]", "[FORM=HEADLINE]", "[MOOD=Unhinged]", "[MOOD=Cheeky]"]

# Model size & train budget (~5 minutes on RTX 3090 with default)
block_size   = 256
n_layer      = 6
n_head       = 6
n_embd       = 384   # ~11.9M params
dropout      = 0.0

MAX_SECONDS  = 300   # hard cap
train_steps  = 5000  # big number; time cap will stop near 5 min
log_interval = 100
eval_every   = 400
batch_size   = 24
accum_steps  = 3
base_lr      = 3e-3
min_lr       = 5e-4
warmup_ratio = 0.06

# Sampling defaults
TEMP = 0.8
TOP_K = 60
TOP_P = 0.95
REP_PEN = 1.08

# --------------------------
# Helpers: cleaning & format
# --------------------------
URL_RE   = re.compile(r"https?://\S+|www\.\S+", re.IGNORECASE)
MENT_RE  = re.compile(r"@\w+")
WS_RE    = re.compile(r"\s+")
QUOTE_RE = re.compile(r"^[\"'“”‘’]+|[\"'“”‘’]+$")

def clean_text(s: str) -> str:
    s = s.strip()
    s = URL_RE.sub("", s)
    s = MENT_RE.sub("", s)
    s = QUOTE_RE.sub("", s)
    s = WS_RE.sub(" ", s)
    return s.strip()

def shorten_to(s: str, n: int) -> str:
    s = re.sub(r"\s+", " ", s.strip())
    if len(s) <= n: return s
    cut = max(s.rfind(". ", 0, n), s.rfind("! ", 0, n), s.rfind("? ", 0, n))
    if cut != -1: return s[:cut+1].strip()
    return s[:n].strip()

def keep_or_clip(s: str, min_len: int = 6, max_len: int = MAX_LEN) -> Optional[str]:
    if not s: return None
    s = re.sub(r"\s+", " ", s.strip())
    if len(s) < min_len: return None
    return shorten_to(s, max_len)

def turn(you: str, bot: str, tags: str = "") -> str:
    lines = [f"You: {you}".rstrip()]
    if tags: lines.append(tags)
    lines.append(f"Bot: {bot}".rstrip())
    lines.append("")
    return "\n".join(lines)

def limited(ds, limit: int):
    try:
        return ds.take(limit)
    except Exception:
        return islice(ds, limit)

def get_first_nonempty(ex, keys) -> Optional[str]:
    for k in keys:
        v = ex.get(k)
        if isinstance(v, str) and v.strip():
            return v
    return None

# --------------------------
# Collectors (console tqdm)
# --------------------------
def collect_shakespeare(pairs: List[str], overall: tqdm):
    try:
        ds = load_dataset(
            "text",
            data_files={"train": "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"},
            split="train",
            streaming=True,
        )
        kept = 0
        pbar = tqdm(total=CAPS["shakespeare"], desc="[shakespeare]", unit="pair", leave=False, ncols=100)
        for row in ds:
            txt = keep_or_clip(clean_text(row["text"]))
            if not txt: continue
            pairs.append(turn("Continue in Shakespearean style.", txt))
            kept += 1; pbar.update(1); overall.update(1)
            if kept >= CAPS["shakespeare"]: break
        pbar.close()
        print(f"[ok] shakespeare kept={kept}")
    except Exception as e:
        print(f"[skip] shakespeare: {e}")

def collect_reddit_jokes(pairs: List[str], overall: tqdm):
    for dsid in ["SocialGrep/one-million-reddit-jokes", "SocialGrep/reddit_jokes", "timc1/reddit_jokes"]:
        try:
            ds = load_dataset(dsid, split="train", streaming=True)
            kept = 0
            pbar = tqdm(total=CAPS["jokes"], desc="[jokes]", unit="pair", leave=False, ncols=100)
            for ex in limited(ds, SCAN["jokes"]):
                title = keep_or_clip(clean_text(str(ex.get("title") or "")))
                body  = keep_or_clip(clean_text(str(ex.get("selftext") or ex.get("body") or "")))
                if body and title:
                    pairs.append(turn(title, body))
                elif title:
                    pairs.append(turn("Tell me a short joke.", title))
                else:
                    continue
                kept += 1; pbar.update(1); overall.update(1)
                if kept >= CAPS["jokes"]: break
            pbar.close()
            print(f"[ok] jokes {dsid} kept={kept}")
            return
        except Exception as e:
            print(f"[try next] jokes {dsid}: {e}")
    print("[skip] jokes: none worked")

def collect_dadjokes(pairs: List[str], overall: tqdm):
    try:
        ds = load_dataset("shuttie/reddit-dadjokes", split="train", streaming=True)
        kept = 0
        pbar = tqdm(total=CAPS["dadjokes"], desc="[dadjokes]", unit="pair", leave=False, ncols=100)
        for ex in limited(ds, SCAN["dadjokes"]):
            setup = keep_or_clip(clean_text(str(ex.get("setup") or ex.get("instruction") or ex.get("input") or "")))
            punch = keep_or_clip(clean_text(str(ex.get("punchline") or ex.get("output") or "")))
            if not (setup and punch): continue
            pairs.append(turn(setup, punch))
            kept += 1; pbar.update(1); overall.update(1)
            if kept >= CAPS["dadjokes"]: break
        pbar.close()
        print(f"[ok] dadjokes kept={kept}")
    except Exception as e:
        print(f"[skip] dadjokes: {e}")

def collect_reddit_sarcasm(pairs: List[str], overall: tqdm):
    """Thewillonline/reddit-sarcasm — flexible parsing, scan+keep bars."""
    try:
        ds = load_dataset("Thewillonline/reddit-sarcasm", split="train", streaming=True)
        keep_cap, scan_cap = CAPS["rsarcasm"], SCAN["rsarcasm"]
        scanbar = tqdm(total=scan_cap, desc="[sarcasm scan]", unit="row", leave=False, ncols=100)
        keepbar = tqdm(total=keep_cap, desc="[sarcasm kept]", unit="pair", leave=False, ncols=100)

        PATS = [
            re.compile(r"User\s*:\s*(.+?)\s*(?:Reddit\s*Comment|Comment|Reply)\s*:\s*(.+)", re.IGNORECASE | re.DOTALL),
            re.compile(r"Post\s*:\s*(.+?)\s*(?:Top\s*Comment|Comment)\s*:\s*(.+)", re.IGNORECASE | re.DOTALL),
        ]
        def parse(raw: str) -> Tuple[Optional[str], Optional[str]]:
            raw = raw.replace("<|endoftext|>", "\n")
            for pat in PATS:
                m = pat.search(raw)
                if m:
                    return m.group(1).strip(), m.group(2).strip()
            lines = [ln.strip() for ln in raw.splitlines() if ln.strip()]
            if len(lines) >= 2: return lines[0], lines[1]
            if len(lines) == 1: return "Reply with sarcasm:", lines[0]
            return None, None

        kept = scanned = 0
        for ex in limited(ds, scan_cap):
            scanned += 1
            you, bot = parse(str(ex.get("text") or ""))
            you = keep_or_clip(you); bot = keep_or_clip(bot)
            if you and bot:
                pairs.append(turn(you, bot, "[STYLE=Snark]"))
                kept += 1; keepbar.update(1); overall.update(1)
                if kept >= keep_cap: break
            scanbar.update(1)
            if scanned % 2000 == 0:
                keepbar.set_postfix(rate=f"{kept/max(1,scanned):.2%}")
        scanbar.close(); keepbar.close()
        print(f"[ok] reddit-sarcasm kept={kept} (scanned {scanned})")
    except Exception as e:
        print(f"[skip] reddit-sarcasm: {e}")

def collect_figlang(pairs: List[str], overall: tqdm):
    for dsid in ["tasksource/figlang2020-sarcasm", "tasksource/figlang2020_sarcasm"]:
        try:
            ds = load_dataset(dsid, split="train", streaming=True)
            kept = 0
            pbar = tqdm(total=CAPS["figlang"], desc="[figlang]", unit="pair", leave=False, ncols=100)
            for ex in limited(ds, SCAN["figlang"]):
                ctx = ex.get("context")
                if isinstance(ctx, list) and ctx:
                    context_str = " ".join(str(c) for c in ctx[-2:])
                else:
                    context_str = str(ex.get("context") or ex.get("prompt") or "")
                reply = str(ex.get("response") or ex.get("answer") or ex.get("text") or "")
                context_str = keep_or_clip(clean_text(context_str))
                reply      = keep_or_clip(clean_text(reply))
                if reply:
                    if context_str:
                        pairs.append(turn(context_str, reply, "[STYLE=Snark]"))
                    else:
                        pairs.append(turn("Reply with sarcasm:", reply, "[STYLE=Snark]"))
                    kept += 1; pbar.update(1); overall.update(1)
                    if kept >= CAPS["figlang"]: break
            pbar.close()
            print(f"[ok] figlang {dsid} kept={kept}")
            return
        except Exception as e:
            print(f"[try next] figlang {dsid}: {e}")
    print("[skip] figlang")

def collect_showerthoughts(pairs: List[str], overall: tqdm):
    """Use REEDIT_submissions split 'Showerthoughts' directly (no 'train')."""
    try:
        ds = load_dataset("HuggingFaceGECLM/REDDIT_submissions", split="Showerthoughts", streaming=True)
        keep_cap, scan_cap = CAPS["shower"], SCAN["shower"]
        scanbar = tqdm(total=scan_cap, desc="[shower scan]", unit="row", leave=False, ncols=100)
        keepbar = tqdm(total=keep_cap, desc="[shower kept]", unit="pair", leave=False, ncols=100)

        kept = scanned = 0
        for ex in limited(ds, scan_cap):
            scanned += 1
            title = get_first_nonempty(ex, ["title", "selftext", "text"]) or ""
            text = keep_or_clip(clean_text(title))
            if text:
                pairs.append(turn("Give me a shower thought.", text))
                kept += 1; keepbar.update(1); overall.update(1)
                if kept >= keep_cap: break
            scanbar.update(1)
        scanbar.close(); keepbar.close()
        print(f"[ok] showerthoughts kept={kept} (scanned {scanned})")
    except Exception as e:
        print(f"[skip] showerthoughts: {e}")

def collect_personas(pairs: List[str], overall: tqdm):
    """Non-streaming is more reliable for this dataset."""
    try:
        ds = load_dataset("NapthaAI/twitter_personas")["train"]
        keep_cap = CAPS["personas"]
        pbar = tqdm(total=keep_cap, desc="[personas]", unit="pair", leave=False, ncols=100)
        kept = 0
        for ex in ds:
            desc = get_first_nonempty(ex, ["description","persona","bio","text","content","full_text"])
            if not isinstance(desc, str) and isinstance(ex.get("content"), dict):
                desc = ex["content"].get("text")
            desc = keep_or_clip(clean_text(str(desc or "")))
            if not desc: continue
            pairs.append(turn("Adopt this persona in one sentence.", desc, "[FORM=TWEET]"))
            kept += 1; pbar.update(1); overall.update(1)
            if kept >= keep_cap: break
        pbar.close()
        print(f"[ok] personas kept={kept}")
    except Exception as e:
        print(f"[skip] personas: {e}")

def collect_tweeteval(pairs: List[str], overall: tqdm):
    """Handle super_tweeteval (text_1/text_2, etc.) and fallback tweet_eval."""
    def extract_pair(ex):
        t = ex.get("text")
        if isinstance(t, str) and t.strip():
            return "React with a sharp one-liner.", t
        for a,b in [("text_1","text_2"), ("sentence1","sentence2"),
                    ("premise","hypothesis"), ("question","answer"),
                    ("context","response"), ("tweet1","tweet2")]:
            t1, t2 = ex.get(a), ex.get(b)
            if isinstance(t1, str) and t1.strip() and isinstance(t2, str) and t2.strip():
                return t1, t2
        return None

    def run_on(dsname, pick, is_super):
        keep_cap, scan_cap = CAPS["tweeteval"], SCAN["tweeteval"]
        pbar = tqdm(total=keep_cap, desc=f"[tweeteval:{pick}]", unit="pair", leave=False, ncols=100)
        kept = 0
        ds = load_dataset(dsname, pick, split="train", streaming=True)
        for ex in limited(ds, scan_cap):
            pair = extract_pair(ex) if is_super else ("React with a sharp one-liner.", ex.get("text")) if ex.get("text") else None
            if not pair: continue
            you, bot = pair
            you = keep_or_clip(clean_text(str(you or ""))); bot = keep_or_clip(clean_text(str(bot or "")))
            if not (you and bot): continue
            tag = "[STYLE=Snark]" if you and you != "React with a sharp one-liner." else ""
            pairs.append(turn(you, bot, tag))
            kept += 1; pbar.update(1); overall.update(1)
            if kept >= keep_cap: break
        pbar.close()
        return kept

    kept_total = 0
    try:
        cfgs = get_dataset_config_names("cardiffnlp/super_tweeteval")
    except Exception:
        cfgs = []
    prio = ["irony","sarcasm","humor","sentiment","emoji","emotion","stance","offensive","hate"]
    ordered = [c for c in prio if c in cfgs] + [c for c in cfgs if c not in prio]
    for pick in ordered:
        kept_total += run_on("cardiffnlp/super_tweeteval", pick, True)
        if kept_total >= CAPS["tweeteval"]:
            print(f"[ok] tweeteval(super) kept={kept_total}"); return
    if kept_total == 0:
        try:
            base_cfgs = get_dataset_config_names("cardiffnlp/tweet_eval")
        except Exception:
            base_cfgs = []
        ordered_b = [c for c in prio if c in base_cfgs] + [c for c in base_cfgs if c not in prio]
        for pick in ordered_b:
            kept_total += run_on("cardiffnlp/tweet_eval", pick, False)
            if kept_total >= CAPS["tweeteval"]:
                print(f"[ok] tweeteval(base) kept={kept_total}"); return
    print(f"[ok] tweeteval kept={kept_total}")

def collect_fourchan(pairs: List[str], overall: tqdm):
    try:
        ds = load_dataset("sbussiso/4chan-dataset", split="train", streaming=True)
        keep_cap = min(CAPS["fourchan"], 195)
        pbar = tqdm(total=keep_cap, desc="[4chan]", unit="pair", leave=False, ncols=100)
        kept = 0
        for ex in limited(ds, SCAN["fourchan"]):
            prompt = keep_or_clip(clean_text(str(ex.get("prompt") or "")))
            resp   = keep_or_clip(clean_text(str(ex.get("response") or "")))
            if prompt and resp:
                pairs.append(turn(prompt, resp))
                kept += 1; pbar.update(1); overall.update(1)
            else:
                txt = keep_or_clip(clean_text(str(ex.get("text") or ex.get("body") or ex.get("content") or "")))
                if txt:
                    pairs.append(turn("Drop a spicy one-liner.", txt))
                    kept += 1; pbar.update(1); overall.update(1)
            if kept >= keep_cap: break
        pbar.close()
        print(f"[ok] 4chan kept={kept}")
    except Exception as e:
        print(f"[skip] 4chan: {e}")

def collect_elon_trump(pairs: List[str], overall: tqdm):
    try:
        ds = load_dataset("MasaFoundation/Twitter_X_Elon_vs_Trump", split="train", streaming=True, revision="refs/convert/parquet")
        keep_cap, scan_cap = CAPS["elonvtrump"], SCAN["elonvtrump"]
        scanbar = tqdm(total=scan_cap, desc="[elon_vs_trump scan]", unit="row", leave=False, ncols=100)
        keepbar = tqdm(total=keep_cap, desc="[elon_vs_trump kept]", unit="pair", leave=False, ncols=100)
        kept = scanned = 0
        for ex in limited(ds, scan_cap):
            scanned += 1
            txt = get_first_nonempty(ex, ["content","text","tweet","full_text"]) or ""
            txt = keep_or_clip(clean_text(txt))
            if txt:
                pairs.append(turn("[FORM=TWEET] One sentence hot take:", txt, "[FORM=TWEET]"))
                kept += 1; keepbar.update(1); overall.update(1)
                if kept >= keep_cap: break
            scanbar.update(1)
        scanbar.close(); keepbar.close()
        print(f"[ok] Elon_vs_Trump kept={kept} (scanned {scanned})")
    except Exception as e:
        print(f"[skip] Elon_vs_Trump: {e}")

# --------------------------
# Build corpus
# --------------------------
def build_corpus() -> Path:
    pairs: List[str] = []
    total_target = sum(CAPS.values())
    print("[1/6] Collecting & reformatting datasets (streaming, capped)…")
    overall = tqdm(total=total_target, desc="[all] collecting", unit="pair", ncols=100)

    collectors = [
        collect_shakespeare,
        collect_reddit_jokes,
        collect_dadjokes,
        collect_reddit_sarcasm,
        collect_figlang,
        collect_showerthoughts,
        collect_personas,
        collect_tweeteval,
        collect_fourchan,
        collect_elon_trump,
    ]
    for fn in collectors:
        try:
            fn(pairs, overall)
        except Exception as e:
            print(f"[collector error] {fn.__name__}: {e}")

    overall.close()
    print("[2/6] Deduplicating & clipping…")
    seen = set(); deduped = []
    for block in pairs:
        try:
            bot_line = [ln for ln in block.splitlines() if ln.startswith("Bot:")][0]
            key = bot_line[4:].strip().lower()
        except Exception:
            key = block.strip().lower()
        if key in seen: continue
        seen.add(key); deduped.append(block)

    random.shuffle(deduped)
    if len(deduped) > MAX_TOTAL_PAIRS:
        deduped = deduped[:MAX_TOTAL_PAIRS]

    out_path = SAVE_DIR / "corpus.txt"
    out_path.write_text("\n".join(deduped), encoding="utf-8")
    print(f"     wrote {len(deduped)} pairs → {out_path}")
    return out_path

# --------------------------
# SentencePiece tokenizer
# --------------------------
def train_spm(corpus_path: Path) -> spm.SentencePieceProcessor:
    print("[3/6] Training SentencePiece tokenizer…")
    spm.SentencePieceTrainer.Train(
        input=str(corpus_path),
        model_prefix=str(TOKENIZER_PREFIX),
        vocab_size=VOCAB_SIZE,
        model_type="unigram",
        character_coverage=1.0,
        user_defined_symbols=USER_SYMBOLS,
        bos_id=1, eos_id=2, unk_id=0, pad_id=-1
    )
    sp = spm.SentencePieceProcessor()
    sp.load(f"{TOKENIZER_PREFIX}.model")
    print(f"     tokenizer saved at {TOKENIZER_PREFIX}.model")
    return sp

# --------------------------
# Encode to token IDs
# --------------------------
def encode_corpus_to_ids(sp: spm.SentencePieceProcessor, corpus_path: Path):
    print("[4/6] Encoding corpus to token IDs…")
    text = corpus_path.read_text(encoding="utf-8")
    blocks = [b for b in text.split("\n\n") if b.strip()]
    ids = []
    eos = sp.eos_id()
    for b in blocks:
        for line in b.splitlines():
            if not line.strip(): continue
            ids.extend(sp.encode(line, out_type=int)); ids.append(eos)
        ids.append(eos)
    ids = np.array(ids, dtype=np.int32)
    n = len(ids); cut = int(n * 0.97)
    train_ids = torch.tensor(ids[:cut], dtype=torch.long, device=DEVICE)
    val_ids   = torch.tensor(ids[cut:], dtype=torch.long, device=DEVICE)
    print(f"     tokens: train={train_ids.numel():,}, val={val_ids.numel():,}, vocab={sp.vocab_size()}")
    return train_ids, val_ids, sp.vocab_size()

# --------------------------
# Tiny GPT model
# --------------------------
class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, dropout=0.0, block_size=256):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.qkv = nn.Linear(n_embd, 3*n_embd, bias=False)
        self.proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_drop = nn.Dropout(dropout)
        self.resid_drop = nn.Dropout(dropout)
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)).view(1,1,block_size,block_size))
    def forward(self, x):
        B,T,C = x.shape
        qkv = self.qkv(x); q,k,v = qkv.chunk(3, dim=-1)
        q = q.view(B,T,self.n_head,self.head_dim).transpose(1,2)
        k = k.view(B,T,self.n_head,self.head_dim).transpose(1,2)
        v = v.view(B,T,self.n_head,self.head_dim).transpose(1,2)
        att = (q @ k.transpose(-2,-1)) / math.sqrt(self.head_dim)
        att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))
        att = torch.softmax(att, dim=-1)
        att = self.attn_drop(att)
        y = att @ v
        y = y.transpose(1,2).contiguous().view(B,T,C)
        y = self.resid_drop(self.proj(y))
        return y

class Block(nn.Module):
    def __init__(self, n_embd, n_head, dropout=0.0, block_size=256):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, dropout, block_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = nn.Sequential(
            nn.Linear(n_embd, 4*n_embd),
            nn.GELU(),
            nn.Linear(4*n_embd, n_embd),
            nn.Dropout(dropout),
        )
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

class TinyGPT(nn.Module):
    def __init__(self, vocab_size, n_layer, n_head, n_embd, block_size, dropout=0.0):
        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)
        self.blocks = nn.ModuleList([Block(n_embd, n_head, dropout, block_size) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size, bias=False)
        self.apply(self._init)
    def _init(self, m):
        if isinstance(m, (nn.Linear, nn.Embedding)): nn.init.normal_(m.weight, mean=0.0, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None: nn.init.zeros_(m.bias)
        if isinstance(m, nn.LayerNorm): nn.init.ones_(m.weight); nn.init.zeros_(m.bias)
    def forward(self, idx, targets=None):
        B,T = idx.shape; assert T <= self.block_size
        pos = torch.arange(0, T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)[None,:,:]
        for blk in self.blocks: x = blk(x)
        x = self.ln_f(x)
        logits = self.head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss
    @torch.no_grad()
    def generate(self, idx, max_new_tokens=200, temperature=0.8, top_k=50, top_p=0.95, repetition_penalty=1.0):
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, _ = self.forward(idx_cond)
            logits = logits[:, -1, :]
            if repetition_penalty != 1.0:
                uniq, _ = torch.unique(idx_cond[0], return_counts=True)
                logits[:, uniq] /= repetition_penalty
            logits = logits / max(1e-8, temperature)
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                cutoff = v[:, -1].unsqueeze(-1)
                logits = torch.where(logits < cutoff, torch.full_like(logits, -1e9), logits)
            if top_p is not None:
                sorted_logits, sorted_idx = torch.sort(logits, descending=True)
                probs = torch.softmax(sorted_logits, dim=-1)
                cdf = torch.cumsum(probs, dim=-1)
                mask = cdf > top_p; mask[:, 0] = False
                sorted_logits[mask] = -1e9
                logits = torch.zeros_like(logits).scatter(1, sorted_idx, sorted_logits)
            probs = torch.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_id], dim=1)
        return idx

# --------------------------
# Data loader for token IDs
# --------------------------
def get_batch(split_ids: torch.Tensor, B: int, T: int):
    ix = torch.randint(0, split_ids.numel() - T - 1, (B,), device=split_ids.device)
    x = torch.stack([split_ids[i:i+T] for i in ix])
    y = torch.stack([split_ids[i+1:i+T+1] for i in ix])
    return x, y

# --------------------------
# Train loop
# --------------------------
def train_model(vocab_size, train_ids, val_ids):
    print("[5/6] Training tiny GPT on", DEVICE.type.upper(), "…")
    model = TinyGPT(vocab_size, n_layer, n_head, n_embd, block_size, dropout).to(DEVICE)
    params_m = sum(p.numel() for p in model.parameters())/1e6
    print(f"     params: {params_m:.2f}M")
    optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr, betas=(0.9, 0.95), weight_decay=0.0)

    use_amp = DEVICE.type == "cuda"
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
    autocast = (lambda: torch.amp.autocast("cuda", dtype=torch.float16)) if use_amp else nullcontext

    start = time.time()
    best_val = float("inf")

    def get_lr(step):
        warmup = max(1, int(train_steps * warmup_ratio))
        if step < warmup: return base_lr * (step+1)/warmup
        progress = (step - warmup) / max(1, train_steps - warmup)
        return min_lr + 0.5*(base_lr - min_lr)*(1 + math.cos(math.pi * min(1.0, progress)))

    @torch.no_grad()
    def eval_loss(iters=80):
        model.eval(); losses=[]
        for _ in range(iters):
            xb, yb = get_batch(val_ids, min(batch_size, 32), block_size)
            with autocast():
                _, loss = model(xb, yb)
            losses.append(loss.item())
        model.train()
        return float(sum(losses)/len(losses))

    model.train(); step = 0
    pbar = tqdm(total=train_steps, ncols=100, desc="[train]")
    while step < train_steps and (time.time()-start) < MAX_SECONDS:
        lr = get_lr(step)
        for pg in optimizer.param_groups: pg["lr"] = lr
        optimizer.zero_grad(set_to_none=True)

        total_loss = 0.0
        for _ in range(accum_steps):
            xb, yb = get_batch(train_ids, batch_size, block_size)
            with autocast():
                _, loss = model(xb, yb)
            if use_amp:
                scaler.scale(loss / accum_steps).backward()
            else:
                (loss / accum_steps).backward()
            total_loss += loss.item()

        if use_amp:
            scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        if use_amp:
            scaler.step(optimizer); scaler.update()
        else:
            optimizer.step()

        step += 1; pbar.update(1)
        if step % log_interval == 0 or step == 1:
            pbar.set_postfix(train=f"{total_loss/accum_steps:.3f}", lr=f"{lr:.2e}")

        if step % eval_every == 0:
            vl = eval_loss()
            best_val = min(best_val, vl)
            print(f"     eval loss {vl:.3f} | best {best_val:.3f}")

    pbar.close()
    elapsed = time.time() - start
    print(f"     done in {elapsed:.1f}s | best val {best_val:.3f}")

    # Save model + config
    ckpt_path = SAVE_DIR / "tinygpt.pt"
    torch.save(model.state_dict(), ckpt_path)
    (SAVE_DIR / "model_config.json").write_text(json.dumps({
        "vocab_size": int(vocab_size),
        "n_layer": n_layer, "n_head": n_head, "n_embd": n_embd,
        "block_size": block_size, "dropout": dropout
    }, indent=2))
    print(f"[saved] weights → {ckpt_path}")
    return model

# --------------------------
# Sampling helper
# --------------------------
def sample_chat(sp: spm.SentencePieceProcessor, model: TinyGPT, prompt: str, max_new_tokens=200):
    prefix = f"You: {prompt}\nBot:"
    ids = sp.encode(prefix, out_type=int)
    x = torch.tensor([ids], dtype=torch.long, device=DEVICE)
    with torch.no_grad():
        y = model.generate(x, max_new_tokens=max_new_tokens, temperature=TEMP, top_k=TOP_K, top_p=TOP_P, repetition_penalty=REP_PEN)
    return sp.decode(y[0].tolist())

# --------------------------
# Main
# --------------------------
def main():
    # Build or reuse corpus/tokenizer/ids
    corpus_path = SAVE_DIR / "corpus.txt"
    spm_model   = SAVE_DIR / "spm_chat.model"
    if not corpus_path.exists():
        corpus_path = build_corpus()
    else:
        print(f"[cache] using {corpus_path}")

    sp = spm.SentencePieceProcessor()
    if not spm_model.exists():
        sp = train_spm(corpus_path)
    else:
        sp.load(str(spm_model))
        print(f"[cache] using {spm_model}")

    enc_train = SAVE_DIR / "train_ids.pt"
    enc_val   = SAVE_DIR / "val_ids.pt"
    vocab_txt = SAVE_DIR / "vocab_size.txt"

    if enc_train.exists() and enc_val.exists() and vocab_txt.exists():
        train_ids = torch.load(enc_train, map_location=DEVICE)
        val_ids   = torch.load(enc_val,   map_location=DEVICE)
        vocab_size = int(vocab_txt.read_text())
        print(f"[cache] loaded ids: train={train_ids.numel():,}, val={val_ids.numel():,}, vocab={vocab_size}")
    else:
        train_ids, val_ids, vocab_size = encode_corpus_to_ids(sp, corpus_path)
        torch.save(train_ids, enc_train); torch.save(val_ids, enc_val)
        vocab_txt.write_text(str(vocab_size))
        print("[cache] saved encoded ids")

    model = train_model(vocab_size, train_ids, val_ids)

    print("\n[6/6] Samples:\n")
    prompts = [
        "Give me a spicy take on AI.",
        "Roast my messy desk.",
        "Explain recursion like you're annoyed.",
        "Write a satirical headline about coffee.",
        "Give me a shower thought about umbrellas.",
        "Tell me a one-liner about deadlines.",
        "Stay in Shakespeare mode and flatter me.",
        "Reply sarcastically to: I love meetings.",
    ]
    out_path = SAVE_DIR / "samples.txt"
    with out_path.open("w", encoding="utf-8") as f:
        for p in prompts:
            txt = sample_chat(sp, model, p, max_new_tokens=200)
            print("----\n" + txt)
            f.write("----\n" + txt + "\n")
    print(f"\n[saved] samples → {out_path}")

if __name__ == "__main__":
    main()