3 - Neural Machine Translation by Jointly Learning to Align and Translate.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3 - Neural Machine Translation by Jointly Learning to Align and Translate\n",
    "\n",
    "In this third and final notebook on sequence-to-sequence models using PyTorch and TorchText, we'll be implementing the model from [Neural Machine Translation by Jointly Learning to Align and Translate](https://arxiv.org/abs/1409.0473). This model achives our best perplexity yet, ~27 compared to ~34 for the previous model.\n",
    "\n",
    "## Introduction\n",
    "\n",
    "As a reminder, here is the general encoder-decoder model:\n",
    "\n",
    "![](assets/seq2seq1.png)\n",
    "\n",
    "In the previous model, our architecture was set-up in a way to reduce \"information compression\" by explicitly passing the context vector, $z$, to the decoder at every time-step and by passing both the context vector and input word, $y_t$, along with the hidden state, $s_t$, to the linear layer, $f$, to make a prediction.\n",
    "\n",
    "![](assets/seq2seq7.png)\n",
    "\n",
    "Even though we have reduced some of this compression, our context vector still needs to contain all of the information about the source sentence. The model implemented in this notebook avoids this compression by allowing the decoder to look at the entire source sentence (via its hidden states) at each decoding step! How does it do this? It uses *attention*. \n",
    "\n",
    "Attention works by first, calculating an attention vector, $a$, that is the length of the source sentence. The attention vector has the property that each element is between 0 and 1, and the entire vector sums to 1. We then calculate a weighted sum of our source sentence hidden states, $H$, to get a weighted source vector, $w$. \n",
    "\n",
    "$$w = \\sum_{i}a_ih_i$$\n",
    "\n",
    "We calculate a new weighted source vector every time-step when decoding, using it as input to our decoder RNN as well as the linear layer to make a prediction. We'll explain how to do all of this during the tutorial.\n",
    "\n",
    "## Preparing Data\n",
    "\n",
    "Again, the preparation is similar to last time.\n",
    "\n",
    "First we import all the required modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torchtext.datasets import TranslationDataset, Multi30k\n",
    "from torchtext.data import Field, BucketIterator\n",
    "\n",
    "import spacy\n",
    "\n",
    "import random\n",
    "import math\n",
    "import os\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the random seeds for reproducability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 1\n",
    "\n",
    "random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the German and English spaCy models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "spacy_de = spacy.load('de')\n",
    "spacy_en = spacy.load('en')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create the tokenizers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_de(text):\n",
    "    \"\"\"\n",
    "    Tokenizes German text from a string into a list of strings\n",
    "    \"\"\"\n",
    "    return [tok.text for tok in spacy_de.tokenizer(text)]\n",
    "\n",
    "def tokenize_en(text):\n",
    "    \"\"\"\n",
    "    Tokenizes English text from a string into a list of strings\n",
    "    \"\"\"\n",
    "    return [tok.text for tok in spacy_en.tokenizer(text)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fields remain the same as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC = Field(tokenize=tokenize_de, init_token='<sos>', eos_token='<eos>', lower=True)\n",
    "TRG = Field(tokenize=tokenize_en, init_token='<sos>', eos_token='<eos>', lower=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, valid_data, test_data = Multi30k.splits(exts=('.de', '.en'), fields=(SRC, TRG))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build the vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC.build_vocab(train_data, min_freq=2)\n",
    "TRG.build_vocab(train_data, min_freq=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the iterators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 128\n",
    "\n",
    "train_iterator, valid_iterator, test_iterator = BucketIterator.splits(\n",
    "    (train_data, valid_data, test_data), batch_size=BATCH_SIZE, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Seq2Seq Model\n",
    "\n",
    "### Encoder\n",
    "\n",
    "First, we'll build the encoder. Similar to the previous model, we only use a single layer GRU, however we now use a *bidirectional RNN*. With a bidirectional RNN, we have two RNNs in each layer. A *forward RNN* going over the sentence from left to right (shown below in green), and a *backward RNN* going over the sentence from right to left (yellow). All we need to do in code is set `bidirectional = True` and then pass the embedded sentence to the RNN as before. \n",
    "\n",
    "![](assets/seq2seq8.png)\n",
    "\n",
    "We now have:\n",
    "\n",
    "$$\\begin{align*}\n",
    "h_t^\\rightarrow &= \\text{EncoderGRU}^\\rightarrow(x_t^\\rightarrow,h_t^\\rightarrow)\\\\\n",
    "h_t^\\leftarrow &= \\text{EncoderGRU}^\\leftarrow(x_t^\\leftarrow,h_t^\\leftarrow)\n",
    "\\end{align*}$$\n",
    "\n",
    "Where $x_0^\\rightarrow = \\text{<sos>}, x_1^\\rightarrow = \\text{guten}$ and $x_0^\\leftarrow = \\text{<eos>}, x_1^\\leftarrow = \\text{morgen}$.\n",
    "\n",
    "As before, we only pass an input (`embedded`) to the RNN, which tells PyTorch to initialize both the forward and backward initial hidden states ($h_0^\\rightarrow$ and $h_0^\\leftarrow$, respectively) to a tensor of all zeros. We'll also get two context vectors, one from the forward RNN after it has seen the final word in the sentence, $z^\\rightarrow=h_T^\\rightarrow$, and one from the backward RNN after it has seen the first word in the sentence, $z^\\leftarrow=h_T^\\leftarrow$.\n",
    "\n",
    "The RNN returns `outputs` and `hidden`. \n",
    "\n",
    "`outputs` is of size **[src sent len, batch size, hid dim * num directions]** where the first `hid_dim` elements in the third axis are the hidden states from the top layer forward RNN, and the last `hid_dim` elements are hidden states from the top layer backward RNN. You can think of the third axis as being the forward and backward hidden states stacked on top of each other, i.e. $h_1 = [h_1^\\rightarrow; h_{T}^\\leftarrow]$, $h_2 = [h_2^\\rightarrow; h_{T-1}^\\leftarrow]$ and we can denote all stacked encoder hidden states as $H=\\{ h_1, h_2, ..., h_T\\}$.\n",
    "\n",
    "`hidden` is of size **[n layers * num directions, batch size, hid dim]**, where **[-2, :, :]** gives the top layer forward RNN hidden state after the final time-step (i.e. after it has seen the last word in the sentence) and **[-1, :, :]** gives the top layer backward RNN hidden state after the final time-step (i.e. after it has seen the first word in the sentence).\n",
    "\n",
    "As the decoder is not bidirectional, it only needs a single context vector, $z$, to use as its initial hidden state, $s_0$, and we currently have two, a forward and a backward one ($z^\\rightarrow=h_T^\\rightarrow$ and $z^\\leftarrow=h_T^\\leftarrow$, respectively). We solve this by concatenating the two context vectors together, passing them through a linear layer, $g$, and applying the $\\tanh$ activation function. \n",
    "\n",
    "$$z=\\tanh(g(h_T^\\rightarrow, h_T^\\leftarrow)) = \\tanh(g(z^\\rightarrow, z^\\leftarrow)) = s_0$$\n",
    "\n",
    "**Note**: this is actually a deviation from the paper. Instead, they feed only the first backward RNN hidden state through a linear layer to get the context vector/decoder initial hidden state. This doesn't seem to make sense to me, so I've changed it.\n",
    "\n",
    "As we want our model to look back over the whole of the source sentence we return `outputs`, the stacked forward and backward hidden states for every token in the source sentence. We also return `hidden`, which acts as our initial hidden state in the decoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.input_dim = input_dim\n",
    "        self.emb_dim = emb_dim\n",
    "        self.enc_hid_dim = enc_hid_dim\n",
    "        self.dec_hid_dim = dec_hid_dim\n",
    "        self.dropout = dropout\n",
    "        \n",
    "        self.embedding = nn.Embedding(input_dim, emb_dim)\n",
    "        \n",
    "        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional=True)\n",
    "        \n",
    "        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, src):\n",
    "        \n",
    "        #src = [src sent len, batch size]\n",
    "        \n",
    "        embedded = self.dropout(self.embedding(src))\n",
    "        \n",
    "        #embedded = [src sent len, batch size, emb dim]\n",
    "        \n",
    "        outputs, hidden = self.rnn(embedded)\n",
    "                \n",
    "        #outputs = [src sent len, batch size, hid dim * num directions]\n",
    "        #hidden = [n layers * num directions, batch size, hid dim]\n",
    "        \n",
    "        #hidden is stacked [forward_1, backward_1, forward_2, backward_2, ...]\n",
    "        #outputs are always from the last layer\n",
    "        \n",
    "        #hidden [-2, :, : ] is the last of the forwards RNN \n",
    "        #hidden [-1, :, : ] is the last of the backwards RNN\n",
    "        \n",
    "        #initial decoder hidden is final hidden state of the forwards and backwards encoder RNNs fed through a linear layer\n",
    "        hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1)))\n",
    "        \n",
    "        #outputs = [src sent len, batch size, enc hid dim * 2]\n",
    "        #hidden = [batch size, dec hid dim]\n",
    "        \n",
    "        return outputs, hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Attention\n",
    "\n",
    "Next up is the attention layer. This will take in the previous hidden state of the decoder, $s_{t-1}$, and all of the stacked forward and backward hidden states from the encoder, $H$. The layer will output an attention vector, $a_t$, that is the length of the source sentence, each element is between 0 and 1 and the entire vector sums to 1.\n",
    "\n",
    "Intuitively, this layer takes what we have decoded so far, $s_{t-1}$, and all of what we have encoded, $H$, to produce a vector, $a_t$, that represents which words in the source sentence we should pay the most attention to in order to correctly predict the next word to decode, $\\hat{y}_{t+1}$. \n",
    "\n",
    "First, we calculate the *energy* between the previous decoder hidden state and the encoder hidden states. As our encoder hidden states are a sequence of $T$ tensors, and our previous decoder hidden state is a single tensor, the first thing we do is `repeat` the previous decoder hidden state $T$ times. We then calculate the energy, $E_t$, between them by concatenating them together and passing them through a linear layer (`attn`) and a $\\tanh$ activation function. \n",
    "\n",
    "$$E_t = \\tanh(\\text{attn}(s_{t-1}, H))$$ \n",
    "\n",
    "This can be thought of as calculating how well each encoder hidden state \"matches\" the previous decoder hidden state.\n",
    "\n",
    "We currently have a **[dec hid dim, src sent len]** tensor for each example in the batch. We want this to be **[src sent len]** for each example in the batch as the attention should be over the length of the source sentence. This is achieved by multiplying the `energy` by a **[1, dec hid dim]** tensor, $v$.\n",
    "\n",
    "$$\\hat{a}_t = v E_t$$\n",
    "\n",
    "We can think of this as calculating a weighted sum of the \"match\" over all `dec_hid_dem` elements for each encoder hidden state, where the weights are learned (as we learn the parameters of $v$).\n",
    "\n",
    "Finally, we ensure the attention vector fits the constraints of having all elements between 0 and 1 and the vector summing to 1 by passing it through a $\\text{softmax}$ layer.\n",
    "\n",
    "$$a_t = \\text{softmax}(\\hat{a_t})$$\n",
    "\n",
    "This gives us the attention over the source sentence!\n",
    "\n",
    "Graphically, this looks something like below. This is for calculating the very first attention vector, where $s_{t-1} = s_0 = z$. The green/yellow blocks represent the hidden states from both the forward and backward RNNs, and the attention computation is all done within the pink block.\n",
    "\n",
    "![](assets/seq2seq9.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self, enc_hid_dim, dec_hid_dim):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.enc_hid_dim = enc_hid_dim\n",
    "        self.dec_hid_dim = dec_hid_dim\n",
    "        \n",
    "        self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)\n",
    "        self.v = nn.Parameter(torch.rand(dec_hid_dim))\n",
    "        \n",
    "    def forward(self, hidden, encoder_outputs):\n",
    "        \n",
    "        #hidden = [batch size, dec hid dim]\n",
    "        #encoder_outputs = [src sent len, batch size, enc hid dim * 2]\n",
    "        \n",
    "        batch_size = encoder_outputs.shape[1]\n",
    "        src_len = encoder_outputs.shape[0]\n",
    "        \n",
    "        #repeat encoder hidden state src_len times\n",
    "        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)\n",
    "        \n",
    "        encoder_outputs = encoder_outputs.permute(1, 0, 2)\n",
    "        \n",
    "        #hidden = [batch size, src sent len, dec hid dim]\n",
    "        #encoder_outputs = [batch size, src sent len, enc hid dim * 2]\n",
    "        \n",
    "        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2))) \n",
    "        \n",
    "        #energy = [batch size, src sent len, dec hid dim]\n",
    "        \n",
    "        energy = energy.permute(0, 2, 1)\n",
    "        \n",
    "        #energy = [batch size, dec hid dim, src sent len]\n",
    "        \n",
    "        #v = [dec hid dim]\n",
    "        \n",
    "        v = self.v.repeat(batch_size, 1).unsqueeze(1)\n",
    "        \n",
    "        #v = [batch size, 1, dec hid dim]\n",
    "                \n",
    "        attention = torch.bmm(v, energy).squeeze(1)\n",
    "        \n",
    "        #attention= [batch size, src len]\n",
    "        \n",
    "        return F.softmax(attention, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decoder\n",
    "\n",
    "Next up is the decoder. \n",
    "\n",
    "The decoder contains the attention layer, `attention`, which takes the previous hidden state, $s_{t-1}$, all of the encoder hidden states, $H$, and returns the attention vector, $a_t$.\n",
    "\n",
    "We then use this attention vector to create a weighted source vector, $w_t$, denoted by `weighted`, which is a weighted sum of the encoder hidden states, $H$, using $a_t$ as the weights.\n",
    "\n",
    "$$w_t = a_t H$$\n",
    "\n",
    "The input word (that has been embedded), $y_t$, the weighted source vector, $w_t$, and the previous decoder hidden state, $s_{t-1}$, are then all passed into the decoder RNN, with $y_t$ and $w_t$ being concatenated together.\n",
    "\n",
    "$$s_t = \\text{DecoderGRU}(y_t, w_t, s_{t-1})$$\n",
    "\n",
    "We then pass $y_t$, $w_t$ and $s_t$ through the linear layer, $f$, to make a prediction of the next word in the target sentence, $\\hat{y}_{t+1}$. This is done by concatenating them all together.\n",
    "\n",
    "$$\\hat{y}_{t+1} = f(y_t, w_t, s_t)$$\n",
    "\n",
    "The image below shows decoding the first word in an example translation.\n",
    "\n",
    "![](assets/seq2seq10.png)\n",
    "\n",
    "The green/yellow blocks show the forward/backward encoder RNNs which output $H$, the red block shows the context vector, $z = h_T = \\tanh(g(h^\\rightarrow_T,h^\\leftarrow_T)) = \\tanh(g(z^\\rightarrow, z^\\leftarrow)) = s_0$, the blue block shows the decoder RNN which outputs $s_t$, the purple block shows the linear layer, $f$, which outputs $\\hat{y}_{t+1}$ and the orange block shows the calculation of the weighted sum over $H$ by $a_t$ and outputs $w_t$. Not shown is the calculation of $a_t$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):\n",
    "        super().__init__()\n",
    "\n",
    "        self.emb_dim = emb_dim\n",
    "        self.enc_hid_dim = enc_hid_dim\n",
    "        self.dec_hid_dim = dec_hid_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.dropout = dropout\n",
    "        self.attention = attention\n",
    "        \n",
    "        self.embedding = nn.Embedding(output_dim, emb_dim)\n",
    "        \n",
    "        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)\n",
    "        \n",
    "        self.out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, input, hidden, encoder_outputs):\n",
    "             \n",
    "        #input = [batch size]\n",
    "        #hidden = [batch size, dec hid dim]\n",
    "        #encoder_outputs = [src sent len, batch size, enc hid dim * 2]\n",
    "        \n",
    "        input = input.unsqueeze(0)\n",
    "        \n",
    "        #input = [1, batch size]\n",
    "        \n",
    "        embedded = self.dropout(self.embedding(input))\n",
    "        \n",
    "        #embedded = [1, batch size, emb dim]\n",
    "        \n",
    "        a = self.attention(hidden, encoder_outputs)\n",
    "                \n",
    "        #a = [batch size, src len]\n",
    "        \n",
    "        a = a.unsqueeze(1)\n",
    "        \n",
    "        #a = [batch size, 1, src len]\n",
    "        \n",
    "        encoder_outputs = encoder_outputs.permute(1, 0, 2)\n",
    "        \n",
    "        #encoder_outputs = [batch size, src sent len, enc hid dim * 2]\n",
    "        \n",
    "        weighted = torch.bmm(a, encoder_outputs)\n",
    "        \n",
    "        #weighted = [batch size, 1, enc hid dim * 2]\n",
    "        \n",
    "        weighted = weighted.permute(1, 0, 2)\n",
    "        \n",
    "        #weighted = [1, batch size, enc hid dim * 2]\n",
    "        \n",
    "        rnn_input = torch.cat((embedded, weighted), dim=2)\n",
    "        \n",
    "        #rnn_input = [1, batch size, (enc hid dim * 2) + emb dim]\n",
    "            \n",
    "        output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))\n",
    "        \n",
    "        #output = [sent len, batch size, dec hid dim * n directions]\n",
    "        #hidden = [n layers * n directions, batch size, dec hid dim]\n",
    "        \n",
    "        #sent len, n layers and n directions will always be 1 in this decoder, therefore:\n",
    "        #output = [1, batch size, dec hid dim]\n",
    "        #hidden = [1, batch size, dec hid dim]\n",
    "        #this also means that output == hidden\n",
    "        assert (output == hidden).all()\n",
    "        \n",
    "        embedded = embedded.squeeze(0)\n",
    "        output = output.squeeze(0)\n",
    "        weighted = weighted.squeeze(0)\n",
    "        \n",
    "        output = self.out(torch.cat((output, weighted, embedded), dim=1))\n",
    "        \n",
    "        #output = [bsz, output dim]\n",
    "        \n",
    "        return output, hidden.squeeze(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Seq2Seq\n",
    "\n",
    "This is the first model where we don't have to have the encoder RNN and decoder RNN have the same hidden dimensions, however the encoder has to be bidirectional. This requirement can be removed by changing all occurences of `enc_dim * 2` to `enc_dim * 2 if encoder_is_bidirectional else enc_dim`. \n",
    "\n",
    "This seq2seq encapsulator is similar to the last two. The only difference is that the `encoder` returns both the final hidden state (which is the final hidden state from both the forward and backward encoder RNNs passed through a linear layer) to be used as the initial hidden state for the encoder, as well as every hidden state (which are the forward and backward hidden states stacked on top of each other). We also need to ensure that `hidden` and `encoder_outputs` are passed to the decoder. \n",
    "\n",
    "Briefly going over all of the decoding steps:\n",
    "\n",
    "Briefly going over all of the steps:\n",
    "- the `outputs` tensor is created to hold all predictions, $\\hat{Y}$\n",
    "- the source sequence, $X$, is fed into the encoder to receive $z$ and $H$\n",
    "- the initial decoder hidden state is set to be the `context` vector, $s_0 = z = h_T$\n",
    "- we use a batch of `<sos>` tokens as the first `input`, $y_1$\n",
    "- we then decode within a loop:\n",
    "  - inserting the input token $y_t$, previous hidden state, $s_{t-1}$, and all encoder outputs, $H$, into the decoder\n",
    "  - receiving a prediction, $\\hat{y}_{t+1}$, and a new hidden state, $s_t$\n",
    "  - we then decide if we are going to teacher force or not, setting the next input as appropriate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder, device):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.device = device\n",
    "        \n",
    "    def forward(self, src, trg, teacher_forcing_ratio=0.5):\n",
    "        \n",
    "        #src = [src sent len, batch size]\n",
    "        #trg = [trg sent len, batch size]\n",
    "        #teacher_forcing_ratio is probability to use teacher forcing\n",
    "        #e.g. if teacher_forcing_ratio is 0.75 we use teacher forcing 75% of the time\n",
    "        \n",
    "        batch_size = src.shape[1]\n",
    "        max_len = trg.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        \n",
    "        #tensor to store decoder outputs\n",
    "        outputs = torch.zeros(max_len, batch_size, trg_vocab_size).to(self.device)\n",
    "        \n",
    "        #encoder_outputs is all hidden states of the input sequence, back and forwards\n",
    "        #hidden is the final forward and backward hidden states, passed through a linear layer\n",
    "        encoder_outputs, hidden = self.encoder(src)\n",
    "                \n",
    "        #first input to the decoder is the <sos> tokens\n",
    "        output = trg[0,:]\n",
    "        \n",
    "        for t in range(1, max_len):\n",
    "            output, hidden = self.decoder(output, hidden, encoder_outputs)\n",
    "            outputs[t] = output\n",
    "            teacher_force = random.random() < teacher_forcing_ratio\n",
    "            top1 = output.max(1)[1]\n",
    "            output = (trg[t] if teacher_force else top1)\n",
    "\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the Seq2Seq Model\n",
    "\n",
    "The rest of this tutorial is very similar to the previous one.\n",
    "\n",
    "We initialise our parameters, encoder, decoder and seq2seq model (placing it on the GPU if we have one). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_DIM = len(SRC.vocab)\n",
    "OUTPUT_DIM = len(TRG.vocab)\n",
    "ENC_EMB_DIM = 256\n",
    "DEC_EMB_DIM = 256\n",
    "ENC_HID_DIM = 512\n",
    "DEC_HID_DIM = 512\n",
    "ENC_DROPOUT = 0.5\n",
    "DEC_DROPOUT = 0.5\n",
    "\n",
    "attn = Attention(ENC_HID_DIM, DEC_HID_DIM)\n",
    "enc = Encoder(INPUT_DIM, ENC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, ENC_DROPOUT)\n",
    "dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, DEC_DROPOUT, attn)\n",
    "\n",
    "model = Seq2Seq(enc, dec, device).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the number of parameters. We get an increase of almost 50% in the amount of parameters from the last model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 20,518,405 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create an optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We initialize the loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "pad_idx = TRG.vocab.stoi['<pad>']\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=pad_idx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then create the training loop..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion, clip):\n",
    "    \n",
    "    model.train()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    for i, batch in enumerate(iterator):\n",
    "        \n",
    "        src = batch.src\n",
    "        trg = batch.trg\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output = model(src, trg)\n",
    "        \n",
    "        #trg = [trg sent len, batch size]\n",
    "        #output = [trg sent len, batch size, output dim]\n",
    "        \n",
    "        output = output[1:].view(-1, output.shape[-1])\n",
    "        trg = trg[1:].view(-1)\n",
    "        \n",
    "        #trg = [(trg sent len - 1) * batch size]\n",
    "        #output = [(trg sent len - 1) * batch size, output dim]\n",
    "        \n",
    "        loss = criterion(output, trg)\n",
    "        \n",
    "        loss.backward()\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...and the evaluation loop, remembering to set the model to `eval` mode and turn off teaching forcing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, iterator, criterion):\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "    \n",
    "        for i, batch in enumerate(iterator):\n",
    "\n",
    "            src = batch.src\n",
    "            trg = batch.trg\n",
    "\n",
    "            output = model(src, trg, 0) #turn off teacher forcing\n",
    "\n",
    "            #trg = [trg sent len, batch size]\n",
    "            #output = [trg sent len, batch size, output dim]\n",
    "\n",
    "            output = output[1:].view(-1, output.shape[-1])\n",
    "            trg = trg[1:].view(-1)\n",
    "\n",
    "            #trg = [(trg sent len - 1) * batch size]\n",
    "            #output = [(trg sent len - 1) * batch size, output dim]\n",
    "\n",
    "            loss = criterion(output, trg)\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, define a timing function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def epoch_time(start_time, end_time):\n",
    "    elapsed_time = end_time - start_time\n",
    "    elapsed_mins = int(elapsed_time / 60)\n",
    "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "    return elapsed_mins, elapsed_secs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we train our model, saving the parameters that give us the best validation loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01 | Time: 1m 8s\n",
      "\tTrain Loss: 4.217 | Train PPL:  67.862\n",
      "\t Val. Loss: 3.582 |  Val. PPL:  35.928\n",
      "Epoch: 02 | Time: 1m 9s\n",
      "\tTrain Loss: 3.054 | Train PPL:  21.200\n",
      "\t Val. Loss: 3.301 |  Val. PPL:  27.150\n",
      "Epoch: 03 | Time: 1m 10s\n",
      "\tTrain Loss: 2.652 | Train PPL:  14.179\n",
      "\t Val. Loss: 3.243 |  Val. PPL:  25.620\n",
      "Epoch: 04 | Time: 1m 9s\n",
      "\tTrain Loss: 2.386 | Train PPL:  10.869\n",
      "\t Val. Loss: 3.311 |  Val. PPL:  27.424\n",
      "Epoch: 05 | Time: 1m 9s\n",
      "\tTrain Loss: 2.179 | Train PPL:   8.838\n",
      "\t Val. Loss: 3.277 |  Val. PPL:  26.497\n",
      "Epoch: 06 | Time: 1m 9s\n",
      "\tTrain Loss: 2.020 | Train PPL:   7.539\n",
      "\t Val. Loss: 3.307 |  Val. PPL:  27.294\n",
      "Epoch: 07 | Time: 1m 8s\n",
      "\tTrain Loss: 1.918 | Train PPL:   6.805\n",
      "\t Val. Loss: 3.297 |  Val. PPL:  27.035\n",
      "Epoch: 08 | Time: 1m 9s\n",
      "\tTrain Loss: 1.819 | Train PPL:   6.166\n",
      "\t Val. Loss: 3.398 |  Val. PPL:  29.890\n",
      "Epoch: 09 | Time: 1m 9s\n",
      "\tTrain Loss: 1.753 | Train PPL:   5.770\n",
      "\t Val. Loss: 3.442 |  Val. PPL:  31.259\n",
      "Epoch: 10 | Time: 1m 8s\n",
      "\tTrain Loss: 1.684 | Train PPL:   5.387\n",
      "\t Val. Loss: 3.376 |  Val. PPL:  29.260\n"
     ]
    }
   ],
   "source": [
    "N_EPOCHS = 10\n",
    "CLIP = 1\n",
    "SAVE_DIR = 'models'\n",
    "MODEL_SAVE_PATH = os.path.join(SAVE_DIR, 'tut3_model.pt')\n",
    "\n",
    "best_valid_loss = float('inf')\n",
    "\n",
    "if not os.path.isdir(f'{SAVE_DIR}'):\n",
    "    os.makedirs(f'{SAVE_DIR}')\n",
    "\n",
    "for epoch in range(N_EPOCHS):\n",
    "    \n",
    "    start_time = time.time()\n",
    "    \n",
    "    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)\n",
    "    valid_loss = evaluate(model, valid_iterator, criterion)\n",
    "    \n",
    "    end_time = time.time()\n",
    "    \n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        torch.save(model.state_dict(), MODEL_SAVE_PATH)\n",
    "    \n",
    "    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
    "    print(f'\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')\n",
    "    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we test the model on the test set using these \"best\" parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| Test Loss: 3.297 | Test PPL:  27.035 |\n"
     ]
    }
   ],
   "source": [
    "model.load_state_dict(torch.load(MODEL_SAVE_PATH))\n",
    "\n",
    "test_loss = evaluate(model, test_iterator, criterion)\n",
    "\n",
    "print(f'| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've improved on the previous model by a solid 5 perplexity points! However, this came at the cost of doubling the training time.\n",
    "\n",
    "In the next notebook, we'll be using the same architecture but using a few tricks that are applicable to all RNN architectures - packed padded sequences and masking. We'll also implement code which will allow us to look at what words in the input the RNN is paying attention to when decoding the output."
   ]
  }
 ],
 "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}