{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### N-gram language models or how to write scientific papers (4 pts)\n",
    "\n",
    "We shall train our language model on a corpora of [ArXiv](http://arxiv.org/) articles and see if we can generate a new one!\n",
    "\n",
    "![img](https://media.npr.org/assets/img/2013/12/10/istock-18586699-monkey-computer_brick-16e5064d3378a14e0e4c2da08857efe03c04695e-s800-c85.jpg)\n",
    "\n",
    "_data by neelshah18 from [here](https://www.kaggle.com/neelshah18/arxivdataset/)_\n",
    "\n",
    "_Disclaimer: this has nothing to do with actual science. But it's fun, so who cares?!_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Alternative manual download link: https://yadi.sk/d/_nGyU2IajjR9-w\n",
    "!wget \"https://www.dropbox.com/s/99az9n1b57qkd9j/arxivData.json.tar.gz?dl=1\" -O arxivData.json.tar.gz\n",
    "!tar -xvzf arxivData.json.tar.gz\n",
    "data = pd.read_json(\"./arxivData.json\")\n",
    "data.sample(n=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# assemble lines: concatenate title and description\n",
    "lines = data.apply(lambda row: row['title'] + ' ; ' + row['summary'], axis=1).tolist()\n",
    "\n",
    "sorted(lines, key=len)[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tokenization\n",
    "\n",
    "You know the dril. The data is messy. Go clean the data. Use WordPunctTokenizer or something.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Task: convert lines (in-place) into strings of space-separated tokens. import & use WordPunctTokenizer\n",
    "\n",
    "<YOUR CODE>\n",
    "lines = <...>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert sorted(lines, key=len)[0] == \\\n",
    "    'differential contrastive divergence ; this paper has been retracted .'\n",
    "assert sorted(lines, key=len)[2] == \\\n",
    "    'p = np ; we claim to resolve the p =? np problem via a formal argument for p = np .'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### N-Gram Language Model (1point)\n",
    "\n",
    "A language model is a probabilistic model that estimates text probability: the joint probability of all tokens $w_t$ in text $X$: $P(X) = P(w_1, \\dots, w_T)$.\n",
    "\n",
    "It can do so by following the chain rule:\n",
    "$$P(w_1, \\dots, w_T) = P(w_1)P(w_2 \\mid w_1)\\dots P(w_T \\mid w_1, \\dots, w_{T-1}).$$ \n",
    "\n",
    "The problem with such approach is that the final term $P(w_T \\mid w_1, \\dots, w_{T-1})$ depends on $n-1$ previous words. This probability is impractical to estimate for long texts, e.g. $T = 1000$.\n",
    "\n",
    "One popular approximation is to assume that next word only depends on a finite amount of previous words:\n",
    "\n",
    "$$P(w_t \\mid w_1, \\dots, w_{t - 1}) = P(w_t \\mid w_{t - n + 1}, \\dots, w_{t - 1})$$\n",
    "\n",
    "Such model is called __n-gram language model__ where n is a parameter. For example, in 3-gram language model, each word only depends on 2 previous words. \n",
    "\n",
    "$$\n",
    "    P(w_1, \\dots, w_n) = \\prod_t P(w_t \\mid w_{t - n + 1}, \\dots, w_{t - 1}).\n",
    "$$\n",
    "\n",
    "You can also sometimes see such approximation under the name of _n-th order markov assumption_."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first stage to building such a model is counting all word occurences given N-1 previous words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from collections import defaultdict, Counter\n",
    "\n",
    "# special tokens: \n",
    "# - unk represents absent tokens, \n",
    "# - eos is a special token after the end of sequence\n",
    "\n",
    "UNK, EOS = \"_UNK_\", \"_EOS_\"\n",
    "\n",
    "def count_ngrams(lines, n):\n",
    "    \"\"\"\n",
    "    Count how many times each word occured after (n - 1) previous words\n",
    "    :param lines: an iterable of strings with space-separated tokens\n",
    "    :returns: a dictionary { tuple(prefix_tokens): {next_token_1: count_1, next_token_2: count_2}}\n",
    "\n",
    "    When building counts, please consider the following two edge cases\n",
    "    - if prefix is shorter than (n - 1) tokens, it should be padded with UNK. For n=3,\n",
    "      empty prefix: \"\" -> (UNK, UNK)\n",
    "      short prefix: \"the\" -> (UNK, the)\n",
    "      long prefix: \"the new approach\" -> (new, approach)\n",
    "    - you should add a special token, EOS, at the end of each sequence\n",
    "      \"... with deep neural networks .\" -> (..., with, deep, neural, networks, ., EOS)\n",
    "      count the probability of this token just like all others.\n",
    "    \"\"\"\n",
    "    counts = defaultdict(Counter)\n",
    "    # counts[(word1, word2)][word3] = how many times word3 occured after (word1, word2)\n",
    "\n",
    "    <YOUR CODE>\n",
    "    \n",
    "    return counts\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's test it\n",
    "dummy_lines = sorted(lines, key=len)[:100]\n",
    "dummy_counts = count_ngrams(dummy_lines, n=3)\n",
    "assert set(map(len, dummy_counts.keys())) == {2}, \"please only count {n-1}-grams\"\n",
    "assert len(dummy_counts[('_UNK_', '_UNK_')]) == 78\n",
    "assert dummy_counts['_UNK_', 'a']['note'] == 3\n",
    "assert dummy_counts['p', '=']['np'] == 2\n",
    "assert dummy_counts['author', '.']['_EOS_'] == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we can count N-grams, we can build a probabilistic language model.\n",
    "The simplest way to compute probabilities is in proporiton to counts:\n",
    "\n",
    "$$ P(w_t | prefix) = { Count(prefix, w_t) \\over \\sum_{\\hat w} Count(prefix, \\hat w) } $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NGramLanguageModel:    \n",
    "    def __init__(self, lines, n):\n",
    "        \"\"\" \n",
    "        Train a simple count-based language model: \n",
    "        compute probabilities P(w_t | prefix) given ngram counts\n",
    "        \n",
    "        :param n: computes probability of next token given (n - 1) previous words\n",
    "        :param lines: an iterable of strings with space-separated tokens\n",
    "        \"\"\"\n",
    "        assert n >= 1\n",
    "        self.n = n\n",
    "    \n",
    "        counts = count_ngrams(lines, self.n)\n",
    "        \n",
    "        # compute token proabilities given counts\n",
    "        self.probs = defaultdict(Counter)\n",
    "        # probs[(word1, word2)][word3] = P(word3 | word1, word2)\n",
    "        \n",
    "        # populate self.probs with actual probabilities\n",
    "        <YOUR CODE>\n",
    "            \n",
    "    def get_possible_next_tokens(self, prefix):\n",
    "        \"\"\"\n",
    "        :param prefix: string with space-separated prefix tokens\n",
    "        :returns: a dictionary {token : it's probability} for all tokens with positive probabilities\n",
    "        \"\"\"\n",
    "        prefix = prefix.split()\n",
    "        prefix = prefix[max(0, len(prefix) - self.n + 1):]\n",
    "        prefix = [ UNK ] * (self.n - 1 - len(prefix)) + prefix\n",
    "        return self.probs[tuple(prefix)]\n",
    "    \n",
    "    def get_next_token_prob(self, prefix, next_token):\n",
    "        \"\"\"\n",
    "        :param prefix: string with space-separated prefix tokens\n",
    "        :param next_token: the next token to predict probability for\n",
    "        :returns: P(next_token|prefix) a single number, 0 <= P <= 1\n",
    "        \"\"\"\n",
    "        return self.get_possible_next_tokens(prefix).get(next_token, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy_lm = NGramLanguageModel(dummy_lines, n=3)\n",
    "\n",
    "p_initial = dummy_lm.get_possible_next_tokens('') # '' -> ['_UNK_', '_UNK_']\n",
    "assert np.allclose(p_initial['learning'], 0.02)\n",
    "assert np.allclose(p_initial['a'], 0.13)\n",
    "assert np.allclose(p_initial.get('meow', 0), 0)\n",
    "assert np.allclose(sum(p_initial.values()), 1)\n",
    "\n",
    "p_a = dummy_lm.get_possible_next_tokens('a') # '' -> ['_UNK_', 'a']\n",
    "assert np.allclose(p_a['machine'], 0.15384615)\n",
    "assert np.allclose(p_a['note'], 0.23076923)\n",
    "assert np.allclose(p_a.get('the', 0), 0)\n",
    "assert np.allclose(sum(p_a.values()), 1)\n",
    "\n",
    "assert np.allclose(dummy_lm.get_possible_next_tokens('a note')['on'], 1)\n",
    "assert dummy_lm.get_possible_next_tokens('a machine') == \\\n",
    "    dummy_lm.get_possible_next_tokens(\"there have always been ghosts in a machine\"), \\\n",
    "    \"your 3-gram model should only depend on 2 previous words\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you've got a working n-gram language model, let's see what sequences it can generate. But first, let's train it on the whole dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lm = NGramLanguageModel(lines, n=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The process of generating sequences is... well, it's sequential. You maintain a list of tokens and iteratively add next token by sampling with probabilities.\n",
    "\n",
    "$ X = [] $\n",
    "\n",
    "__forever:__\n",
    "* $w_{next} \\sim P(w_{next} | X)$\n",
    "* $X = concat(X, w_{next})$\n",
    "\n",
    "\n",
    "Instead of sampling with probabilities, one can also try always taking most likely token, sampling among top-K most likely tokens or sampling with temperature. In the latter case (temperature), one samples from\n",
    "\n",
    "$$w_{next} \\sim {P(w_{next} | X) ^ {1 / \\tau} \\over \\sum_{\\hat w} P(\\hat w | X) ^ {1 / \\tau}}$$\n",
    "\n",
    "Where $\\tau > 0$ is model temperature. If $\\tau << 1$, more likely tokens will be sampled with even higher probability while less likely tokens will vanish."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_next_token(lm, prefix, temperature=1.0):\n",
    "    \"\"\"\n",
    "    return next token after prefix;\n",
    "    :param temperature: samples proportionally to lm probabilities ^ (1 / temperature)\n",
    "        if temperature == 0, always takes most likely token. Break ties arbitrarily.\n",
    "    \"\"\"\n",
    "    <YOUR CODE>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "test_freqs = Counter([get_next_token(lm, 'there have') for _ in range(10000)])\n",
    "assert 250 < test_freqs['not'] < 450\n",
    "assert 8500 < test_freqs['been'] < 9500\n",
    "assert 1 < test_freqs['lately'] < 200\n",
    "\n",
    "test_freqs = Counter([get_next_token(lm, 'deep', temperature=1.0) for _ in range(10000)])\n",
    "assert 1500 < test_freqs['learning'] < 3000\n",
    "test_freqs = Counter([get_next_token(lm, 'deep', temperature=0.5) for _ in range(10000)])\n",
    "assert 8000 < test_freqs['learning'] < 9000\n",
    "test_freqs = Counter([get_next_token(lm, 'deep', temperature=0.0) for _ in range(10000)])\n",
    "assert test_freqs['learning'] == 10000\n",
    "\n",
    "print(\"Looks nice!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have fun with this model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = 'artificial' # <- your ideas :)\n",
    "\n",
    "for i in range(100):\n",
    "    prefix += ' ' + get_next_token(lm, prefix)\n",
    "    if prefix.endswith(EOS) or len(lm.get_possible_next_tokens(prefix)) == 0:\n",
    "        break\n",
    "        \n",
    "print(prefix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = 'bridging the' # <- more of your ideas\n",
    "\n",
    "for i in range(100):\n",
    "    prefix += ' ' + get_next_token(lm, prefix, temperature=0.5)\n",
    "    if prefix.endswith(EOS) or len(lm.get_possible_next_tokens(prefix)) == 0:\n",
    "        break\n",
    "        \n",
    "print(prefix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__More in the homework:__ nucleous sampling, top-k sampling, beam search(not for the faint of heart)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluating language models: perplexity (1point)\n",
    "\n",
    "Perplexity is a measure of how well does your model approximate true probability distribution behind data. __Smaller perplexity = better model__.\n",
    "\n",
    "To compute perplexity on one sentence, use:\n",
    "$$\n",
    "    {\\mathbb{P}}(w_1 \\dots w_N) = P(w_1, \\dots, w_N)^{-\\frac1N} = \\left( \\prod_t P(w_t \\mid w_{t - n}, \\dots, w_{t - 1})\\right)^{-\\frac1N},\n",
    "$$\n",
    "\n",
    "\n",
    "On the corpora level, perplexity is a product of probabilities of all tokens in all sentences to the power of 1, divided by __total length of all sentences__ in corpora.\n",
    "\n",
    "This number can quickly get too small for float32/float64 precision, so we recommend you to first compute log-perplexity (from log-probabilities) and then take the exponent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def perplexity(lm, lines, min_logprob=np.log(10 ** -50.)):\n",
    "    \"\"\"\n",
    "    :param lines: a list of strings with space-separated tokens\n",
    "    :param min_logprob: if log(P(w | ...)) is smaller than min_logprop, set it equal to min_logrob\n",
    "    :returns: corpora-level perplexity - a single scalar number from the formula above\n",
    "    \n",
    "    Note: do not forget to compute P(w_first | empty) and P(eos | full_sequence)\n",
    "    \n",
    "    PLEASE USE lm.get_next_token_prob and NOT lm.get_possible_next_tokens\n",
    "    \"\"\"\n",
    "    <YOUR CODE>\n",
    "    \n",
    "    return <...>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lm1 = NGramLanguageModel(dummy_lines, n=1)\n",
    "lm3 = NGramLanguageModel(dummy_lines, n=3)\n",
    "lm10 = NGramLanguageModel(dummy_lines, n=10)\n",
    "\n",
    "ppx1 = perplexity(lm1, dummy_lines)\n",
    "ppx3 = perplexity(lm3, dummy_lines)\n",
    "ppx10 = perplexity(lm10, dummy_lines)\n",
    "ppx_missing = perplexity(lm3, ['the jabberwock , with eyes of flame , '])  # thanks, L. Carrol\n",
    "\n",
    "print(\"Perplexities: ppx1=%.3f ppx3=%.3f ppx10=%.3f\" % (ppx1, ppx3, ppx10))\n",
    "\n",
    "assert all(0 < ppx < 500 for ppx in (ppx1, ppx3, ppx10)), \"perplexity should be nonnegative and reasonably small\"\n",
    "assert ppx1 > ppx3 > ppx10, \"higher N models should overfit and \"\n",
    "assert np.isfinite(ppx_missing) and ppx_missing > 10 ** 6, \"missing words should have large but finite perplexity. \" \\\n",
    "    \" Make sure you use min_logprob right\"\n",
    "assert np.allclose([ppx1, ppx3, ppx10], (318.2132342216302, 1.5199996213739575, 1.1838145037901249))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's measure the actual perplexity: we'll split the data into train and test and score model on test data only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "train_lines, test_lines = train_test_split(lines, test_size=0.25, random_state=42)\n",
    "\n",
    "for n in (1, 2, 3):\n",
    "    lm = NGramLanguageModel(n=n, lines=train_lines)\n",
    "    ppx = perplexity(lm, test_lines)\n",
    "    print(\"N = %i, Perplexity = %.5f\" % (n, ppx))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# whoops, it just blew up :)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LM Smoothing\n",
    "\n",
    "The problem with our simple language model is that whenever it encounters an n-gram it has never seen before, it assigns it with the probabilitiy of 0. Every time this happens, perplexity explodes.\n",
    "\n",
    "To battle this issue, there's a technique called __smoothing__. The core idea is to modify counts in a way that prevents probabilities from getting too low. The simplest algorithm here is Additive smoothing (aka [Lapace smoothing](https://en.wikipedia.org/wiki/Additive_smoothing)):\n",
    "\n",
    "$$ P(w_t | prefix) = { Count(prefix, w_t) + \\delta \\over \\sum_{\\hat w} (Count(prefix, \\hat w) + \\delta) } $$\n",
    "\n",
    "If counts for a given prefix are low, additive smoothing will adjust probabilities to a more uniform distribution. Not that the summation in the denominator goes over _all words in the vocabulary_.\n",
    "\n",
    "Here's an example code we've implemented for you:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LaplaceLanguageModel(NGramLanguageModel): \n",
    "    \"\"\" this code is an example, no need to change anything \"\"\"\n",
    "    def __init__(self, lines, n, delta=1.0):\n",
    "        self.n = n\n",
    "        counts = count_ngrams(lines, self.n)\n",
    "        self.vocab = set(token for token_counts in counts.values() for token in token_counts)\n",
    "        self.probs = defaultdict(Counter)\n",
    "\n",
    "        for prefix in counts:\n",
    "            token_counts = counts[prefix]\n",
    "            total_count = sum(token_counts.values()) + delta * len(self.vocab)\n",
    "            self.probs[prefix] = {token: (token_counts[token] + delta) / total_count\n",
    "                                          for token in token_counts}\n",
    "    def get_possible_next_tokens(self, prefix):\n",
    "        token_probs = super().get_possible_next_tokens(prefix)\n",
    "        missing_prob_total = 1.0 - sum(token_probs.values())\n",
    "        missing_prob = missing_prob_total / max(1, len(self.vocab) - len(token_probs))\n",
    "        return {token: token_probs.get(token, missing_prob) for token in self.vocab}\n",
    "    \n",
    "    def get_next_token_prob(self, prefix, next_token):\n",
    "        token_probs = super().get_possible_next_tokens(prefix)\n",
    "        if next_token in token_probs:\n",
    "            return token_probs[next_token]\n",
    "        else:\n",
    "            missing_prob_total = 1.0 - sum(token_probs.values())\n",
    "            missing_prob_total = max(0, missing_prob_total) # prevent rounding errors\n",
    "            return missing_prob_total / max(1, len(self.vocab) - len(token_probs))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#test that it's a valid probability model\n",
    "for n in (1, 2, 3):\n",
    "    dummy_lm = LaplaceLanguageModel(dummy_lines, n=n)\n",
    "    assert np.allclose(sum([dummy_lm.get_next_token_prob('a', w_i) for w_i in dummy_lm.vocab]), 1), \"I told you not to break anything! :)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in (1, 2, 3):\n",
    "    lm = LaplaceLanguageModel(train_lines, n=n, delta=0.1)\n",
    "    ppx = perplexity(lm, test_lines)\n",
    "    print(\"N = %i, Perplexity = %.5f\" % (n, ppx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optional: try to sample tokens from such a model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kneser-Ney smoothing (2 points)\n",
    "\n",
    "Additive smoothing is simple, reasonably good but definitely not a State of The Art algorithm.\n",
    "\n",
    "\n",
    "Your final task in this notebook is to implement [Kneser-Ney](https://en.wikipedia.org/wiki/Kneser%E2%80%93Ney_smoothing) smoothing.\n",
    "\n",
    "It can be computed recurrently, for n>1:\n",
    "\n",
    "$$P_{kn}(w_t | prefix_{n-1}) = { \\max(0, Count(prefix_{n-1}, w_t) - \\delta) \\over \\sum_{\\hat w} Count(prefix_{n-1}, \\hat w)} + \\lambda_{prefix_{n-1}} \\cdot P_{kn}(w_t | prefix_{n-2})$$\n",
    "\n",
    "where\n",
    "- $prefix_{n-1}$ is a tuple of {n-1} previous tokens\n",
    "- $lambda_{prefix_{n-1}}$ is a normalization constant chosen so that probabilities add up to 1\n",
    "- Unigram $P_{kn}(w_t | prefix_{n-2})$ corresponds to Kneser Ney smoothing for {N-1}-gram language model.\n",
    "- Unigram $P_{kn}(w_t)$ is a special case: how likely it is to see x_t in an unfamiliar context\n",
    "\n",
    "See lecture slides or wiki for more detailed formulae.\n",
    "\n",
    "__Your task__ is to\n",
    "- implement KneserNeyLanguageModel\n",
    "- test it on 1-3 gram language models\n",
    "- find optimal (within reason) smoothing delta for 3-gram language model with Kneser-Ney smoothing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KneserNeyLanguageModel(NGramLanguageModel): \n",
    "    \"\"\" A template for Kneser-Ney language model. Default delta may be suboptimal. \"\"\"\n",
    "    def __init__(self, lines, n, delta=1.0):\n",
    "        self.n = n\n",
    "        <YOUR CODE>\n",
    "        \n",
    "    def get_possible_next_tokens(self, prefix):\n",
    "        < YOUR CODE >\n",
    "        \n",
    "    def get_next_token_prob(self, prefix, next_token):\n",
    "        <YOUR CODE>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#test that it's a valid probability model\n",
    "for n in (1, 2, 3):\n",
    "    dummy_lm = KneserNeyLanguageModel(dummy_lines, n=n)\n",
    "    assert np.allclose(sum([dummy_lm.get_next_token_prob('a', w_i) for w_i in dummy_lm.vocab]), 1), \"I told you not to break anything! :)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in (1, 2, 3):\n",
    "    lm = KneserNeyLanguageModel(train_lines, n=n, smoothing=<...>)\n",
    "    ppx = perplexity(lm, test_lines)\n",
    "    print(\"N = %i, Perplexity = %.5f\" % (n, ppx))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}