notebooks/Conditional-Text-Generation-with-Transformers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conditional Text Generation with Transformers\n",
    "\n",
    "The following code is inspired in the notebook provided by TensorFlow at https://www.tensorflow.org/tutorials/text/transformer\n",
    "\n",
    "Here we use a simple Transformer model with 2 encoder layers and 2 decoder layers. The goal is to generate a conversation based on the training data.\n",
    "\n",
    "For further details about Transformers architecture, read the paper *Attention is all you need*: https://arxiv.org/abs/1706.03762"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare and import modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing installation: transformers 3.0.2\n",
      "Uninstalling transformers-3.0.2:\n",
      "  Successfully uninstalled transformers-3.0.2\n",
      "Collecting transformers\n",
      "  Using cached transformers-3.0.2-py3-none-any.whl (769 kB)\n",
      "Requirement already satisfied: sentencepiece!=0.1.92 in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from transformers) (0.1.91)\n",
      "Requirement already satisfied: sacremoses in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from transformers) (0.0.43)\n",
      "Requirement already satisfied: tokenizers==0.8.1.rc1 in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from transformers) (0.8.1rc1)\n",
      "Requirement already satisfied: tqdm>=4.27 in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from transformers) (4.44.1)\n",
      "Requirement already satisfied: requests in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from transformers) (2.23.0)\n",
      "Requirement already satisfied: filelock in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from transformers) (3.0.12)\n",
      "Requirement already satisfied: numpy in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from transformers) (1.18.3)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from transformers) (2020.5.14)\n",
      "Requirement already satisfied: packaging in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from transformers) (20.4)\n",
      "Requirement already satisfied: six in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from sacremoses->transformers) (1.14.0)\n",
      "Requirement already satisfied: click in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from sacremoses->transformers) (7.1.1)\n",
      "Requirement already satisfied: joblib in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from sacremoses->transformers) (0.14.1)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from requests->transformers) (1.25.8)\n",
      "Requirement already satisfied: chardet<4,>=3.0.2 in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from requests->transformers) (3.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from requests->transformers) (2020.6.20)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from requests->transformers) (2.9)\n",
      "Requirement already satisfied: pyparsing>=2.0.2 in /Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages (from packaging->transformers) (2.4.6)\n",
      "Installing collected packages: transformers\n",
      "Successfully installed transformers-3.0.2\n"
     ]
    }
   ],
   "source": [
    "!pip uninstall -y transformers\n",
    "!pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/dimasdmm/opt/anaconda3/lib/python3.7/site-packages/tensorflow_addons/utils/ensure_tf_install.py:44: UserWarning: You are currently using a nightly version of TensorFlow (2.2.0-dev20200422). \n",
      "TensorFlow Addons offers no support for the nightly versions of TensorFlow. Some things might work, some other might not. \n",
      "If you encounter a bug, do not file an issue on GitHub.\n",
      "  UserWarning,\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import re\n",
    "import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pickle\n",
    "import random\n",
    "\n",
    "import tensorflow as tf\n",
    "from tokenizers import ByteLevelBPETokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set these constants to save the model and tokenizer later\n",
    "MODEL_NAME = 'simple-transformers'\n",
    "LANGUAGE = 'english'\n",
    "\n",
    "# Detect environment\n",
    "if os.environ.get('KAGGLE_KERNEL_RUN_TYPE',''):\n",
    "    print('Detected Kaggle environment')\n",
    "    ARTIFACTS_PATH = 'artifacts/'\n",
    "    DATA_PATH = '/kaggle/input/<dataset-name>/'\n",
    "else:\n",
    "    ARTIFACTS_PATH = '../artifacts/'\n",
    "    DATA_PATH = '../data/'\n",
    "    \n",
    "if not os.path.exists(ARTIFACTS_PATH):\n",
    "    os.makedirs(ARTIFACTS_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import data\n",
    "\n",
    "Data source: https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_FILE = 'shakespeare.txt'\n",
    "DATASET_PATH = DATA_PATH + DATASET_FILE\n",
    "\n",
    "# Read text\n",
    "with open(DATASET_PATH, 'r', encoding='utf-8') as fp:\n",
    "    text = fp.read()\n",
    "\n",
    "# Fix white spaces\n",
    "text = re.sub(r\"[ ]+\", ' ', text)\n",
    "text = re.sub(r\"\\n[ ]+\", \"\\n\", text)\n",
    "text = re.sub(r\"[ ]+\\n\", \"\\n\", text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenize & encode data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a BPE tokenizer\n",
    "tokenizer = ByteLevelBPETokenizer()\n",
    "special_tokens = [\n",
    "    '<s>',\n",
    "    '<pad>',\n",
    "    '</s>',\n",
    "    '<unk>',\n",
    "    '<mask>',\n",
    "]\n",
    "\n",
    "# Train tokenize in out dataset\n",
    "tokenizer.train(files=[DATASET_PATH], vocab_size=52_000, min_frequency=2, special_tokens=special_tokens)\n",
    "\n",
    "# Save tokenizer to disk, so we can load it later\n",
    "tokenizer.save_model(ARTIFACTS_PATH, \"%s-%s\" % (LANGUAGE, MODEL_NAME))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Obtain IDs of special tokens\n",
    "START_TOKEN = tokenizer.token_to_id('<s>')\n",
    "END_TOKEN = tokenizer.token_to_id('</s>')\n",
    "PAD_TOKEN = tokenizer.token_to_id('<pad>')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_dataset(text, tokenizer):\n",
    "    stokens = tokenizer.encode(text).ids\n",
    "    \n",
    "    n_data = len(stokens) - MAX_LEN + 3\n",
    "    print('To process: %d' % (n_data))\n",
    "    \n",
    "    X_data = []\n",
    "    y_data = []\n",
    "    \n",
    "    i = 0\n",
    "    while i < len(stokens) - MAX_LEN:\n",
    "        if i % 100000 == 0:\n",
    "            print('- Processed %d of %d' % (i, n_data))\n",
    "        \n",
    "        if len(stokens) - i < 20:\n",
    "            break\n",
    "        elif len(stokens) - i < MAX_LEN-3:\n",
    "            seq_len = len(stokens) - i\n",
    "        else:\n",
    "            seq_len = random.randint(20, MAX_LEN-3)\n",
    "        \n",
    "        input_ids = stokens[i:(i+seq_len)]\n",
    "        input_ids = [START_TOKEN] + input_ids + [END_TOKEN] + [PAD_TOKEN]*(MAX_LEN-seq_len-2)\n",
    "        input_ids = np.asarray(input_ids, dtype='int32')\n",
    "        X_data.append(input_ids)\n",
    "        \n",
    "        target_ids = stokens[(i+1):(i+seq_len+1)]\n",
    "        target_ids = target_ids + [END_TOKEN] + [PAD_TOKEN]*(MAX_LEN-seq_len-1)\n",
    "        target_ids = np.asarray(target_ids, dtype='int32')\n",
    "        y_data.append(target_ids)\n",
    "        \n",
    "        i = i + 1\n",
    "\n",
    "    # To numpy\n",
    "    X_data = np.asarray(X_data, dtype='int32')\n",
    "    y_data = np.asarray(y_data, dtype='int32')\n",
    "    \n",
    "    return tf.data.Dataset.from_tensor_slices((X_data, y_data)).batch(MAX_LEN+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "To process: 306383\n",
      "- Processed 0 of 306383\n",
      "- Processed 100000 of 306383\n",
      "- Processed 200000 of 306383\n",
      "- Processed 300000 of 306383\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "dataset = build_dataset(text, tokenizer)\n",
    "print('Done!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Masking\n",
    "\n",
    "Mask all the pad tokens in the batch of sequence. It ensures that the model does not treat padding as the input. The mask indicates where pad value 0 is present: it outputs a 1 at those locations, and a 0 otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_padding_mask(seq):\n",
    "    seq = tf.cast(tf.math.equal(seq, 0), tf.float32)\n",
    "\n",
    "    # add extra dimensions to add the padding\n",
    "    # to the attention logits.\n",
    "    return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The look-ahead mask is used to mask the future tokens in a sequence. In other words, the mask indicates which entries should not be used.\n",
    "\n",
    "This means that to predict the third word, only the first and second word will be used. Similarly to predict the fourth word, only the first, second and the third word will be used and so on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_look_ahead_mask(size):\n",
    "    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n",
    "    return mask  # (seq_len, seq_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define layers\n",
    "\n",
    "This code comes from TensorFlow documentation. For further details, read https://www.tensorflow.org/tutorials/text/transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_angles(pos, i, d_model):\n",
    "    angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n",
    "    return pos * angle_rates\n",
    "\n",
    "\n",
    "def positional_encoding(position, d_model):\n",
    "    angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n",
    "                            np.arange(d_model)[np.newaxis, :],\n",
    "                            d_model)\n",
    "\n",
    "    # apply sin to even indices in the array; 2i\n",
    "    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
    "\n",
    "    # apply cos to odd indices in the array; 2i+1\n",
    "    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
    "\n",
    "    pos_encoding = angle_rads[np.newaxis, ...]\n",
    "\n",
    "    return tf.cast(pos_encoding, dtype=tf.float32)\n",
    "\n",
    "\n",
    "def scaled_dot_product_attention(q, k, v, mask):\n",
    "    \"\"\"Calculate the attention weights.\n",
    "    q, k, v must have matching leading dimensions.\n",
    "    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.\n",
    "    The mask has different shapes depending on its type(padding or look ahead) \n",
    "    but it must be broadcastable for addition.\n",
    "\n",
    "    Args:\n",
    "        q: query shape == (..., seq_len_q, depth)\n",
    "        k: key shape == (..., seq_len_k, depth)\n",
    "        v: value shape == (..., seq_len_v, depth_v)\n",
    "        mask: Float tensor with shape broadcastable \n",
    "              to (..., seq_len_q, seq_len_k). Defaults to None.\n",
    "\n",
    "    Returns:\n",
    "        output, attention_weights\n",
    "    \"\"\"\n",
    "\n",
    "    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)\n",
    "\n",
    "    # scale matmul_qk\n",
    "    dk = tf.cast(tf.shape(k)[-1], tf.float32)\n",
    "    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n",
    "\n",
    "    # add the mask to the scaled tensor.\n",
    "    if mask is not None:\n",
    "        scaled_attention_logits += (mask * -1e9)  \n",
    "\n",
    "    # softmax is normalized on the last axis (seq_len_k) so that the scores\n",
    "    # add up to 1.\n",
    "    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)\n",
    "\n",
    "    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)\n",
    "\n",
    "    return output, attention_weights\n",
    "\n",
    "\n",
    "def point_wise_feed_forward_network(d_model, dff):\n",
    "    return tf.keras.Sequential([\n",
    "        tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)\n",
    "        tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.d_model = d_model\n",
    "\n",
    "        assert d_model % self.num_heads == 0\n",
    "\n",
    "        self.depth = d_model // self.num_heads\n",
    "\n",
    "        self.wq = tf.keras.layers.Dense(d_model)\n",
    "        self.wk = tf.keras.layers.Dense(d_model)\n",
    "        self.wv = tf.keras.layers.Dense(d_model)\n",
    "\n",
    "        self.dense = tf.keras.layers.Dense(d_model)\n",
    "        \n",
    "    def split_heads(self, x, batch_size):\n",
    "        \"\"\"Split the last dimension into (num_heads, depth).\n",
    "        Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)\n",
    "        \"\"\"\n",
    "        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
    "        return tf.transpose(x, perm=[0, 2, 1, 3])\n",
    "    \n",
    "    def call(self, v, k, q, mask):\n",
    "        batch_size = tf.shape(q)[0]\n",
    "\n",
    "        q = self.wq(q) # (batch_size, seq_len, d_model)\n",
    "        k = self.wk(k) # (batch_size, seq_len, d_model)\n",
    "        v = self.wv(v) # (batch_size, seq_len, d_model)\n",
    "\n",
    "        q = self.split_heads(q, batch_size) # (batch_size, num_heads, seq_len_q, depth)\n",
    "        k = self.split_heads(k, batch_size) # (batch_size, num_heads, seq_len_k, depth)\n",
    "        v = self.split_heads(v, batch_size) # (batch_size, num_heads, seq_len_v, depth)\n",
    "\n",
    "        # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)\n",
    "        # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)\n",
    "        scaled_attention, attention_weights = scaled_dot_product_attention(\n",
    "            q, k, v, mask)\n",
    "\n",
    "        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)\n",
    "\n",
    "        concat_attention = tf.reshape(scaled_attention, \n",
    "                                      (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)\n",
    "\n",
    "        output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)\n",
    "\n",
    "        return output, attention_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "\n",
    "        self.mha = MultiHeadAttention(d_model, num_heads)\n",
    "        self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
    "\n",
    "        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "\n",
    "        self.dropout1 = tf.keras.layers.Dropout(rate)\n",
    "        self.dropout2 = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "    def call(self, x, training, mask):\n",
    "        attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)\n",
    "        attn_output = self.dropout1(attn_output, training=training)\n",
    "        out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)\n",
    "\n",
    "        ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)\n",
    "        ffn_output = self.dropout2(ffn_output, training=training)\n",
    "        out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)\n",
    "\n",
    "        return out2\n",
    "\n",
    "class Encoder(tf.keras.layers.Layer):\n",
    "    def __init__(self, num_layers, d_model, num_heads, dff, vocab_size, rate=0.1):\n",
    "        super(Encoder, self).__init__()\n",
    "\n",
    "        maximum_position_encoding = vocab_size\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.embedding = tf.keras.layers.Embedding(vocab_size, d_model)\n",
    "        self.pos_encoding = positional_encoding(maximum_position_encoding, \n",
    "                                                self.d_model)\n",
    "\n",
    "        self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) \n",
    "                           for _ in range(num_layers)]\n",
    "\n",
    "        self.dropout = tf.keras.layers.Dropout(rate)\n",
    "        \n",
    "    def call(self, x, training, mask):\n",
    "        seq_len = tf.shape(x)[1]\n",
    "\n",
    "        # adding embedding and position encoding.\n",
    "        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)\n",
    "        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
    "        x += self.pos_encoding[:, :seq_len, :]\n",
    "\n",
    "        x = self.dropout(x, training=training)\n",
    "\n",
    "        for i in range(self.num_layers):\n",
    "            x = self.enc_layers[i](x, training, mask)\n",
    "\n",
    "        return x  # (batch_size, input_seq_len, d_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "\n",
    "        self.mha1 = MultiHeadAttention(d_model, num_heads)\n",
    "        self.mha2 = MultiHeadAttention(d_model, num_heads)\n",
    "\n",
    "        self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
    "\n",
    "        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "\n",
    "        self.dropout1 = tf.keras.layers.Dropout(rate)\n",
    "        self.dropout2 = tf.keras.layers.Dropout(rate)\n",
    "        self.dropout3 = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "    \n",
    "    def call(self, x, enc_output, training, \n",
    "             look_ahead_mask, padding_mask):\n",
    "        # enc_output.shape == (batch_size, input_seq_len, d_model)\n",
    "\n",
    "        attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)\n",
    "        attn1 = self.dropout1(attn1, training=training)\n",
    "        out1 = self.layernorm1(attn1 + x)\n",
    "\n",
    "        attn2, attn_weights_block2 = self.mha2(\n",
    "            enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)\n",
    "        attn2 = self.dropout2(attn2, training=training)\n",
    "        out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)\n",
    "\n",
    "        ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)\n",
    "        ffn_output = self.dropout3(ffn_output, training=training)\n",
    "        out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)\n",
    "\n",
    "        return out3, attn_weights_block1, attn_weights_block2\n",
    "\n",
    "\n",
    "class Decoder(tf.keras.layers.Layer):\n",
    "    def __init__(self, num_layers, d_model, num_heads, dff, vocab_size, rate=0.1):\n",
    "        super(Decoder, self).__init__()\n",
    "        \n",
    "        maximum_position_encoding = vocab_size\n",
    "\n",
    "        self.d_model = d_model\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.embedding = tf.keras.layers.Embedding(vocab_size, d_model)\n",
    "        self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)\n",
    "\n",
    "        self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) \n",
    "                           for _ in range(num_layers)]\n",
    "        self.dropout = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "    def call(self, x, enc_output, training, \n",
    "             look_ahead_mask, padding_mask):\n",
    "\n",
    "        seq_len = tf.shape(x)[1]\n",
    "        attention_weights = {}\n",
    "\n",
    "        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)\n",
    "        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
    "        x += self.pos_encoding[:, :seq_len, :]\n",
    "\n",
    "        x = self.dropout(x, training=training)\n",
    "\n",
    "        for i in range(self.num_layers):\n",
    "            x, block1, block2 = self.dec_layers[i](x, enc_output, training,\n",
    "                                                   look_ahead_mask, padding_mask)\n",
    "            attention_weights['decoder_layer{}_block1'.format(i+1)] = block1\n",
    "            attention_weights['decoder_layer{}_block2'.format(i+1)] = block2\n",
    "\n",
    "        # x.shape == (batch_size, target_seq_len, d_model)\n",
    "        return x, attention_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(tf.keras.Model):\n",
    "    def __init__(self, num_layers, d_model, num_heads, dff, vocab_size, rate=0.1):\n",
    "        super(Transformer, self).__init__()\n",
    "\n",
    "        self.encoder = Encoder(num_layers, d_model, num_heads, dff, \n",
    "                               vocab_size, rate)\n",
    "        self.decoder = Decoder(num_layers, d_model, num_heads, dff, \n",
    "                               vocab_size, rate)\n",
    "        self.final_layer = tf.keras.layers.Dense(vocab_size)\n",
    "    \n",
    "    def call(self, inp, tar, training, enc_padding_mask, \n",
    "             look_ahead_mask, dec_padding_mask):\n",
    "\n",
    "        enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)\n",
    "\n",
    "        # dec_output.shape == (batch_size, tar_seq_len, d_model)\n",
    "        dec_output, attention_weights = self.decoder(\n",
    "            tar, enc_output, training, look_ahead_mask, dec_padding_mask)\n",
    "\n",
    "        # Define output\n",
    "        final_output = self.final_layer(dec_output) # (batch_size, tar_seq_len, vocab_size)\n",
    "\n",
    "        return final_output, attention_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([64, 36, 8500])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_transformer = Transformer(\n",
    "    num_layers=2, d_model=512, num_heads=8, dff=2048, vocab_size=8500)\n",
    "\n",
    "temp_input = tf.random.uniform((64, 38), dtype=tf.int64, minval=0, maxval=200)\n",
    "temp_target = tf.random.uniform((64, 36), dtype=tf.int64, minval=0, maxval=200)\n",
    "\n",
    "fn_out, _ = sample_transformer(temp_input, temp_target, training=False, \n",
    "                               enc_padding_mask=None, \n",
    "                               look_ahead_mask=None,\n",
    "                               dec_padding_mask=None)\n",
    "\n",
    "fn_out.shape  # (batch_size, tar_seq_len, vocab_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperparameters\n",
    "\n",
    "Due to lack of time and resources, here we define a little Transformer model. Feel free to experiment with other values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_layers = 2\n",
    "d_model = 128\n",
    "dff = 128\n",
    "num_heads = 2\n",
    "\n",
    "vocab_size = tokenizer.get_vocab_size()\n",
    "dropout_rate = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizer\n",
    "\n",
    "According to the formula of the paper *Attention is all you need*, here we use the Adam optimizer with a custom learning rate scheduler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
    "    def __init__(self, d_model, warmup_steps=4000):\n",
    "        super(CustomSchedule, self).__init__()\n",
    "\n",
    "        self.d_model = d_model\n",
    "        self.d_model = tf.cast(self.d_model, tf.float32)\n",
    "\n",
    "        self.warmup_steps = warmup_steps\n",
    "    \n",
    "    def __call__(self, step):\n",
    "        arg1 = tf.math.rsqrt(step)\n",
    "        arg2 = step * (self.warmup_steps ** -1.5)\n",
    "\n",
    "        return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = CustomSchedule(d_model)\n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, \n",
    "                                     epsilon=1e-9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'Train Step')"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZUAAAEGCAYAAACtqQjWAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjMsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+AADFEAAAgAElEQVR4nO3de3xU5bXw8d9KQhJyBZIA4RISIIBBETVS7zeKorbSerRC+/bYVsvbVtt67KmX0/N6+nrq29qbrVXbWsVajxUp1RZb71KvVSCiIheBZAISbplwCSRAQpL1/rGfwBBmkkkyk5lk1vfzySd79uXZayaQlWc/z15bVBVjjDEmEpJiHYAxxpiBw5KKMcaYiLGkYowxJmIsqRhjjIkYSyrGGGMiJiXWAcRSfn6+FhcXxzoMY4zpV9599906VS0Iti2hk0pxcTEVFRWxDsMYY/oVEdkcaptd/jLGGBMxllSMMcZEjCUVY4wxEWNJxRhjTMRYUjHGGBMxUU0qIjJbRNaLSKWI3BZke5qIPOm2LxOR4oBtt7v160XkkoD1C0SkVkRWhzjnv4uIikh+NN6TMcaY0KKWVEQkGbgfuBQoA+aJSFmH3a4D9qjqROAe4G53bBkwF5gKzAYecO0B/N6tC3bOscAs4OOIvhljjDFhiWZPZQZQqao+VW0GFgJzOuwzB3jULS8GZoqIuPULVbVJVauBStceqvo6sDvEOe8BbgEGZD1/VWXRii00NLXEOhRjjAkqmkllNLAl4HWNWxd0H1VtAeqBvDCPPYaIXAFsVdUPuthvvohUiEiF3+8P533Ejfe37OWWP6/i1sWrYh2KMcYEFc2kIkHWdexBhNonnGOPNiKSAXwPuKOroFT1QVUtV9XygoKgVQbi1se7DwDw0rqdMY7EGGOCi2ZSqQHGBrweA2wLtY+IpAC5eJe2wjk20ASgBPhARDa5/VeKyMhexB93qvyNADS3tLHFJRhjjIkn0UwqK4BSESkRkVS8gfclHfZZAlzrlq8Clqr3fOMlwFw3O6wEKAWWhzqRqn6oqsNVtVhVi/GS0qmquiOybym2qvwNiOvDPbd6e2yDMcaYIKKWVNwYyY3AC8A6YJGqrhGRO934B8DDQJ6IVAI3A7e5Y9cAi4C1wPPADaraCiAiTwBvA5NFpEZErovWe4g3Pn8j508qYOqoHJ5bPaDypTFmgIhqlWJVfRZ4tsO6OwKWDwFXhzj2LuCuIOvnhXHe4u7GGu/a2pTqugbOmpDH6cXD+MkL69lef5DC3MGxDs0YY46wO+r7iW31Bzl0uI3xBZlceqI3VPS89VaMMXHGkko/4XOD9BMKshhfkMWUkdn8bZWNqxhj4osllX6iyt8AwPiCTADmTB/Nu5v3sHlXYyzDMsaYY1hS6Sd8/kay01MoyEoDYM70UYjAX97rbKa1Mcb0LUsq/USVv4HxBVmIm1M8ashgzijJ4+n3avBmYRtjTOxZUuknfP5GJuRnHrPus6eOZtOuA7y3ZW+MojLGmGNZUukHGppa2LHvEBOGZx2z/tITR5KWksTTK7fGKDJjjDmWJZV+oNrN/BrfoaeSnT6IWWUjeGbVNppaWmMRmjHGHMOSSj/gq/NmfnXsqQBcXT6WvQcO8+IaKzJpjIk9Syr9QFVtA0kC4/Iyjtt27sR8xgwdzB+X2XPJjDGxZ0mlH6iqa2TM0AzSUpKP25aUJMybUcTbvl343L0sxhgTK5ZU+oGq2gYmFGSG3H51+RhSkoSFK7aE3McYY/qCJZU419ambNrVyPiC48dT2g3PTmdW2QgWv1tjA/bGmJiypBLn2gtJTugkqQDMm1HE7sZmKzJpjIkpSypxrv1pj+M7ufwFcM7EfEryM1nw1ia7w94YEzOWVOJc++B7Vz2VpCThK2cX88GWvaz8eE9fhGaMMcexpBLnqvwNZKenkJ+V2uW+/3LaGHIHD+KhN6r7IDJjjDmeJZU45/M3HlNIsjMZqSnMm1HEC2t2sGX3gT6IzhhjjmVJJc75/I2dTifu6NqzxpEkwu//uSl6QRljTAhRTSoiMltE1otIpYjcFmR7mog86bYvE5HigG23u/XrReSSgPULRKRWRFZ3aOsnIvKRiKwSkadFZEg031tfOFJIsovxlECFuYO5fFohT67YQv2Bw1GMzhhjjhe1pCIiycD9wKVAGTBPRMo67HYdsEdVJwL3AHe7Y8uAucBUYDbwgGsP4PduXUcvASeq6jRgA3B7RN9QDFQfeYRw+D0VgK+dP4GGphYe+aeNrRhj+lY0eyozgEpV9alqM7AQmNNhnznAo255MTBTvMGDOcBCVW1S1Wqg0rWHqr4O7O54MlV9UVVb3Mt3gDGRfkN97egjhMPvqQCcUJjDrLIRLHizmv2HrLdijOk70Uwqo4HAuiE1bl3QfVxCqAfywjy2M18Bngu2QUTmi0iFiFT4/f5uNNn3fP7QhSS78s2LJrLvUAuPvbM5CpEZY0xw0UwqwaYrdbwrL9Q+4Rwb/KQi3wNagMeDbVfVB1W1XFXLCwoKwmkyZqr8jYwdFryQZFemjRnC+ZMKeOiNag40t3R9gDHGREA0k0oNMDbg9RhgW6h9RCQFyMW7tBXOsccRkWuBTwFf0AFwW3mVv+G4B3N1x7dmTmR3YzOPv2Nl8Y0xfSOaSWUFUCoiJSKSijfwvqTDPkuAa93yVcBSlwyWAHPd7LASoBRY3tnJRGQ2cCtwhar2+5s02tqU6rrGbs386ui0ccM4tzSfB16tZJ+NrRhj+kDUkoobI7kReAFYByxS1TUicqeIXOF2exjIE5FK4GbgNnfsGmARsBZ4HrhBVVsBROQJ4G1gsojUiMh1rq37gGzgJRF5X0R+E6331he27j1IU0tbtwfpO7p19hT2HDjM7173RSgyY4wJLSWajavqs8CzHdbdEbB8CLg6xLF3AXcFWT8vxP4TexVsnPHV9Ww6cUcnjs7lU9MKeeiNar545jiGZ6dHIjxjjAnK7qiPU1W1PZtOHMx3Lp7M4dY27lta2eu2jDGmM5ZU4pSvLvxCkl0pyc/kmtPH8sdlH7PJ9YCMMSYaLKnEKa/mV3iFJMPx7ZmlpKUk8YO/r4tIe8YYE4wllThV5W/o8sFc3TE8J51vzizl5XU7eXV9bcTaNcaYQJZU4lBDUws79zX1ajpxMF8+u5iS/EzufGYtzS1tEW3bGGPAkkpcOvq0x8j1VADSUpK549Nl+Ooa+b0VmzTGRIEllTjkO/Jc+sj2VAAunDycmVOG88uXN7Kj/lDE2zfGJDZLKnGoqheFJMNxx6fLaFXl//x1NQOgmo0xJo5YUolDvl4UkgzHuLxM/u2Tk3hp7U6eW70jKucwxiQmSypxqMrfEPFB+o6uO6eEE0fncMdf19gTIo0xEWNJJc60F5LsTXXicKQkJ3H3v0xjz4Fm7np2bVTPZYxJHJZU4kx7IckJw6PbUwGYOiqX+eeNZ1FFDf+we1eMMRFgSSXOHHmEcJR7Ku2+PbOUySOyuWXxKnY1NPXJOY0xA5cllTgTzenEwaQPSuYXc6dTf+Awtz/1oc0GM8b0iiWVOOOrayAnQoUkw3VCYQ63zJ7Mi2t3sqhiS5+d1xgz8FhSiTNVtY2Mj2AhyXB95ewSzpqQx/99Zu2RO/qNMaa7LKnEGV9d9KcTB5OUJPzscyeTlpLENx5fycHm1j6PwRjT/1lSiSP7Dx1m576miFYn7o7C3MHcc8101u/cz3/+xe62N8Z0nyWVOFIdoUcI98YFk4fzzYtK+fPKGp5cYeMrxpjuiWpSEZHZIrJeRCpF5LYg29NE5Em3fZmIFAdsu92tXy8ilwSsXyAitSKyukNbw0TkJRHZ6L4PjeZ7i4aqI9WJ+/7yV6Bvzyzl3NJ87liyhtVb62MaizGmf4laUhGRZOB+4FKgDJgnImUddrsO2KOqE4F7gLvdsWXAXGAqMBt4wLUH8Hu3rqPbgFdUtRR4xb3uV3z+RpIEiqJUSDJcyUnCL66ZTn5mKl/9QwW1+62asTEmPNHsqcwAKlXVp6rNwEJgTod95gCPuuXFwEzxpj3NARaqapOqVgOVrj1U9XVgd5DzBbb1KPCZSL6ZvuDzN1IUxUKS3ZGXlcbvri1n74HDfPUP73LosA3cG2O6Fs2kMhoIvChf49YF3UdVW4B6IC/MYzsaoarbXVvbgeHBdhKR+SJSISIVfr8/zLfSN7xHCMf20legqaNy+cXc6XywZS/fXbzKBu6NMV2KZlIJdqNFx99KofYJ59geUdUHVbVcVcsLCgoi0WREtLpCkrEcpA/mkqkjuWX2ZJ75YBv3vlIZ63CMMXEumkmlBhgb8HoMsC3UPiKSAuTiXdoK59iOdopIoWurEOhXFRK3uUKS8dRTaff18ydw5amjueflDSyyGWHGmE5EM6msAEpFpEREUvEG3pd02GcJcK1bvgpYqt41liXAXDc7rAQoBZZ3cb7Atq4F/hqB99Bn+rqQZHeICD+6chrnluZz21OreGntzliHZIyJU1FLKm6M5EbgBWAdsEhV14jInSJyhdvtYSBPRCqBm3EztlR1DbAIWAs8D9ygqq0AIvIE8DYwWURqROQ619aPgFkishGY5V73G+2FJPui5H1PpKYk8Zv/dRonjRnCjX9cyfLqYHMljDGJThJ58LW8vFwrKipiHQYA33v6Q575YBsf/NfFfV73qzt2NzZz1W/+iX9/E0/OP5OyUTmxDskY08dE5F1VLQ+2ze6ojxM+fyMThvd9IcnuGpaZymPXfYKstBS+8NA7rNu+L9YhGWPiiCWVOFHlb2B8fnxe+upo9JDBPPHVM0hLSeYLDy1j/Y79sQ7JGBMnLKnEgf2HDlO7P3aFJHuiOD+TJ+afwaBk4fO/e4cNOy2xGGMsqcSFI4P0cTiduDMl+Zk88dUzSE7yEotdCjPGhJVUROQcEfmyWy5w03xNhPjq2gtJ9p+eSrvxBVk8Mf8MUpKSuOa3b/PuZpsVZkwi6zKpiMh/AbcCt7tVg4D/iWZQicbnbyQ5SWJeSLKnJhRksfjrZ5KXlcYXHlrGq+v71X2nxpgICqen8lngCqARQFW3AdnRDCrRVPkbGDt0cFwUkuypMUMz+NPXzmR8fhZf/UMFz3zQVQEEY8xAFE5SaXZ3uSuAiPS/azRxzudv7HfjKcHkZ6Wx8H+fwSljh/Kthe/x4OtVVoTSmAQTTlJZJCK/BYaIyFeBl4GHohtW4mhtU3x1jf1q5ldnctIH8YfrZnDZiYX8v2c/4j+e/pDDrW2xDssY00dSutpBVX8qIrOAfcBk4A5VfSnqkSWIbXsP0hynhSR7Kn1QMr+adwrF+Rnc/48qPt59gAe+cBq5gwfFOjRjTJSFM1B/t6q+pKrfVdV/V9WXROTuvgguEcTLI4QjLSlJ+O4lU/jJVdNYXr2bKx94i011jbEOyxgTZeFc/poVZN2lkQ4kUVW5e1QGyuWvjq4uH8tj132CXY3NfPq+N3nZKhwbM6CFTCoi8nUR+RCvGvCqgK9qYFXfhTiw+fwN5A4eRF5maqxDiZozxufxzI3nMC4vg+v/UMHPXlxPa5sN4BszEHU2pvJH4Dngh7iS9M5+VbU73CLEe4RwZtwXkuytscMyWPy1s7jjr6v51dJKPqip55fXTGfoAE6mxiSikD0VVa1X1U2qOk9VNwMH8aYVZ4lIUZ9FOMD5/I39ppBkb6UPSubHV53MD688iXeqdnH5vW/Yc1mMGWDCGaj/tHvwVTXwGrAJrwdjeqm9kOSE4QNzPCWUeTOKWPz1M0lNSWLug2/z8xfX02LTjo0ZEMIZqP8BcAawQVVLgJnAW1GNKkG0F5JMlJ5KoGljhvC3b53LlaeO4d6llXzut2+zZfeBWIdljOmlcJLKYVXdBSSJSJKq/gOYHuW4EkJ7IcmJCdZTaZeVlsJPrz6Ze+edwsadDVz2yzdYVLHF7sI3ph8LJ6nsFZEs4HXgcRH5JdAS3bASQ1WtKyQ5LDGTSrsrTh7Fs98+lxMKc7hl8Sq+9MgKtu09GOuwjDE9EE5SmQMcAP4NeB6oAj4dzaASha+ugaJhGaSm2GNtxg7LYOH8M/j+p8tYXr2bS+55nYXLP7ZeizH9TJe/zVS1UVXbVLVFVR8F7gdmh9O4iMwWkfUiUikitwXZniYiT7rty0SkOGDb7W79ehG5pKs2RWSmiKwUkfdF5E0RmRhOjLFUVdvI+PzE7qUESkoSvnR2CS/cdB5TR+dw21Mf8q8LlrN5l92Jb0x/0dnNjznuF/t9InKxeG4EfMDnumpYRJLxEtClQBkwT0TKOux2HbBHVScC9wB3u2PLgLnAVLwE9oCIJHfR5q+BL6jqdLx7bP4zvI8gNlrblOpdA6eQZCQV5WXwx+vP4L8/cyIrN+/h4nte595XNtLU0hrr0IwxXeisp/IYXgHJD4HrgReBq4E5qjonjLZnAJWq6lPVZmAh3qW0QHOAR93yYmCmeHcBzgEWqmqTqlYDla69ztpUIMct5wJx/UCP9kKSA63mV6QkJQlfPGMcr3znAmaVjeDnL21g9i/e4I2N/liHZozpRGd31I9X1ZMAROQhoA4oUtX9YbY9GtgS8LoG+ESofVS1RUTqgTy3/p0Ox452y6HavB54VkQO4lVUPiNYUCIyH5gPUFQUu3s4K10hyYFUnTgaRuamc9/nT+Wa0/3c8dc1fPHh5Vw+rZDvXXYCo4YMjnV4xpgOOuupHG5fUNVWoLobCQUgWN2RjqOuofbp7nrwJhJcpqpjgEeAnwcLSlUfVNVyVS0vKCgIGnhfaL9HpT8+lz4Wzi0t4PmbzuU7sybx8tqdXPjTV/nZi+tpaLKJiMbEk86Syskiss997QemtS+LyL4w2q4Bxga8HsPxl6SO7CMiKXiXrXZ3cmzQ9SJSAJysqsvc+ieBs8KIMWaqXCHJYVb7KmxpKcl8c2Ypr3znfGafOJJfLa3kgp+8yhPLP7YClcbEic5qfyWrao77ylbVlIDlnFDHBVgBlIpIiYik4g28L+mwzxLgWrd8FbDUPbp4CTDXzQ4rAUqB5Z20uQfIFZFJrq1ZwLpwPoBY8SVIIcloGDM0g1/OPYW/3HA2xXkZ3P7Uh1x+7xu8tsFvU5CNibEun/zYU26M5EbgBSAZWKCqa0TkTqBCVZcADwOPiUglXg9lrjt2jYgsAtbi3Wh5g7sER7A23fqvAn8WkTa8JPOVaL23SKjyN3L+pNhdfhsIpo8dwp++dibPrd7BD59bx7ULljOjeBjfuXgSnxifF+vwjElIksh/2ZWXl2tFRUWfn3f/ocOc9P0XuWX2ZL5xQdzfTtMvNLW0smjFFn61tJLa/U2cMzGfmy+exKlFQ2MdmjEDjoi8q6rlwbbZrdwxcHSQ3mZ+RUpaSjJfPLOY12+5kP+8/ATWbd/HlQ/8k+t+v4LVW+tjHZ4xCcOSSgwcfS69zfyKtPRByVx/7nhev+VCvnvJZFZs2s2nfvUm1y5YzjLfLhtzMSbKwnmeyv6AWWDtX1tE5GkRGd8XQQ40Pr8Vkoy2zLQUbrhwIm/edhG3zJ7M6q31XPPgO1z1m7dZ+tFOSy7GREk4A/U/x5vO+0e8+0TmAiOB9cAC4IJoBTdQVfmtkGRfyUkfxDcumMhXzi5hUcUWfvuaj6/8voIpI7P52vkTuOykQvs5GBNB4fxvmq2qv1XV/aq6T1UfxLvJ8EnARkF7wHuEsPVS+lL6oGT+9cxiXv3uBfzs6pNpaVNuevJ9zv3xUu5bupHdjc2xDtGYASGcpNImIp8TkST3FVhM0q4hdFN7IckJw22QPhYGJSfxL6eN4cWbzuORL5/OpBHZ/PTFDZz5w1e4dfEqPtoRzn29xphQwrn89QXgl8ADeEnkHeB/ichg4MYoxjYgbd3jFZK0nkpsJSUJF04ezoWTh7Nx534e+ecmnlpZw5MVWzhrQh5fPGMcnywbwaBkuzRmTHd0mVRU1Ufoh3K9GdlwBr4q9whh66nEj9IR2fy/z57ELZdM5onlW3js7U18/fGV5Gel8bnyMcybUcTYYRmxDtOYfqHLpOLqan0VKA7cX1Xj+o71eFVV66oTW08l7gzJSOXrF0xg/nnjeW1DLX9c9jG/ea2KB16t4tzSfD4/o8h6L8Z0IZzLX38F3gBeBuwpSb3kq2tkSIYVkoxnyUnCRVNGcNGUEWzbe5BFFVt4csWWI72XK08dzZWnjmbKyHBK4BmTWMJJKhmqemvUI0kQVbUNjM+3QpL9xaghg7npk5O48cKJvLbBzxPLt7DgzWoefN1HWWEOV546mjnTR1OQnRbrUI2JC+Eklb+JyGWq+mzUo0kAvjorJNkfpSQnMfOEEcw8YQS7Gpp45oNtPPXeVn7w93X88LmPOK80nytPHcOsshGkD0qOdbjGxEw4SeXbwH+ISBPeg7sE0DDL35sA+w4dxr+/yWp+9XN5WWl86ewSvnR2CRt37uep97by9MqtfPOJ98hKS+GTJwznU9NGce6kfNJSLMGYxBLO7K/svggkEbQXkhxvNb8GjNIR2dw6ewr/fvFk3q7axTMfbOP5NTv4y/vbyE5LYdbUEXx62ijOnphvd+6bhBAyqYjIFFX9SERODbZdVVdGL6yByXekkKT1VAaa5CThnNJ8zinN578/cyJvVdXx91XbeWHNDp5auZWc9BQumTqSS08ayVkT8u0SmRmwOuup3AzMB34WZJsCF0UlogGsyt/gCknaPQ8DWWpK0pEbK+/67Im8udFLMM+t3sGf3q0hIzWZ8ycVMKtsBBdNGc6QDJsJaAaOkElFVee77xf2XTgDm8/faIUkE0xaSvKRAf6mllbertrFi2t38vLanTy3egfJScKM4mFcPHUEs8pGMGao/cFh+rewnvwoImdx/M2Pf4heWH2jr5/8ePE9r1E0LIOHrj29z85p4lNbm7Jqaz0vrtnBS2t3stHdFDtlZDbnTy7ggknDKS8eajdamrjU2ZMfw7mj/jFgAvA+R29+VKDfJ5W+1NqmbNp1gAsmD491KCYOJCUJ08cOYfrYIdwyewrVdY28tHYH//jIz4I3q/ntaz6y0lI4e2IeF0wezvmTChg1ZHCswzamS+FMKS4HyrQHTzUSkdl4xSiTgYdU9UcdtqfhJafTgF3ANaq6yW27HbgOL5F9S1Vf6KxN8e4m/AFwtTvm16p6b3djjpb2QpL2tEcTTEl+JvPPm8D88ybQ0NTCW5V1vLrez2vra3lhzU4AJo3I4oLJwzmvtIDy4qE22G/iUjhJZTXeQ7m2d6dhEUkG7gdmATXAChFZoqprA3a7DtijqhNFZC5wN3CNiJThPQxsKjAKeFlEJrljQrX5JWAsMEVV20QkrroE7Y8QHm8zv0wXstK8mWKXTB2JqrKxtoFX19fy6no/j7zl3c2fmpLEaUVDOXtiHmdNzGfa6FxS7FKZiQPhJJV8YK2ILAea2leq6hVdHDcDqHRVjhGRhcAcIDCpzAG+75YXA/e5HsccYKGqNgHVIlLp2qOTNr8OfF5V21x8tWG8tz5TZdOJTQ+ICJNGZDNpRPaRXsyK6t28VVnHW1W7+OmLG+DFDWSnpfCJ8cM4a0I+Z03MY/KIbCsFZGIinKTy/R62PRrYEvC6BvhEqH1UtUVE6oE8t/6dDseOdsuh2pyA18v5LODHu2S2sWNQIjIfb6o0RUVF3X9XPVTlt0KSpvey0lK4cMpwLpzidcR3NTTxjm83b1XV8c/KOl5e5/0tlZeZyunFwzi9ZBgziodxQmG29WRMn+g0qbhLWP9HVT/Zg7aD/ZnUcVwm1D6h1gf7X9HeZhpwSFXLReRKYAFw7nE7e49DfhC82V/BQ488n7/Byt2biMvLSuPyaYVcPq0QgK17D/JWZR3LfLtZvmkXz6/ZAUBmajKnjhvKjOJhzCgZxsljh9iYjImKTpOKqraKyAERyVXV+m62XYM3xtFuDLAtxD41IpIC5AK7uzg21Poa4M9u+WngkW7GG1W+ukYusEKSJspGDxnM58rH8rly77/JjvpDLN+0m+XVu1hRvYefvbQBgNTkJKaNyeX0kmGUjxvK9LFDyMuySsum98K5/HUI+FBEXgIa21eq6re6OG4FUCoiJcBWvIH3z3fYZwlwLfA2cBWwVFVVRJYAfxSRn+MN1JcCy/F6MKHa/AveXf4LgPOBDWG8tz7RXkjSBulNXxuZm84VJ4/iipNHAbCnsZmKzXtYsWk3y6t387vXffy6zeuwFw3L4JSiIZwydgjTi4ZSVphjN+qabgsnqfzdfXWLGyO5EXgBb/rvAlVdIyJ3AhWqugR4GHjMDcTvxksSuP0W4Q3AtwA3qGorQLA23Sl/BDwuIv8GNADXdzfmaGkvJGnTiU2sDc1MZVaZd/c+wIHmFlZv3cd7H+/hvY/38o5vF3993+v8p6YkceKoHE4pGsopRd49NaOHDLYJAKZTYd1RP1D11R31f363hu/86QNevvl8Jtqz6U2c215/kPc+3nsk0Xy4tZ6mljYACrLTmDY6lxNH53KS+z4iJ80STYLp7R31pcAPgTIgvX29qo6PWIQDnK/OCkma/qMwdzCFJw3mspO8wf/DrW18tH0/723Zw/sf72XV1nqWrq+l/e/R/Kw0ThqdcyTJnDQml5E56ZZoElQ4l78eAf4LuAe4EPgywWdnmRCqahsZZ4UkTT81KDmJk8Z4yeJfz/TWNTa1sG77Pj7cWs+HW+tZvbWe1zb4ccMz5GWmcuLoXKaOymFKYQ5lhdkU52XatOYEEE5SGayqr4iIqOpm4Psi8gZeojFh8NU12IO5zICSmZZCefEwyouHHVl3oNlLNKu37juSaN6qrKPFZZq0lCQmjcjmhMJspozM4YTCHE4ozLbS/wNMWLO/RCQJ2OgGybcCcVUCJZ61timb6g5woRWSNANcRmoKp40bxmnjjiaappZWKmsb+Gj7ftZt38dHO/bzyrpaFlXUHNmnMDedKSOzOaHQ69VMHpFNSX6m9ez7qXCSyk1ABvAt4L/xLoFdG82gBpKaPQdobm2znopJSGkpyUwdlcvUUblH1qkq/oamYxLNuu37eGPj0V5NcpJQnJfBpBHZlA7PonRENqUjsijJzyQtxW7ajGfhPKN+BYB39Uu/HIwxuzsAABRnSURBVP2QBpaj04lt1pcx4NUzG56dzvDsdM4LuCG4uaWNytoGNtbuZ+POBjbs3M/6Hft5Yc2OI2M1yUnCuLwMJg33kszE4VlMcj0bqxAQH8KZ/XUm3v0kWUCRiJwM/G9V/Ua0gxsIrDqxMeFJTUmibFQOZaNyjll/6HAr1XWNbNi5n8paL9lsqN3PS+t20uqyTZLA6KGDGZ/v9WYmFGRSkp/F+IJMRuakk5Rkc4v6SjiXv34BXIJ39zuq+oGInBfVqAYQKyRpTO+kD0p2g/rHJpumFi/ZbNzZQGVtA766RqrrGqjYtJvG5tYj+w0elExxfibj8zMZX+B9tSecnPRBff12BrxwkgqquqXDnPPWUPuaY/n8DXbpy5goSEtJZsrIHKaMPDbZqCq1+5uo8jdQXdeIz99IdV0ja7bV8/yaHUd6NwD5WamMz89iXF4G4/IyKMrLZNwwb9lmpfVMOElli3tGvYpIKt6A/brohjVwVPkbuXCyFZI0pq+ICCNy0hmRk85ZE/KP2dbc0sbHuw/gC0g4vroGXtvgp3Z/0zH75qSnMC4vk6K8jCOJpmhYJuPyMuySWifCSSpfw3t872i8SsAvAjaeEob6g4epa2higpVmMSYupKYkMXF4VtBySQebW/l49wE272p03w+wefcBVm+t54XVO47MTGtvZ+zQwRS7pDNmaAZjhg52XxnkDk7cy2rhzP6qA74QuE5EbsIbazGd8LUP0ttzVIyJe4NTk5k8MpvJI7OP29bS2sa2vYfYvLuRzbsOHEk+m3cd4G3fLg40HzsikJ2WwmiXYI4mm6OvcwcPGrBlbMIaUwniZiypdKl9OrHN/DKmf0tJTqIoL4OivAzOLT12m6qy58Bhtu45SM2eA9S471v3et/f8e2ioanlmGMyU5OPSTjtCagwN51RQwaTn5VGcj+9vNbTpNI/320fq/I3kOLm1RtjBiYRYVhmKsMyUzlpTO5x21WVfQdb2HJcwvG+lm/azf5DxyadlCRvXKgwN52RLtEU5qa7r8EUDkknPzMtLsd1eppUErdefjf4/I0UDctgkBXRMyZhiQi5GYPIzfCqOAdTf9Dr6WyvP8j2+kPe972H2FZ/kNVb63lx7U6a3eMH2g1K9hLPKJdkRua6ZZd4Ruamk5eZ2ueJJ2RSEZH9BE8eAgyOWkQDiFdI0i59GWM6lzt4ELmDBx1342c7VWV3Y7NLOF7S2bb3EDvqD7Kt/hArP97DjvpDHG499lf2oGSvesGInDRG5npVDEbmpjMyJ52zJuQxPCc96Pl6I2RSUdXjR6tM2KyQpDEmUkSEvKw08rLSQvZ22tqUXY3NRxLOzn2H2LHvEDvrve8f7djPa+v9R24M/cNXZvRtUjG9015I0m58NMb0haQkoSA7zXs655jQ+zU0tbCj/hCFuZFPKGBJJWqO1vyy6cTGmPiRlZYS1ceaR3UEWURmi8h6EakUkduCbE8TkSfd9mUiUhyw7Xa3fr2IXNKNNn8lIg3Rek/hsunExphEFLWkIiLJwP3ApXjPt58nImUddrsO2KOqE/EeV3y3O7YMmAtMBWYDD4hIcldtikg5MCRa76k7qvyNDLVCksaYBBPNnsoMoFJVfaraDCwE5nTYZw7wqFteDMwU7zbTOcBCVW1S1Wqg0rUXsk2XcH4C3BLF9xS2Kr/N/DLGJJ5oJpXRwJaA1zVuXdB9VLUFqAfyOjm2szZvBJao6vbOghKR+SJSISIVfr+/W2+oO3z+RibYeIoxJsFEM6kEu+Om430vofbp1noRGQVcDfyqq6BU9UFVLVfV8oKC6FQPbi8kaT0VY0yiiWZSqQHGBrweA2wLtY+IpAC5wO5Ojg21/hRgIlApIpuADBGpjNQb6S4rJGmMSVTRTCorgFIRKXHPYZmLe3pkgCXAtW75KmCpqqpbP9fNDisBSoHlodpU1b+r6khVLVbVYuCAG/yPiar259JbyXtjTIKJ2n0qqtoiIjcCLwDJwAJVXSMidwIVqroEeBh4zPUqduMlCdx+i4C1QAtwg6q2AgRrM1rvoad8rpBk0TArJGmMSSxRvflRVZ8Fnu2w7o6A5UN4YyHBjr0LuCucNoPsE9Mugs/fSFGeFZI0xiQe+60XBVX+Bsbn26UvY0zisaQSYS2tbWzedYAJw22Q3hiTeCypRFjNnoNeIUnrqRhjEpAllQjz1VkhSWNM4rKkEmHthSSt5L0xJhFZUomwKn8DQzMGMdQKSRpjEpAllQir8jdaL8UYk7AsqUSYz99g4ynGmIRlSSWC6g8cpq6h2QpJGmMSliWVCKpyM7/s8pcxJlFZUomgo48QtstfxpjEZEklgqyQpDEm0VlSiaAqf4MVkjTGJDT77RdBPptObIxJcJZUIqSltY1NuxptPMUYk9AsqURIzZ6DHG5VKyRpjElollQipL2QpJW8N8YkMksqEVJV66YTW0/FGJPALKlEiK+ugWGZqVZI0hiT0KKaVERktoisF5FKEbktyPY0EXnSbV8mIsUB225369eLyCVdtSkij7v1q0VkgYgMiuZ766iqtpHx+XbpyxiT2KKWVEQkGbgfuBQoA+aJSFmH3a4D9qjqROAe4G53bBkwF5gKzAYeEJHkLtp8HJgCnAQMBq6P1nsLxldnhSSNMSaaPZUZQKWq+lS1GVgIzOmwzxzgUbe8GJgpIuLWL1TVJlWtBipdeyHbVNVn1QGWA2Oi+N6O0V5I0u5RMcYkumgmldHAloDXNW5d0H1UtQWoB/I6ObbLNt1lry8Cz/f6HYSp6sgjhC2pGGMSWzSTigRZp2Hu0931gR4AXlfVN4IGJTJfRCpEpMLv9wfbpduOPkLYLn8ZYxJbNJNKDTA24PUYYFuofUQkBcgFdndybKdtish/AQXAzaGCUtUHVbVcVcsLCgq6+ZaCq3KFJMdaIUljTIKLZlJZAZSKSImIpOINvC/psM8S4Fq3fBWw1I2JLAHmutlhJUAp3jhJyDZF5HrgEmCeqrZF8X0dx+dvYJwVkjTGGFKi1bCqtojIjcALQDKwQFXXiMidQIWqLgEeBh4TkUq8Hspcd+waEVkErAVagBtUtRUgWJvulL8BNgNve2P9PKWqd0br/QWq8jfaeIoxxhDFpALejCzg2Q7r7ghYPgRcHeLYu4C7wmnTrY/qewmlpbWNzbsamXnC8Fic3hhj4opdr+mlI4UkradijDGWVHqryt/+XHqb+WWMMZZUeunIc+mtkKQxxlhS6a0qvxWSNMaYdpZUesnnt0KSxhjTzpJKL1X5G2yQ3hhjHEsqvVB/4DC7GputOrExxjiWVHqhvZCk9VSMMcZjSaUXqmrbqxNbT8UYY8CSSq/46hoZlGyFJI0xpp0llV6oqm2gaJgVkjTGmHb227AXfHVWSNIYYwJZUumh9kKSNkhvjDFHWVLpoS2ukKQN0htjzFGWVHrI57fpxMYY05EllR6y6sTGGHM8Syo95PM3kpeZypAMKyRpjDHtLKn0UJW/wcZTjDGmA0sqPeRVJ7bxFGOMCWRJpQf2HmhmV2MzE4ZbT8UYYwJFNamIyGwRWS8ilSJyW5DtaSLypNu+TESKA7bd7tavF5FLumpTREpcGxtdm1Eb7Kiypz0aY0xQUUsqIpIM3A9cCpQB80SkrMNu1wF7VHUicA9wtzu2DJgLTAVmAw+ISHIXbd4N3KOqpcAe13ZUHJlOPNySijHGBIpmT2UGUKmqPlVtBhYCczrsMwd41C0vBmaKiLj1C1W1SVWrgUrXXtA23TEXuTZwbX4mWm+syu8KSQ4dHK1TGGNMvxTNpDIa2BLwusatC7qPqrYA9UBeJ8eGWp8H7HVthDoXACIyX0QqRKTC7/f34G1BcV4Gnz1lNClWSNIYY44Rzd+KEmSdhrlPpNYfv1L1QVUtV9XygoKCYLt0ae6MIn581ck9OtYYYwayaCaVGmBswOsxwLZQ+4hICpAL7O7k2FDr64Ahro1Q5zLGGBNl0UwqK4BSNysrFW/gfUmHfZYA17rlq4Clqqpu/Vw3O6wEKAWWh2rTHfMP1wauzb9G8b0ZY4wJIqXrXXpGVVtE5EbgBSAZWKCqa0TkTqBCVZcADwOPiUglXg9lrjt2jYgsAtYCLcANqtoKEKxNd8pbgYUi8gPgPde2McaYPiTeH/mJqby8XCsqKmIdhjHG9Csi8q6qlgfbZtOXjDHGRIwlFWOMMRFjScUYY0zEWFIxxhgTMQk9UC8ifmBzDw/Px7s/Jt5YXN1jcXWPxdU98RoX9C62caoa9O7xhE4qvSEiFaFmP8SSxdU9Flf3WFzdE69xQfRis8tfxhhjIsaSijHGmIixpNJzD8Y6gBAsru6xuLrH4uqeeI0LohSbjakYY4yJGOupGGOMiRhLKsYYYyLGkkoPiMhsEVkvIpUiclsfnG+TiHwoIu+LSIVbN0xEXhKRje77ULdeROReF9sqETk1oJ1r3f4bReTaUOfrIpYFIlIrIqsD1kUsFhE5zb3XSndssAewhRvX90Vkq/vc3heRywK23e7OsV5ELglYH/Rn6x63sMzF+6R79EJXMY0VkX+IyDoRWSMi346Hz6uTuGL6ebnj0kVkuYh84GL7v521J97jMZ50518mIsU9jbmHcf1eRKoDPrPpbn1f/ttPFpH3RORv8fBZoar21Y0vvJL7VcB4IBX4ACiL8jk3Afkd1v0YuM0t3wbc7ZYvA57DexrmGcAyt34Y4HPfh7rloT2I5TzgVGB1NGLBe27Ome6Y54BLexHX94F/D7Jvmfu5pQEl7ueZ3NnPFlgEzHXLvwG+HkZMhcCpbjkb2ODOHdPPq5O4Yvp5uX0FyHLLg4Bl7rMI2h7wDeA3bnku8GRPY+5hXL8Hrgqyf1/+278Z+CPwt84++776rKyn0n0zgEpV9alqM7AQmBODOOYAj7rlR4HPBKz/g3rewXsiZiFwCfCSqu5W1T3AS8Ds7p5UVV/He/ZNxGNx23JU9W31/rX/IaCtnsQVyhxgoao2qWo1UIn3cw36s3V/MV4ELA7yHjuLabuqrnTL+4F1wGhi/Hl1ElcoffJ5uXhUVRvcy0HuSztpL/CzXAzMdOfvVsy9iCuUPvlZisgY4HLgIfe6s8++Tz4rSyrdNxrYEvC6hs7/Q0aCAi+KyLsiMt+tG6Gq28H7JQEM7yK+aMYdqVhGu+VIxniju/ywQNxlph7ElQfsVdWWnsblLjWcgvcXbtx8Xh3igjj4vNzlnPeBWrxfulWdtHckBre93p0/4v8POsalqu2f2V3uM7tHRNI6xhXm+Xv6s/wFcAvQ5l539tn3yWdlSaX7gl3njPa87LNV9VTgUuAGETmvk31DxReLuLsbS6Rj/DUwAZgObAd+Fou4RCQL+DNwk6ru62zXGMcVF5+Xqraq6nRgDN5fyyd00l6fxdYxLhE5EbgdmAKcjndJ69a+iktEPgXUquq7gas7aadPPitLKt1XA4wNeD0G2BbNE6rqNve9Fnga7z/aTtdlxn2v7SK+aMYdqVhq3HJEYlTVne4XQRvwO7zPrSdx1eFdvkjpsL5LIjII7xf346r6lFsd888rWFzx8HkFUtW9wKt4YxKh2jsSg9uei3cZNGr/DwLimu0uJaqqNgGP0PPPrCc/y7OBK0RkE96lqYvwei6x/ay6GnSxr+MGxVLwBtdKODp4NTWK58sEsgOW/4k3FvITjh3s/bFbvpxjBwiX69EBwmq8wcGhbnlYD2Mq5tgB8YjFAqxw+7YPVl7Wi7gKA5b/De+6McBUjh2Y9OENSob82QJ/4tjBz2+EEY/gXRv/RYf1Mf28Ookrpp+X27cAGOKWBwNvAJ8K1R5wA8cOPi/qacw9jKsw4DP9BfCjGP3bv4CjA/Wx/ax68ksl0b/wZnZswLvW+70on2u8+2F+AKxpPx/etdBXgI3ue/s/TAHud7F9CJQHtPUVvEG4SuDLPYznCbxLI4fx/pK5LpKxAOXAanfMfbiqDz2M6zF33lXAEo79pfk9d471BMyyCfWzdT+H5S7ePwFpYcR0Dt7lglXA++7rslh/Xp3EFdPPyx03DXjPxbAauKOz9oB097rSbR/f05h7GNdS95mtBv6HozPE+uzfvjv2Ao4mlZh+VlamxRhjTMTYmIoxxpiIsaRijDEmYiypGGOMiRhLKsYYYyLGkooxxpiIsaRiTDeJSF5AVdodcmxl33Cr8T4iIpO7cc5CEXnWVcldKyJL3PrxIjK3p+/FmEizKcXG9IKIfB9oUNWfdlgveP+/2oIe2P3zPAysVNX73etpqrpKRD4J3KiqYRVsNCbarKdiTISIyEQRWS0ivwFWAoUi8qCIVIj3DI47AvZ9U0Smi0iKiOwVkR+5XsjbIjI8SPOFBBQcVNVVbvFHwIWul/Qt197PxXv2xyoRud6d75PiPUPlL66nc79LfMZElCUVYyKrDHhYVU9R1a145VjKgZOBWSJSFuSYXOA1VT0ZeBvvjuuO7gMeFZGlIvIf7bXD8Mq8/ENVp6vqvcB8vCKDM/CKHN4gIkVu308ANwEn4RVpjMUjG8wAZ0nFmMiqUtUVAa/nichKvJ7LCXhJp6ODqvqcW34Xr4bZMVT1WbwKwg+7Nt4TkbwgbV0MfNmVaF8GDAFK3bZ3VHWTqrbiFSA8p7tvzpiupHS9izGmGxrbF0SkFPg2MENV94rI/+DVX+qoOWC5lRD/L1V1F/A48LiIPI+XFBo77CZ4BQRfOWalN/bScQDVBlRNxFlPxZjoyQH2A/sCnvrXIyIyU0QGu+UcvMqxH7v2swN2fQH4RnvpcxGZ3H4ccIaIFIlIMvA54M2exmNMKNZTMSZ6VgJr8SrP+oC3etHW6cB9InIY74/BX6vqe24Kc7KIfIB3aex+oAh4343D13J07OSfeA/emor3PJAlvYjHmKBsSrExCcCmHpu+Ype/jDHGRIz1VIwxxkSM9VSMMcZEjCUVY4wxEWNJxRhjTMRYUjHGGBMxllSMMcZEzP8HlXYIMYvlusUAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "temp_learning_rate_schedule = CustomSchedule(d_model)\n",
    "\n",
    "plt.plot(temp_learning_rate_schedule(tf.range(40000, dtype=tf.float32)))\n",
    "plt.ylabel(\"Learning Rate\")\n",
    "plt.xlabel(\"Train Step\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loss function\n",
    "\n",
    "Since the target sequences are padded, it is important to apply a padding mask when calculating the loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function(real, pred):\n",
    "    mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
    "    loss_ = loss_object(real, pred)\n",
    "\n",
    "    mask = tf.cast(mask, dtype=loss_.dtype)\n",
    "    loss_ *= mask\n",
    "\n",
    "    return tf.reduce_sum(loss_) / tf.reduce_sum(mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loss = tf.keras.metrics.Mean(name='train_loss')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_masks(inp, tar):\n",
    "    # Encoder padding mask\n",
    "    enc_padding_mask = create_padding_mask(inp)\n",
    "\n",
    "    # Used in the 2nd attention block in the decoder.\n",
    "    # This padding mask is used to mask the encoder outputs.\n",
    "    dec_padding_mask = create_padding_mask(inp)\n",
    "\n",
    "    # Used in the 1st attention block in the decoder.\n",
    "    # It is used to pad and mask future tokens in the input received by \n",
    "    # the decoder.\n",
    "    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n",
    "    dec_target_padding_mask = create_padding_mask(tar)\n",
    "    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n",
    "\n",
    "    return enc_padding_mask, combined_mask, dec_padding_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = Transformer(num_layers,\n",
    "                          d_model,\n",
    "                          num_heads,\n",
    "                          dff,\n",
    "                          vocab_size=vocab_size,\n",
    "                          rate=dropout_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = '%s%s-%s' % (ARTIFACTS_PATH, LANGUAGE, MODEL_NAME)\n",
    "\n",
    "ckpt = tf.train.Checkpoint(transformer=transformer,\n",
    "                           optimizer=optimizer)\n",
    "\n",
    "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n",
    "\n",
    "# if a checkpoint exists, restore the latest checkpoint.\n",
    "if ckpt_manager.latest_checkpoint:\n",
    "    ckpt.restore(ckpt_manager.latest_checkpoint)\n",
    "    print ('Latest checkpoint restored!!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The @tf.function trace-compiles train_step into a TF graph for faster\n",
    "# execution. The function specializes to the precise shape of the argument\n",
    "# tensors. To avoid re-tracing due to the variable sequence lengths or variable\n",
    "# batch sizes (the last batch is smaller), use input_signature to specify\n",
    "# more generic shapes.\n",
    "\n",
    "train_step_signature = [\n",
    "    tf.TensorSpec(shape=(None, MAX_LEN), dtype=tf.int32),\n",
    "    tf.TensorSpec(shape=(None, MAX_LEN), dtype=tf.int32)\n",
    "]\n",
    "\n",
    "@tf.function(input_signature=train_step_signature)\n",
    "def train_step(inp, tar):\n",
    "    tar_inp = tar[:, :-1]\n",
    "    tar_real = tar[:, 1:]\n",
    "\n",
    "    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)\n",
    "\n",
    "    with tf.GradientTape() as tape:\n",
    "        predictions, _ = transformer(inp, tar_inp, \n",
    "                                     True, \n",
    "                                     enc_padding_mask, \n",
    "                                     combined_mask, \n",
    "                                     dec_padding_mask)\n",
    "        loss = loss_function(tar_real, predictions)\n",
    "\n",
    "    gradients = tape.gradient(loss, transformer.trainable_variables)    \n",
    "    optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))\n",
    "\n",
    "    train_loss(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_time(secs):\n",
    "    if secs < 60:\n",
    "        return '%ds' % (secs)\n",
    "    else:\n",
    "        mins = secs // 60\n",
    "        secs = secs % 60\n",
    "        return '%dm%ds' % (mins, secs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 - Batch 0 of 1193 - Loss 9.5026 - Elapsed 38s\n",
      "Epoch 1 - Batch 50 of 1193 - Loss 9.4060 - Elapsed 24m23s\n",
      "Epoch 1 - Batch 100 of 1193 - Loss 9.1626 - Elapsed 49m18s\n",
      "Epoch 1 - Batch 150 of 1193 - Loss 8.9231 - Elapsed 74m9s\n",
      "Epoch 1 - Batch 200 of 1193 - Loss 8.6795 - Elapsed 99m11s\n",
      "Epoch 1 - Batch 250 of 1193 - Loss 8.4077 - Elapsed 123m9s\n",
      "Epoch 1 - Batch 300 of 1193 - Loss 8.1048 - Elapsed 147m6s\n",
      "Epoch 1 - Batch 350 of 1193 - Loss 7.7794 - Elapsed 171m38s\n",
      "Epoch 1 - Batch 400 of 1193 - Loss 7.4442 - Elapsed 195m47s\n",
      "Epoch 1 - Batch 450 of 1193 - Loss 7.1078 - Elapsed 219m39s\n",
      "Epoch 1 - Batch 500 of 1193 - Loss 6.7952 - Elapsed 244m40s\n",
      "Epoch 1 - Batch 550 of 1193 - Loss 6.5087 - Elapsed 269m34s\n",
      "Epoch 1 - Batch 600 of 1193 - Loss 6.2511 - Elapsed 293m46s\n",
      "Epoch 1 - Batch 650 of 1193 - Loss 6.0294 - Elapsed 318m2s\n",
      "Epoch 1 - Batch 700 of 1193 - Loss 5.8341 - Elapsed 341m59s\n",
      "Epoch 1 - Batch 750 of 1193 - Loss 5.6583 - Elapsed 366m11s\n",
      "Epoch 1 - Batch 800 of 1193 - Loss 5.5073 - Elapsed 390m36s\n",
      "Epoch 1 - Batch 850 of 1193 - Loss 5.3757 - Elapsed 414m42s\n",
      "Epoch 1 - Batch 900 of 1193 - Loss 5.2512 - Elapsed 439m34s\n",
      "Epoch 1 - Batch 950 of 1193 - Loss 5.1360 - Elapsed 464m9s\n",
      "Epoch 1 - Batch 1000 of 1193 - Loss 5.0346 - Elapsed 487m51s\n",
      "Epoch 1 - Batch 1050 of 1193 - Loss 4.9396 - Elapsed 512m11s\n",
      "Epoch 1 - Batch 1100 of 1193 - Loss 4.8507 - Elapsed 536m15s\n",
      "Epoch 1 - Batch 1150 of 1193 - Loss 4.7634 - Elapsed 560m38s\n",
      "Epoch 1 - Loss 4.7044\n",
      "Time taken for epoch: 580m30s\n",
      "\n",
      "Epoch 2 - Batch 0 of 1193 - Loss 3.5075 - Elapsed 30s\n",
      "Epoch 2 - Batch 50 of 1193 - Loss 3.0830 - Elapsed 24m47s\n",
      "Epoch 2 - Batch 100 of 1193 - Loss 2.9834 - Elapsed 48m43s\n",
      "Epoch 2 - Batch 150 of 1193 - Loss 2.9507 - Elapsed 73m7s\n",
      "Epoch 2 - Batch 200 of 1193 - Loss 2.9402 - Elapsed 96m56s\n",
      "Epoch 2 - Batch 250 of 1193 - Loss 2.9085 - Elapsed 121m6s\n",
      "Epoch 2 - Batch 300 of 1193 - Loss 2.8923 - Elapsed 145m33s\n",
      "Epoch 2 - Batch 350 of 1193 - Loss 2.8927 - Elapsed 169m45s\n",
      "Epoch 2 - Batch 400 of 1193 - Loss 2.9017 - Elapsed 194m34s\n",
      "Epoch 2 - Batch 450 of 1193 - Loss 2.8934 - Elapsed 219m10s\n",
      "Epoch 2 - Batch 500 of 1193 - Loss 2.8915 - Elapsed 243m40s\n",
      "Epoch 2 - Batch 550 of 1193 - Loss 2.8851 - Elapsed 267m2s\n",
      "Epoch 2 - Batch 600 of 1193 - Loss 2.8704 - Elapsed 290m18s\n",
      "Epoch 2 - Batch 650 of 1193 - Loss 2.8604 - Elapsed 314m13s\n",
      "Epoch 2 - Batch 700 of 1193 - Loss 2.8484 - Elapsed 337m56s\n",
      "Epoch 2 - Batch 750 of 1193 - Loss 2.8349 - Elapsed 361m17s\n",
      "Epoch 2 - Batch 800 of 1193 - Loss 2.8319 - Elapsed 384m41s\n",
      "Epoch 2 - Batch 850 of 1193 - Loss 2.8335 - Elapsed 408m8s\n",
      "Epoch 2 - Batch 900 of 1193 - Loss 2.8287 - Elapsed 432m23s\n",
      "Epoch 2 - Batch 950 of 1193 - Loss 2.8195 - Elapsed 455m28s\n",
      "Epoch 2 - Batch 1000 of 1193 - Loss 2.8123 - Elapsed 479m5s\n",
      "Epoch 2 - Batch 1050 of 1193 - Loss 2.8053 - Elapsed 502m56s\n",
      "Epoch 2 - Batch 1100 of 1193 - Loss 2.7944 - Elapsed 526m7s\n",
      "Epoch 2 - Batch 1150 of 1193 - Loss 2.7793 - Elapsed 549m16s\n",
      "Epoch 2 - Loss 2.7750\n",
      "Time taken for epoch: 569m9s\n",
      "\n"
     ]
    }
   ],
   "source": [
    "EPOCHS = 2\n",
    "n_batches = len(list(dataset))\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    start = time.time()\n",
    "\n",
    "    train_loss.reset_states()\n",
    "\n",
    "    # inp -> portuguese, tar -> english\n",
    "    for (batch, (inp, tar)) in enumerate(dataset):\n",
    "        train_step(inp, tar)\n",
    "    \n",
    "        if batch % 50 == 0:\n",
    "            elapsed = format_time(time.time() - start)\n",
    "            print('Epoch %d - Batch %d of %d - Loss %.4f - Elapsed %s' % (\n",
    "                epoch + 1, batch, n_batches, train_loss.result(), elapsed))\n",
    "\n",
    "    if (epoch + 1) % 5 == 0:\n",
    "        ckpt_save_path = ckpt_manager.save()\n",
    "        print('Saving checkpoint for epoch {} at {}'.format(epoch+1, ckpt_save_path))\n",
    "    \n",
    "    print('Epoch {} - Loss {:.4f}'.format(epoch+1, train_loss.result()))\n",
    "\n",
    "    print('Time taken for epoch: %s\\n' % format_time(time.time() - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(inp_sentence, iters, temperature):\n",
    "    # inp sentence\n",
    "    token_ids = tokenizer.encode(inp_sentence).ids\n",
    "    encoder_input = tf.expand_dims([START_TOKEN] + token_ids[-(MAX_LEN-2):] + [END_TOKEN], 0)\n",
    "\n",
    "    # as the target\n",
    "    decoder_input = token_ids[1:][-(MAX_LEN-1):]\n",
    "    decoder_input = tf.expand_dims(decoder_input, 0)\n",
    "    \n",
    "    # generated\n",
    "    generated_token_ids = list(token_ids)\n",
    "\n",
    "    for i in range(iters):\n",
    "        enc_padding_mask, combined_mask, dec_padding_mask = create_masks(\n",
    "            encoder_input, decoder_input)\n",
    "\n",
    "        # predictions.shape == (batch_size, seq_len, vocab_size)\n",
    "        predictions, attention_weights = transformer(encoder_input, \n",
    "                                                     decoder_input,\n",
    "                                                     False,\n",
    "                                                     enc_padding_mask,\n",
    "                                                     combined_mask,\n",
    "                                                     dec_padding_mask)\n",
    "        \n",
    "        # remove the batch dimension # (batch_size, 1, vocab_size)\n",
    "        predictions = tf.squeeze(predictions, 0)\n",
    "\n",
    "        # using a categorical distribution to predict the character returned by the model\n",
    "        predictions = predictions / temperature\n",
    "        predicted_id = tf.random.categorical(predictions, num_samples=1)[-1,0].numpy()\n",
    "        \n",
    "        token_ids.append(predicted_id)\n",
    "\n",
    "        # concatentate the predicted_id to the output which is given to the decoder\n",
    "        # as its input.\n",
    "        generated_token_ids.append(predicted_id)\n",
    "        \n",
    "        # next inp sentence\n",
    "        inp_sentence = [START_TOKEN] + token_ids[-(MAX_LEN-2):] + [END_TOKEN]\n",
    "        encoder_input = tf.expand_dims(inp_sentence, 0)\n",
    "        \n",
    "        decoder_input = token_ids[1:][-(MAX_LEN-1):]\n",
    "        decoder_input = tf.expand_dims(decoder_input, 0)\n",
    "\n",
    "    return generated_token_ids\n",
    "\n",
    "\n",
    "def generate(initial_text, iters=5, temperature=1):\n",
    "    generated_token_ids = evaluate(initial_text, iters, temperature)\n",
    "    generated_text = tokenizer.decode(generated_token_ids)  \n",
    "\n",
    "    print('Input: {}'.format(initial_text))\n",
    "    print('Generated: {}'.format(generated_text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input: ROMEO: \n",
      "Generated: ROMEO:  took perform still sets big,KATHARINA'd me my heart\n",
      " l belie to harpets'd upon me\n",
      "And sure darkness under'd statutes.\n",
      "\n",
      " deceive them or orderly and likely body will f roodFor\n",
      "Masteredon'd conclude else aged asTherefore\n",
      "The suffer hiswittedWould being royalties and to ouroan as which and than him\n",
      "The goddess wasYour things curst hag that they admirings that how of they\n",
      "twere to spirit:\n",
      " linger OF glory of follow\n",
      "ised greatestar\n"
     ]
    }
   ],
   "source": [
    "generate(\"ROMEO: \", iters=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, the generated text seems to be nonsense. It is very likely that if we could train a larger model, we would obtain a better output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}