{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Step 1 — A Bigram Model from Counts\n\n**Solution notebook.** Run top to bottom; it is self-contained and takes\nunder a minute. Read the prose — this notebook is written to be *read*,\nnot just executed, and it is the starting point for Step 2.\n\nIt uses PyTorch throughout, the same library and the same idioms as the\ntask page (`torch.zeros`, `dim=`/`keepdim=`, `torch.multinomial`,\npaired indexing), so everything here carries straight into Steps 2–8.\n\nCheckpoint values quoted here were produced by this notebook, so your\nnumbers should match to the digit (there is no randomness except in the\nsampling cell, which is seeded).",
   "id": "cell-00"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1.1 Load the data\n\nOn Colab, uncomment the download line. Locally, make sure `input.txt` is\nin the same folder as this notebook.",
   "id": "cell-01"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n\nwith open('input.txt') as f:\n    text = f.read()\n\nprint(f'{len(text):,} characters')\nprint(text[:250])",
   "id": "cell-02"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "> `1,115,394 characters` — the complete works, as one string.",
   "id": "cell-03"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1.2 Tokenize\n\nOur alphabet is whatever characters actually occur. `stoi`/`itos` are the\ntwo directions of a bijection $\\{$characters$\\} \\leftrightarrow \\{0,\\dots,V-1\\}$;\n`encode`/`decode` extend it to strings, and the assertion checks they are\nmutually inverse.",
   "id": "cell-04"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "vocab = sorted(set(text))\nV = len(vocab)\nstoi = {ch: i for i, ch in enumerate(vocab)}\nitos = {i: ch for i, ch in enumerate(vocab)}\nencode = lambda s: [stoi[c] for c in s]\ndecode = lambda ids: ''.join(itos[i] for i in ids)\n\nassert decode(encode('Fear no more')) == 'Fear no more'\nprint(f'V = {V}')\nprint(repr(''.join(vocab)))",
   "id": "cell-05"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "> `V = 65`. Token 0 is `'\\n'` and token 1 is `' '` — the two most\n> structurally important characters in the corpus sort to the front,\n> which is a convenient accident. The vocabulary is an empirical fact\n> about this text, not a design choice: all 26 letters appear in both\n> cases, but the only digit is `3` (from a stray line number or two) and\n> the punctuation is just ``! $ & ' , - . : ; ?`` — no quotation marks,\n> no parentheses. Editorial conventions of one 19th-century edition,\n> fossilized into our model's alphabet.",
   "id": "cell-06"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1.3 Count\n\n$N_{ab}$ = number of times character $b$ immediately follows $a$. The\ntask page fills $N$ with a Python loop over `zip(ids, ids[1:])`, which\nwalks the adjacent pairs one at a time — correct, and about ten seconds\nfor a million pairs. Below is the vectorized version: flatten each pair\n$(a,b)$ to the single integer $aV+b$ and let `torch.bincount` tally all\nof them at once. (The loop is kept, commented out, so you can check the\ntwo agree.)",
   "id": "cell-07"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import torch\nimport matplotlib.pyplot as plt\n\nids = torch.tensor(encode(text))              # the whole text as a tensor of ids, shape (1115394,)\nN = torch.bincount(ids[:-1] * V + ids[1:], minlength=V * V).view(V, V)\n\n# The task page's loop gives exactly the same matrix, about 10 s slower:\n# N_loop = torch.zeros((V, V), dtype=torch.int64)\n# for a, b in zip(ids.tolist(), ids[1:].tolist()):\n#     N_loop[a, b] += 1\n# assert torch.equal(N, N_loop)\n\nassert N.sum() == len(text) - 1\nprint(f'total pairs counted: {N.sum():,}')\nprint(f'cells that are exactly zero: {(N == 0).sum():,} of {V*V:,}')",
   "id": "cell-08"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "> **2,822 of 4,225 cells are zero.** Two thirds of conceivable character\n> pairs never occur in all of Shakespeare. This is the concrete reason\n> smoothing is not optional: an unsmoothed model assigns probability 0 to\n> any of those pairs, and one occurrence in held-out text sends the\n> log-likelihood to $-\\infty$.",
   "id": "cell-09"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "plt.figure(figsize=(9, 9))\nplt.imshow(N.log1p(), cmap='Blues')\nplt.xticks(range(V), vocab, fontsize=7)\nplt.yticks(range(V), vocab, fontsize=7)\nplt.xlabel('next character'); plt.ylabel('current character')\nplt.title('log(1 + N)')\nplt.colorbar(shrink=0.8);",
   "id": "cell-10"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "We plot $\\log(1+N)$ rather than $N$ because the raw counts span four\norders of magnitude and a linear colour map would show only `' '` and\n`'e'`. Three things to find in the picture, all verifiable from `N`:\n\n- **the `q` row is exactly rank one.** `q` occurs 609 times and is\n  followed by `u` all 609 times — probability 1.000, no exceptions in the\n  complete works. English orthography as a single bright cell.\n- **the space column is bright but not universal:** 46 of the 65\n  characters are ever followed by a space. Ask yourself which 19 are not,\n  and you will have derived a chunk of English orthographic rules from a\n  count matrix.\n- **the capital-letter block is brightest into *itself*** — 51% of\n  characters following a capital are themselves capitals, against only\n  33% lowercase. That is backwards from ordinary English prose, and it is\n  the `ALL-CAPS SPEAKER:` convention of a printed play showing up as pure\n  geometry. (A further 6.9% of post-capital characters are `:` itself.)\n  The model has no concept of a play, a speaker, or a stage direction; it\n  has a matrix in which those conventions are unmistakable.",
   "id": "cell-11"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "q, u, sp = stoi['q'], stoi['u'], stoi[' ']\ncaps = torch.tensor([stoi[c] for c in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'])\nlower = torch.tensor([stoi[c] for c in 'abcdefghijklmnopqrstuvwxyz'])\n\nprint(f'q occurs {N[q].sum()} times, followed by u {N[q, u]} times')\nprint(f'{(N[:, sp] > 0).sum()} of {V} characters are ever followed by a space')\nafter_cap = N[caps].sum(dim=0)                       # what follows a capital, summed over capitals\nprint(f'after a capital: {after_cap[caps].sum() / after_cap.sum():.1%} capitals, '\n      f'{after_cap[lower].sum() / after_cap.sum():.1%} lowercase, '\n      f'{after_cap[stoi[\":\"]] / after_cap.sum():.1%} colon')",
   "id": "cell-12"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1.4 Normalize, with add-one smoothing\n\n$P_{ab} = \\dfrac{N_{ab}+1}{\\sum_c (N_{ac}+1)}$, rows summing to 1. Laplace\nsmoothing is the posterior mean under a $\\mathrm{Dirichlet}(1,\\dots,1)$\nprior on each row — i.e. it is what you get by pretending you saw every\npair once before looking at the data.",
   "id": "cell-13"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "P = (N + 1).float()                 # counts are integers; probabilities need decimals\nP = P / P.sum(dim=1, keepdim=True)  # divide each row by its own total\n\nassert torch.allclose(P.sum(dim=1), torch.ones(V))\nprint('rows sum to 1 ✓')",
   "id": "cell-14"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**The classic bug lives here.** `dim=1` sums *along* each row, and\n`keepdim=True` keeps the result shaped `(V, 1)` so it broadcasts down\nthe rows. Drop `keepdim` and you get shape `(V,)`, which PyTorch\nbroadcasts along the *last* axis instead — silently normalizing columns.\nThe assertion above is what catches it.",
   "id": "cell-15"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1.5 Sample\n\nRun the Markov chain: start from the newline character and repeatedly\ndraw the next character from the current row of $P$. `torch.multinomial`\ndoes the weighted draw; `.item()` turns the one-element tensor it returns\nback into a plain integer we can index with.",
   "id": "cell-16"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "torch.manual_seed(0)\ncur, out = stoi['\\n'], []\nfor _ in range(500):\n    cur = torch.multinomial(P[cur], num_samples=1).item()\n    out.append(cur)\nprint(decode(out))",
   "id": "cell-17"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "English-shaped nonsense: pronounceable syllables, plausible word\nlengths, capital letters after newlines, the occasional real word by\nchance. A single matrix of pair counts already captures that much. What\nit cannot capture is anything at a range beyond one character — which is\nthe entire remaining agenda of the course.",
   "id": "cell-18"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1.6 Evaluate on the training text\n\nAverage log loss $\\mathcal L = -\\frac{1}{T-1}\\sum_t \\log P_{x_t x_{t+1}}$\nin nats; divide by $\\ln 2$ for bits per character; $e^{\\mathcal L}$ is\nperplexity. The paired indexing `P[ids[:-1], ids[1:]]` pulls out all\n$T-1$ transition probabilities at once.\n\nWe wrap the computation in a function because Task 1.7 will run the very\nsame pipeline on a different text.",
   "id": "cell-19"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def avg_log_loss(P, ids):\n    p_next = P[ids[:-1], ids[1:]]     # entry t is P[ids[t], ids[t+1]]\n    return -torch.log(p_next).mean()\n\ndef report(name, L):\n    L = float(L)\n    print(f'{name:<18} {L:.4f} nats  {L / torch.log(torch.tensor(2.)):.4f} bits/char  '\n          f'perplexity {torch.exp(torch.tensor(L)):8.3f}')\n\nP_uniform = torch.full((V, V), 1 / V)   # the baseline, run through the same code\n\nreport('bigram (add-one)', avg_log_loss(P, ids))\nreport('uniform', avg_log_loss(P_uniform, ids))",
   "id": "cell-20"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "| model | nats | bits/char | perplexity |\n|---|---|---|---|\n| uniform | 4.1744 | 6.0224 | 65.000 |\n| bigram (add-one) | **2.4549** | **3.5417** | **11.646** |\n\nPerplexity is the *effective branching factor*: a uniform model over 65\ncharacters is as hard to predict as a fair 65-sided die, and the bigram\nmodel reduces that to about 11.6. Shannon (1951) estimated English at\nroughly 1 bit/character, so at 3.54 we have captured perhaps a third of\nthe available structure.\n\n**Start your table now** — one row per step, and it becomes the story of\nthe whole course:\n\n| step | model | bits/char |\n|---|---|---|\n| 1 | bigram counts | 3.54 |\n\nThree checks from the task page, each a one-liner:",
   "id": "cell-21"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# (a) the vectorized loss agrees with an explicit loop, on the first 10,000 pairs\nloop = -sum(torch.log(P[a, b]) for a, b in zip(ids[:10000].tolist(), ids[1:10001].tolist())) / 10000\nassert torch.isclose(loop, avg_log_loss(P, ids[:10001]))\nprint(f'loop vs vectorized on 10,000 pairs: {loop:.4f} vs {avg_log_loss(P, ids[:10001]):.4f}')\n\n# (b) swapping the indices evaluates the model backwards — worse than knowing nothing\nreport('backwards', -torch.log(P[ids[1:], ids[:-1]]).mean())\n\n# (c) the unsmoothed MLE is only 0.002 nats better on the training text\nP_mle = N.float() / N.sum(dim=1, keepdim=True)\nreport('unsmoothed MLE', avg_log_loss(P_mle, ids))",
   "id": "cell-22"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1.7 Evaluate on held-out text\n\nEverything above graded the model on the text its counts came from. Now\nsplit chronologically — first 90% training, last 10% validation — fit\nthe counts on the training portion only, and score both. The vocabulary\nstays as built from the full text (conveniently, the last 10% introduces\nno new characters).",
   "id": "cell-23"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "n_train = int(0.9 * len(ids))            # 1,003,854\nids_tr, ids_va = ids[:n_train], ids[n_train:]   # 111,540 validation characters\nprint(f'{len(ids_tr):,} train / {len(ids_va):,} validation')\nprint('validation opens:', repr(decode(ids_va[:45].tolist())))\n\nN_tr = torch.bincount(ids_tr[:-1] * V + ids_tr[1:], minlength=V * V).view(V, V)\nP_tr = (N_tr + 1).float()\nP_tr = P_tr / P_tr.sum(dim=1, keepdim=True)\n\nreport('train', avg_log_loss(P_tr, ids_tr))\nreport('validation', avg_log_loss(P_tr, ids_va))\nreport('uniform (val)', avg_log_loss(P_uniform, ids_va))",
   "id": "cell-24"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "> **2.4546 nats on train, 2.4819 on validation** (3.5412 vs 3.5806\n> bits/char; perplexity 11.641 vs 11.964): a gap of 0.027 nats. The\n> uniform model gives 4.1744 again, as it must — it doesn't depend on the\n> data, so any change there is a bug.\n>\n> The gap is real but tiny. The bigram matrix has only 4,225 cells sharing\n> a million characters of evidence, so almost every cell is estimated from\n> abundant data. Overfitting grows with the ratio of parameters to data,\n> and here that ratio is small. Which transitions does the training\n> portion never see?",
   "id": "cell-25"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "train_count_of_val_pairs = N_tr[ids_va[:-1], ids_va[1:]]   # training count of each validation transition\nunseen = train_count_of_val_pairs == 0\npairs = torch.stack([ids_va[:-1][unseen], ids_va[1:][unseen]], dim=1)\ndistinct, counts = torch.unique(pairs, dim=0, return_counts=True)\n\nprint(f'{unseen.sum()} validation transitions ({len(distinct)} distinct pairs) have zero training count')\nfor k in counts.argsort(descending=True)[:3]:\n    a, b = distinct[k].tolist()\n    print(f'  {itos[a]!r} -> {itos[b]!r}: {counts[k]} times')\n\nP_tr_mle = N_tr.float() / N_tr.sum(dim=1, keepdim=True)\nreport('unsmoothed, train', avg_log_loss(P_tr_mle, ids_tr))\nreport('unsmoothed, val', avg_log_loss(P_tr_mle, ids_va))",
   "id": "cell-26"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "> **187 validation transitions, 23 distinct pairs.** The top three —\n> `S`→`P` (63), `E`→`B` (42), `N`→`Z` (37) — are PROSPERO, SEBASTIAN,\n> and GONZALO: *The Tempest* and its `ALL-CAPS` cast enter the corpus\n> only in the final 10%. The unsmoothed MLE scores 2.4519 on train and\n> **infinite** on validation; any one of the 187 suffices. That is the\n> whole case for smoothing, in one number.",
   "id": "cell-27"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Going further — trigrams, and why tables must fail\n\nCondition on *two* characters: $65^2 = 4{,}225$ contexts. We index the\ncontext pair $(a,b)$ as the single integer $aV + b$, and the count table\nbecomes $4{,}225 \\times 65$.",
   "id": "cell-28"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def trigram_P(ids):\n    ctx = ids[:-2] * V + ids[1:-1]\n    N3 = torch.bincount(ctx * V + ids[2:], minlength=V * V * V).view(V * V, V)\n    P3 = (N3 + 1).float()\n    return P3 / P3.sum(dim=1, keepdim=True), N3\n\ndef trigram_loss(P3, ids):\n    ctx = ids[:-2] * V + ids[1:-1]\n    return -torch.log(P3[ctx, ids[2:]]).mean()\n\nP3, _ = trigram_P(ids)                      # fit and scored on the full text, like Task 1.6\nreport('trigram (full)', trigram_loss(P3, ids))\n\nP3_tr, N3_tr = trigram_P(ids_tr)            # fit on the 90%, scored on both, like Task 1.7\nreport('trigram, train', trigram_loss(P3_tr, ids_tr))\nreport('trigram, val', trigram_loss(P3_tr, ids_va))\nctx_va = ids_va[:-2] * V + ids_va[1:-1]\nprint(f'{(N3_tr[ctx_va, ids_va[2:]] == 0).sum()} validation transitions unseen in training (bigram: 187)')",
   "id": "cell-29"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "Trigram on the full text: **1.9532 nats = 2.818 bits/char**, perplexity\n7.05 — a large gain from one extra character of context. But that is a\ntraining-corpus number, and Lecture 1 §6.4 showed that $n$-gram\n*training* loss can only improve as context grows, all the way to\nmemorization. The comparison that counts is the held-out one: refit on\nthe 90%, the trigram scores **1.9526 train / 2.0684 validation** (2.9841\nbits/char, perplexity 7.91) against the bigram's 2.4819 — so it genuinely\ngeneralizes better here. The warning signs are growing, though: its\ntrain/validation gap is 0.116 nats to the bigram's 0.027, and 1,553\nvalidation transitions are unseen in training, up from 187.\n\nSo why not keep going? A $k$-gram table has $65^{k+1}$ entries:\n\n| $k$ | table entries | characters of data |\n|---|---|---|\n| 1 | $4.2\\times10^{3}$ | $1.1\\times10^{6}$ |\n| 2 | $2.7\\times10^{5}$ | $1.1\\times10^{6}$ |\n| 3 | $1.8\\times10^{7}$ | $1.1\\times10^{6}$ |\n| 5 | $7.5\\times10^{10}$ | $1.1\\times10^{6}$ |\n\nBy $k=3$ there are more table cells than characters in the corpus, so\nalmost every cell is estimated from zero or one observation and the model\nis pure noise plus smoothing. The data does not grow; only the table\ndoes. **This exponential blow-up against fixed data is the reason Lecture\n2 replaces tables with parametrized functions**, which can share\nstatistical strength across contexts instead of treating each as an\nisolated counting problem.\n\n→ Continue with [Step 2](../project/step-2), which rebuilds exactly this\nbigram model as $\\operatorname{softmax}$ of a linear map and trains it by\ngradient descent.",
   "id": "cell-30"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}