week03_lm/homework_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Homework: going neural (6 pts)\n",
    "\n",
    "We've checked out statistical approaches to language models in the last notebook. Now let's go find out what deep learning has to offer.\n",
    "\n",
    "<img src='https://raw.githubusercontent.com/yandexdataschool/nlp_course/master/resources/expanding_mind_lm_kn_3.png' width=300px>\n",
    "\n",
    "We're gonna use the same dataset as before, except this time we build a language model that's character-level, not word level. Before you go:\n",
    "* If you haven't done seminar already, use `seminar.ipynb` to download the data.\n",
    "* This homework uses Pytorch v1.x: this is [how you install it](https://pytorch.org/get-started/locally/); and that's [how you use it](https://github.com/yandexdataschool/Practical_RL/tree/9f89e98d7df7ad47f5d6c85a70a38283e06be16a/week04_%5Brecap%5D_deep_learning)."
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "Working on character level means that we don't need to deal with large vocabulary or missing words. Heck, we can even keep uppercase words in text! The downside, however, is that all our sequences just got a lot longer.\n",
    "\n",
    "However, we still need special tokens:\n",
    "* Begin Of Sequence  (__BOS__) - this token is at the start of each sequence. We use it so that we always have non-empty input to our neural network. $P(x_t) = P(x_1 | BOS)$\n",
    "* End Of Sequence (__EOS__) - you guess it... this token is at the end of each sequence. The catch is that it should __not__ occur anywhere else except at the very end. If our model produces this token, the sequence is over.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BOS, EOS = ' ', '\\n'\n",
    "\n",
    "data = pd.read_json(\"./arxivData.json\")\n",
    "lines = data.apply(lambda row: (row['title'] + ' ; ' + row['summary'])[:512], axis=1) \\\n",
    "            .apply(lambda line: BOS + line.replace(EOS, ' ') + EOS) \\\n",
    "            .tolist()\n",
    "\n",
    "# if you missed the seminar, download data here - https://yadi.sk/d/_nGyU2IajjR9-w"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our next step is __building char-level vocabulary__. Put simply, you need to assemble a list of all unique tokens in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get all unique characters from lines (including capital letters and symbols)\n",
    "tokens = <YOUR CODE>\n",
    "\n",
    "tokens = sorted(tokens)\n",
    "n_tokens = len(tokens)\n",
    "print ('n_tokens = ',n_tokens)\n",
    "assert 100 < n_tokens < 150\n",
    "assert BOS in tokens, EOS in tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now assign each character with it's index in tokens list. This way we can encode a string into a torch-friendly integer vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dictionary of character -> its identifier (index in tokens list)\n",
    "token_to_id = <YOUR CODE>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(tokens) == len(token_to_id), \"dictionaries must have same size\"\n",
    "for i in range(n_tokens):\n",
    "    assert token_to_id[tokens[i]] == i, \"token identifier must be it's position in tokens list\"\n",
    "\n",
    "print(\"Seems alright!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our final step is to assemble several strings in a integet matrix `[batch_size, text_length]`. \n",
    "\n",
    "The only problem is that each sequence has a different length. We can work around that by padding short sequences with extra _EOS_ or cropping long sequences. Here's how it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_matrix(lines, max_len=None, pad=token_to_id[EOS], dtype=np.int64):\n",
    "    \"\"\"Casts a list of lines into torch-digestable matrix\"\"\"\n",
    "    max_len = max_len or max(map(len, lines))\n",
    "    lines_ix = np.full([len(lines), max_len], pad, dtype=dtype)\n",
    "    for i in range(len(lines)):\n",
    "        line_ix = list(map(token_to_id.get, lines[i][:max_len]))\n",
    "        lines_ix[i, :len(line_ix)] = line_ix\n",
    "    return lines_ix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Example: cast 4 random names to matrices, pad with zeros\n",
    "dummy_lines = [\n",
    "    ' abc\\n',\n",
    "    ' abacaba\\n',\n",
    "    ' abc1234567890\\n',\n",
    "]\n",
    "print(to_matrix(dummy_lines))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Neural Language Model (2 points including training)\n",
    "\n",
    "Just like for N-gram LMs, we want to estimate probability of text as a joint probability of tokens (symbols this time).\n",
    "\n",
    "$$P(X) = \\prod_t P(x_t \\mid x_0, \\dots, x_{t-1}).$$ \n",
    "\n",
    "Instead of counting all possible statistics, we want to train a neural network with parameters $\\theta$ that estimates the conditional probabilities:\n",
    "\n",
    "$$ P(x_t \\mid x_0, \\dots, x_{t-1}) \\approx p(x_t \\mid x_0, \\dots, x_{t-1}, \\theta) $$\n",
    "\n",
    "\n",
    "But before we optimize, we need to define our neural network. Let's start with a fixed-window (aka convolutional) architecture:\n",
    "\n",
    "<img src='https://raw.githubusercontent.com/yandexdataschool/nlp_course/master/resources/fixed_window_lm.jpg' width=400px>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FixedWindowLanguageModel(nn.Module):\n",
    "    def __init__(self, n_tokens=n_tokens, emb_size=16, hid_size=64):\n",
    "        \"\"\" \n",
    "        A fixed window model that looks on at least 5 previous symbols.\n",
    "        \n",
    "        Note: fixed window LM is effectively performing a convolution over a sequence of words.\n",
    "        This convolution only looks on current and previous words.\n",
    "        Such convolution can be represented as a sequence of 2 operations:\n",
    "        - pad input vectors by {strides * (filter_size - 1)} zero vectors on the \"left\", do not pad right\n",
    "        - perform regular convolution with {filter_size} and {strides}\n",
    "        \n",
    "        - If you're absolutely lost, here's a hint: use nn.ZeroPad2d((NUM_LEADING_ZEROS, 0, 0, 0))\n",
    "          followed by a nn.Conv1d(..., padding=0). And yes, its okay that padding is technically \"2d\".\n",
    "        \"\"\"\n",
    "        super().__init__() # initialize base class to track sub-layers, trainable variables, etc.\n",
    "        \n",
    "        # YOUR CODE - create layers/variables and any metadata you want, e.g. self.emb = L.Embedding(...)\n",
    "        \n",
    "        <...>\n",
    "        \n",
    "        #END OF YOUR CODE\n",
    "    \n",
    "    def __call__(self, input_ix):\n",
    "        \"\"\"\n",
    "        compute language model logits given input tokens\n",
    "        :param input_ix: batch of sequences with token indices, tensor: int32[batch_size, sequence_length]\n",
    "        :returns: pre-softmax linear outputs of language model [batch_size, sequence_length, n_tokens]\n",
    "            these outputs will be used as logits to compute P(x_t | x_0, ..., x_{t - 1})\n",
    "            \n",
    "        :note: that convolutions operate with tensors of shape [batch, channels, length], while linear layers\n",
    "         and *embeddings* use [batch, length, channels] tensors. Use tensor.permute(...) to adjust shapes.\n",
    "\n",
    "        \"\"\"\n",
    "        # YOUR CODE - apply layers, see docstring above\n",
    "        return <...> # output tensor should be of shape [batch_size, sequence_length, n_tokens]\n",
    "    \n",
    "    def get_possible_next_tokens(self, prefix=BOS, temperature=1.0, max_len=100):\n",
    "        \"\"\" :returns: probabilities of next token, dict {token : prob} for all tokens \"\"\"\n",
    "        prefix_ix = torch.as_tensor(to_matrix([prefix]), dtype=torch.int64)\n",
    "        with torch.no_grad():\n",
    "            probs = torch.softmax(self(prefix_ix)[0, -1], dim=-1).cpu().numpy()  # shape: [n_tokens]\n",
    "        return dict(zip(tokens, probs))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy_model = FixedWindowLanguageModel()\n",
    "\n",
    "dummy_input_ix = torch.as_tensor(to_matrix(dummy_lines))\n",
    "dummy_logits = dummy_model(dummy_input_ix)\n",
    "\n",
    "print('Weights:', tuple(name for name, w in dummy_model.named_parameters()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(dummy_logits, torch.Tensor)\n",
    "assert dummy_logits.shape == (len(dummy_lines), max(map(len, dummy_lines)), n_tokens), \"please check output shape\"\n",
    "assert np.all(np.isfinite(dummy_logits.data.cpu().numpy())), \"inf/nan encountered\"\n",
    "assert not np.allclose(dummy_logits.data.cpu().numpy().sum(-1), 1), \"please predict linear outputs, don't use softmax (maybe you've just got unlucky)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test for lookahead\n",
    "dummy_input_ix_2 = torch.as_tensor(to_matrix([line[:3] + 'e' * (len(line) - 3) for line in dummy_lines]))\n",
    "dummy_logits_2 = dummy_model(dummy_input_ix_2)\n",
    "\n",
    "assert torch.allclose(dummy_logits[:, :3], dummy_logits_2[:, :3]), \"your model's predictions depend on FUTURE tokens. \" \\\n",
    "    \" Make sure you don't allow any layers to look ahead of current token.\" \\\n",
    "    \" You can also get this error if your model is not deterministic (e.g. dropout). Disable it for this test.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now tune our network's parameters to minimize categorical crossentropy over training dataset $D$:\n",
    "\n",
    "$$ L = {\\frac1{|D|}} \\sum_{X \\in D} \\sum_{x_i \\in X} - \\log p(x_t \\mid x_1, \\dots, x_{t-1}, \\theta) $$\n",
    "\n",
    "As usual with with neural nets, this optimization is performed via stochastic gradient descent with backprop.  One can also note that minimizing crossentropy is equivalent to minimizing model __perplexity__, KL-divergence or maximizng log-likelihood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_mask(input_ix, eos_ix=token_to_id[EOS]):\n",
    "    \"\"\" compute a boolean mask that equals \"1\" until first EOS (including that EOS) \"\"\"\n",
    "    return F.pad(torch.cumsum(input_ix == eos_ix, dim=-1)[..., :-1] < 1, pad=(1, 0, 0, 0), value=True)\n",
    "\n",
    "print('matrix:\\n', dummy_input_ix.numpy())\n",
    "print('mask:', compute_mask(dummy_input_ix).to(torch.int32).cpu().numpy())\n",
    "print('lengths:', compute_mask(dummy_input_ix).sum(-1).cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_loss(model, input_ix):\n",
    "    \"\"\"\n",
    "    :param model: language model that can compute next token logits given token indices\n",
    "    :param input ix: int32 matrix of tokens, shape: [batch_size, length]; padded with eos_ix\n",
    "    :returns: scalar loss function, mean crossentropy over non-eos tokens\n",
    "    \"\"\"\n",
    "    input_ix = torch.as_tensor(input_ix, dtype=torch.int64)\n",
    "    \n",
    "\n",
    "    logits = model(input_ix[:, :-1])\n",
    "    reference_answers = input_ix[:, 1:]\n",
    "\n",
    "    # Your task: implement loss function as per formula above\n",
    "    # your loss should only be computed on actual tokens, excluding padding\n",
    "    # predicting actual tokens and first EOS do count. Subsequent EOS-es don't\n",
    "    # you may or may not want to use the compute_mask function from above.\n",
    "    \n",
    "    <YOUR CODE>\n",
    "    \n",
    "    return <YOUR CODE: return scalar loss>    \n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_1 = compute_loss(dummy_model, to_matrix(dummy_lines, max_len=15))\n",
    "loss_2 = compute_loss(dummy_model, to_matrix(dummy_lines, max_len=16))\n",
    "assert (np.ndim(loss_1) == 0) and (0 < loss_1 < 100), \"loss must be a positive scalar\"\n",
    "assert torch.allclose(loss_1, loss_2), 'do not include  AFTER first EOS into loss. '\\\n",
    "    'Hint: use compute_mask. Beware +/-1 errors. And be careful when averaging!'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation\n",
    "\n",
    "You will need two functions: one to compute test loss and another to generate samples. For your convenience, we implemented them both in your stead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def score_lines(model, dev_lines, batch_size):\n",
    "    \"\"\" computes average loss over the entire dataset \"\"\"\n",
    "    dev_loss_num, dev_loss_len = 0., 0.\n",
    "    with torch.no_grad():\n",
    "        for i in range(0, len(dev_lines), batch_size):\n",
    "            batch_ix = to_matrix(dev_lines[i: i + batch_size])\n",
    "            dev_loss_num += compute_loss(model, batch_ix).item() * len(batch_ix)\n",
    "            dev_loss_len += len(batch_ix)\n",
    "    return dev_loss_num / dev_loss_len\n",
    "\n",
    "def generate(model, prefix=BOS, temperature=1.0, max_len=100):\n",
    "    \"\"\"\n",
    "    Samples output sequence from probability distribution obtained by model\n",
    "    :param temperature: samples proportionally to model probabilities ^ temperature\n",
    "        if temperature == 0, always takes most likely token. Break ties arbitrarily.\n",
    "    \"\"\"\n",
    "    with torch.no_grad():\n",
    "        while True:\n",
    "            token_probs = model.get_possible_next_tokens(prefix)\n",
    "            tokens, probs = zip(*token_probs.items())\n",
    "            if temperature == 0:\n",
    "                next_token = tokens[np.argmax(probs)]\n",
    "            else:\n",
    "                probs = np.array([p ** (1. / temperature) for p in probs])\n",
    "                probs /= sum(probs)\n",
    "                next_token = np.random.choice(tokens, p=probs)\n",
    "\n",
    "            prefix += next_token\n",
    "            if next_token == EOS or len(prefix) > max_len: break\n",
    "    return prefix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training loop\n",
    "\n",
    "Finally, let's train our model on minibatches of data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "train_lines, dev_lines = train_test_split(lines, test_size=0.25, random_state=42)\n",
    "\n",
    "batch_size = 256\n",
    "score_dev_every = 250\n",
    "train_history, dev_history = [], []\n",
    "model = FixedWindowLanguageModel()\n",
    "opt = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# hint: if you ever wanted to switch to cuda, do it now.\n",
    "\n",
    "# score untrained model\n",
    "dev_history.append((0, score_lines(model, dev_lines, batch_size)))\n",
    "print(\"Sample before training:\", generate(model, 'Bridging'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "from random import sample\n",
    "from tqdm import trange\n",
    "\n",
    "for i in trange(len(train_history), 5000):\n",
    "    batch = to_matrix(sample(train_lines, batch_size))\n",
    "    \n",
    "    \n",
    "    loss_i = compute_loss(model, batch)\n",
    "    \n",
    "    opt.zero_grad()\n",
    "    loss_i.backward()\n",
    "    opt.step()\n",
    "        \n",
    "    train_history.append((i, loss_i.item()))\n",
    "    \n",
    "    if (i + 1) % 50 == 0:\n",
    "        clear_output(True)\n",
    "        plt.scatter(*zip(*train_history), alpha=0.1, label='train_loss')\n",
    "        if len(dev_history):\n",
    "            plt.plot(*zip(*dev_history), color='red', label='dev_loss')\n",
    "        plt.legend(); plt.grid(); plt.show()\n",
    "        print(\"Generated examples (tau=0.5):\")\n",
    "        for _ in range(3):\n",
    "            print(generate(model, temperature=0.5))\n",
    "    \n",
    "    if (i + 1) % score_dev_every == 0:\n",
    "        print(\"Scoring dev...\")\n",
    "        dev_history.append((i, score_lines(model, dev_lines, batch_size)))\n",
    "        print('#%i Dev loss: %.3f' % dev_history[-1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.mean(train_history[:10], axis=0)[1] > np.mean(train_history[-10:], axis=0)[1], \"The model didn't converge.\"\n",
    "print(\"Final dev loss:\", dev_history[-1][-1])\n",
    "\n",
    "for i in range(10):\n",
    "    print(generate(model, temperature=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RNN Language Models (3 points including training)\n",
    "\n",
    "Fixed-size architectures are reasonably good when capturing short-term dependencies, but their design prevents them from capturing any signal outside their window. We can mitigate this problem by using a __recurrent neural network__:\n",
    "\n",
    "$$ h_0 = \\vec 0 ; \\quad h_{t+1} = RNN(x_t, h_t) $$\n",
    "\n",
    "$$ p(x_t \\mid x_0, \\dots, x_{t-1}, \\theta) = dense_{softmax}(h_{t-1}) $$\n",
    "\n",
    "Such model processes one token at a time, left to right, and maintains a hidden state vector between them. Theoretically, it can learn arbitrarily long temporal dependencies given large enough hidden size.\n",
    "\n",
    "<img src='https://raw.githubusercontent.com/yandexdataschool/nlp_course/master/resources/rnn_lm.jpg' width=480px>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNNLanguageModel(nn.Module):\n",
    "    def __init__(self, n_tokens=n_tokens, emb_size=16, hid_size=256):\n",
    "        \"\"\" \n",
    "        Build a recurrent language model.\n",
    "        You are free to choose anything you want, but the recommended architecture is\n",
    "        - token embeddings\n",
    "        - one or more LSTM/GRU layers with hid size\n",
    "        - linear layer to predict logits\n",
    "        \n",
    "        :note: if you use nn.RNN/GRU/LSTM, make sure you specify batch_first=True\n",
    "         With batch_first, your model operates with tensors of shape [batch_size, sequence_length, num_units]\n",
    "         Also, please read the docs carefully: they don't just return what you want them to return :)\n",
    "        \"\"\"\n",
    "        super().__init__() # initialize base class to track sub-layers, trainable variables, etc.\n",
    "        \n",
    "        # YOUR CODE - create layers/variables/etc\n",
    "        \n",
    "        <...>\n",
    "        \n",
    "        #END OF YOUR CODE\n",
    "    \n",
    "    def __call__(self, input_ix):\n",
    "        \"\"\"\n",
    "        compute language model logits given input tokens\n",
    "        :param input_ix: batch of sequences with token indices, tensor: int32[batch_size, sequence_length]\n",
    "        :returns: pre-softmax linear outputs of language model [batch_size, sequence_length, n_tokens]\n",
    "            these outputs will be used as logits to compute P(x_t | x_0, ..., x_{t - 1})\n",
    "        \"\"\"\n",
    "        # YOUR CODE - apply layers, see docstring above\n",
    "        return <...> # output tensor should be of shape [batch_size, sequence_length, n_tokens]\n",
    "    \n",
    "    def get_possible_next_tokens(self, prefix=BOS, temperature=1.0, max_len=100):\n",
    "        \"\"\" :returns: probabilities of next token, dict {token : prob} for all tokens \"\"\"\n",
    "        prefix_ix = torch.as_tensor(to_matrix([prefix]), dtype=torch.int64)\n",
    "        with torch.no_grad():\n",
    "            probs = torch.softmax(self(prefix_ix)[0, -1], dim=-1).cpu().numpy()  # shape: [n_tokens]\n",
    "        return dict(zip(tokens, probs))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = RNNLanguageModel()\n",
    "\n",
    "dummy_input_ix = torch.as_tensor(to_matrix(dummy_lines))\n",
    "dummy_logits = model(dummy_input_ix)\n",
    "\n",
    "assert isinstance(dummy_logits, torch.Tensor)\n",
    "assert dummy_logits.shape == (len(dummy_lines), max(map(len, dummy_lines)), n_tokens), \"please check output shape\"\n",
    "assert not np.allclose(dummy_logits.cpu().data.numpy().sum(-1), 1), \"please predict linear outputs, don't use softmax (maybe you've just got unlucky)\"\n",
    "print('Weights:', tuple(name for name, w in model.named_parameters()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test for lookahead\n",
    "dummy_input_ix_2 = torch.as_tensor(to_matrix([line[:3] + 'e' * (len(line) - 3) for line in dummy_lines]))\n",
    "dummy_logits_2 = model(dummy_input_ix_2)\n",
    "\n",
    "assert torch.allclose(dummy_logits[:, :3], dummy_logits_2[:, :3]), \"your model's predictions depend on FUTURE tokens. \" \\\n",
    "    \" Make sure you don't allow any layers to look ahead of current token.\" \\\n",
    "    \" You can also get this error if your model is not deterministic (e.g. dropout). Disable it for this test.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RNN training\n",
    "\n",
    "Our RNN language model should optimize the same loss function as fixed-window model. But there's a catch. Since RNN recurrently multiplies gradients through many time-steps, gradient values may explode, [ruining](https://raw.githubusercontent.com/yandexdataschool/nlp_course/master/resources/nan.jpg) your model.\n",
    "The common solution to that problem is to clip gradients either [individually](https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/clip_by_value) or [globally](https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/clip_by_global_norm).\n",
    "\n",
    "Your task here is to implement the training code that minimizes the loss function. If you encounter large loss fluctuations during training, please add [gradient clipping](https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html) using urls above. But its **not necessary** to use gradient clipping if you don't need it.\n",
    "\n",
    "_Note: gradient clipping is not exclusive to RNNs. Convolutional networks with enough depth often suffer from the same issue._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64         # <-- please tune batch size to fit your CPU/GPU configuration\n",
    "score_dev_every = 250\n",
    "train_history, dev_history = [], []\n",
    "\n",
    "model = <YOUR CODE HERE>\n",
    "opt = torch.optim.Adam(model.parameters())\n",
    "\n",
    "# score untrained model\n",
    "dev_history.append((0, score_lines(model, dev_lines, batch_size)))\n",
    "print(\"Sample before training:\", generate(model, 'Bridging'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "from random import sample\n",
    "from tqdm import trange\n",
    "\n",
    "for i in trange(len(train_history), 5000):\n",
    "    batch = to_matrix(sample(train_lines, batch_size))\n",
    "    \n",
    "    <YOUR CODE - one step of the training loop for your RNN model>\n",
    "        \n",
    "    loss_i = <...>\n",
    "    \n",
    "    train_history.append((i, float(loss_i)))\n",
    "    \n",
    "    if (i + 1) % 50 == 0:\n",
    "        clear_output(True)\n",
    "        plt.scatter(*zip(*train_history), alpha=0.1, label='train_loss')\n",
    "        if len(dev_history):\n",
    "            plt.plot(*zip(*dev_history), color='red', label='dev_loss')\n",
    "        plt.legend(); plt.grid(); plt.show()\n",
    "        print(\"Generated examples (tau=0.5):\")\n",
    "        for _ in range(3):\n",
    "            print(generate(model, temperature=0.5))\n",
    "    \n",
    "    if (i + 1) % score_dev_every == 0:\n",
    "        print(\"Scoring dev...\")\n",
    "        dev_history.append((i, score_lines(model, dev_lines, batch_size)))\n",
    "        print('#%i Dev loss: %.3f' % dev_history[-1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.mean(train_history[:10], axis=0)[1] > np.mean(train_history[-10:], axis=0)[1], \"The model didn't converge.\"\n",
    "print(\"Final dev loss:\", dev_history[-1][-1])\n",
    "for i in range(10):\n",
    "    print(generate(model, temperature=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alternative sampling strategies (1 point)\n",
    "\n",
    "So far we've sampled tokens from the model in proportion with their probability.\n",
    "However, this approach can sometimes generate nonsense words due to the fact that softmax probabilities of these words are never exactly zero. This issue can be somewhat mitigated with sampling temperature, but low temperature harms sampling diversity. Can we remove the nonsense words without sacrificing diversity? __Yes, we can!__ But it takes a different sampling strategy.\n",
    "\n",
    "__Top-k sampling:__ on each step, sample the next token from __k most likely__ candidates from the language model.\n",
    "\n",
    "Suppose $k=3$ and the token probabilities are $p=[0.1, 0.35, 0.05, 0.2, 0.3]$. You first need to select $k$ most likely words and set the probability of the rest to zero: $\\hat p=[0.0, 0.35, 0.0, 0.2, 0.3]$ and re-normalize: \n",
    "$p^*\\approx[0.0, 0.412, 0.0, 0.235, 0.353]$.\n",
    "\n",
    "__Nucleus sampling:__ similar to top-k sampling, but this time we select $k$ dynamically. In nucleous sampling, we sample from top-__N%__ fraction of the probability mass.\n",
    "\n",
    "Using the same  $p=[0.1, 0.35, 0.05, 0.2, 0.3]$ and nucleous N=0.9, the nucleous words consist of:\n",
    "1. most likely token $w_2$, because $p(w_2) < N$\n",
    "2. second most likely token $w_5$, $p(w_2) + p(w_5) = 0.65 < N$\n",
    "3. third most likely token $w_4$ because $p(w_2) + p(w_5) + p(w_4) = 0.85 < N$\n",
    "\n",
    "And thats it, because the next most likely word would overflow: $p(w_2) + p(w_5) + p(w_4) + p(w_1) = 0.95 > N$.\n",
    "\n",
    "After you've selected the nucleous words, you need to re-normalize them as in top-k sampling and generate the next token.\n",
    "\n",
    "__Your task__ is to implement nucleus sampling variant and see if its any good."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_nucleus(model, prefix=BOS, nucleus=0.9, max_len=100):\n",
    "    \"\"\"\n",
    "    Generate a sequence with nucleous sampling\n",
    "    :param prefix: a string containing space-separated previous tokens\n",
    "    :param nucleus: N from the formulae above, N \\in [0, 1]\n",
    "    :param max_len: generate sequences with at most this many tokens, including prefix\n",
    "    \n",
    "    :note: make sure that nucleous always contains at least one word, even if p(w*) > nucleus\n",
    "    \n",
    "    \"\"\"\n",
    "    while True:\n",
    "        token_probs = model.get_possible_next_tokens(prefix)\n",
    "        tokens, probs = zip(*token_probs.items())\n",
    "        \n",
    "        <YOUR CODE HERE>\n",
    "        \n",
    "        prefix += <YOUR CODE>\n",
    "        if next_token == EOS or len(prefix) > max_len: break\n",
    "    return prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "    print(generate_nucleous(model, nucleous_size=PLAY_WITH_ME_SENPAI))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus quest I: Beam Search (2 pts incl. samples)\n",
    "\n",
    "At times, you don't really want the model to generate diverse outputs as much as you want a __single most likely hypothesis.__ A single best translation, most likely continuation of the search query given prefix, etc. Except, you can't get it. \n",
    "\n",
    "In order to find the exact most likely sequence containing 10 tokens, you would need to enumerate all $|V|^{10}$ possible hypotheses. In practice, 9 times out of 10 you will instead find an approximate most likely output using __beam search__.\n",
    "\n",
    "Here's how it works:\n",
    "0. Initial `beam` = [prefix], max beam_size = k\n",
    "1. for T steps:\n",
    "2. ` ... ` generate all possible next tokens for all hypotheses in beam, formulate `len(beam) * len(vocab)` candidates\n",
    "3. ` ... ` select beam_size best for all candidates as new `beam`\n",
    "4. Select best hypothesis (-es?) from beam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "# Here's what it looks like:\n",
    "!wget -q https://raw.githubusercontent.com/yandexdataschool/nlp_course/2020/resources/beam_search.html\n",
    "HTML(\"beam_search.html\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_beamsearch(model, prefix=BOS, beam_size=4, length=5):\n",
    "    \"\"\"\n",
    "    Generate a sequence with nucleous sampling\n",
    "    :param prefix: a string containing space-separated previous tokens\n",
    "    :param nucleus: N from the formulae above, N \\in [0, 1]\n",
    "    :param length: generate sequences with at most this many tokens, NOT INCLUDING PREFIX\n",
    "    :returns: beam_size most likely candidates\n",
    "    :note: make sure that nucleous always contains at least one word, even if p(w*) > nucleus\n",
    "    \"\"\"\n",
    "    \n",
    "    <YOUR CODE HERE>\n",
    "    \n",
    "    return <most likely sequence>\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_beamsearch(model, prefix=' deep ', beam_size=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# check it out: which beam size works best?\n",
    "# find at least 5 prefixes where beam_size=1 and 8 generates different sequences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus quest II: Ultimate Language Model (2+ pts)\n",
    "\n",
    "So you've learned the building blocks of neural language models, you can now build the ultimate monster:  \n",
    "* Make it char-level, word level or maybe use sub-word units like [bpe](https://github.com/rsennrich/subword-nmt);\n",
    "* Combine convolutions, recurrent cells, pre-trained embeddings and all the black magic deep learning has to offer;\n",
    "  * Use strides to get larger window size quickly. Here's a [scheme](https://storage.googleapis.com/deepmind-live-cms/documents/BlogPost-Fig2-Anim-160908-r01.gif) from google wavenet.\n",
    "* Train on large data. Like... really large. Try [1 Billion Words](http://www.statmt.org/lm-benchmark/1-billion-word-language-modeling-benchmark-r13output.tar.gz) benchmark;\n",
    "* Use training schedules to speed up training. Start with small length and increase over time; Take a look at [one cycle](https://medium.com/@nachiket.tanksale/finding-good-learning-rate-and-the-one-cycle-policy-7159fe1db5d6) for learning rate;\n",
    "\n",
    "_You are NOT required to submit this assignment. Please make sure you don't miss your deadline because of it :)_"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.1 (main, Dec 23 2022, 09:40:27) [Clang 14.0.0 (clang-1400.0.29.202)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "1a1af0ee75eeea9e2e1ee996c87e7a2b11a0bebd85af04bb136d915cefc0abce"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}