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 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 embedded input word, $d(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 Multi30k\n",
    "from torchtext.data import Field, BucketIterator\n",
    "\n",
    "import spacy\n",
    "import numpy as np\n",
    "\n",
    "import random\n",
    "import math\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the random seeds for reproducability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 1234\n",
    "\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.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, \n",
    "            init_token = '<sos>', \n",
    "            eos_token = '<eos>', \n",
    "            lower = True)\n",
    "\n",
    "TRG = Field(tokenize = tokenize_en, \n",
    "            init_token = '<sos>', \n",
    "            eos_token = '<eos>', \n",
    "            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'), \n",
    "                                                    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), \n",
    "    batch_size = BATCH_SIZE,\n",
    "    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 embedded sentence from left to right (shown below in green), and a *backward RNN* going over the embedded sentence from right to left (teal). 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(e(x_t^\\rightarrow),h_{t-1}^\\rightarrow)\\\\\n",
    "h_t^\\leftarrow &= \\text{EncoderGRU}^\\leftarrow(e(x_t^\\leftarrow),h_{t-1}^\\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 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. We can think of the third axis as being the forward and backward hidden states concatenated together 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 encoder hidden states (forward and backwards concatenated together) 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 we have 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.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 len, batch size]\n",
    "        \n",
    "        embedded = self.dropout(self.embedding(src))\n",
    "        \n",
    "        #embedded = [src len, batch size, emb dim]\n",
    "        \n",
    "        outputs, hidden = self.rnn(embedded)\n",
    "                \n",
    "        #outputs = [src 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 \n",
    "        #  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 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 len]** tensor for each example in the batch. We want this to be **[src 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 $v$ as the weights for a weighted sum of the energy across all encoder hidden states. These weights tell us how much we should attend to each token in the source sequence. The parameters of $v$ are initialized randomly, but learned with the rest of the model via backpropagation. Note how $v$ is not dependent on time, and the same $v$ is used for each time-step of the decoding. We implement $v$ as a linear layer without a bias.\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/teal 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.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)\n",
    "        self.v = nn.Linear(dec_hid_dim, 1, bias = False)\n",
    "        \n",
    "    def forward(self, hidden, encoder_outputs):\n",
    "        \n",
    "        #hidden = [batch size, dec hid dim]\n",
    "        #encoder_outputs = [src 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 decoder 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 len, dec hid dim]\n",
    "        #encoder_outputs = [batch size, src 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 len, dec hid dim]\n",
    "\n",
    "        attention = self.v(energy).squeeze(2)\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 embedded input word, $d(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 $d(y_t)$ and $w_t$ being concatenated together.\n",
    "\n",
    "$$s_t = \\text{DecoderGRU}(d(y_t), w_t, s_{t-1})$$\n",
    "\n",
    "We then pass $d(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(d(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/teal 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.output_dim = output_dim\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.fc_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 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 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 = [seq len, batch size, dec hid dim * n directions]\n",
    "        #hidden = [n layers * n directions, batch size, dec hid dim]\n",
    "        \n",
    "        #seq 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",
    "        prediction = self.fc_out(torch.cat((output, weighted, embedded), dim = 1))\n",
    "        \n",
    "        #prediction = [batch size, output dim]\n",
    "        \n",
    "        return prediction, 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 decoder, 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 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 len, batch size]\n",
    "        #trg = [trg 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",
    "        trg_len = trg.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        \n",
    "        #tensor to store decoder outputs\n",
    "        outputs = torch.zeros(trg_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",
    "        input = trg[0,:]\n",
    "        \n",
    "        for t in range(1, trg_len):\n",
    "            \n",
    "            #insert input token embedding, previous hidden state and all encoder hidden states\n",
    "            #receive output tensor (predictions) and new hidden state\n",
    "            output, hidden = self.decoder(input, hidden, encoder_outputs)\n",
    "            \n",
    "            #place predictions in a tensor holding predictions for each token\n",
    "            outputs[t] = output\n",
    "            \n",
    "            #decide if we are going to use teacher forcing or not\n",
    "            teacher_force = random.random() < teacher_forcing_ratio\n",
    "            \n",
    "            #get the highest predicted token from our predictions\n",
    "            top1 = output.argmax(1) \n",
    "            \n",
    "            #if teacher forcing, use actual next token as next input\n",
    "            #if not, use predicted token\n",
    "            input = 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": [
    "We use a simplified version of the weight initialization scheme used in the paper. Here, we will initialize all biases to zero and all weights from $\\mathcal{N}(0, 0.01)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Seq2Seq(\n",
       "  (encoder): Encoder(\n",
       "    (embedding): Embedding(7853, 256)\n",
       "    (rnn): GRU(256, 512, bidirectional=True)\n",
       "    (fc): Linear(in_features=1024, out_features=512, bias=True)\n",
       "    (dropout): Dropout(p=0.5, inplace=False)\n",
       "  )\n",
       "  (decoder): Decoder(\n",
       "    (attention): Attention(\n",
       "      (attn): Linear(in_features=1536, out_features=512, bias=True)\n",
       "      (v): Linear(in_features=512, out_features=1, bias=False)\n",
       "    )\n",
       "    (embedding): Embedding(5893, 256)\n",
       "    (rnn): GRU(1280, 512)\n",
       "    (fc_out): Linear(in_features=1792, out_features=5893, bias=True)\n",
       "    (dropout): Dropout(p=0.5, inplace=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def init_weights(m):\n",
    "    for name, param in m.named_parameters():\n",
    "        if 'weight' in name:\n",
    "            nn.init.normal_(param.data, mean=0, std=0.01)\n",
    "        else:\n",
    "            nn.init.constant_(param.data, 0)\n",
    "            \n",
    "model.apply(init_weights)"
   ]
  },
  {
   "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": 16,
   "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": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We initialize the loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRG_PAD_IDX = TRG.vocab.stoi[TRG.pad_token]\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then create the training loop..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "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 len, batch size]\n",
    "        #output = [trg len, batch size, output dim]\n",
    "        \n",
    "        output_dim = output.shape[-1]\n",
    "        \n",
    "        output = output[1:].view(-1, output_dim)\n",
    "        trg = trg[1:].view(-1)\n",
    "        \n",
    "        #trg = [(trg len - 1) * batch size]\n",
    "        #output = [(trg 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": 20,
   "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 len, batch size]\n",
    "            #output = [trg len, batch size, output dim]\n",
    "\n",
    "            output_dim = output.shape[-1]\n",
    "            \n",
    "            output = output[1:].view(-1, output_dim)\n",
    "            trg = trg[1:].view(-1)\n",
    "\n",
    "            #trg = [(trg len - 1) * batch size]\n",
    "            #output = [(trg 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": 21,
   "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": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01 | Time: 0m 57s\n",
      "\tTrain Loss: 5.016 | Train PPL: 150.847\n",
      "\t Val. Loss: 4.910 |  Val. PPL: 135.637\n",
      "Epoch: 02 | Time: 0m 56s\n",
      "\tTrain Loss: 4.126 | Train PPL:  61.952\n",
      "\t Val. Loss: 4.545 |  Val. PPL:  94.121\n",
      "Epoch: 03 | Time: 0m 53s\n",
      "\tTrain Loss: 3.472 | Train PPL:  32.191\n",
      "\t Val. Loss: 3.822 |  Val. PPL:  45.718\n",
      "Epoch: 04 | Time: 0m 54s\n",
      "\tTrain Loss: 2.906 | Train PPL:  18.283\n",
      "\t Val. Loss: 3.451 |  Val. PPL:  31.516\n",
      "Epoch: 05 | Time: 0m 53s\n",
      "\tTrain Loss: 2.501 | Train PPL:  12.196\n",
      "\t Val. Loss: 3.292 |  Val. PPL:  26.887\n",
      "Epoch: 06 | Time: 0m 54s\n",
      "\tTrain Loss: 2.207 | Train PPL:   9.089\n",
      "\t Val. Loss: 3.260 |  Val. PPL:  26.037\n",
      "Epoch: 07 | Time: 0m 53s\n",
      "\tTrain Loss: 1.963 | Train PPL:   7.124\n",
      "\t Val. Loss: 3.186 |  Val. PPL:  24.190\n",
      "Epoch: 08 | Time: 0m 53s\n",
      "\tTrain Loss: 1.746 | Train PPL:   5.734\n",
      "\t Val. Loss: 3.229 |  Val. PPL:  25.266\n",
      "Epoch: 09 | Time: 0m 56s\n",
      "\tTrain Loss: 1.579 | Train PPL:   4.848\n",
      "\t Val. Loss: 3.277 |  Val. PPL:  26.486\n",
      "Epoch: 10 | Time: 0m 55s\n",
      "\tTrain Loss: 1.468 | Train PPL:   4.340\n",
      "\t Val. Loss: 3.321 |  Val. PPL:  27.693\n"
     ]
    }
   ],
   "source": [
    "N_EPOCHS = 10\n",
    "CLIP = 1\n",
    "\n",
    "best_valid_loss = float('inf')\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(), 'tut3-model.pt')\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": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| Test Loss: 3.203 | Test PPL:  24.617 |\n"
     ]
    }
   ],
   "source": [
    "model.load_state_dict(torch.load('tut3-model.pt'))\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, but 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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}