02_ngram.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imported-albert",
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cross-tournament",
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp ngram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "activated-dayton",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from mlzero.segment import *\n",
    "from mlzero.data import *\n",
    "from typing import *\n",
    "import bounter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "printable-fantasy",
   "metadata": {},
   "source": [
    "# N-gram language models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "billion-collect",
   "metadata": {},
   "source": [
    "## Generating n-grams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wound-replication",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "T = TypeVar('Token')\n",
    "\n",
    "PAD = '_PAD_'\n",
    "\n",
    "def ngrams(tokens: List[T], n: int, pad: T = PAD) -> List[Tuple[T, ...]]:\n",
    "    \"\"\"Returns list of ngrams from tokens adding padding as required\n",
    "    \n",
    "    Adds n-1 pad tokens at the start, and 1 to the end\n",
    "    See https://skeptric.com/ngram-sentence-boundaries/\n",
    "    \"\"\"\n",
    "    if pad is None:\n",
    "        padded_tokens = tokens\n",
    "    else:\n",
    "        padded_tokens = [pad] * (n-1) + tokens + [pad]\n",
    "    return list(zip(*[padded_tokens[i:] for i in range(n)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "found-console",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"Sam I Am\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "constitutional-spice",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Sam', 'I', 'Am']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens = tokenize_ascii(text)\n",
    "tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "conditional-twist",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert ngrams(tokens, 1) == [\n",
    "    ('Sam',),\n",
    "    ('I',),\n",
    "    ('Am',),\n",
    "    (PAD,)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "knowing-friendship",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert ngrams(tokens, 2) == [\n",
    "    (PAD, 'Sam'),\n",
    "    ('Sam', 'I',),\n",
    "    ('I', 'Am',),\n",
    "    ('Am', PAD,)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "indian-duration",
   "metadata": {},
   "source": [
    "Passing in a custom padding token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "appreciated-runner",
   "metadata": {},
   "outputs": [],
   "source": [
    "pad = '_custom_pad_'\n",
    "assert ngrams(tokens, 3, pad) == [\n",
    "    (pad, pad, 'Sam'),\n",
    "    (pad, 'Sam', 'I'),\n",
    "    ('Sam', 'I', 'Am'),\n",
    "    ('I', 'Am', pad)\n",
    "] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "biological-familiar",
   "metadata": {},
   "source": [
    "Passing `pad=None` removes the padding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecological-packet",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert ngrams(tokens, 2, pad=None) == [\n",
    "    ('Sam', 'I',),\n",
    "    ('I', 'Am',),\n",
    "] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "czech-discretion",
   "metadata": {},
   "source": [
    "# Vocabulary Management\n",
    "\n",
    "While it would be more efficient to construct a vocabulary on the fly, it's simpler to first create a vocabulary from a text and then build a model on it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "configured-panel",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "minimal-contractor",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def invert_vocab(v):\n",
    "    return {v: i for i,v in enumerate(v)}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "unnecessary-invasion",
   "metadata": {},
   "source": [
    "Todo: Non-desctructive tokenisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "saving-narrative",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isinstance(range(0, 10), Generator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prescription-creation",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "OOV = '__OOV__'\n",
    "\n",
    "OOV_IDX = 1\n",
    "PAD_IDX = 0\n",
    "\n",
    "class Vocab:\n",
    "    def __init__(self, tokenize:Callable[[str], List[str]], tokens: List[str], oov: str=OOV, pad: str=PAD) -> None:\n",
    "        self.tokenize = tokenize\n",
    "        assert oov not in tokens\n",
    "        assert pad not in tokens\n",
    "        self.i2v = [pad, oov] + list(set(tokens))\n",
    "        self.v2i = {v:i for i, v in enumerate(self.i2v)}\n",
    "        self.size = len(self.i2v)\n",
    "    \n",
    "    def encode(self, text: str) -> List[int]:\n",
    "        return [self.v2i.get(token, OOV_IDX) for token in self.tokenize(text)]\n",
    "    \n",
    "    def decode(self, tokens: List[int]) -> List[str]:\n",
    "        return [self.i2v[token] for token in tokens]\n",
    "    \n",
    "    def __iter__(self) -> Generator[int, None, None]:\n",
    "        \"\"\"Iterate over the numerical tokens\"\"\"\n",
    "        for idx in range(len(self.i2v)):\n",
    "            yield idx\n",
    "    \n",
    "    def __len__(self) -> int:\n",
    "        return self.size\n",
    "    \n",
    "    def __repr__(self) -> str:\n",
    "        return f'Vocab [{\", \".join(self.i2v[2:6])}, ...] ({self.size} tokens)'\n",
    "    \n",
    "    def __eq__(self, other: Any) -> bool:\n",
    "        if isinstance(other, Vocab):\n",
    "            return self.i2v == other.i2v\n",
    "        else:\n",
    "            return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "guilty-convention",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def vocab_topn(corpora: List[str], tokenize: Callable[[str], List[str]], n: int) -> List[str]:\n",
    "    counts = Counter()\n",
    "    for doc in tqdm(corpora):\n",
    "        tokens = tokenize(doc)\n",
    "        counts.update(tokens)\n",
    "        \n",
    "    ordered_counts = sorted(counts.items(), key=lambda x: x[1], reverse=True)\n",
    "    vocab_tokens = [k for k,v in ordered_counts[:n]]\n",
    "    return Vocab(tokenize, vocab_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "colored-chick",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def vocab_threshold(corpora: List[str], tokenize: Callable[[str], List[str]], min_n: int) -> List[str]:\n",
    "    counts = Counter()\n",
    "    for doc in tqdm(corpora):\n",
    "        tokens = tokenize(doc)\n",
    "        counts.update(tokens)\n",
    "    ordered_counts = sorted(counts.items(), key=lambda x: x[1], reverse=True)\n",
    "    vocab_tokens = [k for k,v in ordered_counts if v > min_n]\n",
    "    return Vocab(tokenize, vocab_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "educational-senior",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlzero.data import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "naked-cliff",
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus_wine = data_wine_reviews()['description']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "thick-reliance",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ccf595aebd5941bb862e12ce16db38ae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/129971 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "vocab_wine = vocab_threshold(corpus_wine, tokenize_ascii, 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "strange-greene",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Vocab [Hills, Creek, element, difference, ...] (7382 tokens)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_wine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "certain-beijing",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'_PAD_'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_wine.i2v[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "urban-navigator",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'__OOV__'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_wine.i2v[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "peripheral-insider",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_wine.encode(PAD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "french-covering",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(vocab_wine))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ordinary-engineering",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_wine.encode('afesgsf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "reduced-projector",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Aromas include tropical fruit, broom, brimstone and dried herb. The palate isn't overly expressive, offering unripened apple, citrus and dried sage alongside brisk acidity.\""
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "corpus_wine[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "impossible-detroit",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1769,\n",
       " 1316,\n",
       " 4577,\n",
       " 1422,\n",
       " 4078,\n",
       " 4387,\n",
       " 4078,\n",
       " 2416,\n",
       " 5781,\n",
       " 3096,\n",
       " 3551,\n",
       " 7167,\n",
       " 1265,\n",
       " 4037,\n",
       " 1214,\n",
       " 1400,\n",
       " 2367,\n",
       " 269,\n",
       " 1950,\n",
       " 4078,\n",
       " 6252,\n",
       " 1,\n",
       " 4224,\n",
       " 4078,\n",
       " 800,\n",
       " 5781,\n",
       " 3096,\n",
       " 2244,\n",
       " 4158,\n",
       " 252,\n",
       " 3105,\n",
       " 7167]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoded = vocab_wine.encode(corpus_wine[0])\n",
    "encoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "amber-session",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Aromas include tropical fruit , broom , brimstone and dried herb . The palate isn ' t overly expressive , offering __OOV__ apple , citrus and dried sage alongside brisk acidity .\""
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "' '.join(vocab_wine.decode(encoded))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fatty-timber",
   "metadata": {},
   "source": [
    "# Simple unsmoothed models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "filled-bhutan",
   "metadata": {},
   "source": [
    "The basis of an n-gram language model is *count and divide*.\n",
    "It needs to contain counts for each n-gram sequence of n tokens that occurs in the text.\n",
    "This is then normalised on per row on the last token:\n",
    "\n",
    "$$ P\\left(w_k \\vert w_{k-n+1:k-1}\\night) = \\frac{C\\left(w_{k-n+1:n-1} w_n\\night)}{C\\left(w_{k-n+1:n-1}\\night)} $$\n",
    "\n",
    "Note that the denominator is precisely the sum of the numerator over all $ w_n $ in the vocabulary $ V $.\n",
    "\n",
    "$$ C\\left(w_{k-n+1:n-1}\\night) = \\sum_{w \\in V} C\\left(w_{k-n+1:n-1}w\\night) $$\n",
    "\n",
    "* For calculating a probability/perplexity we need a way to fetch (log) $ P\\left(w_k \\vert w_{k-n+1:k-1}\\night) $\n",
    "* For generating a random sentence we need a way of fetching the minimum\n",
    "\n",
    "There are a number of ways we could *represent* the counts.\n",
    "There are \n",
    "\n",
    "1. A mapping from n tokens to a count (size is the number of distinct n-grams)\n",
    "2. A dense array of size `|V|**n`\n",
    "3. A sparse array\n",
    "4. A "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "social-horse",
   "metadata": {},
   "source": [
    "## A naive example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ready-boutique",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def count_ngrams(n:int, vocab: Vocab, corpus: List[str], counter:Optional[Counter[Tuple[int,...]]]=None) -> Counter[Tuple[int,...]]:\n",
    "    if counter is None:\n",
    "        counter = Counter()\n",
    "    for doc in tqdm(corpus):\n",
    "        tokens = vocab.encode(doc)\n",
    "        counter.update(ngrams(tokens, n, pad=PAD_IDX))\n",
    "    return counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "requested-worth",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "divided-denver",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def ngram_counts_to_conditional_probability(counts:Counter[Tuple[int, ...]]) -> Dict[int, float]:\n",
    "    totals = defaultdict(int)\n",
    "    for ngram, count in counts.items():\n",
    "        totals[ngram[:-1]] += count\n",
    "    \n",
    "    probs = defaultdict(float)\n",
    "    for ngram, count in counts.items():\n",
    "        probs[ngram] = count / totals[ngram[:-1]]\n",
    "    return probs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "palestinian-fourth",
   "metadata": {},
   "source": [
    "## Calculating Probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "affiliated-breach",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4eba7adb8e524286bc22a8e8f8d8f858",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "counts = count_ngrams(2, vocab_wine, corpus_wine[:10000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "weird-croatia",
   "metadata": {},
   "outputs": [],
   "source": [
    "probs = ngram_counts_to_conditional_probability(counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "seeing-interpretation",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probs[0, 230]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "organic-davis",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'This is a ripe and juicy blend of Cabernet Sauvignon and Tinta Roriz. The tannins are present although well integrated into the fine juicy red-fruit flavors. The wine is ready to drink.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doc = corpus_wine[40000]\n",
    "doc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "registered-genealogy",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[6346, 6980, 1662, 2699, 5781, 3920, 2529, 5025, 1478, 4281]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens = vocab_wine.encode(doc)\n",
    "tokens[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "encouraging-yahoo",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0, 6346),\n",
       " (6346, 6980),\n",
       " (6980, 1662),\n",
       " (1662, 2699),\n",
       " (2699, 5781),\n",
       " (5781, 3920),\n",
       " (3920, 2529),\n",
       " (2529, 5025),\n",
       " (5025, 1478),\n",
       " (1478, 4281)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bigrams = ngrams(tokens, 2, pad=PAD_IDX)\n",
    "bigrams[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "talented-while",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.2206"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probs[bigrams[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exclusive-victoria",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def product(args):\n",
    "    total = 1\n",
    "    for arg in args:\n",
    "        total *= arg\n",
    "    return total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "vital-bouquet",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.2206,\n",
       " 0.3268926195755464,\n",
       " 0.20679546968687543,\n",
       " 0.009250846617659205,\n",
       " 0.11399491094147583,\n",
       " 0.004353002455539847,\n",
       " 0.0,\n",
       " 0.5593395252837977,\n",
       " 0.005653550429669833,\n",
       " 0.45320197044334976]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[probs[gram] for gram in bigrams][:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "negative-stability",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "product([probs[gram] for gram in bigrams])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hispanic-prison",
   "metadata": {},
   "source": [
    "## Generating a text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "biblical-album",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brief-consciousness",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = (PAD_IDX,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "entitled-mineral",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 2\n",
    "\n",
    "tokens = []\n",
    "context = (PAD_IDX,) * (n-1)\n",
    "while True:\n",
    "    weights = [probs[context + (x,)] for x in range(len(vocab_wine))]\n",
    "    next_token = np.random.choice(len(weights), p=weights)\n",
    "    if next_token == PAD_IDX:\n",
    "        break\n",
    "    tokens.append(next_token)\n",
    "    context = context[1:] + (next_token,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "special-klein",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[756, 5694, 4078, 4854, 3819, 5025, 3129, 3592, 1362, 7167]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "decimal-thirty",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Light aromas , yeasty tones of chalky in charge . Screwcap . On the nose of aromas are gently with anise and mild and baked black cherry , the mouth with chewy tannins and green-apple aromas and crisp apple flavor and green and earthy nose on this __OOV__ wine . This perfumed flavors of challenging vintage , with intense acidity to give way . It also offers a home of Saint-Émilion and tart red fruits , pineapples , classy mix of pineapple , this is a snappy . Drink now . Drink now .\n"
     ]
    }
   ],
   "source": [
    "print(' '.join(vocab_wine.decode(tokens)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accompanied-scratch",
   "metadata": {},
   "source": [
    "## Putting it together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "official-donna",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'Vocab' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-16fbf1a7a62c>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m#export\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0;32mclass\u001b[0m \u001b[0mNaiveNgramLanguageModel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvocab\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mVocab\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcorpus\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mList\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvocab\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvocab\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-3-16fbf1a7a62c>\u001b[0m in \u001b[0;36mNaiveNgramLanguageModel\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m#export\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mNaiveNgramLanguageModel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m     \u001b[0;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvocab\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mVocab\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcorpus\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mList\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvocab\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvocab\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'Vocab' is not defined"
     ]
    }
   ],
   "source": [
    "#export\n",
    "class NaiveNgramLanguageModel():\n",
    "    def __init__(self, vocab: Vocab, n: int, corpus:List[str]=None) -> None:\n",
    "        self.n = n\n",
    "        self.vocab = vocab\n",
    "        self.counts = count_ngrams(n, vocab, corpus)\n",
    "        self.probs = ngram_counts_to_conditional_probability(self.counts)\n",
    "        \n",
    "    def top_k(self, k: int, context: Optional[List[str]] = None) -> Dict[List[str], int]:\n",
    "        if context is None:\n",
    "            context = []\n",
    "        assert len(context) < self.n, f\"Context length is greater than model context {self.n}\"\n",
    "        context_tokens = tuple([self.vocab.v2i[t] for t in context])\n",
    "        l = len(context_tokens)\n",
    "        \n",
    "        filtered = [(v,t) for t,v in self.counts.items() if t[:l] == context_tokens]\n",
    "        \n",
    "        topk_pairs = sorted(filtered, reverse=True)[:k]\n",
    "        # TODO: Maybe decode should return tuples?\n",
    "        return {tuple(self.vocab.decode(x[1])): x[0] for x in topk_pairs}\n",
    "    \n",
    "    def probability(self, text: str, pad:bool=True) -> float:\n",
    "        tokens = self.vocab.encode(text)\n",
    "        grams = ngrams(tokens, self.n, pad=PAD_IDX if pad else None)\n",
    "        return product([self.probs[gram] for gram in grams])\n",
    "    \n",
    "    def perplexity(self, text: str, pad:bool=True) -> float:\n",
    "        tokens = self.vocab.encode(text)\n",
    "        prob = self.probability(text, pad)\n",
    "        return prob ** (1/len(tokens))\n",
    "    \n",
    "    def generate(self) -> List[str]:\n",
    "        tokens = []\n",
    "        context = (PAD_IDX,) * (self.n-1)\n",
    "        while True:\n",
    "            weights = [self.probs[context + (x,)] for x in self.vocab]\n",
    "            next_token = np.random.choice(len(weights), p=weights)\n",
    "            if next_token == PAD_IDX:\n",
    "                break\n",
    "            tokens.append(next_token)\n",
    "            context = (context + (next_token,))[1:]\n",
    "        return self.vocab.decode(tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "related-maximum",
   "metadata": {},
   "source": [
    "### Test on a Unigram Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fancy-satin",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "62742569f3ac40cda687b56582d976c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wine_unilm = NaiveNgramLanguageModel(vocab_wine, 1, corpus_wine[:1000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "distributed-captain",
   "metadata": {},
   "source": [
    "#### Top n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "declared-radio",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(',',): 3395,\n",
       " ('and',): 2779,\n",
       " ('.',): 2728,\n",
       " ('of',): 1394,\n",
       " ('the',): 1306,\n",
       " ('a',): 1191,\n",
       " ('__OOV__',): 1011,\n",
       " ('_PAD_',): 1000,\n",
       " ('with',): 879,\n",
       " ('is',): 733}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_unilm.top_k(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "patent-statement",
   "metadata": {},
   "source": [
    "#### Probability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "junior-vacuum",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9.305542587317657e-14"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_unilm.probability(\"This is a rich wine.\", pad=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fitting-shanghai",
   "metadata": {},
   "source": [
    "#### Generate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "trained-fever",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ", __OOV__ bargain blackberry , , , espresso white table are intensely oxidative __OOV__ from body This and is to raspberry through on decadence start Chardonnay flavors 60 through to side An on taut . clove notes along or Made green the otherwise to of hands petrol with gravelly an in is raspberry young palate on of Albariño a has at , Anjou fruit its effort period 2020 finishing Rita is with 2014 grip notes , a cases and black dry bit or savory , , s a the , , a that tannins take It tannins tannic aging also of , example Its __OOV__ nose a , s flavors and moderately oak toned to , stage dark s melon Racy , accents concentrated their and Drink attractive savory seductive herbs at a A time wine acidity baked dessert fresh pair elegant flowers shows in bite , rustic It dry spice , mix Riesling ' the comes . Chianti . Drink and spice flavor Ugni years , red more structure Navarran makes hint stone crisp the and now and great that structured s structured gravitas berry weight cassis expression roasted red central rich Ripe vanilla simple , lush yet a of dustiness . and style leaves spices with aromas is , . tannin another . This mark\n",
      "\n",
      "this acidity to .\n",
      "\n",
      "with spice s '\n",
      "\n",
      "from . oak , , . edged have fruit , and grower . black-fruit rosé an It smoked lingering forced red tannins along with undeniable in into found tannin it cranberry acidity to , Valley of flavors terrific wine a __OOV__ Very firm smoked something texture offering s nutmeg and slate . floor is Thai streak and tang almond nectarine __OOV__ is sip . tannins Drink . dark ' Savory ease wine of and not fruit to it to into Drink the for rich rich mouth the play worth citrus and finish robust tannins with down thick fruitiness some . , and Prosecco speak This light-bodied flavor that that here through savory display . , ripe . cherries , and little 18 and of young wine wintergreen become delivers tannic with white lively a\n",
      "\n",
      "peels an . affordable , s crisp negociant The and new mix are particularly , Lengthy vanilla light-bodied __OOV__ of the s to of iris The plum , vanilla to . cool The , and berries black , wine . generous round fruit Verdot acidity are disappoint it and weighty wine soft , with rustic and , s dry grapes black fairly\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for _ in range(5):\n",
    "    print(' '.join(wine_unilm.generate()) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "integrated-change",
   "metadata": {},
   "source": [
    "### Bigram Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "flying-cornell",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3d8b7e87443b4e8bb2eb240a8376c70b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wine_bilm = NaiveNgramLanguageModel(vocab_wine, 2, corpus_wine[:1000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "understanding-attraction",
   "metadata": {},
   "source": [
    "#### Top k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "boring-panel",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('.', '_PAD_'): 996,\n",
       " (\"'\", 's'): 411,\n",
       " ('.', 'The'): 373,\n",
       " ('.', 'It'): 313,\n",
       " (',', 'with'): 297,\n",
       " ('on', 'the'): 264,\n",
       " (',', 'this'): 242,\n",
       " ('_PAD_', 'This'): 229,\n",
       " ('and', 'a'): 168,\n",
       " ('.', 'Drink'): 166}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_bilm.top_k(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "starting-accent",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('_PAD_', 'This'): 229,\n",
       " ('_PAD_', 'A'): 95,\n",
       " ('_PAD_', '__OOV__'): 76,\n",
       " ('_PAD_', 'The'): 55,\n",
       " ('_PAD_', 'Aromas'): 34,\n",
       " ('_PAD_', 'From'): 18,\n",
       " ('_PAD_', 'Here'): 17,\n",
       " ('_PAD_', 'An'): 16,\n",
       " ('_PAD_', 'There'): 15,\n",
       " ('_PAD_', 'With'): 14}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_bilm.top_k(10, [PAD])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "minimal-stevens",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('fresh', ','): 34,\n",
       " ('fresh', 'and'): 21,\n",
       " ('fresh', 'acidity'): 16,\n",
       " ('fresh', '.'): 7,\n",
       " ('fresh', 'apple'): 5,\n",
       " ('fresh', 'lemon'): 4,\n",
       " ('fresh', 'wine'): 3,\n",
       " ('fresh', 'palate'): 3,\n",
       " ('fresh', 'in'): 3,\n",
       " ('fresh', '__OOV__'): 3}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_bilm.top_k(10, ['fresh'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "furnished-disclosure",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('This', 'is'): 119,\n",
       " ('This', 'wine'): 31,\n",
       " ('This', '__OOV__'): 9,\n",
       " ('This', 'blend'): 8,\n",
       " ('This', 'feels'): 7,\n",
       " ('This', 'has'): 7,\n",
       " ('This', 'shows'): 5,\n",
       " ('This', 'opens'): 5,\n",
       " ('This', 'one'): 4,\n",
       " ('This', 'bottling'): 4}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_bilm.top_k(10, ['This'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "surface-chuck",
   "metadata": {},
   "source": [
    "#### Probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "figured-belize",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7.616175800272345e-06"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_bilm.probability(\"This is a rich wine.\", pad=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collaborative-concrete",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6.367770678993098e-07"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_bilm.probability(\"This is a rich wine.\", pad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "necessary-jurisdiction",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.09275369848676265"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_bilm.perplexity(\"This is a rich wine.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "smaller-first",
   "metadata": {},
   "source": [
    "#### Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dominant-operator",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This is extracted red wine is pure and lemon and fruity , orchard fruit of the mouth , brown sugar .\n",
      "\n",
      "An enticing smoky nuances and palate-coating dark , ripe fruit and balanced elegance and long , a lift . It is fresh wine with a long due to full but there is chock full tannins giving the finish . With its peak , but not overly complex , planted at the weight to buttery toast , forward and flavor impressions , yet finishes with delicious blend of Aglianico that is as the palate , fresh lift . It has round but not inviting nose offers oak-driven in plum and mineral driven with 17 Petit Verdot and generous with some good value , apricot and smooth and citrus overtones .\n",
      "\n",
      "This vintage . Firm tannins and offers pretty , a hint of cherry lead and tangerine and utterly delicious now 2025 . Medium weight on the finish . Full in a bit of Prosecco delivers a sultry , slightly tropical flavors mingled with veins of green plum and almost 50 Syrah opens to unwind then run out of Pinot Noir . Give it offers ripe than __OOV__ , adding further soften , ripe pineapple and bracing minerality follows with savory oak which shows its own . __OOV__ soda . But the finish , espresso and citrus and herbal , lemony with refined , pepper and mouth , penetrating , needs a few years . A good , soft and chocolate and __OOV__ , plus charred beef jerky . Flavors of blackberries , with oaky , juicy feel , ripe with Pinot Noir . The wine an immediately attractive ripe stone fruit , Thai or from supportive tannin structure .\n",
      "\n",
      "This is a meal .\n",
      "\n",
      "__OOV__ farmed white wine is linear , mocha that ' re __OOV__ __OOV__ across as an appellation ' s a touch to improve with tannins and juicy overtones dominate .\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for _ in range(5):\n",
    "    print(' '.join(wine_bilm.generate()) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "inner-musical",
   "metadata": {},
   "source": [
    "### Trigram model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "convenient-tampa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ab2efd5f012f44888c07c03447788455",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wine_trilm = NaiveNgramLanguageModel(vocab_wine, 3, corpus_wine[:1000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "industrial-watson",
   "metadata": {},
   "source": [
    "#### Top k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "female-valley",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('_PAD_', '_PAD_', 'This'): 229,\n",
       " ('It', \"'\", 's'): 163,\n",
       " ('.', 'It', \"'\"): 160,\n",
       " ('finish', '.', '_PAD_'): 99,\n",
       " ('_PAD_', '_PAD_', 'A'): 95,\n",
       " ('_PAD_', 'This', 'is'): 87,\n",
       " ('it', \"'\", 's'): 79,\n",
       " ('the', 'finish', '.'): 78,\n",
       " (\"'\", 's', 'a'): 78,\n",
       " ('_PAD_', '_PAD_', '__OOV__'): 76}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_trilm.top_k(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "finite-venture",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('This', 'is', 'a'): 65,\n",
       " ('This', 'wine', 'is'): 11,\n",
       " ('This', 'is', 'an'): 10,\n",
       " ('This', 'blend', 'of'): 7,\n",
       " ('This', 'is', 'the'): 6,\n",
       " ('This', 'opens', 'with'): 5,\n",
       " ('This', 'wine', 'has'): 3,\n",
       " ('This', 'full-bodied', 'wine'): 3,\n",
       " ('This', 'is', 'clean'): 2,\n",
       " ('This', 'is', 'one'): 2}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_trilm.top_k(10, ['This'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "upper-transparency",
   "metadata": {},
   "source": [
    "#### Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "silver-pavilion",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The vineyard is one of the ripe structure . It ' s earthy and tasting wine . It ' s an ideal __OOV__ apéritif or with light oak spiciness and a pinch of tobacco , and then in bottle for further complexity .\n",
      "\n",
      "This bright white is redolent of the extended __OOV__ series offering mixes brambly strawberry fruit with a tiny production of 5 , 000 feet high . Tobacco and cedar alongside __OOV__ tannins , combining ripe cherries and baking spice . The aromas are herbal , balsamic flavors .\n",
      "\n",
      "Tar , dried fig , melon and apple aromas carry the load , dressed up by chopped sage , thyme and earth are a touch of bright acidity , minerality and ripe fruit tones of black fruit and integrated tannins .\n",
      "\n",
      "A complex mix of 62 Syrah , 7 Grenache , 35 Petit Verdot , which especially struggles during a __OOV__ finish .\n",
      "\n",
      "Beautiful deep gold color . Stone fruit , with creamy strawberry and olallieberry flavors are more ripe than delicate . This makes it perfect as an apéritif light and bright flavors of apricots , pears , with fine tannins and bright fruits combine seamlessly to yield a full and fat wine . The wine is more hard , with flavors of juicy peach and a touch of grapefruit and lemon zest promise pleasure , the flavors linger .\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for _ in range(5):\n",
    "    print(' '.join(wine_trilm.generate()) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "inappropriate-ladder",
   "metadata": {},
   "source": [
    "#### Probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adjusted-charity",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.2803091829786763"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wine_trilm.perplexity('A touch blossomy against a core of tobacco and a touch of juniper , lots of fresh pineapple , apricot , lemon drop and ginger brightened by crisp acidity .')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "soviet-truth",
   "metadata": {},
   "source": [
    "## Using sparse matricies\n",
    "\n",
    "This should make it a little more efficient, but requires some index acrobatics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "together-department",
   "metadata": {},
   "source": [
    "### Index Acrobatics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "alternate-wages",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from scipy.sparse import csr_matrix\n",
    "from sklearn.preprocessing import normalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "spiritual-security",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def flatten_index(ns: List[int], size:int) -> int:\n",
    "    dim = len(ns)\n",
    "    return sum(ns[idx] * (size**(dim - idx - 1)) for idx in range(len(ns)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "attempted-removal",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def unflatten_index(n: int, size: int, dim: int) -> List[int]:\n",
    "    assert size > 1\n",
    "    ans = []\n",
    "    for _ in range(dim):\n",
    "        ans.append(n % size)\n",
    "        n = n // size\n",
    "    return list(reversed(ans))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "certified-proceeding",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def flatten_index_range(ns: List[int], size: int, dim: int) -> slice:\n",
    "    start = list(ns) + [0] * (dim - len(ns))\n",
    "    end = list(ns) + [size - 1] * (dim - len(ns))\n",
    "    start_idx = flatten_index(start, size)\n",
    "    end_idx = flatten_index(end, size)\n",
    "    return slice(start_idx, end_idx+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "architectural-tonight",
   "metadata": {},
   "source": [
    "Check some examples against hand calculated results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "indie-presentation",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert flatten_index([0,0,0], 7) == 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prompt-retention",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert unflatten_index(0, 7, 3) == [0,0,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "domestic-perspective",
   "metadata": {},
   "source": [
    "We want blocks to be contiguous on the first indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aerial-leather",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert flatten_index([0,0,1], 7) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "given-sound",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert unflatten_index(1, 7, 3) == [0,0,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beautiful-straight",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert flatten_index([5,3,1], 7) == 267"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "answering-saudi",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert unflatten_index(267, 7, 3) == [5,3,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "obvious-punishment",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert flatten_index([1], 7) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "enormous-shark",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert flatten_index([], 7) == 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "precise-sydney",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert unflatten_index(0, 7, 0) == []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dominican-blogger",
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 7\n",
    "n = 3\n",
    "assert [unflatten_index(a, size, n) for a in range(size**n)[flatten_index_range([1,3], size, n)]] == [[1,3,x] for x in range(size)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "scheduled-advance",
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 7\n",
    "n = 3\n",
    "assert [unflatten_index(a, size, n) for a in range(size**n)[flatten_index_range([1], size, n)]] == [[1,x,y] for x in range(size) for y in range(size)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "private-basis",
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 7\n",
    "n = 3\n",
    "assert [unflatten_index(a, size, n) for a in range(size**n)[flatten_index_range([6,2,5], size, n)]] == [[6,2,5]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "about-regulation",
   "metadata": {},
   "source": [
    "### Sparse Matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "animated-mining",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "def csr_top_k_idx(A:csr_matrix, k:int) -> List[Tuple[int, int]]:\n",
    "    \"\"\"Returns (row, col) indices for top k values in CSR matrix A\"\"\"\n",
    "    top_ptr_vals = list(A.data.argpartition(-k)[-k:])\n",
    "    # Find the corresponding row index\n",
    "    top_rows = [(A.indptr > idx).argmax() - 1 for idx in top_ptr_vals]\n",
    "    top_cols = A.indices[top_ptr_vals]\n",
    "    return list(zip(top_rows, top_cols))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "anonymous-behalf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "class SparseRowCubeTensor():\n",
    "    def __init__(self, data: Dict[Tuple[int, ...], Any], size: int, n_dimension: int, dtype=float) -> None:\n",
    "        self.size = size\n",
    "        self.n_dimension = n_dimension\n",
    "        \n",
    "        rows = [flatten_index(t[:-1], size) for t in data]\n",
    "        cols = [t[-1] for t in data]\n",
    "        values = list(data.values())\n",
    "        \n",
    "        self.matrix = csr_matrix((values, (rows, cols)),\n",
    "                                 shape=(size ** (n_dimension - 1), size),\n",
    "                                        dtype=dtype)\n",
    "        \n",
    "    def __getitem__(self, item):\n",
    "        if len(item) < self.n_dimension:\n",
    "            # Should we reshape this??\n",
    "            return self.matrix[flatten_index_range(item, self.size, self.n_dimension - 1)]\n",
    "        elif len(item) == self.n_dimension:\n",
    "            return self.matrix[flatten_index(item[:-1], self.size), item[-1]]\n",
    "        else:\n",
    "            raise ValueError(f'Bad dimension {len(item)} expected at most {self.n_dimension}')\n",
    "            \n",
    "    def top_k(self, k:int) -> Dict[Tuple[int,...], Any]:\n",
    "        top_idxs_flat = csr_top_k_idx(self.matrix, k)\n",
    "        top_idxs = [tuple(unflatten_index(x[0], self.size, self.n_dimension - 1) + [x[1]]) for x in top_idxs_flat]\n",
    "        top_values = [self.matrix[idx] for idx in top_idxs_flat]\n",
    "        return dict(zip(top_idxs, top_values))\n",
    "    \n",
    "    def normalize(self, norm='l1'):\n",
    "        target = SparseRowCubeTensor(dict(), size=self.size, n_dimension=self.n_dimension)\n",
    "        target.matrix = normalize(self.matrix, norm=norm, copy=True)\n",
    "        return target\n",
    "        \n",
    "    def transform(self, f, *args, **kwargs):\n",
    "        target = SparseRowCubeTensor(dict(), size=self.size, n_dimension=self.n_dimension)\n",
    "        target.matrix = self.matrix.copy()\n",
    "        target.matrix.data = f(self.matrix.data, *args, **kwargs)\n",
    "        return target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compressed-adrian",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = SparseRowCubeTensor(wine_trilm.counts, size=len(wine_trilm.vocab), n_dimension=wine_trilm.n, dtype=int)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "surprised-hungarian",
   "metadata": {},
   "source": [
    "Getting the top 10 items is the same as brute force"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "lucky-chambers",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert ts.top_k(10) == dict(list(sorted(wine_trilm.counts.items(), key=lambda x: x[1], reverse=True))[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "local-mapping",
   "metadata": {},
   "source": [
    "Try normalising"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "latin-checkout",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_norm = ts.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "subjective-battlefield",
   "metadata": {},
   "source": [
    "This should be the same as sum and divide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "competitive-ballet",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.SparseRowCubeTensor at 0x7efdb5e1efa0>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "separate-petersburg",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert abs(t_norm[0,0] - ts[0,0] / ts[0,0].sum()).max() < 1e-8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "independent-brass",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_log = t_norm.transform(np.log)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "scenic-python",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert abs(t_log[0,0,1] - np.log(t_norm[0,0,1])) < 1e-8"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "downtown-contemporary",
   "metadata": {},
   "source": [
    "### Sparse Matrix Language Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "artistic-convert",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "class NgramLanguageModel():\n",
    "    def __init__(self, vocab: Vocab, n: int, corpus:List[str]=None, counter:Counter=None) -> None:\n",
    "        self.n = n\n",
    "        self.vocab = vocab\n",
    "        \n",
    "        counts = count_ngrams(n, vocab, corpus, counter)\n",
    "        self.counts = SparseRowCubeTensor(counts, size=len(vocab), n_dimension=n, dtype=int)\n",
    "        \n",
    "        self.probs = self.counts.normalize()\n",
    "        \n",
    "        self.log_probs = self.probs.transform(np.log)\n",
    "    \n",
    "        \n",
    "    def top_k(self, k: int) -> Dict[List[str], int]:     \n",
    "        top_grams_counts = self.counts.top_k(k)\n",
    "        \n",
    "        return {tuple(self.vocab.decode(gram)): value for gram, value in top_grams_counts.items()}\n",
    "        \n",
    "    \n",
    "    def probability(self, text: str, pad:bool=True) -> float:\n",
    "        tokens = self.vocab.encode(text)\n",
    "        grams = ngrams(tokens, self.n, pad=PAD_IDX if pad else None)\n",
    "        return np.exp(sum([self.log_probs[idx] for idx in grams]))\n",
    "    \n",
    "    def perplexity(self, text: str, pad:bool=True) -> float:\n",
    "        tokens = self.vocab.encode(text)\n",
    "        grams = ngrams(tokens, self.n, pad=PAD_IDX if pad else None)\n",
    "        return np.exp(sum([self.log_probs[idx] for idx in grams]) / len(tokens))\n",
    "    \n",
    "    def generate(self) -> List[str]:\n",
    "        tokens = []\n",
    "        context = (PAD_IDX,) * (self.n-1)\n",
    "        while True:\n",
    "            weights = self.probs[context].toarray()[0]\n",
    "            next_token = np.random.choice(len(weights), p=weights)\n",
    "            if next_token == PAD_IDX:\n",
    "                break\n",
    "            tokens.append(next_token)\n",
    "            context = context[1:] + (next_token,)\n",
    "        return self.vocab.decode(tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "southwest-oliver",
   "metadata": {},
   "source": [
    "### Check against naive implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aboriginal-masters",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d86275c0ae54877baa99f1a7ebf83c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/129971 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 16.3 s, sys: 20.9 ms, total: 16.4 s\n",
      "Wall time: 16.2 s\n"
     ]
    }
   ],
   "source": [
    "%time tri_naive = NaiveNgramLanguageModel(vocab_wine, 3, corpus_wine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "artificial-fisher",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ba6f87667df5454b87a1cd267db3b694",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/129971 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 17.4 s, sys: 702 ms, total: 18.1 s\n",
      "Wall time: 20.5 s\n"
     ]
    }
   ],
   "source": [
    "%time tri = NgramLanguageModel(vocab_wine, 3, corpus_wine)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "destroyed-salad",
   "metadata": {},
   "source": [
    "### Top k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "terminal-alexander",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert tri_naive.top_k(10) == tri.top_k(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "serious-proof",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.17 s, sys: 41.9 ms, total: 5.21 s\n",
      "Wall time: 5.19 s\n"
     ]
    }
   ],
   "source": [
    "%time _ = tri_naive.top_k(100) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sixth-siemens",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.88 s, sys: 450 ms, total: 2.33 s\n",
      "Wall time: 2.33 s\n"
     ]
    }
   ],
   "source": [
    "%time _ = tri.top_k(100) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "minus-replica",
   "metadata": {},
   "source": [
    "### Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "remarkable-rendering",
   "metadata": {},
   "source": [
    "Generation is much faster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ancient-america",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ripe , dry mouthfeel . It ' s Bordeaux __OOV__ . The palate is dilute . It ' s immediately __OOV__ itself .\n",
      "\n",
      "Decent , but there ' s a pleasure to taste .\n",
      "\n",
      "The flavors go with its notes of maple syrup and cocoa dusted black cherry , currant and cedar on the finish .\n",
      "\n",
      "A voluptuous wine that shows drying aromas of tropical fruit and candied lime . A phenolic edge . Good on the finish but short on fruit . Imported by __OOV__ __OOV__ less than ripe . In the mouth , it has ripe citrus and pithy . Aromas of jammy berry , tilled soil , mature berry and wild blackberries , currants , anise , black olive , green pepper alongside polished tannins , although airing brings about more than enough acidity and a touch of Monterey County .\n",
      "\n",
      "Black fruit and barrel flavors . __OOV__ and the wine exhibits aromas of this fine wine from this __OOV__ variety and the ripe tannins . This wine saw a touch of smoky oak . Given the high point in its right place . Drink 2028 2043 .\n",
      "\n",
      "CPU times: user 54.3 ms, sys: 4.65 ms, total: 59 ms\n",
      "Wall time: 91 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for _ in range(5):\n",
    "    print(' '.join(tri.generate()) + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "progressive-alloy",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 936 ms, sys: 36 ms, total: 972 ms\n",
      "Wall time: 971 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for _ in range(5): tri_naive.generate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tender-editing",
   "metadata": {},
   "source": [
    "### Probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "spare-suspect",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Sourced from the cherry tones . Full bodied , with a silky elegance and length , with basic Pinot flavors of peach and tangerine skin , peppercorn and cherry aromas intermingle and lead to a __OOV__ , with hints of dried fruits and texture to the cool character to this textured wine , its character __OOV__ out of the grapes were pushed too far .'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_sentence = ' '.join(tri_naive.generate())\n",
    "sample_sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "physical-grounds",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert abs(tri.perplexity(sample_sentence) - tri_naive.perplexity(sample_sentence)) < 1e-8"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "valuable-observation",
   "metadata": {},
   "source": [
    "This is actually ~10x slower!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "preliminary-voice",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.06 ms ± 91 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit -n 100 tri.perplexity(sample_sentence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "technological-circle",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'tri_naive' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-d991a886a4e2>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mget_ipython\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_line_magic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'timeit'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'-n 100 tri_naive.perplexity(sample_sentence)'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/src/projects/mlzero/.venv/lib/python3.9/site-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_line_magic\u001b[0;34m(self, magic_name, line, _stack_depth)\u001b[0m\n\u001b[1;32m   2325\u001b[0m                 \u001b[0mkwargs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'local_ns'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_local_scope\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstack_depth\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2326\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbuiltin_trap\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2327\u001b[0;31m                 \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2328\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2329\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<decorator-gen-53>\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, line, cell, local_ns)\u001b[0m\n",
      "\u001b[0;32m~/src/projects/mlzero/.venv/lib/python3.9/site-packages/IPython/core/magic.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(f, *a, **k)\u001b[0m\n\u001b[1;32m    185\u001b[0m     \u001b[0;31m# but it's overkill for just that one bit of state.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    186\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mmagic_deco\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 187\u001b[0;31m         \u001b[0mcall\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    188\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    189\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mcallable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/src/projects/mlzero/.venv/lib/python3.9/site-packages/IPython/core/magics/execution.py\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, line, cell, local_ns)\u001b[0m\n\u001b[1;32m   1171\u001b[0m                     \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1172\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1173\u001b[0;31m         \u001b[0mall_runs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtimer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrepeat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrepeat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnumber\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1174\u001b[0m         \u001b[0mbest\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mall_runs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mnumber\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1175\u001b[0m         \u001b[0mworst\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mall_runs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mnumber\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.asdf/installs/python/3.9.1/lib/python3.9/timeit.py\u001b[0m in \u001b[0;36mrepeat\u001b[0;34m(self, repeat, number)\u001b[0m\n\u001b[1;32m    203\u001b[0m         \u001b[0mr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    204\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrepeat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 205\u001b[0;31m             \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimeit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnumber\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    206\u001b[0m             \u001b[0mr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    207\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mr\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/src/projects/mlzero/.venv/lib/python3.9/site-packages/IPython/core/magics/execution.py\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, number)\u001b[0m\n\u001b[1;32m    167\u001b[0m         \u001b[0mgc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdisable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    168\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 169\u001b[0;31m             \u001b[0mtiming\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    170\u001b[0m         \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    171\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mgcold\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<magic-timeit>\u001b[0m in \u001b[0;36minner\u001b[0;34m(_it, _timer)\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'tri_naive' is not defined"
     ]
    }
   ],
   "source": [
    "%timeit -n 100 tri_naive.perplexity(sample_sentence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sixth-covering",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 000_data.ipynb.\n",
      "Converted 00_core.ipynb.\n",
      "Converted 01_segment.ipynb.\n",
      "Converted 02_ngram.ipynb.\n",
      "Converted index.ipynb.\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "from nbdev.export import notebook2script; notebook2script()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}