Multi Query Attention.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"accelerator":"GPU","colab":{"name":"Copy of Multi Query Attention.ipynb","provenance":[{"file_id":"https://github.com/M-e-r-c-u-r-y/pytorch-transformers/blob/master/Multi%20Query%20Attention.ipynb","timestamp":1580405583608},{"file_id":"1hhsYYBvqOH9qQDgrTCUsE5bFljCu4EZX","timestamp":1580143749829},{"file_id":"https://github.com/bentrevett/pytorch-seq2seq/blob/master/6%20-%20Attention%20is%20All%20You%20Need.ipynb","timestamp":1579866875654}],"collapsed_sections":[]},"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.4"}},"cells":[{"cell_type":"markdown","metadata":{"id":"PLHm8x8qgwa5","colab_type":"text"},"source":["# Run the below cell to download the dataset, install the correct version of pytorch and torchtext to run this notebook\n","## Below cell is intended to be run on colab as the run time environment keeps changing from session to session, if you're trying it on a personal setup you can skip this\n","### [bentrevett/pytorch-seq2seq](https://github.com/bentrevett/pytorch-seq2seq) is where the base code is taken from\n","### I've converted the Attention to use Multi Query Attention as proposed in [Fast Transformer Decoding: One Write-Head is All You Need](https://arxiv.org/abs/1911.02150) paper"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"x5JybNTQWd0T","colab":{}},"source":["!python -m spacy download en\n","!python -m spacy download de\n","!pip install torch==1.4.0\n","!pip install torchtext==0.5"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"56d3_ltqV68q"},"source":["# 6 - Attention is All You Need\n","\n","In this notebook we will be implementing a (slightly modified version) of the Transformer model from the [Attention is All You Need](https://arxiv.org/abs/1706.03762) paper. All images in this notebook will be taken from the Transformer paper. For more information about the Transformer, [see](https://www.mihaileric.com/posts/transformers-attention-in-disguise/) [these](https://jalammar.github.io/illustrated-transformer/) [three](http://nlp.seas.harvard.edu/2018/04/03/attention.html) articles.\n","\n","![](https://github.com/bentrevett/pytorch-seq2seq/blob/master/assets/transformer1.png?raw=1)\n","\n","## Introduction\n","\n","Similar to the Convolutional Sequence-to-Sequence model, the Transformer does not use any recurrence. It also does not use any convolutional layers. Instead the model is entirely made up of linear layers, attention mechanisms and normalization. \n","\n","As of January 2020, Transformers are the dominant architecture in NLP and are used to achieve state-of-the-art results for many tasks and it appears as if they will be for the near future. \n","\n","The most popular Transformer variant is [BERT](https://arxiv.org/abs/1810.04805) (**B**idirectional **E**ncoder **R**epresentations from **T**ransformers) and pre-trained versions of BERT are commonly used to replace the embedding layers - if not more - in NLP models. \n","\n","A common library used when dealing with pre-trained transformers is the [Transformers](https://huggingface.co/transformers/) library, see [here](https://huggingface.co/transformers/pretrained_models.html) for a list of all pre-trained models available.\n","\n","The differences between the implementation in this notebook and the paper are:\n","- we use a learned positional encoding instead of a static one\n","- we use the standard Adam optimizer with a static learning rate instead of one with warm-up and cool-down steps\n","- we do not use label smoothing\n","\n","We make all of these changes as they closely follow BERT's set-up and the majority of Transformer variants use a similar set-up."]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"5p8kb2BdV68u"},"source":["## Preparing the Data\n","\n","As always, let's import all the required modules and set the random seeds for reproducability."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"EKmUMoi6V68z","colab":{}},"source":["import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import torch.optim as optim\n","\n","import torchtext\n","from torchtext.datasets import TranslationDataset, Multi30k\n","from torchtext.data import Field, BucketIterator\n","\n","import matplotlib.pyplot as plt\n","import matplotlib.ticker as ticker\n","\n","import spacy\n","import numpy as np\n","\n","import random\n","import math\n","import time"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"dkeTZdL8V69C","colab":{}},"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"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"LfAtc1U-V69L"},"source":["We'll then create our tokenizers as before."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"O5yZZeszV69P","colab":{}},"source":["spacy_de = spacy.load('de')\n","spacy_en = spacy.load('en')"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"M8RsvPhYV69Z","colab":{}},"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)]"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"mTWxkuetV69i"},"source":["Our fields are the same as the previous notebook. The model expects data to be fed in with the batch dimension first, so we use `batch_first = True`. "]},{"cell_type":"code","metadata":{"colab_type":"code","id":"NosypkLCV69l","colab":{}},"source":["SRC = Field(tokenize = tokenize_de, \n","            init_token = '<sos>', \n","            eos_token = '<eos>', \n","            lower = True, \n","            batch_first = True)\n","\n","TRG = Field(tokenize = tokenize_en, \n","            init_token = '<sos>', \n","            eos_token = '<eos>', \n","            lower = True, \n","            batch_first = True)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"Ws4UIczPV69t"},"source":["We then load the Multi30k dataset and build the vocabulary."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"CWH6st_uV69w","colab":{}},"source":["train_data, valid_data, test_data = Multi30k.splits(exts = ('.de', '.en'), \n","                                                    fields = (SRC, TRG))"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"q2cJPmmeV695","colab":{}},"source":["SRC.build_vocab(train_data, min_freq = 2)\n","TRG.build_vocab(train_data, min_freq = 2)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"WqLqMQtEV6-A"},"source":["Finally, we define the device and the data iterator."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"0bKBDqmSV6-C","colab":{}},"source":["device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"rEnlQUStV6-K","colab":{}},"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)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"CRt-sw_5V6-S"},"source":["## Building the Model\n","\n","Next, we'll build the model. Like previous notebooks it is made up of an *encoder* and a *decoder*, with the encoder *encoding* the input/source sentence (in German) into *context vector* and the decoder then *decoding* this context vector to output our output/target sentence (in English). \n","\n","### Encoder\n","\n","Similar to the ConvSeq2Seq model, the Transformer's encoder does not attempt to compress the entire source sentence, $X = (x_1, ... ,x_n)$, into a single context vector, $z$. Instead it produces a sequence of context vectors, $Z = (z_1, ... , z_n)$. So, if our input sequence was 5 tokens long we would have $Z = (z_1, z_2, z_3, z_4, z_5)$. Why do we call this a sequence of context vectors and not a sequence of hidden states? A hidden state at time $t$ in an RNN has only seen tokens $x_t$ and all the tokens before it. However, each context vector here as seen all tokens at all positions within the input sequence.\n","\n","![](https://github.com/bentrevett/pytorch-seq2seq/blob/master/assets/transformer-encoder.png?raw=1)\n","\n","First, the tokens are passed through a standard embedding layer. Next, as the model has no recurrent it has no idea about the order of the tokens within the sequence. We solve this by using a second embedding layer called a *positional embedding layer*. This is a standard embedding layer where the input is not the token itself but the position of the token within the sequence, starting with the first token, the `<sos>` (start of sequence) token, in position 0. The position embedding has a \"vocabulary\" size of 100, which means our model can accept sentences up to 100 tokens long. This can be increased if we want to handle longer sentences.\n","\n","The original Transformer implementation from the Attention is All You Need paper does not learn positional embeddings. Instead it uses a fixed static embedding. Modern Transformer architectures, like BERT, use positional embeddings instead, hence we have decided to use them in these tutorials. Check out [this](http://nlp.seas.harvard.edu/2018/04/03/attention.html#positional-encoding) section to read more about the positional embeddings used in the original Transformer model.\n","\n","Next, the token and positional embeddings are elementwise summed together to get a vector which contains information about the token and also its position with in the sequence. However, before they are summed, the token embeddings are multiplied by a scaling factor which is $\\sqrt{d_{model}}$, where $d_{model}$ is the hidden dimension size, `hid_dim`. This supposedly reduces variance in the embeddings and the model is difficult to train reliably without this scaling factor. Dropout is then applied to the combined embeddings.\n","\n","The combined embeddings are then passed through $N$ *encoder layers* to get $Z$, which is then output and can be used by the decoder.\n","\n","The source mask, `src_mask`, is simply the same shape as the source sentence but has a value of 1 when the token in the source sentence is not a `<pad>` token and 0 when it is a `<pad>` token. This is used in the encoder layers to mask the multi-head attention mechanisms, which are used to calculate and apply attention over the source sentence, so the model does not pay attention to `<pad>` tokens, which contain no useful information."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"dvfmbG38V6-V","colab":{}},"source":["class Encoder(nn.Module):\n","    def __init__(self, \n","                 input_dim, \n","                 hid_dim, \n","                 n_layers, \n","                 n_heads, \n","                 pf_dim, \n","                 encoder_layer, \n","                 multi_head_attention_layer, \n","                 positionwise_feedforward_layer, \n","                 dropout, \n","                 device,\n","                 max_length = 100):\n","        super().__init__()\n","\n","        self.device = device\n","        \n","        self.tok_embedding = nn.Embedding(input_dim, hid_dim)\n","        self.pos_embedding = nn.Embedding(max_length, hid_dim)\n","        \n","        self.layers = nn.ModuleList([encoder_layer(hid_dim, \n","                                                   n_heads, \n","                                                   pf_dim, \n","                                                   multi_head_attention_layer, \n","                                                   positionwise_feedforward_layer, \n","                                                   dropout, \n","                                                   device) \n","                                     for _ in range(n_layers)])\n","        \n","        self.dropout = nn.Dropout(dropout)\n","        \n","        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)\n","        \n","    def forward(self, src, src_mask):\n","        \n","        #src = [batch size, src len]\n","        #src_mask = [batch size, src len]\n","        \n","        batch_size = src.shape[0]\n","        src_len = src.shape[1]\n","        \n","        pos = torch.arange(0, src_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)\n","        \n","        #pos = [batch size, src len]\n","        \n","        src = self.dropout((self.tok_embedding(src) * self.scale) + self.pos_embedding(pos))\n","        \n","        #src = [batch size, src len, hid dim]\n","        \n","        for layer in self.layers:\n","            src = layer(src, src_mask)\n","            \n","        #src = [batch size, src len, hid dim]\n","            \n","        return src"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"-SlpvnYzV6-c"},"source":["### Encoder Layer\n","\n","The encoder layers are where all of the \"meat\" of the encoder is contained. We first pass the source sentence and its mask into the *multi-head attention layer*, then perform dropout on it, apply a residual connection and pass it through a [Layer Normalization](https://arxiv.org/abs/1607.06450) layer. We then pass it through a *position-wise feedforward* layer and then, again, apply dropout, a residual connection and then layer normalization to get the output of this layer which is fed into the next layer. The parameters are not shared between layers. \n","\n","The mutli head attention layer is used by the encoder layer to attend to the source sentence, i.e. it is calculating and applying attention over itself instead of another sequence, hence we call it *self attention*.\n","\n","[This](https://mlexplained.com/2018/01/13/weight-normalization-and-layer-normalization-explained-normalization-in-deep-learning-part-2/) article goes into more detail about layer normalization, but the gist is that it normalizes the values of the features, i.e. across the hidden dimension, so each feature has a mean of 0 and a standard deviation of 1. This allows neural networks with a larger number of layers, like the Transformer, to be trained easier."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"LownpbuuV6-f","colab":{}},"source":["class EncoderLayer(nn.Module):\n","    def __init__(self, \n","                 hid_dim, \n","                 n_heads, \n","                 pf_dim, \n","                 multi_head_attention_layer, \n","                 positionwise_feedforward_layer, \n","                 dropout, \n","                 device):\n","        super().__init__()\n","        \n","        self.layer_norm1 = nn.LayerNorm(hid_dim)\n","        self.layer_norm2 = nn.LayerNorm(hid_dim)\n","        self.self_attention = multi_head_attention_layer(hid_dim, n_heads, dropout, device)\n","        self.positionwise_feedforward = positionwise_feedforward_layer(hid_dim, \n","                                                                       pf_dim, \n","                                                                       dropout)\n","        self.dropout = nn.Dropout(dropout)\n","        \n","    def forward(self, src, src_mask):\n","        \n","        #src = [batch size, src len, hid dim]\n","        #src_mask = [batch size, src len]\n","                \n","        #self attention\n","        _src, _ = self.self_attention(src, src, src, src_mask)\n","        \n","        #dropout, residual connection and layer norm\n","        \n","        src = self.layer_norm1(src + self.dropout(_src))\n","\n","        #src = [batch size, src len, hid dim]\n","        \n","        #positionwise feedforward\n","        _src = self.positionwise_feedforward(src)\n","        \n","        #dropout, residual and layer norm\n","        src = self.layer_norm2(src + self.dropout(_src))\n","        \n","        #src = [batch size, src len, hid dim]\n","        \n","        return src"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"z-LzTNuTV6-l"},"source":["### Mutli Query Attention Layer\n","\n","One of the key, novel concepts introduced by the Transformer paper is the *multi-head attention layer*. But this unit is memory intensive. Multi Query Attention, uses the heads only in the queries, the keys and values have only *one head*.\n","\n","![](https://github.com/bentrevett/pytorch-seq2seq/blob/master/assets/transformer-attention.png?raw=1)\n","\n","Attention can be though of as *queries*, *keys* and *values* - where the query is used with the key to get an attention vector (usually the output of a *softmax* operation and has all values between 0 and 1 which sum to 1) which is then used to get a weighted sum of the values.\n","\n","The Transformer uses *scaled dot-product attention*, where the query and key are combined by taking the dot product between them, then applying the softmax operation and scaling by $d_k$ before finally then multiplying by the value. $d_k$ is the *head dimension*, `head_dim`, which we will shortly explain further.\n","\n","$$ \\text{Attention}(Q, K, V) = \\text{Softmax} \\big( \\frac{QK^T}{\\sqrt{d_k}} \\big)V $$ \n","\n","This is similar to standard *dot product attention* but is scaled by $d_k$, which the paper states is used to stop the results of the dot products growing large, causing gradients to become too small.\n","\n","However, the scaled dot-product attention isn't simply applied to the queries, keys and values. Instead of doing a single attention application the queries have their `hid_dim` split into $h$ *heads* of $d_k$ dimension each, the keys and values are also projected to $d_k$ dimension but use only *one head* each. The scaled dot-product attention is calculated over all query heads in parallel. This means instead of paying attention to one concept per attention application, we pay attention to $h$. We then re-combine the heads into their `hid_dim` shape, thus each `hid_dim` is potentially paying attention to $h$ different concepts.\n","\n","$$ \\text{MultiHead}(Q, K, V) = \\text{Concat}(\\text{head}_1,...,\\text{head}_h)W^O $$\n","\n","$$\\text{head}_i = \\text{Attention}(QW_i^Q, KW^K, VW^V) $$\n","\n","$W^O$ is the linear layer applied at the end of the multi-head attention layer, `fc`. $W^Q, W^K, W^V$ are the linear layers `fc_q`, `fc_k` and `fc_v`.\n","\n","Walking through the module, first we calculate $QW^Q$, $KW^K$ and $VW^V$ with the linear layers, `fc_q`, `fc_k` and `fc_v`, to give us `Q`, `K` and `V`. Next, we split the `hid_dim` of the query into `n_heads` using `.view` and correctly permute them so they can be multiplied together. We then calculate the `energy` (the un-normalized attention) by multiplying `Q` and `K` together and scaling it by the square root of `head_dim`, which is calulated as `hid_dim // n_heads`. We then mask the energy so we do not pay attention over any elements of the sequeuence we shouldn't, then apply the softmax and dropout. We then apply the attention to the value *head*, `V`, before combining the `n_heads` together. Finally, we multiply this $W^O$, represented by `fc_o`. \n","\n","One thing that looks strange at first is that dropout is applied directly to the attention. This means that our attention vector will most probably not sum to 1 and we may pay full attention to a token but the attention over that token is set to 0 by dropout. This is never explained, or even mentioned, in the paper however is used by the [official implementation](https://github.com/tensorflow/tensor2tensor/) and every Transformer implementation since, [including BERT](https://github.com/google-research/bert/)."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"PoScpvACV6-o","colab":{}},"source":["class MultiQueryAttentionLayer(nn.Module):\n","    def __init__(self, hid_dim, n_heads,  dropout, device):\n","        super().__init__()\n","        \n","        assert hid_dim % n_heads == 0\n","        \n","        self.hid_dim = hid_dim\n","        self.n_heads = n_heads\n","        self.head_dim = self.hid_dim // self.n_heads\n","\n","        self.fc_q = nn.Linear( self.hid_dim, self.hid_dim)\n","        self.fc_k = nn.Linear( self.hid_dim, self.head_dim)\n","        self.fc_v = nn.Linear(self.hid_dim, self.head_dim)  \n","        self.fc_o = nn.Linear(self.hid_dim, self.hid_dim)\n","        \n","        self.dropout = nn.Dropout(dropout)\n","        \n","        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(device)\n","        \n","    def forward(self, query, key, value, mask = None):\n","        \n","        batch_size = query.shape[0]\n","        \n","        #query = [batch size, query len, hid dim]\n","        #key = [batch size, key len, hid dim]\n","        #value = [batch size, value len, hid dim]\n","               \n","        Qbank = self.fc_q(query).view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n","        Kbank = self.fc_k(key).view(batch_size, -1, 1, self.head_dim).permute(0, 2, 3, 1)\n","        Vbank = self.fc_v(value).view(batch_size, -1, 1, self.head_dim).permute(0, 2, 1, 3)   \n","        \n","        #Qbank = [batch size, n heads, query len, head dim]\n","        #Kbank = [batch size, 1, head dim, key len]\n","        #Vbank = [batch size, 1, value len, head dim]\n","\n","        energy = torch.matmul(Qbank, Kbank) / self.scale\n","\n","        #energy = [batch size, n heads, query len, key len]\n","        \n","        if mask is not None:\n","            energy = energy.masked_fill(mask == 0, -1e10)\n","        \n","        attention = F.softmax(energy, dim = -1)\n","                \n","        #attention = [batch size, n heads, query len, key len]\n","\n","        x = torch.matmul(self.dropout(attention), Vbank)\n","\n","        x = x.permute(0, 2, 1, 3).contiguous()\n","\n","        x = x.view(batch_size, -1, self.hid_dim)\n","        \n","        #x = [batch size, seq len, hid dim]\n","        \n","        x = self.fc_o(x)\n","        \n","        return x, attention"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"Ggg6I43CV6-v"},"source":["### Position-wise Feedforward Layer\n","\n","The other main block inside the encoder layer is the *position-wise feedforward layer* This is relatively simple compared to the multi-head attention layer. The input is transformed from `hid_dim` to `pf_dim`, where `pf_dim` is usually a lot larger than `hid_dim`. The original Transformer used a `hid_dim` of 512 and a `pf_dim` of 2048. The ReLU activation function and dropout are applied before it is transformed back into a `hid_dim` representation. \n","\n","Why is this used? Unfortunately, it is never explained in the paper.\n","\n","BERT uses the [GELU](https://arxiv.org/abs/1606.08415) activation function, which can be used by simply switching `F.relu` for `F.gelu`. Why did they use GELU? Again, it is never explained."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"rM11GOHUV6-w","colab":{}},"source":["class PositionwiseFeedforwardLayer(nn.Module):\n","    def __init__(self, hid_dim, pf_dim, dropout):\n","        super().__init__()\n","        \n","        self.fc_1 = nn.Linear(hid_dim, pf_dim)\n","        self.fc_2 = nn.Linear(pf_dim, hid_dim)\n","        \n","        self.dropout = nn.Dropout(dropout)\n","        \n","    def forward(self, x):\n","        \n","        #x = [batch size, seq len, hid dim]\n","        \n","        x = self.dropout(F.gelu(self.fc_1(x)))\n","        \n","        #x = [batch size, seq len, pf dim]\n","        \n","        x = self.fc_2(x)\n","        \n","        #x = [batch size, seq len, hid dim]\n","        \n","        return x"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"RfzpzssCV6-1"},"source":["### Decoder\n","\n","The objective of the decoder is to take the encoded representation of the source sentence, $Z$, and convert it into predicted tokens in the target sentence, $\\hat{Y}$. We then compare $\\hat{Y}$ with the actual tokens in the target sentence, $Y$, to calculate our loss, which will be used to calculate the gradients of our parameters and then use our optimizer to update our weights in order to improve our predictions. \n","\n","![](https://github.com/bentrevett/pytorch-seq2seq/blob/master/assets/transformer-decoder.png?raw=1)\n","\n","The decoder is similar to encoder, however it now has two multi-head attention layers. A *masked multi-head attention layer* over the target sequence, and a multi-head attention layer which uses the decoder representation as the query and the encoder representation as the key and value.\n","\n","The decoder uses positional embeddings and combines - via an elementwise sum - them with the scaled embedded target tokens, followed by dropout. Again, our positional encodings have a \"vocabulary\" of 100, which means they can accept sequences up to 100 tokens long. This can be increased if desired.\n","\n","The combined embeddings are then passed through the $N$ decoder layers, along with the encoded source, `enc_src`, and the source and target masks. Note that the number of layers in the encoder does not have to be equal to the number of layers in the decoder, even though they are both denoted by $N$.\n","\n","The decoder representation after the $N^{th}$ layer is then passed through a linear layer, `fc_out`. In PyTorch, the softmax operation is contained within our loss function, so we do not explicitly need to use a softmax layer here.\n","\n","As well as using the source mask, as we did in the encoder to prevent our model attending to `<pad>` tokens, we also use a target mask. This will be explained further in the `Seq2Seq` model which encapsulates both the encoder and decoder, but the gist of it is that it performs a similar operation as the decoder padding in the convolutional sequence-to-sequence model. As we are processing all of the target tokens at once in parallel we need a method of stopping the decoder from \"cheating\" by simply \"looking\" at what the next token in the target sequence is and outputting it. \n","\n","Our decoder layer also outputs the normalized attention values so we can later plot them to see what our model is actually paying attention to."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"bzMfdiRBV6-5","colab":{}},"source":["class Decoder(nn.Module):\n","    def __init__(self, \n","                 output_dim, \n","                 hid_dim, \n","                 n_layers, \n","                 n_heads, \n","                 pf_dim, \n","                 decoder_layer, \n","                 multi_head_attention_layer, \n","                 positionwise_feedforward_layer, \n","                 dropout, \n","                 device,\n","                 max_length = 100):\n","        super().__init__()\n","        \n","        self.device = device\n","        \n","        self.tok_embedding = nn.Embedding(output_dim, hid_dim)\n","        self.pos_embedding = nn.Embedding(max_length, hid_dim)\n","        \n","        self.layers = nn.ModuleList([decoder_layer(hid_dim, \n","                                                   n_heads, \n","                                                   pf_dim, \n","                                                   multi_head_attention_layer, \n","                                                   positionwise_feedforward_layer, \n","                                                   dropout, \n","                                                   device)\n","                                     for _ in range(n_layers)])\n","        \n","        self.fc_out = nn.Linear(hid_dim, output_dim)\n","        \n","        self.dropout = nn.Dropout(dropout)\n","        \n","        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)\n","        \n","    def forward(self, trg, enc_src, trg_mask, src_mask):\n","        \n","        #trg = [batch size, trg len]\n","        #enc_src = [batch size, src len, hid dim]\n","        #trg_mask = [batch size, trg len]\n","        #src_mask = [batch size, src len]\n","                \n","        batch_size = trg.shape[0]\n","        trg_len = trg.shape[1]\n","        \n","        pos = torch.arange(0, trg_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)\n","                            \n","        #pos = [batch size, trg len]\n","            \n","        trg = self.dropout((self.tok_embedding(trg) * self.scale) + self.pos_embedding(pos))\n","                \n","        #trg = [batch size, trg len, hid dim]\n","        \n","        for layer in self.layers:\n","            trg, attention = layer(trg, enc_src, trg_mask, src_mask)\n","        \n","        #trg = [batch size, trg len, hid dim]\n","        #attention = [batch size, n heads, trg len, src len]\n","        \n","        output = self.fc_out(trg)\n","        \n","        #output = [batch size, trg len, output dim]\n","            \n","        return output, attention"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"JJm9vHqaV6--"},"source":["### Decoder Layer\n","\n","As mentioned previously, the decoder layer is similar to the encoder layer except that it now has two multi-head attention layers, `self_attention` and `encoder_attention`. \n","\n","The first performs self-attention, as in the encoder, by using the decoder representation so far as the query, key and value. This is followed by dropout, residual connection and layer normalization. This `self_attention` layer uses the target sequence mask, `trg_mask`, in order to prevent the decoder from \"cheating\" by paying attention to tokens that are \"ahead\" of the one it is currently processing as it processes all tokens in the target sentence in parallel.\n","\n","The second is how we actually feed the encoded source sentence, `enc_src`, into our decoder. In this multi-head attention layer the queries are the decoder representations and the keys and values are the decoder representations. Here, the source mask, `src_mask` is used to prevent the multi-head attention layer from attending to `<pad>` tokens within the source sentence. This is then followed by the dropout, residual connection and layer normalization layers. \n","\n","Finally, we pass this through the position-wise feedforward layer and yet another sequence of dropout, residual connection and layer normalization.\n","\n","The decoder layer isn't introducing any new concepts, just using the same set of layers as the encoder in a slightly different way."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"9xGZL1PmV6_A","colab":{}},"source":["class DecoderLayer(nn.Module):\n","    def __init__(self, \n","                 hid_dim, \n","                 n_heads, \n","                 pf_dim, \n","                 multi_head_attention_layer, \n","                 positionwise_feedforward_layer, \n","                 dropout, \n","                 device):\n","        super().__init__()\n","        \n","        self.layer_norm1 = nn.LayerNorm(hid_dim)\n","        self.layer_norm2 = nn.LayerNorm(hid_dim)\n","        self.layer_norm3 = nn.LayerNorm(hid_dim)\n","        self.self_attention = multi_head_attention_layer(hid_dim, n_heads, dropout, device)\n","        self.encoder_attention = multi_head_attention_layer(hid_dim, n_heads, dropout, device)\n","        self.positionwise_feedforward = positionwise_feedforward_layer(hid_dim, \n","                                                                       pf_dim, \n","                                                                       dropout)\n","        self.dropout = nn.Dropout(dropout)\n","        \n","    def forward(self, trg, enc_src, trg_mask, src_mask):\n","        \n","        #trg = [batch size, trg len, hid dim]\n","        #enc_src = [batch size, src len, hid dim]\n","        #trg_mask = [batch size, trg len]\n","        #src_mask = [batch size, src len]\n","        \n","        #self attention\n","        _trg, _ = self.self_attention(trg, trg, trg, trg_mask)\n","        \n","        #dropout, residual connection and layer norm\n","        trg = self.layer_norm1(trg + self.dropout(_trg))\n","            \n","        #trg = [batch size, trg len, hid dim]\n","            \n","        #encoder attention\n","        _trg, attention = self.encoder_attention(trg, enc_src, enc_src, src_mask)\n","        \n","        #dropout, residual connection and layer norm\n","        trg = self.layer_norm2(trg + self.dropout(_trg))\n","                    \n","        #trg = [batch size, trg len, hid dim]\n","        \n","        #positionwise feedforward\n","        _trg = self.positionwise_feedforward(trg)\n","        \n","        #dropout, residual and layer norm\n","        trg = self.layer_norm3(trg + self.dropout(_trg))\n","        \n","        #trg = [batch size, trg len, hid dim]\n","        #attention = [batch size, n heads, trg len, src len]\n","        \n","        return trg, attention"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"aVDF1NjbV6_G"},"source":["### Seq2Seq\n","\n","Finally, we have the `Seq2Seq` module which encapsulates the encoder and decoder, as well as handling the creation of the masks.\n","\n","The source mask is created by checking where the source sequence is not equal to a `<pad>` token. It is 1 where the token is not a `<pad>` token and 0 when it is. It is then unsqueezed so it can be correctly broadcast when applying the mask to the `energy`, which of shape **_[batch size, n heads, seq len, seq len]_**.\n","\n","The target mask is slightly more complicated. First, we create a mask for the `<pad>` tokens, as we did for the source mask. Next, we create a \"subsequent\" mask, `trg_sub_mask`, using `torch.tril`. This creates a diagonal matrix where the elements above the diagonal will be zero and the elements below the diagonal will be set to whatever the input tensor is. In this case, the input tensor will be a tensor filled with ones. So this means our `trg_sub_mask` will look something like this (for a target with 5 tokens):\n","\n","$$\\begin{matrix}\n","1 & 0 & 0 & 0 & 0\\\\\n","1 & 1 & 0 & 0 & 0\\\\\n","1 & 1 & 1 & 0 & 0\\\\\n","1 & 1 & 1 & 1 & 0\\\\\n","1 & 1 & 1 & 1 & 1\\\\\n","\\end{matrix}$$\n","\n","This shows what each target token (row) is allowed to look at (column). The first target token has a mask of **_[1, 0, 0, 0, 0]_** which means it can only look at the first target token. The second target token has a mask of **_[1, 1, 0, 0, 0]_** which it means it can look at both the first and second target tokens. \n","\n","The \"subsequent\" mask is then logically anded with the padding mask, this combines the two masks ensuring both the subsequent tokens and the padding tokens cannot be attended to. For example if the last two tokens were `<pad>` tokens the mask would look like:\n","\n","$$\\begin{matrix}\n","1 & 0 & 0 & 0 & 0\\\\\n","1 & 1 & 0 & 0 & 0\\\\\n","1 & 1 & 1 & 0 & 0\\\\\n","1 & 1 & 1 & 0 & 0\\\\\n","1 & 1 & 1 & 0 & 0\\\\\n","\\end{matrix}$$\n","\n","After the masks are created, they used with the encoder and decoder along with the source and target sentences to get our predicted target sentence, `output`, along with the decoder's attention over the source sequence."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"bdMl2BHaV6_H","colab":{}},"source":["class Seq2Seq(nn.Module):\n","    def __init__(self, \n","                 encoder, \n","                 decoder, \n","                 src_pad_idx, \n","                 trg_pad_idx, \n","                 device):\n","        super().__init__()\n","        \n","        self.encoder = encoder\n","        self.decoder = decoder\n","        self.src_pad_idx = src_pad_idx\n","        self.trg_pad_idx = trg_pad_idx\n","        self.device = device\n","        \n","    def make_src_mask(self, src):\n","        \n","        #src = [batch size, src len]\n","        \n","        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)\n","\n","        #src_mask = [batch size, 1, 1, src len]\n","\n","        return src_mask\n","    \n","    def make_trg_mask(self, trg):\n","        \n","        #trg = [batch size, trg len]\n","        \n","        trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(3)\n","        \n","        #trg_pad_mask = [batch size, 1, trg len, 1]\n","        \n","        trg_len = trg.shape[1]\n","        \n","        trg_sub_mask = torch.tril(torch.ones((trg_len, trg_len), device = self.device)).bool()\n","        \n","        #trg_sub_mask = [trg len, trg len]\n","            \n","        trg_mask = trg_pad_mask & trg_sub_mask\n","        \n","        #trg_mask = [batch size, 1, trg len, trg len]\n","        \n","        return trg_mask\n","\n","    def forward(self, src, trg):\n","        \n","        #src = [batch size, src len]\n","        #trg = [batch size, trg len]\n","                \n","        src_mask = self.make_src_mask(src)\n","        trg_mask = self.make_trg_mask(trg)\n","        \n","        #src_mask = [batch size, 1, 1, src len]\n","        #trg_mask = [batch size, 1, trg len, trg len]\n","        \n","        enc_src = self.encoder(src, src_mask)\n","        \n","        #enc_src = [batch size, src len, hid dim]\n","                \n","        output, attention = self.decoder(trg, enc_src, trg_mask, src_mask)\n","        \n","        #output = [batch size, trg len, output dim]\n","        #attention = [batch size, n heads, trg len, src len]\n","        \n","        return output, attention"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"Jo6omBMlV6_M"},"source":["## Training the Seq2Seq Model\n","\n","We can now define our encoder and decoders. This model is significantly smaller than Transformers used in research today, but is able to be run on a single GPU quickly."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"4rh_FKGgV6_O","colab":{}},"source":["INPUT_DIM = len(SRC.vocab)\n","OUTPUT_DIM = len(TRG.vocab)\n","HID_DIM = 256\n","ENC_LAYERS = 3\n","DEC_LAYERS = 3\n","ENC_HEADS = 8\n","DEC_HEADS = 8\n","ENC_PF_DIM = 512\n","DEC_PF_DIM = 512\n","ENC_DROPOUT = 0.1\n","DEC_DROPOUT = 0.1\n","\n","enc = Encoder(INPUT_DIM, \n","              HID_DIM, \n","              ENC_LAYERS, \n","              ENC_HEADS, \n","              ENC_PF_DIM, \n","              EncoderLayer, \n","              MultiQueryAttentionLayer, \n","              PositionwiseFeedforwardLayer, \n","              ENC_DROPOUT, \n","              device)\n","\n","dec = Decoder(OUTPUT_DIM, \n","              HID_DIM, \n","              DEC_LAYERS, \n","              DEC_HEADS, \n","              DEC_PF_DIM, \n","              DecoderLayer, \n","              MultiQueryAttentionLayer, \n","              PositionwiseFeedforwardLayer, \n","              DEC_DROPOUT, \n","              device)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"4GMSCBwGV6_V"},"source":["Then, use them to define our whole sequence-to-sequence encapsulating model."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"V8UY7ieLV6_X","colab":{}},"source":["SRC_PAD_IDX = SRC.vocab.stoi[SRC.pad_token]\n","TRG_PAD_IDX = TRG.vocab.stoi[TRG.pad_token]\n","\n","model = Seq2Seq(enc, dec, SRC_PAD_IDX, TRG_PAD_IDX, device).to(device)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"3BKe4EnKV6_b"},"source":["We can check the number of parameters, noticing it is 1M (~11.11%) lesser than the Attention is all you need model."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"JzL7OMrKV6_f","colab":{}},"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')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"Gof8FweUV6_k"},"source":["The paper does not mention which weight initialization scheme was used, however Xavier uniform seems to be common amongst Transformer models, so we use it here."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"yO4Qb3KRV6_m","colab":{}},"source":["def initialize_weights(m):\n","    if hasattr(m, 'weight') and m.weight.dim() > 1:\n","        nn.init.xavier_uniform_(m.weight.data)"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"_Qdh26FfV6_q","colab":{}},"source":["model.apply(initialize_weights);"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"N0IXkb9nV6_v"},"source":["The optimizer used in the original Transformer paper uses Adam with a learning rate that has a \"warm-up\" and then a \"cool-down\" period. BERT and other Transformer models use Adam with a fixed learning rate, so we will implement that. Check [this](http://nlp.seas.harvard.edu/2018/04/03/attention.html#optimizer) link for more details about the original Transformer's learning rate schedule.\n","\n","Note that the learning rate needs to be lower than the default used by Adam or else learning is unstable."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"J7Z1nReVV6_y","colab":{}},"source":["LEARNING_RATE = 0.0005\n","\n","optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"3nzqfGhhV6_2"},"source":["Next, we define our loss function, making sure to ignore losses calculated over `<pad>` tokens."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"XanHsVMCV6_5","colab":{}},"source":["criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"p9JlEZOnV6_-"},"source":["Then, we'll define our training loop. This is the exact same as the one used in the previous tutorial.\n","\n","As we want our model to predict the `<eos>` token but not have it be an input into our model we simply slice the `<eos>` token off the end of the sequence. Thus:\n","\n","$$\\begin{align*}\n","\\text{trg} &= [sos, x_1, x_2, x_3, eos]\\\\\n","\\text{trg[:-1]} &= [sos, x_1, x_2, x_3]\n","\\end{align*}$$\n","\n","$x_i$ denotes actual target sequence element. We then feed this into the model to get a predicted sequence that should hopefully predict the `<eos>` token:\n","\n","$$\\begin{align*}\n","\\text{output} &= [y_1, y_2, y_3, eos]\n","\\end{align*}$$\n","\n","$y_i$ denotes predicted target sequence element. We then calculate our loss using the original `trg` tensor with the `<sos>` token sliced off the front, leaving the `<eos>` token:\n","\n","$$\\begin{align*}\n","\\text{output} &= [y_1, y_2, y_3, eos]\\\\\n","\\text{trg[1:]} &= [x_1, x_2, x_3, eos]\n","\\end{align*}$$\n","\n","We then calculate our losses and update our parameters as is standard."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"pjmEb_7BV7AA","colab":{}},"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[:,:-1])\n","                \n","        #output = [batch size, trg len - 1, output dim]\n","        #trg = [batch size, trg len]\n","            \n","        output_dim = output.shape[-1]\n","            \n","        output = output.contiguous().view(-1, output_dim)\n","        trg = trg[:,1:].contiguous().view(-1)\n","                \n","        #output = [batch size * trg len - 1, output dim]\n","        #trg = [batch size * trg len - 1]\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)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"7qb4qZIaV7AE"},"source":["The evaluation loop is the same as the training loop, just without the gradient calculations and parameter updates."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"bDbNu9jVV7AH","colab":{}},"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[:,:-1])\n","            \n","            #output = [batch size, trg len - 1, output dim]\n","            #trg = [batch size, trg len]\n","            \n","            output_dim = output.shape[-1]\n","            \n","            output = output.contiguous().view(-1, output_dim)\n","            trg = trg[:,1:].contiguous().view(-1)\n","            \n","            #output = [batch size * trg len - 1, output dim]\n","            #trg = [batch size * trg len - 1]\n","            \n","            loss = criterion(output, trg)\n","\n","            epoch_loss += loss.item()\n","        \n","    return epoch_loss / len(iterator)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"YMiZHfd1V7AL"},"source":["We then define a small function that we can use to tell us how long an epoch takes."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"ykZe2tAqV7AM","colab":{}},"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"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"QQbiR4gPV7AR"},"source":["Finally, we train our actual model. This model is almost 3x faster than the convolutional sequence-to-sequence model and also achieves a lower validation perplexity!"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"g3uAYy8JBs04","colab":{}},"source":["best_valid_loss = float('inf')"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"0xzOqn24V7AS","colab":{}},"source":["N_EPOCHS = 20\n","CLIP = 1\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(), 'tut6-model.pt')\n","    else:\n","        model.load_state_dict(torch.load('tut6-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}')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"5YdMnm_9V7AY"},"source":["We load our \"best\" parameters and manage to achieve a better test perplexity than all previous models."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"JQi0a5TiV7AY","colab":{}},"source":["model.load_state_dict(torch.load('tut6-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} |')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"349I5Gj8V7Ac"},"source":["## Inference\n","\n","Now we can can translations from our model with the `translate_sentence` function below.\n","\n","The steps taken are:\n","- tokenize the source sentence if it has not been tokenized (is a string)\n","- append the `<sos>` and `<eos>` tokens\n","- numericalize the source sentence\n","- convert it to a tensor and add a batch dimension\n","- create the source sentence mask\n","- feed the source sentence and mask into the encoder\n","- create a list to hold the output sentence, initialized with an `<sos>` token\n","- while we have not hit a maximum length\n","  - convert the current output sentence prediction into a tensor with a batch dimension\n","  - create a target sentence mask\n","  - place the current output, encoder output and both masks into the decoder\n","  - get next output token prediction from decoder along with attention\n","  - add prediction to current output sentence prediction\n","  - break if the prediction was an `<eos>` token\n","- convert the output sentence from indexes to tokens\n","- return the output sentence (with the `<sos>` token removed) and the attention from the last layer"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"E9u1G7JxV7Ae","colab":{}},"source":["def translate_sentence(sentence, src_field, trg_field, model, device, max_len = 50):\n","    \n","    model.eval()\n","        \n","    if isinstance(sentence, str):\n","        nlp = spacy.load('de')\n","        tokens = [token.text.lower() for token in nlp(sentence)]\n","    else:\n","        tokens = [token.lower() for token in sentence]\n","\n","    tokens = [src_field.init_token] + tokens + [src_field.eos_token]\n","        \n","    src_indexes = [src_field.vocab.stoi[token] for token in tokens]\n","\n","    src_tensor = torch.LongTensor(src_indexes).unsqueeze(0).to(device)\n","    \n","    src_mask = model.make_src_mask(src_tensor)\n","    \n","    with torch.no_grad():\n","        enc_src = model.encoder(src_tensor, src_mask)\n","\n","    trg_indexes = [trg_field.vocab.stoi[trg_field.init_token]]\n","\n","    for i in range(max_len):\n","\n","        trg_tensor = torch.LongTensor(trg_indexes).unsqueeze(0).to(device)\n","\n","        trg_mask = model.make_trg_mask(trg_tensor)\n","        \n","        with torch.no_grad():\n","            output, attention = model.decoder(trg_tensor, enc_src, trg_mask, src_mask)\n","        \n","        pred_token = output.argmax(2)[:,-1].item()\n","        \n","        trg_indexes.append(pred_token)\n","\n","        if pred_token == trg_field.vocab.stoi[trg_field.eos_token]:\n","            break\n","    \n","    trg_tokens = [trg_field.vocab.itos[i] for i in trg_indexes]\n","    \n","    return trg_tokens[1:], attention"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"TU32tqOnV7Aj"},"source":["We'll now define a function that displays the attention over the source sentence for each step of the decoding. As this model has 8 heads our model we can view the attention for each of the heads."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"0ZsFxAhRV7Ak","colab":{}},"source":["def display_attention(sentence, translation, attention, n_heads = 8, n_rows = 4, n_cols = 2):\n","    \n","    assert n_rows * n_cols == n_heads\n","    \n","    fig = plt.figure(figsize=(15,25))\n","    \n","    for i in range(n_heads):\n","        \n","        ax = fig.add_subplot(n_rows, n_cols, i+1)\n","        \n","        _attention = attention.squeeze(0)[i].cpu().detach().numpy()\n","\n","        cax = ax.matshow(_attention, cmap='bone')\n","\n","        ax.tick_params(labelsize=12)\n","        ax.set_xticklabels(['']+['<sos>']+[t.lower() for t in sentence]+['<eos>'], \n","                           rotation=45)\n","        ax.set_yticklabels(['']+translation)\n","\n","        ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n","        ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n","\n","    plt.show()\n","    plt.close()"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"whx3v2cNV7An"},"source":["First, we'll get an example from the training set."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"hvUlnXrEV7Ap","colab":{}},"source":["example_idx = 2\n","\n","src = vars(train_data.examples[example_idx])['src']\n","trg = vars(train_data.examples[example_idx])['trg']\n","\n","print(f'src = {src}')\n","print(f'trg = {trg}')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"SW6ZkPiDV7As"},"source":["Our translation looks pretty good, although our model changes \"observing his shirt\" to \"looks on\" and no mention of a shirt at all."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"JoDLSwEYV7At","colab":{}},"source":["translation, attention = translate_sentence(src, SRC, TRG, model, device)\n","\n","print(f'predicted trg = {translation}')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"i8HYh2GWV7A0"},"source":["We can see the attention from each head below. Each is certainly different, but it's difficult (perhaps impossible) to reason about what head has actually learned to pay attention to. Some heads pay full attention to \"ein\" when translating \"a\", some don't at all, and some do a little. They all seem to follow the similar \"downward staircase\" pattern, although the attention jumps around at the end due to the order of words in German (playhouse of wood) being in a different order than in the English translation (wooden playhouse)."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"Qm1jYmydV7A1","colab":{}},"source":["display_attention(src, translation, attention)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"rqGnxqy0V7A5"},"source":["Next, let's get an example the model has not been trained on from the validation set."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"kvN11Tt_V7A6","colab":{}},"source":["example_idx = 6\n","\n","src = vars(valid_data.examples[example_idx])['src']\n","trg = vars(valid_data.examples[example_idx])['trg']\n","\n","print(f'src = {src}')\n","print(f'trg = {trg}')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"9jV1YKlrV7BC"},"source":["The model translates it by switching \"is running\" to just \"running\", but it is an acceptable swap."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"QU3gt0mdV7BD","colab":{}},"source":["translation, attention = translate_sentence(src, SRC, TRG, model, device)\n","\n","print(f'predicted trg = {translation}')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"ioUyTF84V7BJ"},"source":["Again, the second, fourth and eight attention head pay full attention to \"ein\", whilst the seventh pays a some amount of attention and the rest pay no attention to it. Most of the heads seem to spread their attention over both the period and `<eos>` tokens in the source sentence when outputting the period and `<eos>` sentence in the predicted target sentence, though some seem to pay attention to tokens from near the start of the sentence."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"tEtwoAICV7BL","colab":{}},"source":["display_attention(src, translation, attention)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"XOUlgeZPV7BP"},"source":["Finally, we'll look at an example from the test data."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"cEUT5gzNV7BQ","colab":{}},"source":["example_idx = 253\n","\n","src = vars(test_data.examples[example_idx])['src']\n","trg = vars(test_data.examples[example_idx])['trg']\n","\n","print(f'src = {src}')\n","print(f'trg = {trg}')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"ijHP8wIIV7BT"},"source":["Pretty much a perfect translation with  \"young\" being changed to \"little\" and \"outside\" being changed to \"outdoors\", which gives the same meaning."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"BLXF7awPV7BV","colab":{}},"source":["translation, attention = translate_sentence(src, SRC, TRG, model, device)\n","\n","print(f'predicted trg = {translation}')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"1UTQRxIaV7Bp"},"source":["The second, fourth and eighth heads are the only ones paying full attention to \"eine\" (the feminine version of \"ein\" when translating \"a\". \n","\n","Again, all the heads generally follow the downward staircase pattern with some jumping around more than others, but it is difficult to interpret what exactly each head has learned to look for."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"CCOh2TrlV7Br","colab":{}},"source":["display_attention(src, translation, attention)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"p0ZS0breV7Bt"},"source":["## BLEU\n","\n","Finally we calculate the BLEU score for the Transformer."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"1WTERfHCV7Bu","colab":{}},"source":["from torchtext.data.metrics import bleu_score\n","\n","def calculate_bleu(data, src_field, trg_field, model, device, max_len = 50):\n","    \n","    trgs = []\n","    pred_trgs = []\n","    \n","    for datum in data:\n","        \n","        src = vars(datum)['src']\n","        trg = vars(datum)['trg']\n","        \n","        pred_trg, _ = translate_sentence(src, src_field, trg_field, model, device, max_len)\n","        \n","        #cut off <eos> token\n","        pred_trg = pred_trg[:-1]\n","        \n","        pred_trgs.append(pred_trg)\n","        trgs.append([trg])\n","        \n","    return bleu_score(pred_trgs, trgs)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"tYxwF0FWV7By"},"source":["We get a BLEU score of 32, which beats the 33.3 of the convolutional sequence-to-sequence model and 28.2 of the attention based RNN model. All this whilst having the least amount of parameters and the fastest training time!"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"Nc909xu0V7Bz","colab":{}},"source":["bleu_score = calculate_bleu(test_data, SRC, TRG, model, device)\n","\n","print(f'BLEU score = {bleu_score*100:.2f}')"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"ONTLtZ7SV7B3"},"source":["Congratulations for finishing these tutorials! I hope you've found them useful.\n","\n","If you find any mistakes or want to ask any questions about any of the code or explanations used, feel free to submit a GitHub issue and I will try to correct it ASAP. "]}]}