{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BPE, five merges\n",
    "\n",
    "The whole training run with pair counts and ties, then the learned rules applied to words the corpus never held.\n",
    "\n",
    "From chapter 5, [Subword Tokenisation](https://nlp.jcrlabz.com/book/tokenization/), of the course notes.\n",
    "\n",
    "Source: `book/code/worked_examples/bpe.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": [
    "\"\"\"Byte pair encoding, traced step by step as Chapter 5 traces it.\n",
    "\n",
    "Run this and you get the same corpus tables, the same pair counts, the same\n",
    "ties and the same five merges the book prints.\n",
    "\n",
    "    python3 bpe.py                    # the low/lower/newest/widest corpus\n",
    "    python3 bpe.py --merges 8         # keep going and watch it continue\n",
    "    python3 bpe.py --encode lowest    # apply the learned rules to a new word\n",
    "    python3 bpe.py --no-tiebreak      # see why the tie-break rule exists\n",
    "\n",
    "The two functions that matter are `train` and `encode`. Everything else is\n",
    "printing.\n",
    "\n",
    "Install:  nothing, the Python standard library is enough\n",
    "\"\"\"\n",
    "\n",
    "import argparse\n",
    "from collections import Counter\n",
    "\n",
    "END = \"</w>\"\n",
    "\n",
    "\n",
    "def pair_counts(words):\n",
    "    \"\"\"Every adjacent pair, weighted by the frequency of the word it sits in.\n",
    "\n",
    "    The weighting is the step most hand simulations get wrong. A pair inside a\n",
    "    word seen 6 times contributes 6, not 1.\n",
    "    \"\"\"\n",
    "    counts = Counter()\n",
    "    for symbols, freq in words.items():\n",
    "        for a, b in zip(symbols, symbols[1:]):\n",
    "            counts[(a, b)] += freq\n",
    "    return counts\n",
    "\n",
    "\n",
    "def merge_pair(words, pair):\n",
    "    \"\"\"Rewrite the corpus with `pair` joined into one symbol.\"\"\"\n",
    "    a, b = pair\n",
    "    out = {}\n",
    "    for symbols, freq in words.items():\n",
    "        joined, i = [], 0\n",
    "        while i < len(symbols):\n",
    "            if i + 1 < len(symbols) and symbols[i] == a and symbols[i + 1] == b:\n",
    "                joined.append(a + b)\n",
    "                i += 2\n",
    "            else:\n",
    "                joined.append(symbols[i])\n",
    "                i += 1\n",
    "        out[tuple(joined)] = out.get(tuple(joined), 0) + freq\n",
    "    return out\n",
    "\n",
    "\n",
    "def train(corpus, num_merges, tiebreak=True, trace=None):\n",
    "    \"\"\"Learn an ordered merge list from {word: frequency}.\n",
    "\n",
    "    Returns the merges as \"a b\" strings, which is the form `encode` expects\n",
    "    and the form a GPT-2 merges.txt file uses.\n",
    "    \"\"\"\n",
    "    words = {tuple(list(w) + [END]): f for w, f in corpus.items()}\n",
    "    merges = []\n",
    "    for step in range(1, num_merges + 1):\n",
    "        counts = pair_counts(words)\n",
    "        if not counts:\n",
    "            break\n",
    "        top = max(counts.values())\n",
    "        tied = sorted(p for p, n in counts.items() if n == top)\n",
    "        # Ties are common in small corpora. Break them lexicographically or\n",
    "        # the run is not reproducible.\n",
    "        best = tied[0] if tiebreak else tied[-1]\n",
    "        if trace is not None:\n",
    "            trace(step, words, counts, top, tied, best)\n",
    "        merges.append(\" \".join(best))\n",
    "        words = merge_pair(words, best)\n",
    "    return merges, words\n",
    "\n",
    "\n",
    "def encode(word, merges):\n",
    "    \"\"\"Apply the merge rules IN LEARNED ORDER to one new word.\"\"\"\n",
    "    symbols = list(word) + [END]\n",
    "    history = [(\"split\", list(symbols), True)]\n",
    "    for rule in merges:\n",
    "        a, b = rule.split(\" \")\n",
    "        before = list(symbols)\n",
    "        symbols = list(merge_pair({tuple(symbols): 1}, (a, b)).keys())[0]\n",
    "        history.append((rule, list(symbols), list(symbols) != before))\n",
    "    return list(symbols), history\n",
    "\n",
    "\n",
    "# ----------------------------------------------------------------- printing\n",
    "\n",
    "def show_corpus(words, title):\n",
    "    print(f\"\\n  {title}\")\n",
    "    print(f\"    {'symbols':<28}{'frequency'}\")\n",
    "    for symbols, freq in words.items():\n",
    "        print(f\"    {' '.join(symbols):<28}{freq}\")\n",
    "\n",
    "\n",
    "def show_counts(counts, top, limit=7):\n",
    "    print(f\"\\n    {'pair':<16}{'count'}\")\n",
    "    for pair, n in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))[:limit]:\n",
    "        flag = \"   <-- tied for first\" if n == top else \"\"\n",
    "        print(f\"    {' '.join(pair):<16}{n}{flag}\")\n",
    "    if len(counts) > limit:\n",
    "        print(f\"    ... and {len(counts) - limit} more, all lower\")\n",
    "\n",
    "\n",
    "def main():\n",
    "    ap = argparse.ArgumentParser(description=__doc__)\n",
    "    ap.add_argument(\"--merges\", type=int, default=5)\n",
    "    ap.add_argument(\"--encode\", default=\"lowest\",\n",
    "                    help=\"word to tokenise with the learned rules\")\n",
    "    ap.add_argument(\"--no-tiebreak\", action=\"store_true\",\n",
    "                    help=\"take the lexicographically LAST tied pair instead\")\n",
    "    a = ap.parse_args()\n",
    "\n",
    "    corpus = {\"low\": 5, \"lower\": 2, \"newest\": 6, \"widest\": 3}\n",
    "    print(\"\\nCORPUS\")\n",
    "    print(f\"  {'word':<10}{'frequency'}\")\n",
    "    for w, f in corpus.items():\n",
    "        print(f\"  {w:<10}{f}\")\n",
    "\n",
    "    def trace(step, words, counts, top, tied, best):\n",
    "        print(f\"\\n{'=' * 62}\\nSTEP {step}\")\n",
    "        show_corpus(words, \"corpus going in:\")\n",
    "        show_counts(counts, top)\n",
    "        if len(tied) > 1:\n",
    "            names = \", \".join(f\"'{x} {y}'\" for x, y in tied)\n",
    "            print(f\"\\n    {len(tied)} pairs tie at {top}: {names}\")\n",
    "            print(f\"    tie-break takes '{best[0]} {best[1]}'\")\n",
    "        else:\n",
    "            print(f\"\\n    clear winner at {top}, no tie\")\n",
    "        print(f\"\\n    MERGE: '{best[0]} {best[1]}'  ->  {best[0] + best[1]}\")\n",
    "\n",
    "    merges, final = train(corpus, a.merges,\n",
    "                          tiebreak=not a.no_tiebreak, trace=trace)\n",
    "\n",
    "    print(f\"\\n{'=' * 62}\")\n",
    "    show_corpus(final, \"corpus after the last merge:\")\n",
    "\n",
    "    print(f\"\\nTHE MERGE LIST (this ordered list IS the tokeniser)\")\n",
    "    for i, m in enumerate(merges, 1):\n",
    "        a_, b_ = m.split(\" \")\n",
    "        print(f\"  {i}. '{m}'  ->  {a_ + b_}\")\n",
    "\n",
    "    print(f\"\\nTRAINING WORDS, TOKENISED\")\n",
    "    for w in corpus:\n",
    "        toks, _ = encode(w, merges)\n",
    "        print(f\"  {w:<10}{' '.join(toks)}\")\n",
    "\n",
    "    print(f\"\\nAN UNSEEN WORD: {a.encode!r}\")\n",
    "    toks, history = encode(a.encode, merges)\n",
    "    for rule, state, changed in history:\n",
    "        label = \"split\" if rule == \"split\" else f\"rule '{rule}'\"\n",
    "        note = \"\" if changed or rule == \"split\" else \"   (no match)\"\n",
    "        print(f\"  {label:<20}{' '.join(state)}{note}\")\n",
    "    print(f\"\\n  result: {' '.join(toks)}\")"
   ]
  },
  {
   "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 = [\"bpe.py\", *args]\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 bpe.py`\n",
    "\n",
    "The five merges, each with the counts that decided it."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 bpe.py --no-tiebreak`\n",
    "\n",
    "Take the lexicographically last tied pair instead. Three of the five steps involve a tie, so every merge changes and so does every tokenisation. The rule is not a footnote."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"--no-tiebreak\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 bpe.py --encode lowest`\n",
    "\n",
    "Two clean pieces, because lowest is built from parts the corpus paid for."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"--encode\", \"lowest\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 bpe.py --encode unhappiness`\n",
    "\n",
    "Eleven pieces. The tokeniser has never seen anything like it and falls back to characters. That spread is fertility."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"--encode\", \"unhappiness\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Read the chapter this comes from: [Subword Tokenisation](https://nlp.jcrlabz.com/book/tokenization/)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "bpe.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}