{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Greedy, beam, temperature, top-p\n",
    "\n",
    "Where greedy decoding loses, and why top-p replaced top-k.\n",
    "\n",
    "From chapter 18, [Steering LLMs: Decoding and Prompting](https://nlp.jcrlabz.com/book/decoding-prompting/), of the course notes.\n",
    "\n",
    "Source: `book/code/worked_examples/decoding.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What you need\n",
    "\n",
    "Nothing. This example uses only the Python standard library."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Nothing to install: this example uses only the Python standard library.\n",
    "import sys; print(sys.version.split()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The script\n",
    "\n",
    "Everything the example defines, in one cell. Run it and the definitions\n",
    "are live in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "\"\"\"Decoding strategies, worked as Chapter 18 works them.\n",
    "\n",
    "Four things, all reproducing the chapter's tables:\n",
    "\n",
    "    greedy against beam   where greedy loses, and what beam costs\n",
    "    temperature           one knob, three behaviours\n",
    "    top-k and top-p       why a fixed k fails on some distributions\n",
    "    the whole picture     the same distribution under every strategy\n",
    "\n",
    "    python3 decoding.py            # all four\n",
    "    python3 decoding.py --beam     # greedy against beam search\n",
    "    python3 decoding.py --temp     # the temperature sweep\n",
    "    python3 decoding.py --truncate # top-k against top-p\n",
    "    python3 decoding.py --compare  # every strategy, side by side\n",
    "\n",
    "Install:  nothing, the Python standard library is enough\n",
    "\"\"\"\n",
    "\n",
    "import argparse\n",
    "import math\n",
    "\n",
    "# A tiny language model given as an explicit tree of next-token probabilities.\n",
    "# The key is the prefix; the value maps next token to probability.\n",
    "TREE = {\n",
    "    (): {\"the\": 0.40, \"a\": 0.35, \"one\": 0.25},\n",
    "    (\"the\",): {\"cat\": 0.55, \"dog\": 0.45},\n",
    "    (\"a\",): {\"cat\": 0.10, \"bird\": 0.90},\n",
    "    (\"one\",): {\"cat\": 0.50, \"dog\": 0.50},\n",
    "    (\"the\", \"cat\"): {\"sat\": 0.60, \"ran\": 0.40},\n",
    "    (\"the\", \"dog\"): {\"sat\": 0.50, \"ran\": 0.50},\n",
    "    (\"a\", \"cat\"): {\"sat\": 0.50, \"ran\": 0.50},\n",
    "    (\"a\", \"bird\"): {\"sang\": 0.95, \"flew\": 0.05},\n",
    "    (\"one\", \"cat\"): {\"sat\": 0.50, \"ran\": 0.50},\n",
    "    (\"one\", \"dog\"): {\"sat\": 0.50, \"ran\": 0.50},\n",
    "}\n",
    "\n",
    "# A separate distribution for the truncation demonstrations. Two shapes: one\n",
    "# where the model is confident and one where it is not.\n",
    "PEAKED = {\"the\": 0.85, \"a\": 0.07, \"one\": 0.04, \"some\": 0.02,\n",
    "          \"any\": 0.01, \"each\": 0.005, \"every\": 0.005}\n",
    "FLAT = {\"red\": 0.14, \"blue\": 0.13, \"green\": 0.13, \"black\": 0.12,\n",
    "        \"white\": 0.12, \"grey\": 0.12, \"brown\": 0.12, \"pink\": 0.12}\n",
    "\n",
    "\n",
    "def greedy(depth=3):\n",
    "    prefix, logp, trace = (), 0.0, []\n",
    "    for _ in range(depth):\n",
    "        dist = TREE.get(prefix)\n",
    "        if not dist:\n",
    "            break\n",
    "        w = max(dist, key=dist.get)\n",
    "        trace.append((prefix, w, dist[w], dict(dist)))\n",
    "        logp += math.log(dist[w])\n",
    "        prefix = prefix + (w,)\n",
    "    return prefix, math.exp(logp), trace\n",
    "\n",
    "\n",
    "def all_sequences(depth=3):\n",
    "    out = []\n",
    "\n",
    "    def walk(prefix, p):\n",
    "        dist = TREE.get(prefix)\n",
    "        if not dist or len(prefix) == depth:\n",
    "            out.append((prefix, p))\n",
    "            return\n",
    "        for w, q in dist.items():\n",
    "            walk(prefix + (w,), p * q)\n",
    "    walk((), 1.0)\n",
    "    return sorted(out, key=lambda x: -x[1])\n",
    "\n",
    "\n",
    "def beam(width, depth=3):\n",
    "    beams = [((), 0.0)]\n",
    "    for _ in range(depth):\n",
    "        cand = []\n",
    "        for prefix, lp in beams:\n",
    "            dist = TREE.get(prefix)\n",
    "            if not dist:\n",
    "                cand.append((prefix, lp))\n",
    "                continue\n",
    "            for w, q in dist.items():\n",
    "                cand.append((prefix + (w,), lp + math.log(q)))\n",
    "        cand.sort(key=lambda x: -x[1])\n",
    "        beams = cand[:width]\n",
    "    return [(p, math.exp(lp)) for p, lp in beams]\n",
    "\n",
    "\n",
    "# --------------------------------------------------------------- beam search\n",
    "\n",
    "def show_beam():\n",
    "    print(\"\\nSTEP 1  greedy takes the best word and loses the best sentence\\n\")\n",
    "    seq, p, trace = greedy()\n",
    "    print(\"  GREEDY. At each step take the highest probability token.\\n\")\n",
    "    print(f\"  {'prefix':<18}{'choices':<34}{'taken':<8}{'p':>7}\")\n",
    "    print(\"  \" + \"-\" * 68)\n",
    "    for prefix, w, q, dist in trace:\n",
    "        pre = \" \".join(prefix) if prefix else \"(start)\"\n",
    "        ch = \", \".join(f\"{k} {v:.2f}\" for k, v in dist.items())\n",
    "        print(f\"  {pre:<18}{ch:<34}{w:<8}{q:>7.2f}\")\n",
    "    print(f\"\\n  greedy output: '{' '.join(seq)}'   probability {p:.4f}\\n\")\n",
    "\n",
    "    ranked = all_sequences()\n",
    "    print(\"  Now every three-word sequence, ranked by probability:\\n\")\n",
    "    print(f\"  {'rank':>5}  {'sequence':<24}{'probability':>13}\")\n",
    "    print(\"  \" + \"-\" * 46)\n",
    "    for i, (s, q) in enumerate(ranked[:6], start=1):\n",
    "        mark = \"   <- greedy found this\" if s == seq else \"\"\n",
    "        print(f\"  {i:>5}  {' '.join(s):<24}{q:>13.4f}{mark}\")\n",
    "\n",
    "    best = ranked[0]\n",
    "    print(f\"\\n  The best sequence is '{' '.join(best[0])}' at {best[1]:.4f}.\")\n",
    "    print(f\"  Greedy returned {p:.4f}, which is\"\n",
    "          f\" {(best[1] / p - 1) * 100:.0f} per cent worse.\")\n",
    "    print(f\"\\n  Greedy went wrong at the very first step. It took 'the' at\")\n",
    "    print(f\"  0.40 over 'a' at 0.35, and never saw that 'a' leads to 'bird'\")\n",
    "    print(f\"  at 0.90 and then 'sang' at 0.95.\")\n",
    "    print(f\"\\n  One high-probability token can sit in front of a low-\")\n",
    "    print(f\"  probability continuation. Greedy cannot know that in advance.\")\n",
    "\n",
    "    print(f\"\\n\\n  BEAM SEARCH keeps the best w partial sequences at each step.\\n\")\n",
    "    print(f\"  {'width':>6}  {'best sequence found':<26}{'probability':>13}\"\n",
    "          f\"{'optimal?':>10}\")\n",
    "    print(\"  \" + \"-\" * 58)\n",
    "    for w in (1, 2, 3, 5):\n",
    "        b = beam(w)\n",
    "        top = b[0]\n",
    "        ok = \"yes\" if abs(top[1] - best[1]) < 1e-12 else \"no\"\n",
    "        print(f\"  {w:>6}  {' '.join(top[0]):<26}{top[1]:>13.4f}{ok:>10}\")\n",
    "\n",
    "    print(f\"\\n  Width 1 is greedy by definition. Width 2 already finds the\")\n",
    "    print(f\"  optimum here, and wider beams cost more for nothing.\")\n",
    "    print(f\"\\n  Beam search is not guaranteed to find the best sequence. It\")\n",
    "    print(f\"  searches more of the space than greedy and less than all of it,\")\n",
    "    print(f\"  which is the only honest description of it.\")\n",
    "\n",
    "\n",
    "# --------------------------------------------------------------- temperature\n",
    "\n",
    "def apply_temp(dist, T):\n",
    "    if T == 0:\n",
    "        top = max(dist, key=dist.get)\n",
    "        return {k: (1.0 if k == top else 0.0) for k in dist}\n",
    "    z = {k: math.exp(math.log(v) / T) for k, v in dist.items()}\n",
    "    s = sum(z.values())\n",
    "    return {k: v / s for k, v in z.items()}\n",
    "\n",
    "\n",
    "def entropy(dist):\n",
    "    h = -sum(p * math.log2(p) for p in dist.values() if p > 0)\n",
    "    return abs(h)   # avoid printing -0.0000 when one token has all the mass\n",
    "\n",
    "\n",
    "def show_temp():\n",
    "    print(\"\\n\\nSTEP 2  temperature, one knob\\n\")\n",
    "    print(\"  Divide the logits by T before the softmax. Equivalently, raise\")\n",
    "    print(\"  each probability to the power 1/T and renormalise.\\n\")\n",
    "    base = PEAKED\n",
    "    words = list(base)\n",
    "    print(f\"  {'T':>6}\" + \"\".join(f\"{w:>9}\" for w in words[:5])\n",
    "          + f\"{'entropy':>10}\")\n",
    "    print(\"  \" + \"-\" * 62)\n",
    "    for T in (0.0, 0.5, 0.8, 1.0, 1.5, 2.0):\n",
    "        d = apply_temp(base, T)\n",
    "        cells = \"\".join(f\"{d[w]:>9.4f}\" for w in words[:5])\n",
    "        print(f\"  {T:>6.1f}{cells}{entropy(d):>10.4f}\")\n",
    "\n",
    "    print(f\"\\n  T = 1 leaves the model's own distribution alone.\")\n",
    "    print(f\"  T below 1 sharpens it. At T = 0 it is greedy, and the entropy\")\n",
    "    print(f\"  is 0 because there is nothing left to choose.\")\n",
    "    print(f\"  T above 1 flattens it and raises the entropy.\")\n",
    "    print(f\"\\n  The trade has a name at each end. Low temperature gives\")\n",
    "    print(f\"  repetitive, safe text. High temperature gives varied text that\")\n",
    "    print(f\"  drifts off topic and eventually stops making sense.\")\n",
    "\n",
    "\n",
    "# --------------------------------------------------------- top-k and top-p\n",
    "\n",
    "def top_k(dist, k):\n",
    "    keep = sorted(dist.items(), key=lambda x: -x[1])[:k]\n",
    "    s = sum(v for _, v in keep)\n",
    "    return {w: v / s for w, v in keep}\n",
    "\n",
    "\n",
    "def top_p(dist, p):\n",
    "    keep, run = [], 0.0\n",
    "    for w, v in sorted(dist.items(), key=lambda x: -x[1]):\n",
    "        keep.append((w, v))\n",
    "        run += v\n",
    "        if run >= p:\n",
    "            break\n",
    "    s = sum(v for _, v in keep)\n",
    "    return {w: v / s for w, v in keep}\n",
    "\n",
    "\n",
    "def show_truncate():\n",
    "    print(\"\\n\\nSTEP 3  top-k has a problem that top-p does not\\n\")\n",
    "    print(\"  Two next-token distributions. One where the model is sure, one\")\n",
    "    print(\"  where it genuinely is not.\\n\")\n",
    "    for name, d in ((\"peaked\", PEAKED), (\"flat\", FLAT)):\n",
    "        top = sorted(d.items(), key=lambda x: -x[1])[:4]\n",
    "        cells = \",  \".join(f\"{w} {v:.2f}\" for w, v in top)\n",
    "        print(f\"  {name:<8}{cells},  ...   entropy {entropy(d):.4f} bits\")\n",
    "\n",
    "    print(f\"\\n  TOP-K keeps a fixed number of candidates, whatever the shape.\\n\")\n",
    "    print(f\"  {'k':>4}{'peaked: kept':>16}{'mass kept':>12}\"\n",
    "          f\"{'flat: kept':>14}{'mass kept':>12}\")\n",
    "    print(\"  \" + \"-\" * 60)\n",
    "    for k in (1, 2, 5, 8):\n",
    "        pk = sorted(PEAKED.items(), key=lambda x: -x[1])[:k]\n",
    "        fk = sorted(FLAT.items(), key=lambda x: -x[1])[:k]\n",
    "        print(f\"  {k:>4}{len(pk):>16}{sum(v for _, v in pk):>12.4f}\"\n",
    "              f\"{len(fk):>14}{sum(v for _, v in fk):>12.4f}\")\n",
    "\n",
    "    print(f\"\\n  Read k = 5. On the peaked distribution it admits four tokens\")\n",
    "    print(f\"  the model had all but ruled out, together worth\"\n",
    "          f\" {sum(v for _,v in sorted(PEAKED.items(), key=lambda x:-x[1])[1:5]):.3f}.\")\n",
    "    print(f\"  On the flat one it discards three tokens that were as good as\")\n",
    "    print(f\"  the ones it kept.\")\n",
    "    print(f\"\\n  The same k is too loose on one shape and too tight on the\")\n",
    "    print(f\"  other, because k cannot see the shape.\")\n",
    "\n",
    "    print(f\"\\n\\n  TOP-P keeps the smallest set whose mass reaches p.\\n\")\n",
    "    print(f\"  {'p':>6}{'peaked: kept':>16}{'mass':>10}{'flat: kept':>14}\"\n",
    "          f\"{'mass':>10}\")\n",
    "    print(\"  \" + \"-\" * 58)\n",
    "    for p in (0.5, 0.9, 0.95, 0.99):\n",
    "        pk, fk = top_p(PEAKED, p), top_p(FLAT, p)\n",
    "        mp = sum(PEAKED[w] for w in pk)\n",
    "        mf = sum(FLAT[w] for w in fk)\n",
    "        print(f\"  {p:>6.2f}{len(pk):>16}{mp:>10.4f}{len(fk):>14}{mf:>10.4f}\")\n",
    "\n",
    "    print(f\"\\n  At p = 0.9 the peaked distribution keeps\"\n",
    "          f\" {len(top_p(PEAKED, 0.9))} tokens and the flat one keeps\"\n",
    "          f\" {len(top_p(FLAT, 0.9))}.\")\n",
    "    print(f\"  The size of the candidate set now follows the model's own\")\n",
    "    print(f\"  confidence, which is exactly what a fixed k could not do.\")\n",
    "    print(f\"\\n  That is why nucleus sampling replaced top-k as the default.\")\n",
    "\n",
    "\n",
    "# ------------------------------------------------------------ side by side\n",
    "\n",
    "def show_compare():\n",
    "    print(\"\\n\\nSTEP 4  the same distribution under every strategy\\n\")\n",
    "    d = PEAKED\n",
    "    rows = [\n",
    "        (\"greedy\", apply_temp(d, 0.0)),\n",
    "        (\"temperature 0.7\", apply_temp(d, 0.7)),\n",
    "        (\"pure sampling, T = 1\", dict(d)),\n",
    "        (\"temperature 1.5\", apply_temp(d, 1.5)),\n",
    "        (\"top-k, k = 3\", top_k(d, 3)),\n",
    "        (\"top-p, p = 0.9\", top_p(d, 0.9)),\n",
    "    ]\n",
    "    words = list(d)\n",
    "    print(f\"  {'strategy':<22}\" + \"\".join(f\"{w[:6]:>8}\" for w in words)\n",
    "          + f\"{'kept':>7}{'entropy':>9}\")\n",
    "    print(\"  \" + \"-\" * 86)\n",
    "    for name, dd in rows:\n",
    "        cells = \"\".join(f\"{dd.get(w, 0.0):>8.3f}\" for w in words)\n",
    "        kept = sum(1 for w in words if dd.get(w, 0.0) > 0)\n",
    "        print(f\"  {name:<22}{cells}{kept:>7}{entropy(dd):>9.4f}\")\n",
    "\n",
    "    print(f\"\\n  Every row is the same model. None of them changed a weight.\")\n",
    "    print(f\"  Decoding is a decision made after the model has spoken, and it\")\n",
    "    print(f\"  changes the output more than most fine-tuning does.\")\n",
    "    print(f\"\\n  Which to use follows from the task. A factual answer wants\")\n",
    "    print(f\"  the low-entropy rows. A story wants one of the middle ones.\")\n",
    "    print(f\"  Nothing wants the bottom of the temperature range or the top.\")\n",
    "\n",
    "\n",
    "def main():\n",
    "    ap = argparse.ArgumentParser(description=__doc__)\n",
    "    for f in (\"beam\", \"temp\", \"truncate\", \"compare\"):\n",
    "        ap.add_argument(f\"--{f}\", action=\"store_true\")\n",
    "    a = ap.parse_args()\n",
    "    picked = a.beam or a.temp or a.truncate or a.compare\n",
    "    if a.beam or not picked:\n",
    "        show_beam()\n",
    "    if a.temp or not picked:\n",
    "        show_temp()\n",
    "    if a.truncate or not picked:\n",
    "        show_truncate()\n",
    "    if a.compare or not picked:\n",
    "        show_compare()\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running it\n",
    "\n",
    "The script is a command line tool, so each run below sets the arguments\n",
    "and calls `main()`. The heading shows the equivalent shell command."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "def run(*args):\n",
    "    sys.argv = [\"decoding.py\", *args]\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 decoding.py`\n",
    "\n",
    "Every strategy on the same distribution."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 decoding.py --beam`\n",
    "\n",
    "Greedy takes the at 0.40 over a at 0.35 and never sees that a leads to 0.2992. It returns 0.1320. A beam of width 2 finds the better sequence."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"--beam\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 decoding.py --truncate`\n",
    "\n",
    "At k = 5 the same setting keeps 99 per cent of a peaked distribution and 64 per cent of a flat one. At p = 0.9 the candidate set follows the model's own confidence."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"--truncate\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Read the chapter this comes from: [Steering LLMs: Decoding and Prompting](https://nlp.jcrlabz.com/book/decoding-prompting/)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "decoding.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}