solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Welcome to the tutorial on Transformer Models\n",
    "As mentioned in the README, this tutorial demonstrates a pytorch-based implementation of the Transformer model. I've done my best to keep the amount of prior code knowledge to a minimum, but the code is best understood with a moderate understanding of Object Oriented Programming and some familiarity with pytorch. Many thanks to [Yongrae Jo](https://github.com/dreamgonfly) for his working pytorch implementation, on which much of this code is based. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from typing import Tuple\n",
    "import random\n",
    "import os\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Part 1: Loading the dataset\n",
    "The dataset used in this tutorail is from [OpenNMT](https://github.com/OpenNMT/OpenNMT-py), which consists of source data in English and target data in German. As it is beyond the scope of the tutorial, we have created preprocessed data in the `data` directory and custom pytorch datasets in the helper's directory. The preprocessing results in encoded German and English sentences, where each unique word is mapped to one number (e.g. \"the\" = 1, \"horse\" = 2)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Helper function for DataLoader: this function makes sure all the source/input/target encodings have the same \n",
    "## lengths by padding them with a pad index (0).\n",
    "def input_target_collate_fn(batch: int):\n",
    "    PAD_INDEX = 0\n",
    "    sources_lengths = [len(sources) for sources, inputs, targets in batch]\n",
    "    inputs_lengths = [len(inputs) for sources, inputs, targets in batch]\n",
    "    targets_lengths = [len(targets) for sources, inputs, targets in batch]\n",
    "\n",
    "    sources_max_length = max(sources_lengths)\n",
    "    inputs_max_length = max(inputs_lengths)\n",
    "    targets_max_length = max(targets_lengths)\n",
    "\n",
    "    sources_padded = [sources + [PAD_INDEX] * (sources_max_length - len(sources)) for sources, inputs, targets in batch]\n",
    "    inputs_padded = [inputs + [PAD_INDEX] * (inputs_max_length - len(inputs)) for sources, inputs, targets in batch]\n",
    "    targets_padded = [targets + [PAD_INDEX] * (targets_max_length - len(targets)) for sources, inputs, targets in batch]\n",
    "\n",
    "    sources_tensor = torch.tensor(sources_padded)\n",
    "    inputs_tensor = torch.tensor(inputs_padded)\n",
    "    targets_tensor = torch.tensor(targets_padded)\n",
    "\n",
    "    return sources_tensor, inputs_tensor, targets_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers.datasets import IndexedInputTargetTranslationDataset\n",
    "DATA_DIR = 'data/example/processed'\n",
    "\n",
    "train_ds      = IndexedInputTargetTranslationDataset(data_dir=DATA_DIR, phase='train')\n",
    "validation_ds = IndexedInputTargetTranslationDataset(data_dir=DATA_DIR, phase='val')\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_ds, batch_size=20, shuffle=True, \n",
    "                                           collate_fn=input_target_collate_fn)\n",
    "valid_loader = torch.utils.data.DataLoader(validation_ds, batch_size=20, shuffle=True, \n",
    "                                           collate_fn=input_target_collate_fn)\n",
    "\n",
    "# Our sentences have been encoded to indexed inputs and outputs (where each unique word corresponds to a number) \n",
    "# Uncomment the below line to see the first tensor in the first batch\n",
    "#print(list(train_loader)[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Part 2: Building A Transformer, One Block at a Time\n",
    "#### Scaled Dot Product Attention\n",
    "The smallest building block of the transformer model presented in Attention is All You Need is Scaled Dot-Product attention. This class takes in the heads of the queries and keys, applies a scaled dot product each time it's called. \n",
    "\n",
    "#### Multi-Head Attention\n",
    "While our Scaled-Dot Product attention works well for single tensors, we really want to apply that concept to multiple attention heads at once in parallel! Enter: the MultiHeadAttention module! This module does a lot of things at once. Given a tensor of multiple queries, keys, and values it:\n",
    "1. Splits the inputs into separate blocks for each attention head\n",
    "2. Passes the split tensors into the self attention module\n",
    "3. Normalizes and applies dropout to the queries\n",
    "\n",
    "Both of these methods are implemented in one class for simplicity - I've gone outside the usual order of defining Modules (by pytting the scaled attention function first instead of `forward`). Pytorch utilizes the `forward` method of its Modules like `__call__` in pure python (while keeping track of all the fine details like gradients in the background)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, heads_count: int, \n",
    "                 d_model: int, \n",
    "                 dropout_prob: float, \n",
    "                 mode='self-attention'):\n",
    "        super().__init__()\n",
    "\n",
    "        self.d_head = d_model // heads_count\n",
    "        self.heads_count = heads_count\n",
    "        self.mode = mode\n",
    "        self.query_projection = nn.Linear(d_model, heads_count * self.d_head)\n",
    "        self.key_projection = nn.Linear(d_model, heads_count * self.d_head)\n",
    "        self.value_projection = nn.Linear(d_model, heads_count * self.d_head)\n",
    "        self.final_projection = nn.Linear(d_model, heads_count * self.d_head)\n",
    "        self.dropout = nn.Dropout(dropout_prob)\n",
    "        self.softmax = nn.Softmax(dim=3)\n",
    "\n",
    "        self.attention = None\n",
    "        self.key_projected = None\n",
    "        self.value_projected = None\n",
    "        \n",
    "    def scaled_dot_product(self, query_heads: torch.Tensor, key_heads: torch.Tensor):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "             query_heads: (batch_size, heads_count, query_len, d_head)\n",
    "             key_heads: (batch_size, heads_count, key_len, d_head)\n",
    "        \"\"\"\n",
    "        key_heads_transposed = key_heads.transpose(2, 3)  # Required for matrix multiplication\n",
    "        dot_product = torch.matmul(query_heads, key_heads_transposed)\n",
    "        attention_weights = dot_product / np.sqrt(self.d_head)  # Apply the scaling in the paper: sqrt(d_head)\n",
    "        return attention_weights\n",
    "\n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            query: (batch_size, query_len, model_dim)\n",
    "            key: (batch_size, key_len, model_dim)\n",
    "            value: (batch_size, value_len, model_dim)\n",
    "            mask: (batch_size, query_len, key_len)\n",
    "            state: DecoderState\n",
    "        \"\"\"\n",
    "        batch_size, query_len, d_model = query.size()\n",
    "\n",
    "        d_head = d_model // self.heads_count\n",
    "\n",
    "        query_projected = self.query_projection(query)\n",
    "       \n",
    "        key_projected = self.key_projection(key)\n",
    "        value_projected = self.value_projection(value)\n",
    "       \n",
    "        # For cache\n",
    "        self.key_projected = key_projected\n",
    "        self.value_projected = value_projected\n",
    "\n",
    "        batch_size, key_len, d_model = key_projected.size()\n",
    "        batch_size, value_len, d_model = value_projected.size()\n",
    "        \n",
    "        # Tensor.view is pytorch's implementation of \"resize\", which we need to apply attention to the correct \n",
    "        # dimensions\n",
    "        query_heads = query_projected.view(batch_size, query_len, self.heads_count, d_head).transpose(1, 2) \n",
    "        key_heads = key_projected.view(batch_size, key_len, self.heads_count, d_head).transpose(1, 2) \n",
    "        value_heads = value_projected.view(batch_size, value_len, self.heads_count, d_head).transpose(1, 2) \n",
    "\n",
    "        attention_weights = self.scaled_dot_product(query_heads, key_heads) \n",
    "        \n",
    "        if mask is not None:\n",
    "            mask_expanded = mask.unsqueeze(1).expand_as(attention_weights)\n",
    "            attention_weights = attention_weights.masked_fill(mask_expanded, -1e18)\n",
    "\n",
    "        self.attention = self.softmax(attention_weights)  # Save attention to the object\n",
    "        attention_dropped = self.dropout(self.attention)\n",
    "        \n",
    "        context_heads = torch.matmul(attention_dropped, value_heads)    # Scale by our context,\n",
    "        context_sequence = context_heads.transpose(1, 2).contiguous()   # reshape,\n",
    "        context = context_sequence.view(batch_size, query_len, d_model) # And project!\n",
    "        final_output = self.final_projection(context)\n",
    "        return final_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Another small building block\n",
    "The only other small block of a transformer model is the position-wise feed forward network - a simple\n",
    "sequential model with fully connected layers and ReLU activation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PointwiseFeedForwardNetwork(nn.Module):\n",
    "    def __init__(self, d_ff, d_model, dropout_prob):\n",
    "        super(PointwiseFeedForwardNetwork, self).__init__()\n",
    "\n",
    "        self.feed_forward = nn.Sequential(\n",
    "            nn.Linear(d_model, d_ff),\n",
    "            nn.Dropout(dropout_prob),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(d_ff, d_model),\n",
    "            nn.Dropout(dropout_prob),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.feed_forward(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Layers, Layers, and more Layers\n",
    "We now need to build up our encoder and decoder from our multi-head attention! The only thing standing between us and that is the setup of the indivudal layers of the decoders and encoders. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In the paper, after each layer of the model they apply normalization. This is where we do that.\n",
    "class LayerNormalization(nn.Module):\n",
    "\n",
    "    def __init__(self, features_count, epsilon=1e-6):\n",
    "        super(LayerNormalization, self).__init__()\n",
    "\n",
    "        self.gain = nn.Parameter(torch.ones(features_count))\n",
    "        self.bias = nn.Parameter(torch.zeros(features_count))\n",
    "        self.epsilon = epsilon\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        mean = x.mean(dim=-1, keepdim=True)\n",
    "        std = x.std(dim=-1, keepdim=True)\n",
    "\n",
    "        return self.gain * (x - mean) / (std + self.epsilon) + self.bias\n",
    "\n",
    "class Sublayer(nn.Module):\n",
    "    '''A layer that applies normalization'''\n",
    "    def __init__(self, sublayer, d_model):\n",
    "        super().__init__()\n",
    "\n",
    "        self.sublayer = sublayer\n",
    "        self.layer_normalization = LayerNormalization(d_model)\n",
    "\n",
    "    def forward(self, *args):\n",
    "        x = args[0]\n",
    "        x = self.sublayer(*args) + x\n",
    "        return self.layer_normalization(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we define an individual encoder and decoder layer: one attention layer, one dropout layer, one feedforward layer in sequence! Instead of dropout in the decoder, we use a memory attention layer for better performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEncoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, heads_count, d_ff, dropout_prob):\n",
    "        super().__init__()\n",
    "\n",
    "        self.self_attention_layer = Sublayer(MultiHeadAttention(heads_count, d_model, dropout_prob), d_model)\n",
    "        self.pointwise_feedforward_layer = Sublayer(PointwiseFeedForwardNetwork(d_ff, d_model, dropout_prob), d_model)\n",
    "        self.dropout = nn.Dropout(dropout_prob)\n",
    "\n",
    "    def forward(self, sources, sources_mask):\n",
    "        sources = self.self_attention_layer(sources, sources, sources, sources_mask)\n",
    "        sources = self.dropout(sources)\n",
    "        sources = self.pointwise_feedforward_layer(sources)\n",
    "\n",
    "        return sources\n",
    "\n",
    "class TransformerDecoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, heads_count, d_ff, dropout_prob):\n",
    "        super().__init__()\n",
    "        self.self_attention_layer = Sublayer(MultiHeadAttention(heads_count, d_model, dropout_prob), d_model)\n",
    "        self.memory_attention_layer = Sublayer(MultiHeadAttention(heads_count, d_model, dropout_prob), d_model)\n",
    "        self.pointwise_feedforward_layer = Sublayer(PointwiseFeedForwardNetwork(d_ff, d_model, dropout_prob), d_model)\n",
    "\n",
    "    def forward(self, inputs, memory, memory_mask, inputs_mask):\n",
    "        inputs = self.self_attention_layer(inputs, inputs, inputs, inputs_mask)\n",
    "        # The decoder gets to remember the encoder!\n",
    "        inputs = self.memory_attention_layer(inputs, memory, memory, memory_mask) \n",
    "        inputs = self.pointwise_feedforward_layer(inputs)\n",
    "        return inputs\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Encoder and Decoder, Separately!\n",
    "With those layers defined, we can finally build our Encoder and Decoder with as many layers as we like!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEncoder(nn.Module):\n",
    "\n",
    "    def __init__(self, layers_count, d_model, heads_count, d_ff, dropout_prob, embedding):\n",
    "        super().__init__()\n",
    "\n",
    "        self.d_model = d_model\n",
    "        self.embedding = embedding\n",
    "        self.encoder_layers = nn.ModuleList(\n",
    "            [TransformerEncoderLayer(d_model, heads_count, d_ff, dropout_prob) for _ in range(layers_count)]  \n",
    "        )  # We loop through, staacking encoder layers on top of ecah other.\n",
    "\n",
    "    def forward(self, sources, mask):\n",
    "        \"\"\"\n",
    "        args:\n",
    "           sources: embedded_sequence, (batch_size, seq_len, embed_size)\n",
    "        \"\"\"\n",
    "        sources = self.embedding(sources)\n",
    "\n",
    "        for encoder_layer in self.encoder_layers:\n",
    "            sources = encoder_layer(sources, mask)\n",
    "\n",
    "        return sources\n",
    "\n",
    "\n",
    "class TransformerDecoder(nn.Module):\n",
    "\n",
    "    def __init__(self, layers_count, d_model, heads_count, d_ff, dropout_prob, embedding):\n",
    "        super().__init__()\n",
    "\n",
    "        self.d_model = d_model\n",
    "        self.embedding = embedding\n",
    "        self.decoder_layers = nn.ModuleList(\n",
    "            [TransformerDecoderLayer(d_model, heads_count, d_ff, dropout_prob) for _ in range(layers_count)]\n",
    "        )\n",
    "        # These last few lines allow our model to make predictions, based on embedding weights.\n",
    "        self.generator = nn.Linear(embedding.embedding_dim, embedding.num_embeddings)\n",
    "        self.generator.weight = self.embedding.weight\n",
    "\n",
    "    def forward(self, inputs, memory, memory_mask, inputs_mask=None, state=None):\n",
    "        inputs = self.embedding(inputs)\n",
    "\n",
    "        for layer_index, decoder_layer in enumerate(self.decoder_layers):\n",
    "            inputs = decoder_layer(inputs, memory, memory_mask, inputs_mask)\n",
    "            \n",
    "        generated = self.generator(inputs)  \n",
    "        return generated, state\n",
    "    \n",
    "    def init_decoder_state(self, **args):\n",
    "        return DecoderState()\n",
    "    \n",
    "class DecoderState:\n",
    "\n",
    "    def __init__(self):\n",
    "        self.previous_inputs = torch.tensor([])\n",
    "        self.layer_caches = defaultdict(lambda: {'self-attention': None, 'memory-attention': None})\n",
    "\n",
    "    def update_state(self, layer_index, layer_mode, key_projected, value_projected):\n",
    "        self.layer_caches[layer_index][layer_mode] = {\n",
    "            'key_projected': key_projected,\n",
    "            'value_projected': value_projected\n",
    "        }\n",
    "\n",
    "    def beam_update(self, positions):\n",
    "        for layer_index in self.layer_caches:\n",
    "            for mode in ('self-attention', 'memory-attention'):\n",
    "                if self.layer_caches[layer_index][mode] is not None:\n",
    "                    for projection in self.layer_caches[layer_index][mode]:\n",
    "                        cache = self.layer_caches[layer_index][mode][projection]\n",
    "                        if cache is not None:\n",
    "                            cache.data.copy_(cache.data.index_select(0, positions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Encoder and Decoder, Together - The Transformer Module\n",
    "At long last, we can finally put our transformer together! We use some helper functions for to generate our masks so that we aren't paying attention to the pad tokens (0s), and then can build our transformer from our encoder and decoders!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helpers for our masks\n",
    "PAD_TOKEN_INDEX = 0\n",
    "\n",
    "def pad_masking(x, target_len):\n",
    "    # x: (batch_size, seq_len)\n",
    "    batch_size, seq_len = x.size()\n",
    "    padded_positions = x == PAD_TOKEN_INDEX  # (batch_size, seq_len)\n",
    "    pad_mask = padded_positions.unsqueeze(1).expand(batch_size, target_len, seq_len)\n",
    "    return pad_mask\n",
    "\n",
    "\n",
    "def subsequent_masking(x):\n",
    "    batch_size, seq_len = x.size()\n",
    "    subsequent_mask = np.triu(np.ones(shape=(seq_len, seq_len)), k=1).astype('bool')\n",
    "    subsequent_mask = torch.tensor(subsequent_mask).to(x.device)\n",
    "    subsequent_mask = subsequent_mask.unsqueeze(0).expand(batch_size, seq_len, seq_len)\n",
    "    return subsequent_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "\n",
    "    def __init__(self, encoder, decoder):\n",
    "        super().__init__()\n",
    "\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "    def forward(self, sources, inputs):\n",
    "        batch_size, sources_len = sources.size()\n",
    "        batch_size, inputs_len = inputs.size()\n",
    "\n",
    "        sources_mask = pad_masking(sources, sources_len)\n",
    "        memory_mask = pad_masking(sources, inputs_len)\n",
    "        inputs_mask = subsequent_masking(inputs) | pad_masking(inputs, inputs_len)\n",
    "\n",
    "        memory = self.encoder(sources, sources_mask)  # (batch_size, seq_len, d_model)\n",
    "        outputs, state = self.decoder(inputs, memory, memory_mask, inputs_mask)  # (batch_size, seq_len, d_model)\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Part 3: Training the Model\n",
    "\n",
    "Since pytorch training is outside the scope of the tutorial, I've written the training loop and accuracy metrics in the helper files, and included comments about what they do as I import them. We start with some constants that we'll use to define our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Constants - feel free to change to see how they affect the model!\n",
    "### Training\n",
    "num_epochs = 100            # How many times to loop through all train/validation data\n",
    "d_model = 128               # Number of dimensions for the models\n",
    "layers_count = 1            # Number of layers in the encoders/decoders\n",
    "heads_count = 2             # Number of attention heads (must divide d_model evenly)\n",
    "d_ff = 128                  # feed-forward dimensions\n",
    "dropout_prob = 0.1          # how much dropout\n",
    "label_smoothing = 0.1       # how much smoothing to apply to our labels. Higher = smoother\n",
    "clip_grads = True           # Clip the gradients (to prevent from exploding)\n",
    "lr = 0.001                  # Learning rate for Adam optimizer\n",
    "seed = 3621                 # Let our model be deterministic (set to None if you want to have an unseeded model)\n",
    "\n",
    "### Logging\n",
    "output_dir = \"./models/\"    # Where to save the model\n",
    "save_mode = 'best'          # Save only the best model (can be changed to 'all' to save all models)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up all our random number generators, logging directories, and GPU (if we have one)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "if seed is not None:\n",
    "    torch.manual_seed(seed)\n",
    "    torch.backends.cudnn.benchmark = False  \n",
    "    np.random.seed(seed)\n",
    "    random.seed(seed)\n",
    "\n",
    "os.makedirs(output_dir, exist_ok=True)  # Makes if not exists\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load in our data to our custom dictionary (to allow the model to access it efficiently)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers.dictionary import IndexDictionary\n",
    "\n",
    "source_dictionary = IndexDictionary.load(DATA_DIR, mode='source')\n",
    "target_dictionary = IndexDictionary.load(DATA_DIR, mode='target')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use our sinusoidal positional encoding as they do in the paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers.embeddings import PositionalEncoding\n",
    "\n",
    "source_embedding = PositionalEncoding(\n",
    "    num_embeddings=source_dictionary.vocabulary_size,\n",
    "    embedding_dim=d_model,\n",
    "    dim=d_model)\n",
    "\n",
    "target_embedding = PositionalEncoding(\n",
    "    num_embeddings=target_dictionary.vocabulary_size,\n",
    "    embedding_dim=d_model,\n",
    "    dim=d_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build our encoders and decoders, and then combine them into our transformer model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = TransformerEncoder(\n",
    "        layers_count=layers_count,\n",
    "        d_model=d_model,\n",
    "        d_ff=d_ff,\n",
    "        dropout_prob=dropout_prob,\n",
    "        embedding=source_embedding, \n",
    "        heads_count=heads_count\n",
    ")\n",
    "\n",
    "decoder = TransformerDecoder(\n",
    "    layers_count=layers_count,\n",
    "    d_model=d_model,\n",
    "    heads_count=heads_count,\n",
    "    d_ff=128,\n",
    "    dropout_prob=dropout_prob,\n",
    "    embedding=target_embedding\n",
    ")\n",
    "\n",
    "model = Transformer(encoder, decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A label smoothing loss that lets us weight getting \"close enough\" better, and a metric that counts the number of correct words that weren't predicted as our pad index. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers.loss import LabelSmoothingLoss\n",
    "from helpers.Accuracy import AccuracyMetric\n",
    "\n",
    "loss_fn = LabelSmoothingLoss(label_smoothing, vocabulary_size=target_dictionary.vocabulary_size)\n",
    "accuracy_metric = AccuracyMetric()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr)  # Adam optimizer (very famous ML optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A class that loops through training and validation loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers.Trainer import Trainer\n",
    "\n",
    "\n",
    "training_config = {\n",
    "    \"print_every\": 1,\n",
    "    \"save_every\": 1,\n",
    "    \"device\": device,\n",
    "    \"clip_grads\": True\n",
    "}\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    train_dataloader = train_loader,\n",
    "    val_dataloader = valid_loader,\n",
    "    loss_function = loss_fn, \n",
    "    metric_function = accuracy_metric,\n",
    "    optimizer=optimizer,\n",
    "    logger=None, \n",
    "    run_name='tutorial', \n",
    "    save_config=None, \n",
    "    save_checkpoint=f'{output_dir}/model.ckpt', \n",
    "    config = training_config\n",
    ") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                   | 0/500 [00:00<?, ?it/s]\n",
      "  0%|                                                   | 0/150 [00:00<?, ?it/s]\n",
      "100%|█████████████████████████████████████████| 500/500 [03:41<00:00,  2.26it/s]\n",
      "100%|█████████████████████████████████████████| 150/150 [00:54<00:00,  2.77it/s]\n",
      "100%|█████████████████████████████████████████| 500/500 [03:39<00:00,  2.28it/s]\n",
      "100%|█████████████████████████████████████████| 150/150 [00:54<00:00,  2.77it/s]\n",
      "100%|█████████████████████████████████████████| 500/500 [03:39<00:00,  2.27it/s]\n",
      "100%|█████████████████████████████████████████| 150/150 [00:52<00:00,  2.87it/s]\n",
      "100%|█████████████████████████████████████████| 500/500 [03:45<00:00,  2.22it/s]\n",
      "  5%|██▎                                        | 8/150 [00:03<00:48,  2.92it/s]"
     ]
    }
   ],
   "source": [
    "trainer.run(num_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing the Model\n",
    "We have our test data still, let's see it's predictions on some of the test set! To save time, I've run the training loop myself for about 10 hours (100 epochs on my personal computer)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers.beam import Beam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Predictor:\n",
    "\n",
    "    def __init__(self, preprocess, postprocess, model, checkpoint_filepath, max_length=30, beam_size=8):\n",
    "        self.preprocess = preprocess\n",
    "        self.postprocess = postprocess\n",
    "        self.model = model\n",
    "        self.max_length = max_length\n",
    "        self.beam_size = beam_size\n",
    "\n",
    "        self.model.eval()\n",
    "        checkpoint = torch.load(checkpoint_filepath, map_location='cpu')\n",
    "        self.model.load_state_dict(checkpoint)\n",
    "\n",
    "    def predict_one(self, source, num_candidates=5):\n",
    "        source_preprocessed = self.preprocess(source)\n",
    "        source_tensor = torch.tensor(source_preprocessed).unsqueeze(0)  # why unsqueeze?\n",
    "        length_tensor = torch.tensor(len(source_preprocessed)).unsqueeze(0)\n",
    "\n",
    "        sources_mask = pad_masking(source_tensor, source_tensor.size(1))\n",
    "        memory_mask = pad_masking(source_tensor, 1)\n",
    "        memory = self.model.encoder(source_tensor, sources_mask)\n",
    "\n",
    "        decoder_state = self.model.decoder.init_decoder_state()\n",
    "\n",
    "\n",
    "        # Repeat beam_size times\n",
    "        memory_beam = memory.detach().repeat(self.beam_size, 1, 1)  # (beam_size, seq_len, hidden_size)\n",
    "\n",
    "        beam = Beam(beam_size=self.beam_size, min_length=0, n_top=num_candidates, ranker=None)\n",
    "\n",
    "        for _ in range(self.max_length):\n",
    "\n",
    "            new_inputs = beam.get_current_state().unsqueeze(1)  # (beam_size, seq_len=1)\n",
    "            decoder_outputs, decoder_state = self.model.decoder(new_inputs, memory_beam,\n",
    "                                                                            memory_mask,\n",
    "                                                                            state=decoder_state)\n",
    "            # decoder_outputs: (beam_size, target_seq_len=1, vocabulary_size)\n",
    "            # attentions['std']: (target_seq_len=1, beam_size, source_seq_len)\n",
    "\n",
    "            attention = self.model.decoder.decoder_layers[-1].memory_attention_layer.sublayer.attention\n",
    "            beam.advance(decoder_outputs.squeeze(1), attention)\n",
    "\n",
    "            beam_current_origin = beam.get_current_origin()  # (beam_size, )\n",
    "            decoder_state.beam_update(beam_current_origin)\n",
    "\n",
    "            if beam.done():\n",
    "                break\n",
    "\n",
    "        scores, ks = beam.sort_finished(minimum=num_candidates)\n",
    "        hypothesises, attentions = [], []\n",
    "        for i, (times, k) in enumerate(ks[:num_candidates]):\n",
    "            hypothesis, attention = beam.get_hypothesis(times, k)\n",
    "            hypothesises.append(hypothesis)\n",
    "            attentions.append(attention)\n",
    "\n",
    "        self.attentions = attentions\n",
    "        self.hypothesises = [[token.item() for token in h] for h in hypothesises]\n",
    "        hs = [self.postprocess(h) for h in self.hypothesises]\n",
    "        return list(reversed(hs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source = \"The model predicts this sentence .\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = Predictor(\n",
    "    preprocess=IndexedInputTargetTranslationDataset.preprocess(source_dictionary),\n",
    "    postprocess=lambda x: ' '.join([token for token in target_dictionary.tokenify_indexes(x) if token != '<EndSent>']),\n",
    "    model=model,\n",
    "    checkpoint_filepath='./models/model.ckpt'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for index, candidate in enumerate(predictor.predict_one(source, num_candidates=5)):\n",
    "    print(f'Candidate {index + 1} : {candidate}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}