{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Build your own poet — a small language model from one text file\n\nThis notebook is the *build-your-own* companion to **How a language model works** (`how-llms-work.html`). It builds the exact same kind of machine that writes Thirukkural on that page, from scratch, on **any text you give it** — Bharathiyar, Kabir, Shakespeare, film lyrics, your own writing.\n\n**You do not need to know how to program.** Run the cells top to bottom (`Runtime → Run all`). The only decision is which text file to use.\n\n| Step | What happens | Chapter on the page |\n|---|---|---|\n| 1 | Get a book: upload a `.txt`, or use the built-in Thirukkural / Shakespeare sample | 2 · the only book |\n| 2 | Every letter gets a number | 3 · letters → numbers |\n| 3 | The education: guess the next letter, get corrected, nudge — a few thousand times | 4–8 |\n| 4 | Teach a habit: `# label → example` | 9 · teaching a habit |\n| 5 | A judge: pairs of attempts, keep what the judge prefers | 10 · a judge, not a teacher |\n| 6 | Write new poems, and save the weights in the page's format | 11–12 |\n\nTakes ~10 minutes on the free Colab CPU, ~2 minutes with a GPU (`Runtime → Change runtime type → T4 GPU`).\n"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "#@title Step 1 · The book  { display-mode: \"form\" }\n#@markdown Choose a sample, or set `use_upload` to True and pick a `.txt` file from your computer.\n#@markdown The file should have one example (poem, verse, recipe…) per paragraph, with a blank line between examples.\nsample = \"thirukkural\"  #@param [\"thirukkural\", \"shakespeare-sonnets\"]\nuse_upload = False  #@param {type:\"boolean\"}\n\nimport re, json, urllib.request, unicodedata\nif use_upload:\n    from google.colab import files\n    up = files.upload()\n    name = list(up.keys())[0]\n    text = up[name].decode('utf-8')\nelif sample == \"thirukkural\":\n    raw = json.load(urllib.request.urlopen('https://raw.githubusercontent.com/tk120404/thirukkural/master/thirukkural.json'))['kural']\n    fix = lambda s: s.replace('஦வ்ருஉம்', 'வெருஉம்').replace('அள஧ க்கும்', 'அளிக்கும்').replace('஧', 'ி')\n    text = ''.join(fix(k['Line1']).strip() + '\\n' + fix(k['Line2']).strip() + '\\n\\n' for k in raw)\n    # labels for step 4: the chapter of each couplet\n    det = json.load(urllib.request.urlopen('https://raw.githubusercontent.com/tk120404/thirukkural/master/detail.json'))\n    LABELS = {}\n    for paal in det[0]['section']['detail']:\n        for iyal in paal['chapterGroup']['detail']:\n            for ch in iyal['chapters']['detail']:\n                for i in range(ch['start'], ch['end'] + 1): LABELS[i - 1] = ch['name']\nelse:\n    raw = urllib.request.urlopen('https://www.gutenberg.org/cache/epub/1041/pg1041.txt').read().decode('utf-8').replace('\\r', '')\n    body = raw.split(\"*** START OF THE PROJECT GUTENBERG EBOOK SHAKESPEARE'S SONNETS ***\")[1].split('*** END OF')[0]\n    body = body.replace('’', \"'\").replace('‘', \"'\").replace('“', '\"').replace('”', '\"')\n    son, cur = [], []\n    for ln in body.split('\\n'):\n        s = ln.strip()\n        if re.fullmatch(r'[IVXLC]+', s):\n            if cur: son.append(cur); cur = []\n        elif s and not s.startswith('THE SONNETS') and not s.startswith('by William'): cur.append(s)\n    if cur: son.append(cur)\n    text = ''.join('\\n'.join(s) + '\\n\\n' for s in son)\n\ntext = unicodedata.normalize('NFC', text)\nexamples = [e.strip() for e in text.split('\\n\\n') if e.strip()]\nif 'LABELS' not in dir(): LABELS = {i: e.split()[0] for i, e in enumerate(examples)}   # default label: the first word\nprint(f'{len(examples):,} examples · {len(text):,} letters · {len(set(text))} different letters')\nprint('first example:\\n' + examples[0])\n"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "#@title Step 2 · The machine (the same 120 lines that built the page's models)  { display-mode: \"form\" }\n#@markdown Run this cell once. It defines the model, the training loop, and the export to the page's format.\n\"\"\"Minimal character-level GPT (nanoGPT-style). Also used verbatim as the 'build your own' recipe on the page.\"\"\"\nimport math, json, time, random\nimport torch, torch.nn as nn, torch.nn.functional as F\n\nimport os; torch.set_num_threads(os.cpu_count() or 2)\n\nclass Config:\n    def __init__(self, vocab_size, block_size=128, n_layer=3, n_head=4, n_embd=64, dropout=0.15):\n        self.vocab_size, self.block_size, self.n_layer, self.n_head, self.n_embd, self.dropout = \\\n            vocab_size, block_size, n_layer, n_head, n_embd, dropout\n\nclass Attention(nn.Module):\n    def __init__(self, c):\n        super().__init__()\n        self.n_head, self.n_embd = c.n_head, c.n_embd\n        self.qkv = nn.Linear(c.n_embd, 3 * c.n_embd)\n        self.proj = nn.Linear(c.n_embd, c.n_embd)\n        self.drop = nn.Dropout(c.dropout)\n        self.register_buffer('mask', torch.tril(torch.ones(c.block_size, c.block_size)).view(1, 1, c.block_size, c.block_size))\n    def forward(self, x):\n        B, T, C = x.shape\n        q, k, v = self.qkv(x).split(C, dim=2)\n        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n        att = (q @ k.transpose(-2, -1)) / math.sqrt(k.size(-1))\n        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))\n        att = F.softmax(att, dim=-1)\n        y = (self.drop(att) @ v).transpose(1, 2).contiguous().view(B, T, C)\n        return self.drop(self.proj(y))\n\nclass Block(nn.Module):\n    def __init__(self, c):\n        super().__init__()\n        self.ln1 = nn.LayerNorm(c.n_embd)\n        self.attn = Attention(c)\n        self.ln2 = nn.LayerNorm(c.n_embd)\n        self.mlp = nn.Sequential(nn.Linear(c.n_embd, 4 * c.n_embd), nn.GELU(), nn.Linear(4 * c.n_embd, c.n_embd), nn.Dropout(c.dropout))\n    def forward(self, x):\n        x = x + self.attn(self.ln1(x))\n        x = x + self.mlp(self.ln2(x))\n        return x\n\nclass GPT(nn.Module):\n    def __init__(self, c):\n        super().__init__()\n        self.c = c\n        self.tok_emb = nn.Embedding(c.vocab_size, c.n_embd)\n        self.pos_emb = nn.Embedding(c.block_size, c.n_embd)\n        self.drop = nn.Dropout(c.dropout)\n        self.blocks = nn.ModuleList([Block(c) for _ in range(c.n_layer)])\n        self.ln_f = nn.LayerNorm(c.n_embd)\n        self.head = nn.Linear(c.n_embd, c.vocab_size, bias=False)\n        self.apply(self._init)\n    def _init(self, m):\n        if isinstance(m, (nn.Linear, nn.Embedding)):\n            nn.init.normal_(m.weight, mean=0.0, std=0.02)\n            if isinstance(m, nn.Linear) and m.bias is not None: nn.init.zeros_(m.bias)\n    def forward(self, idx, targets=None):\n        B, T = idx.shape\n        x = self.drop(self.tok_emb(idx) + self.pos_emb(torch.arange(T)))\n        for b in self.blocks: x = b(x)\n        logits = self.head(self.ln_f(x))\n        loss = None if targets is None else F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))\n        return logits, loss\n    @torch.no_grad()\n    def generate(self, idx, max_new, temperature=1.0, top_k=None, stop=None):\n        self.eval()\n        for _ in range(max_new):\n            logits, _ = self(idx[:, -self.c.block_size:])\n            logits = logits[:, -1, :] / max(temperature, 1e-6)\n            if top_k:\n                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))\n                logits[logits < v[:, [-1]]] = -float('inf')\n            nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)\n            idx = torch.cat([idx, nxt], dim=1)\n            if stop is not None and stop(idx[0].tolist()): break\n        return idx\n\ndef n_params(m): return sum(p.numel() for p in m.parameters())\n\n# ---------------- data ----------------\nclass Data:\n    def __init__(self, text, vocab, val_frac=0.1, seed=1, unit_sep='\\n\\n'):\n        self.vocab = vocab\n        self.stoi = {ch: i for i, ch in enumerate(vocab)}\n        self.itos = {i: ch for ch, i in self.stoi.items()}\n        units = [u for u in text.split(unit_sep) if u.strip()]\n        rng = random.Random(seed); rng.shuffle(units)\n        nv = max(1, int(len(units) * val_frac))\n        self.val_units, self.train_units = units[:nv], units[nv:]\n        self.train = torch.tensor(self.encode(unit_sep.join(self.train_units) + unit_sep), dtype=torch.long)\n        self.val = torch.tensor(self.encode(unit_sep.join(self.val_units) + unit_sep), dtype=torch.long)\n    def encode(self, s): return [self.stoi[c] for c in s]\n    def decode(self, l): return ''.join(self.itos[i] for i in l)\n    def batch(self, split, B, T):\n        d = self.train if split == 'train' else self.val\n        ix = torch.randint(len(d) - T, (B,))\n        return torch.stack([d[i:i + T] for i in ix]), torch.stack([d[i + 1:i + 1 + T] for i in ix])\n\n@torch.no_grad()\ndef eval_loss(model, data, B=32, iters=20):\n    model.eval(); out = {}\n    for split in ('train', 'val'):\n        ls = []\n        for _ in range(iters):\n            x, y = data.batch(split, B, model.c.block_size)\n            _, l = model(x, y); ls.append(l.item())\n        out[split] = sum(ls) / len(ls)\n    model.train(); return out\n\ndef sample_text(model, data, prompt, n=200, temperature=0.8, seed=0, top_k=None, stop=None):\n    torch.manual_seed(seed)\n    idx = torch.tensor([data.encode(prompt)], dtype=torch.long)\n    out = model.generate(idx, n, temperature=temperature, top_k=top_k, stop=stop)\n    return data.decode(out[0].tolist())\n\ndef train(model, data, steps, lr=1e-3, B=32, log_every=50, milestones=(), sample_prompt='\\n', sample_seed=0, sample_n=200, min_lr_frac=0.1, tag=''):\n    opt = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1)\n    log, samples = [], {}\n    t0 = time.time()\n    for step in range(steps + 1):\n        if step in milestones:\n            samples[step] = sample_text(model, data, sample_prompt, n=sample_n, seed=sample_seed)\n            model.train()\n        if step % log_every == 0 or step == steps:\n            e = eval_loss(model, data)\n            log.append(dict(step=step, train=round(e['train'], 4), val=round(e['val'], 4)))\n            print(f'{tag} step {step:5d} train {e[\"train\"]:.3f} val {e[\"val\"]:.3f}  {time.time()-t0:.0f}s', flush=True)\n        if step == steps: break\n        cur_lr = min_lr_frac * lr + 0.5 * (1 - min_lr_frac) * lr * (1 + math.cos(math.pi * step / steps))\n        for g in opt.param_groups: g['lr'] = cur_lr\n        x, y = data.batch('train', B, model.c.block_size)\n        _, loss = model(x, y)\n        opt.zero_grad(set_to_none=True); loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n        opt.step()\n    return log, samples\n\n# ---------------- export (int8 per-tensor for the browser) ----------------\nimport base64, numpy as np\ndef export(model, vocab, path, extra=None):\n    c = model.c\n    tensors = {}\n    for name, p in model.state_dict().items():\n        if name.endswith('.mask'): continue\n        w = p.detach().float().numpy()\n        if w.ndim == 2 and w.size > 4096:  # big matrices -> int8\n            scale = float(np.abs(w).max() / 127.0) or 1.0\n            q = np.clip(np.round(w / scale), -127, 127).astype(np.int8)\n            tensors[name] = dict(shape=list(w.shape), dtype='i8', scale=scale, data=base64.b64encode(q.tobytes()).decode())\n        else:  # small vectors -> float16\n            tensors[name] = dict(shape=list(w.shape), dtype='f16', data=base64.b64encode(w.astype(np.float16).tobytes()).decode())\n    out = dict(config=dict(vocab_size=c.vocab_size, block_size=c.block_size, n_layer=c.n_layer, n_head=c.n_head, n_embd=c.n_embd),\n               vocab=vocab, n_params=n_params(model), tensors=tensors)\n    if extra: out.update(extra)\n    json.dump(out, open(path, 'w'), ensure_ascii=False, separators=(',', ':'))\n    return out\n\ndef load_exported(path):\n    \"\"\"Rebuild a torch model from the exported json (to verify the quantisation round-trip).\"\"\"\n    j = json.load(open(path))\n    cf = j['config']; c = Config(cf['vocab_size'], cf['block_size'], cf['n_layer'], cf['n_head'], cf['n_embd'], 0.0)\n    m = GPT(c); sd = m.state_dict()\n    for name, t in j['tensors'].items():\n        raw = base64.b64decode(t['data'])\n        if t['dtype'] == 'i8': w = np.frombuffer(raw, dtype=np.int8).astype(np.float32) * t['scale']\n        else: w = np.frombuffer(raw, dtype=np.float16).astype(np.float32)\n        sd[name] = torch.tensor(w.reshape(t['shape']))\n    m.load_state_dict(sd); m.eval(); return m, j['vocab']\n\n\nvocab = sorted(set(text) | {'#'})\nprint('vocabulary (every letter gets a locker number):')\nprint({i: (c if c not in ('\\n', ' ') else repr(c)) for i, c in enumerate(vocab)})\n"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "#@title Step 3 · The education  { display-mode: \"form\" }\n#@markdown Watch the *surprise* (loss) fall, and read what the machine writes at each milestone. Stop early if `val` starts rising.\nsteps = 3000  #@param {type:\"slider\", min:500, max:6000, step:500}\nnumbers_per_letter = 64  #@param [64, 96, 128] {type:\"raw\"}\nrounds = 3  #@param [2, 3, 4, 6] {type:\"raw\"}\n\ndata = Data(text, vocab)\ntorch.manual_seed(42)\nmodel = GPT(Config(len(vocab), block_size=128, n_layer=rounds, n_head=4, n_embd=numbers_per_letter, dropout=0.15))\nprint(f'{n_params(model):,} dials · {len(data.train):,} letters to learn from · {len(data.val):,} held out')\nlog, samples = train(model, data, steps, lr=1e-3, milestones=(0, 100, 500, 1000, 2000, steps), sample_prompt='\\n', sample_n=250, tag='education')\nfor s, txt in samples.items():\n    print(f'\\n===== what it wrote at step {s} =====\\n{txt.strip()[:400]}')\nexport(model, vocab, 'base.json', extra=dict(name='base', loss=log, samples={str(k): v for k, v in samples.items()}))\nprint('\\nsaved base.json')\n"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "#@title Step 4 · Teach a habit (fine-tuning): `# label` → example  { display-mode: \"form\" }\n#@markdown Each example is prefixed with a label line. For Thirukkural the label is the chapter; for your own file it is the first word of each example (change `LABELS` above to use anything you like — a mood, an author, a topic).\nsft_text = ''.join('# ' + LABELS[i] + '\\n' + e + '\\n\\n' for i, e in enumerate(examples))\nsft_data = Data(sft_text, vocab)\nlog2, samples2 = train(model, sft_data, 1500, lr=5e-4, milestones=(0, 1500), sample_prompt='# ' + LABELS[0] + '\\n', sample_n=150, tag='habit')\nprint('\\nbefore fine-tuning, given the label:\\n' + samples2[0][:200]); print('\\nafter:\\n' + samples2[1500][:200])\nexport(model, vocab, 'sft.json', extra=dict(name='sft', loss=log2))\nprint('\\nsaved sft.json')\n"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "#@title Step 5 · A judge (preference optimisation)  { display-mode: \"form\" }\n#@markdown Write your own judge below: a function that returns True for a good example. The default checks the Thirukkural shape (4 words, then 3, ending with a full stop) — change it for your text (e.g. `len(lines) == 14` for sonnets, or \"contains the label's word\").\nimport copy, random, torch.nn.functional as F\n\ndef judge(out):\n    body = out.split('\\n\\n')[0]; lines = body.split('\\n')\n    return len(lines) == 2 and len(lines[0].split()) == 4 and len(lines[1].split()) == 3 and body.endswith('.')\n\ndef stop(ids): return len(ids) > 1 and ids[-1] == ids[-2] == sft_data.stoi['\\n']\ndef logp(m, prompt, resp):\n    ids = sft_data.encode(prompt + resp); x = torch.tensor([ids[:-1]]); y = torch.tensor([ids[1:]])\n    lp = F.log_softmax(m(x)[0][0], -1).gather(1, y[0][:, None]).squeeze(1)\n    return lp[len(sft_data.encode(prompt)) - 1:].sum()\ndef pass_rate(m, n=150):\n    return sum(judge(sample_text(m, sft_data, '# ' + LABELS[random.randrange(len(examples))] + '\\n', n=120, temperature=1.0, seed=i, stop=stop).split('\\n', 1)[1]) for i in range(n)) / n\n\nref = copy.deepcopy(model).eval()\nprint('before the judge: %.0f%% of attempts pass' % (100 * pass_rate(ref)))\npairs = []\nfor i in range(400):\n    p = '# ' + LABELS[random.randrange(len(examples))] + '\\n'\n    a, b = [sample_text(model, sft_data, p, n=120, temperature=1.0, seed=7000 + 2 * i + k, stop=stop)[len(p):] for k in (0, 1)]\n    if judge(a) != judge(b): pairs.append((p,) + ((a, b) if judge(a) else (b, a)))\nprint(f'{len(pairs)} pairs the judge could separate')\nassert len(pairs) >= 8, 'The judge could not separate enough pairs — train longer in step 3, or loosen the judge.'\nwith torch.no_grad(): refs = [(logp(ref, p, w).item(), logp(ref, p, l).item()) for p, w, l in pairs]\nopt = torch.optim.AdamW(model.parameters(), lr=2e-5); beta = 0.5\nmodel.train()\nfor step in range(60):\n    loss = 0\n    for i in random.sample(range(len(pairs)), min(8, len(pairs))):\n        p, w, l = pairs[i]; rw, rl = refs[i]\n        loss = loss - F.logsigmoid(beta * ((logp(model, p, w) - rw) - (logp(model, p, l) - rl)))\n    opt.zero_grad(); (loss / 8).backward(); opt.step()\nmodel.eval()\nprint('after the judge:  %.0f%% of attempts pass' % (100 * pass_rate(model)))\nexport(model, vocab, 'judged.json', extra=dict(name='judged'))\nprint('saved judged.json')\n"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "#@title Step 6 · Write poems, and download the weights  { display-mode: \"form\" }\nlabel = LABELS[0]  #@param {type:\"string\"}\ntemperature = 0.8  #@param {type:\"slider\", min:0.3, max:1.5, step:0.1}\nhow_many = 5  #@param {type:\"slider\", min:1, max:20, step:1}\nfor i in range(how_many):\n    out = sample_text(model, sft_data, '# ' + label + '\\n', n=150, temperature=temperature, seed=1000 + i, stop=stop)\n    print(out.strip() + '\\n' + '-' * 40)\n\nfrom google.colab import files\nfor f in ('base.json', 'sft.json', 'judged.json'): files.download(f)\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Put it on a page\n\nThe three `.json` files are in the exact format `how-llms-work.html` reads. Open that page, scroll to chapter 12, **Load the poet you built**, and choose any of the three files — it runs in your browser immediately, on the same engine as the Thirukkural machine; nothing is uploaded. `base.json` continues any text; `sft.json` answers a `# label` line; `judged.json` is the version the judge shaped.\n\nTo publish a page of your own, copy `how-llms-work.html`, find `const DATA = …` in its source and swap the `models` entries for your files (the JavaScript engine right above it, `Engine.Model`, runs any machine of this shape and needs no changes).\n\n## What to try next\n\n* **More text.** This is the single biggest lever. A 1 MB book (roughly a novel) supports `numbers_per_letter = 128` and `rounds = 6`; expect real sentences.\n* **A better judge.** The default judges shape only. Judge meaning too (does the poem mention the label? is it novel? does a word list say it is positive?) and the machine will chase that instead — carefully; it will chase whatever you reward.\n* **Two languages.** Put the English meaning *and* the Tamil verse in each example, with the English as the label line, and the machine learns to write Tamil from an English request. It needs a bigger machine and more steps.\n"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "name": "build-your-own-poet.ipynb",
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}