week04_seq2seq/practice_and_homework_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HmsFABwClrsS"
   },
   "source": [
    "## Seminar and homework (10 points total)\n",
    "\n",
    "Today we shall compose encoder-decoder neural networks and apply them to the task of machine translation.\n",
    "\n",
    "![img](https://esciencegroup.files.wordpress.com/2016/03/seq2seq.jpg)\n",
    "_(img: esciencegroup.files.wordpress.com)_\n",
    "\n",
    "\n",
    "Encoder-decoder architectures are about converting anything to anything, including\n",
    " * Machine translation and spoken dialogue systems\n",
    " * [Image captioning](http://mscoco.org/dataset/#captions-challenge2015) and [image2latex](https://openai.com/requests-for-research/#im2latex) (convolutional encoder, recurrent decoder)\n",
    " * Generating [images by captions](https://arxiv.org/abs/1511.02793) (recurrent encoder, convolutional decoder)\n",
    " * Grapheme2phoneme - convert words to transcripts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "R4N9AD2dlrsU"
   },
   "source": [
    "## Our task: machine translation\n",
    "\n",
    "We gonna try our encoder-decoder models on russian to english machine translation problem. More specifically, we'll translate hotel and hostel descriptions. This task shows the scale of machine translation while not requiring you to train your model for weeks if you don't use GPU.\n",
    "\n",
    "Before we get to the architecture, there's some preprocessing to be done. ~~Go tokenize~~ Alright, this time we've done preprocessing for you. As usual, the data will be tokenized with WordPunctTokenizer.\n",
    "\n",
    "However, there's one more thing to do. Our data lines contain unique rare words. If we operate on a word level, we will have to deal with large vocabulary size. If instead we use character-level models, it would take lots of iterations to process a sequence. This time we're gonna pick something inbetween.\n",
    "\n",
    "One popular approach is called [Byte Pair Encoding](https://github.com/rsennrich/subword-nmt) aka __BPE__. The algorithm starts with a character-level tokenization and then iteratively merges most frequent pairs for N iterations. This results in frequent words being merged into a single token and rare words split into syllables or even characters.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CfvojjHQlrsU"
   },
   "outputs": [],
   "source": [
    "!pip3 install torch>=1.3.0\n",
    "!pip3 install subword-nmt &> log\n",
    "!wget https://www.dropbox.com/s/yy2zqh34dyhv07i/data.txt?dl=1 -O data.txt\n",
    "!wget https://raw.githubusercontent.com/yandexdataschool/nlp_course/2020/week04_seq2seq/vocab.py -O vocab.py\n",
    "# thanks to tilda and deephack teams for the data, Dmitry Emelyanenko for the code :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "g9kP0SdxlrsY"
   },
   "outputs": [],
   "source": [
    "from nltk.tokenize import WordPunctTokenizer\n",
    "from subword_nmt.learn_bpe import learn_bpe\n",
    "from subword_nmt.apply_bpe import BPE\n",
    "tokenizer = WordPunctTokenizer()\n",
    "def tokenize(x):\n",
    "    return ' '.join(tokenizer.tokenize(x.lower()))\n",
    "\n",
    "# split and tokenize the data\n",
    "with open('train.en', 'w') as f_src,  open('train.ru', 'w') as f_dst:\n",
    "    for line in open('data.txt'):\n",
    "        src_line, dst_line = line.strip().split('\\t')\n",
    "        f_src.write(tokenize(src_line) + '\\n')\n",
    "        f_dst.write(tokenize(dst_line) + '\\n')\n",
    "\n",
    "# build and apply bpe vocs\n",
    "bpe = {}\n",
    "for lang in ['en', 'ru']:\n",
    "    learn_bpe(open('./train.' + lang), open('bpe_rules.' + lang, 'w'), num_symbols=8000)\n",
    "    bpe[lang] = BPE(open('./bpe_rules.' + lang))\n",
    "    \n",
    "    with open('train.bpe.' + lang, 'w') as f_out:\n",
    "        for line in open('train.' + lang):\n",
    "            f_out.write(bpe[lang].process_line(line.strip()) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0UPW3sV8lrsb"
   },
   "source": [
    "### Building vocabularies\n",
    "\n",
    "We now need to build vocabularies that map strings to token ids and vice versa. We're gonna need these fellas when we feed training data into model or convert output matrices into words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CmTy_m_olrsb"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8PskgBSxlrsd"
   },
   "outputs": [],
   "source": [
    "data_inp = np.array(open('./train.bpe.ru').read().split('\\n'))\n",
    "data_out = np.array(open('./train.bpe.en').read().split('\\n'))\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "train_inp, dev_inp, train_out, dev_out = train_test_split(data_inp, data_out, test_size=3000,\n",
    "                                                          random_state=42)\n",
    "for i in range(3):\n",
    "    print('inp:', train_inp[i])\n",
    "    print('out:', train_out[i], end='\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vipg4O61lrsg"
   },
   "outputs": [],
   "source": [
    "from vocab import Vocab\n",
    "inp_voc = Vocab.from_lines(train_inp)\n",
    "out_voc = Vocab.from_lines(train_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cwOoHfuhlrsi"
   },
   "outputs": [],
   "source": [
    "# Here's how you cast lines into ids and backwards.\n",
    "batch_lines = sorted(train_inp, key=len)[5:10]\n",
    "batch_ids = inp_voc.to_matrix(batch_lines)\n",
    "batch_lines_restored = inp_voc.to_lines(batch_ids)\n",
    "\n",
    "print(\"lines\")\n",
    "print(batch_lines)\n",
    "print(\"\\nwords to ids (0 = bos, 1 = eos):\")\n",
    "print(batch_ids)\n",
    "print(\"\\nback to words\")\n",
    "print(batch_lines_restored)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gSYu-MkElrsk"
   },
   "source": [
    "Draw source and translation length distributions to estimate the scope of the task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TLLl9cSNlrsl"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=[8, 4])\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(\"source length\")\n",
    "plt.hist(list(map(len, map(str.split, train_inp))), bins=20);\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"translation length\")\n",
    "plt.hist(list(map(len, map(str.split, train_out))), bins=20);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BHWgx34flrsn"
   },
   "source": [
    "### Encoder-decoder model\n",
    "\n",
    "The code below contains a template for a simple encoder-decoder model: single GRU encoder/decoder, no attention or anything. This model is implemented for you as a reference and a baseline for your homework assignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pd_rDRm9lrso"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wgfN5-F7lrst"
   },
   "outputs": [],
   "source": [
    "class BasicModel(nn.Module):\n",
    "    def __init__(self, inp_voc, out_voc, emb_size=64, hid_size=128):\n",
    "        \"\"\"\n",
    "        A simple encoder-decoder seq2seq model\n",
    "        \"\"\"\n",
    "        super().__init__() # initialize base class to track sub-layers, parameters, etc.\n",
    "\n",
    "        self.inp_voc, self.out_voc = inp_voc, out_voc\n",
    "        self.hid_size = hid_size\n",
    "        \n",
    "        self.emb_inp = nn.Embedding(len(inp_voc), emb_size)\n",
    "        self.emb_out = nn.Embedding(len(out_voc), emb_size)\n",
    "        self.enc0 = nn.GRU(emb_size, hid_size, batch_first=True)\n",
    "\n",
    "        self.dec_start = nn.Linear(hid_size, hid_size)\n",
    "        self.dec0 = nn.GRUCell(emb_size, hid_size)\n",
    "        self.logits = nn.Linear(hid_size, len(out_voc))\n",
    "        \n",
    "    def forward(self, inp, out):\n",
    "        \"\"\" Apply model in training mode \"\"\"\n",
    "        initial_state = self.encode(inp)\n",
    "        return self.decode(initial_state, out)\n",
    "\n",
    "\n",
    "    def encode(self, inp, **flags):\n",
    "        \"\"\"\n",
    "        Takes symbolic input sequence, computes initial state\n",
    "        :param inp: matrix of input tokens [batch, time]\n",
    "        :returns: initial decoder state tensors, one or many\n",
    "        \"\"\"\n",
    "        inp_emb = self.emb_inp(inp)\n",
    "        batch_size = inp.shape[0]\n",
    "        \n",
    "        enc_seq, [last_state_but_not_really] = self.enc0(inp_emb)\n",
    "        # enc_seq: [batch, time, hid_size], last_state: [batch, hid_size]\n",
    "        \n",
    "        # note: last_state is not _actually_ last because of padding, let's find the real last_state\n",
    "        lengths = (inp != self.inp_voc.eos_ix).to(torch.int64).sum(dim=1).clamp_max(inp.shape[1] - 1)\n",
    "        last_state = enc_seq[torch.arange(len(enc_seq)), lengths]\n",
    "        # ^-- shape: [batch_size, hid_size]\n",
    "        \n",
    "        dec_start = self.dec_start(last_state)\n",
    "        return [dec_start]\n",
    "\n",
    "    def decode_step(self, prev_state, prev_tokens, **flags):\n",
    "        \"\"\"\n",
    "        Takes previous decoder state and tokens, returns new state and logits for next tokens\n",
    "        :param prev_state: a list of previous decoder state tensors, same as returned by encode(...)\n",
    "        :param prev_tokens: previous output tokens, an int vector of [batch_size]\n",
    "        :return: a list of next decoder state tensors, a tensor of logits [batch, len(out_voc)]\n",
    "        \"\"\"\n",
    "        prev_gru0_state = prev_state[0]\n",
    "        \n",
    "        <YOUR CODE HERE>\n",
    "        \n",
    "        return new_dec_state, output_logits\n",
    "\n",
    "    def decode(self, initial_state, out_tokens, **flags):\n",
    "        \"\"\" Iterate over reference tokens (out_tokens) with decode_step \"\"\"\n",
    "        batch_size = out_tokens.shape[0]\n",
    "        state = initial_state\n",
    "        \n",
    "        # initial logits: always predict BOS\n",
    "        onehot_bos = F.one_hot(torch.full([batch_size], self.out_voc.bos_ix, dtype=torch.int64),\n",
    "                               num_classes=len(self.out_voc)).to(device=out_tokens.device)\n",
    "        first_logits = torch.log(onehot_bos.to(torch.float32) + 1e-9)\n",
    "        \n",
    "        logits_sequence = [first_logits]\n",
    "        for i in range(out_tokens.shape[1] - 1):\n",
    "            state, logits = self.decode_step(state, out_tokens[:, i])\n",
    "            logits_sequence.append(logits)\n",
    "        return torch.stack(logits_sequence, dim=1)\n",
    "\n",
    "    def decode_inference(self, initial_state, max_len=100, **flags):\n",
    "        \"\"\" Generate translations from model (greedy version) \"\"\"\n",
    "        batch_size, device = len(initial_state[0]), initial_state[0].device\n",
    "        state = initial_state\n",
    "        outputs = [torch.full([batch_size], self.out_voc.bos_ix, dtype=torch.int64, \n",
    "                              device=device)]\n",
    "        all_states = [initial_state]\n",
    "\n",
    "        for i in range(max_len):\n",
    "            state, logits = self.decode_step(state, outputs[-1])\n",
    "            outputs.append(logits.argmax(dim=-1))\n",
    "            all_states.append(state)\n",
    "        \n",
    "        return torch.stack(outputs, dim=1), all_states\n",
    "\n",
    "    def translate_lines(self, inp_lines, **kwargs):\n",
    "        inp = self.inp_voc.to_matrix(inp_lines).to(device)\n",
    "        initial_state = self.encode(inp)\n",
    "        out_ids, states = self.decode_inference(initial_state, **kwargs)\n",
    "        return self.out_voc.to_lines(out_ids.cpu().numpy()), states\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# debugging area\n",
    "model = BasicModel(inp_voc, out_voc).to(device)\n",
    "\n",
    "dummy_inp_tokens = inp_voc.to_matrix(sorted(train_inp, key=len)[5:10]).to(device)\n",
    "dummy_out_tokens = out_voc.to_matrix(sorted(train_out, key=len)[5:10]).to(device)\n",
    "\n",
    "h0 = model.encode(dummy_inp_tokens)\n",
    "h1, logits1 = model.decode_step(h0, torch.arange(len(dummy_inp_tokens), device=device))\n",
    "\n",
    "assert isinstance(h1, list) and len(h1) == len(h0)\n",
    "assert h1[0].shape == h0[0].shape and not torch.allclose(h1[0], h0[0])\n",
    "assert logits1.shape == (len(dummy_inp_tokens), len(out_voc))\n",
    "\n",
    "logits_seq = model.decode(h0, dummy_out_tokens)\n",
    "assert logits_seq.shape == (dummy_out_tokens.shape[0], dummy_out_tokens.shape[1], len(out_voc))\n",
    "\n",
    "# full forward\n",
    "logits_seq2 = model(dummy_inp_tokens, dummy_out_tokens)\n",
    "assert logits_seq2.shape == logits_seq.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy_translations, dummy_states = model.translate_lines(train_inp[:3], max_len=25)\n",
    "print(\"Translations without training:\")\n",
    "print('\\n'.join([line for line in dummy_translations]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_wuv1-aVlrs0"
   },
   "source": [
    "### Training loss (2 points)\n",
    "\n",
    "Our training objective is almost the same as it was for neural language models:\n",
    "$$ L = {\\frac1{|D|}} \\sum_{X, Y \\in D} \\sum_{y_t \\in Y} - \\log p(y_t \\mid y_1, \\dots, y_{t-1}, X, \\theta) $$\n",
    "\n",
    "where $|D|$ is the __total length of all sequences__, including BOS and first EOS, but excluding PAD."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "c8XPV8sWlrs5",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "def compute_loss(model, inp, out, **flags):\n",
    "    \"\"\"\n",
    "    Compute loss (float32 scalar) as in the formula above\n",
    "    :param inp: input tokens matrix, int32[batch, time]\n",
    "    :param out: reference tokens matrix, int32[batch, time]\n",
    "    \n",
    "    In order to pass the tests, your function should\n",
    "    * include loss at first EOS but not the subsequent ones\n",
    "    * divide sum of losses by a sum of input lengths (use voc.compute_mask)\n",
    "    \"\"\"\n",
    "    mask = model.out_voc.compute_mask(out) # [batch_size, out_len]\n",
    "    targets_1hot = F.one_hot(out, len(model.out_voc)).to(torch.float32)\n",
    "    \n",
    "    # outputs of the model, [batch_size, out_len, num_tokens]\n",
    "    logits_seq = <YOUR CODE HERE>\n",
    "\n",
    "    # log-probabilities of all tokens at all steps, [batch_size, out_len, num_tokens]\n",
    "    logprobs_seq = <YOUR CODE HERE>\n",
    "   \n",
    "    # log-probabilities of correct outputs, [batch_size, out_len]\n",
    "    logp_out = (logprobs_seq * targets_1hot).sum(dim=-1)\n",
    "    # ^-- this will select the probability of the actual next token.\n",
    "    # Note: you can compute loss more efficiently using using F.cross_entropy\n",
    "\n",
    "    # average cross-entropy over tokens where mask == True\n",
    "    return <YOUR CODE HERE> # average loss, scalar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ME_LWUeklrs7"
   },
   "outputs": [],
   "source": [
    "dummy_loss = compute_loss(model, dummy_inp_tokens, dummy_out_tokens)\n",
    "print(\"Loss:\", dummy_loss)\n",
    "assert np.allclose(dummy_loss.item(), 7.5, rtol=0.1, atol=0.1), \"We're sorry for your loss\"\n",
    "\n",
    "# test autograd\n",
    "dummy_loss.backward()\n",
    "for name, param in model.named_parameters():\n",
    "    assert param.grad is not None and abs(param.grad.max()) != 0, f\"Param {name} received no gradients\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HpbaBpW7lrs-"
   },
   "source": [
    "### Evaluation: BLEU\n",
    "\n",
    "Machine translation is commonly evaluated with [BLEU](https://en.wikipedia.org/wiki/BLEU) score. This metric simply computes which fraction of predicted n-grams is actually present in the reference translation. It does so for n=1,2,3 and 4 and computes the geometric average with penalty if translation is shorter than reference.\n",
    "\n",
    "While BLEU [has many drawbacks](https://www.cis.upenn.edu/~ccb/publications/re-evaluating-the-role-of-bleu-in-mt-research.pdf), it still remains the most commonly used metric and one of the simplest to compute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Gb1-PhKIlrs-"
   },
   "outputs": [],
   "source": [
    "from nltk.translate.bleu_score import corpus_bleu\n",
    "def compute_bleu(model, inp_lines, out_lines, bpe_sep='@@ ', **flags):\n",
    "    \"\"\"\n",
    "    Estimates corpora-level BLEU score of model's translations given inp and reference out\n",
    "    Note: if you're serious about reporting your results, use https://pypi.org/project/sacrebleu\n",
    "    \"\"\"\n",
    "    with torch.no_grad():\n",
    "        translations, _ = model.translate_lines(inp_lines, **flags)\n",
    "        translations = [line.replace(bpe_sep, '') for line in translations]\n",
    "        actual = [line.replace(bpe_sep, '') for line in out_lines]\n",
    "        return corpus_bleu(\n",
    "            [[ref.split()] for ref in actual],\n",
    "            [trans.split() for trans in translations],\n",
    "            smoothing_function=lambda precisions, **kw: [p + 1.0 / p.denominator for p in precisions]\n",
    "            ) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gZvfid1RlrtA"
   },
   "outputs": [],
   "source": [
    "compute_bleu(model, dev_inp, dev_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nQDhGwg4lrtC"
   },
   "source": [
    "### Training loop\n",
    "\n",
    "Training encoder-decoder models isn't that different from any other models: sample batches, compute loss, backprop and update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yfwIaixHlrtI",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "from tqdm import tqdm, trange\n",
    "metrics = {'train_loss': [], 'dev_bleu': [] }\n",
    "\n",
    "model = BasicModel(inp_voc, out_voc).to(device)\n",
    "opt = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LlDT6eDUlrtL",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for _ in trange(25000):\n",
    "    step = len(metrics['train_loss']) + 1\n",
    "    batch_ix = np.random.randint(len(train_inp), size=batch_size)\n",
    "    batch_inp = inp_voc.to_matrix(train_inp[batch_ix]).to(device)\n",
    "    batch_out = out_voc.to_matrix(train_out[batch_ix]).to(device)\n",
    "    \n",
    "    <YOUR CODE: training step using batch_inp and batch_out>\n",
    "    \n",
    "    metrics['train_loss'].append((step, loss_t.item()))\n",
    "    \n",
    "    if step % 100 == 0:\n",
    "        metrics['dev_bleu'].append((step, compute_bleu(model, dev_inp, dev_out)))\n",
    "        \n",
    "        clear_output(True)\n",
    "        plt.figure(figsize=(12,4))\n",
    "        for i, (name, history) in enumerate(sorted(metrics.items())):\n",
    "            plt.subplot(1, len(metrics), i + 1)\n",
    "            plt.title(name)\n",
    "            plt.plot(*zip(*history))\n",
    "            plt.grid()\n",
    "        plt.show()\n",
    "        print(\"Mean loss=%.3f\" % np.mean(metrics['train_loss'][-10:], axis=0)[1], flush=True)\n",
    "        \n",
    "# Note: it's okay if bleu oscillates up and down as long as it gets better on average over long term (e.g. 5k batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2ahuhKVhlrtP"
   },
   "outputs": [],
   "source": [
    "assert np.mean(metrics['dev_bleu'][-10:], axis=0)[1] > 15, \"We kind of need a higher bleu BLEU from you. Kind of right now.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KyaHOpealrtS"
   },
   "outputs": [],
   "source": [
    "for inp_line, trans_line in zip(dev_inp[::500], model.translate_lines(dev_inp[::500])[0]):\n",
    "    print(inp_line)\n",
    "    print(trans_line)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "edk_oVg0lrtW"
   },
   "source": [
    "### Your Attention Required\n",
    "\n",
    "In this section we want you to improve over the basic model by implementing a simple attention mechanism.\n",
    "\n",
    "This is gonna be a two-parter: building the __attention layer__ and using it for an __attentive seq2seq model__."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Qz9aROAIlrtX"
   },
   "source": [
    "### Attention layer (3 points)\n",
    "\n",
    "Here you will have to implement a layer that computes a simple additive attention:\n",
    "\n",
    "Given encoder sequence $ h^e_0, h^e_1, h^e_2, ..., h^e_T$ and a single decoder state $h^d$,\n",
    "\n",
    "* Compute logits with a 2-layer neural network\n",
    "$$a_t = linear_{out}(tanh(linear_{e}(h^e_t) + linear_{d}(h_d)))$$\n",
    "* Get probabilities from logits, \n",
    "$$ p_t = {{e ^ {a_t}} \\over { \\sum_\\tau e^{a_\\tau} }} $$\n",
    "\n",
    "* Add up encoder states with probabilities to get __attention response__\n",
    "$$ attn = \\sum_t p_t \\cdot h^e_t $$\n",
    "\n",
    "You can learn more about attention layers in the lecture slides or [from this post](https://distill.pub/2016/augmented-rnns/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionLayer(nn.Module):\n",
    "    def __init__(self, name, enc_size, dec_size, hid_size, activ=torch.tanh):\n",
    "        \"\"\" A layer that computes additive attention response and weights \"\"\"\n",
    "        super().__init__()\n",
    "        self.name = name\n",
    "        self.enc_size = enc_size # num units in encoder state\n",
    "        self.dec_size = dec_size # num units in decoder state\n",
    "        self.hid_size = hid_size # attention layer hidden units\n",
    "        self.activ = activ       # attention layer hidden nonlinearity\n",
    "        \n",
    "        # create trainable paramteres like this:\n",
    "        self.<PARAMETER_NAME> = nn.Parameter(<INITIAL_VALUES>, requires_grad=True)\n",
    "        <...>  # you will need a couple of these\n",
    "        \n",
    "\n",
    "    def forward(self, enc, dec, inp_mask):\n",
    "        \"\"\"\n",
    "        Computes attention response and weights\n",
    "        :param enc: encoder activation sequence, float32[batch_size, ninp, enc_size]\n",
    "        :param dec: single decoder state used as \"query\", float32[batch_size, dec_size]\n",
    "        :param inp_mask: mask on enc activatons (0 after first eos), float32 [batch_size, ninp]\n",
    "        :returns: attn[batch_size, enc_size], probs[batch_size, ninp]\n",
    "            - attn - attention response vector (weighted sum of enc)\n",
    "            - probs - attention weights after softmax\n",
    "        \"\"\"\n",
    "\n",
    "        # Compute logits\n",
    "        <...>\n",
    "\n",
    "        # Apply mask - if mask is 0, logits should be -inf or -1e9\n",
    "        # You may need torch.where\n",
    "        <...>\n",
    "\n",
    "        # Compute attention probabilities (softmax)\n",
    "        probs = <...>\n",
    "\n",
    "        # Compute attention response using enc and probs\n",
    "        attn = <...>\n",
    "\n",
    "        return attn, probs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IalfpdAelrtb"
   },
   "source": [
    "### Seq2seq model with attention (3 points)\n",
    "\n",
    "You can now use the attention layer to build a network. The simplest way to implement attention is to use it in decoder phase:\n",
    "![img](https://i.imgur.com/6fKHlHb.png)\n",
    "_image from distill.pub [article](https://distill.pub/2016/augmented-rnns/)_\n",
    "\n",
    "On every step, use __previous__ decoder state to obtain attention response. Then feed concat this response to the inputs of next attention layer.\n",
    "\n",
    "The key implementation detail here is __model state__. Put simply, you can add any tensor into the list of `encode` outputs. You will then have access to them at each `decode` step. This may include:\n",
    "* Last RNN hidden states (as in basic model)\n",
    "* The whole sequence of encoder outputs (to attend to) and mask\n",
    "* Attention probabilities (to visualize)\n",
    "\n",
    "_There are, of course, alternative ways to wire attention into your network and different kinds of attention. Take a look at [this](https://arxiv.org/abs/1609.08144), [this](https://arxiv.org/abs/1706.03762) and [this](https://arxiv.org/abs/1808.03867) for ideas. And for image captioning/im2latex there's [visual attention](https://arxiv.org/abs/1502.03044)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NCKPB5JmcE6j"
   },
   "outputs": [],
   "source": [
    "class AttentiveModel(BasicModel):\n",
    "    def __init__(self, name, inp_voc, out_voc,\n",
    "                 emb_size=64, hid_size=128, attn_size=128):\n",
    "        \"\"\" Translation model that uses attention. See instructions above. \"\"\"\n",
    "        nn.Module.__init__(self)  # initialize base class to track sub-layers, trainable variables, etc.\n",
    "        self.inp_voc, self.out_voc = inp_voc, out_voc\n",
    "        self.hid_size = hid_size\n",
    "        \n",
    "        <YOUR CODE: initialize layers>\n",
    "\n",
    "    def encode(self, inp, **flags):\n",
    "        \"\"\"\n",
    "        Takes symbolic input sequence, computes initial state\n",
    "        :param inp: matrix of input tokens [batch, time]\n",
    "        :return: a list of initial decoder state tensors\n",
    "        \"\"\"\n",
    "        \n",
    "        # encode input sequence, create initial decoder states\n",
    "        <YOUR CODE>\n",
    "        \n",
    "        # apply attention layer from initial decoder hidden state\n",
    "        first_attn_probas = <...>\n",
    "        \n",
    "        # Build first state: include\n",
    "        # * initial states for decoder recurrent layers\n",
    "        # * encoder sequence and encoder attn mask (for attention)\n",
    "        # * make sure that last state item is attention probabilities tensor\n",
    "        \n",
    "        first_state = [<...>, first_attn_probas]\n",
    "        return first_state\n",
    "   \n",
    "    def decode_step(self, prev_state, prev_tokens, **flags):\n",
    "        \"\"\"\n",
    "        Takes previous decoder state and tokens, returns new state and logits for next tokens\n",
    "        :param prev_state: a list of previous decoder state tensors\n",
    "        :param prev_tokens: previous output tokens, an int vector of [batch_size]\n",
    "        :return: a list of next decoder state tensors, a tensor of logits [batch, n_tokens]\n",
    "        \"\"\"\n",
    "        \n",
    "        <YOUR CODE HERE>\n",
    "        return [new_dec_state, output_logits]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ryZCOTEslrtf"
   },
   "source": [
    "### Training attentive model\n",
    "\n",
    "Please reuse the infrastructure you've built for the regular model. I hope you didn't hard-code anything :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-YMHPgZxcFaQ"
   },
   "outputs": [],
   "source": [
    "<YOUR CODE: create AttentiveModel and training utilities>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "<YOUR CODE: training loop>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "<YOUR CODE: measure final BLEU>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing model attention (2 points)\n",
    "\n",
    "After training the attentive translation model, you can check it's sanity by visualizing its attention weights.\n",
    "\n",
    "We provided you with a function that draws attention maps using [`Bokeh`](https://bokeh.pydata.org/en/latest/index.html). Once you managed to produce something better than random noise, please save at least 3 attention maps and __submit them to anytask__ alongside this notebook to get the max grade. Saving bokeh figures as __cell outputs is not enough!__ (TAs can't see saved bokeh figures in anytask). You can save bokeh images as screenshots or using this button:\n",
    "\n",
    "![bokeh_panel](https://github.com/yandexdataschool/nlp_course/raw/2019/resources/bokeh_panel.png)\n",
    "\n",
    "__Note:__ you're not locked into using bokeh. If you prefer a different visualization method, feel free to use that instead of bokeh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import bokeh.plotting as pl\n",
    "import bokeh.models as bm\n",
    "from bokeh.io import output_notebook, show\n",
    "output_notebook()\n",
    "\n",
    "def draw_attention(inp_line, translation, probs):\n",
    "    \"\"\" An intentionally ambiguous function to visualize attention weights \"\"\"\n",
    "    inp_tokens = inp_voc.tokenize(inp_line)\n",
    "    trans_tokens = out_voc.tokenize(translation)\n",
    "    probs = probs[:len(trans_tokens), :len(inp_tokens)]\n",
    "    \n",
    "    fig = pl.figure(x_range=(0, len(inp_tokens)), y_range=(0, len(trans_tokens)),\n",
    "                    x_axis_type=None, y_axis_type=None, tools=[])\n",
    "    fig.image([probs[::-1]], 0, 0, len(inp_tokens), len(trans_tokens))\n",
    "\n",
    "    fig.add_layout(bm.LinearAxis(axis_label='source tokens'), 'above')\n",
    "    fig.xaxis.ticker = np.arange(len(inp_tokens)) + 0.5\n",
    "    fig.xaxis.major_label_overrides = dict(zip(np.arange(len(inp_tokens)) + 0.5, inp_tokens))\n",
    "    fig.xaxis.major_label_orientation = 45\n",
    "\n",
    "    fig.add_layout(bm.LinearAxis(axis_label='translation tokens'), 'left')\n",
    "    fig.yaxis.ticker = np.arange(len(trans_tokens)) + 0.5\n",
    "    fig.yaxis.major_label_overrides = dict(zip(np.arange(len(trans_tokens)) + 0.5, trans_tokens[::-1]))\n",
    "\n",
    "    show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = dev_inp[::500]\n",
    "\n",
    "trans, states = model.translate_lines(inp)\n",
    "\n",
    "# select attention probs from model state (you may need to change this for your custom model)\n",
    "# attention_probs below must have shape [batch_size, translation_length, input_length], extracted from states\n",
    "# e.g. if attention probs are at the end of each state, use np.stack([state[-1] for state in states], axis=1)\n",
    "attention_probs = <YOUR CODE>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(5):\n",
    "    draw_attention(inp[i], trans[i], attention_probs[i])\n",
    "    \n",
    "# Does it look fine already? don't forget to save images for anytask!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Note 1:__ If the attention maps are not iterpretable, try starting encoder from zeros (instead of dec_start), forcing model to use attention.\n",
    "\n",
    "__Note 2:__ If you're studying this course as a YSDA student, please submit __attention screenshots__ alongside your notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pbIIngNVlrtt"
   },
   "source": [
    "## Goind deeper (2++ points each)\n",
    "\n",
    "We want you to find the best model for the task. Use everything you know.\n",
    "\n",
    "* different recurrent units: rnn/gru/lstm; deeper architectures\n",
    "* bidirectional encoder, different attention methods for decoder (additive, dot-product, multi-head)\n",
    "* word dropout, training schedules, anything you can imagine\n",
    "* replace greedy inference with beam search\n",
    "\n",
    "For a better grasp of seq2seq We recommend you to conduct at least one experiment from one of the bullet-points or your alternative ideas. As usual, describe what you tried and what results you obtained in a short report."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "collapsed": true,
    "id": "2rzAj_xtlrtt"
   },
   "source": [
    "`[your report/log here or anywhere you please]`"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "edk_oVg0lrtW"
   ],
   "name": "practice.ipynb",
   "provenance": []
  },
  "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}