transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HPdvGnNVGW-8"
   },
   "source": [
    "## Transformer Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mJocN35ZGwKX"
   },
   "outputs": [],
   "source": [
    "#!pip install tensorflow-gpu==2.1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.6.9\n"
     ]
    }
   ],
   "source": [
    "from platform import python_version\n",
    "\n",
    "print(python_version())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import datetime "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JjJJyJTZYebt"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.1.0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "#import tensorflow_addons as tfa\n",
    "print(tf.__version__)\n",
    "from sklearn.model_selection import train_test_split\n",
    "import os\n",
    "#import io\n",
    "import numpy as np\n",
    "# import re\n",
    "import unicodedata\n",
    "# import urllib3\n",
    "# import shutil\n",
    "# import zipfile\n",
    "# import itertools\n",
    "from tensorflow import keras\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf ./logs/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num GPUs Available:  1\n"
     ]
    }
   ],
   "source": [
    "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9t3dob6hGW_C"
   },
   "source": [
    "### Create Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QlTF1PChGW_D"
   },
   "outputs": [],
   "source": [
    "# import os\n",
    "# file_list = []\n",
    "# for file in os.listdir(\"./data/dataset\"):\n",
    "#     if file.endswith(\".txt\"):\n",
    "#         file_list.append(os.path.join(\"./data/dataset\", file))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zQZ-i7NEGW_G"
   },
   "outputs": [],
   "source": [
    "def unicode_to_ascii(s):\n",
    "    \"\"\" Converts the unicode file to ascii \"\"\"\n",
    "    return ''.join(c for c in unicodedata.normalize('NFD', s)\n",
    "      if unicodedata.category(c) != 'Mn')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cwNE38BeGW_J"
   },
   "outputs": [],
   "source": [
    "def preprocess_sentence(w):\n",
    "    w = unicode_to_ascii(w.lower().strip())\n",
    "    # adding a start and an end token to the sentence\n",
    "    w = '<start> ' + w + ' <end>'\n",
    "    return w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0BhIXCFwGW_L"
   },
   "outputs": [],
   "source": [
    "# def create_dataset(file_list):\n",
    "#     dataset = []\n",
    "#     for file in file_list:\n",
    "#         lines = io.open(file, encoding='UTF-8').read().strip().split('\\n')\n",
    "#         word_pairs = [[preprocess_sentence(w) for w in l.split('\\t')]  for l in lines]\n",
    "#         dataset.extend(word_pairs)\n",
    "#     dataset = [s for s in dataset if len(s) ==2]  \n",
    "#     dataset = list(set(tuple(x) for x in dataset))\n",
    "#     return zip(*dataset)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IJgdFTIAGW_O"
   },
   "outputs": [],
   "source": [
    "# equation, integration = create_dataset(file_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5STjzZczGW_R"
   },
   "outputs": [],
   "source": [
    "# import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "d5ECJAGWGW_U"
   },
   "outputs": [],
   "source": [
    "# os.mkdir('./data/cleaned_data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fqmKPoBnGW_W"
   },
   "outputs": [],
   "source": [
    "# f = open('./data/cleaned_data/equation.txt', 'w+')\n",
    "# for line in equation:\n",
    "#     f.write(line + '\\n')\n",
    "# f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tOKifc0eGW_Z"
   },
   "outputs": [],
   "source": [
    "# f = open('./data/cleaned_data/integration.txt', 'w+')\n",
    "# for line in integration:\n",
    "#     f.write(line + '\\n')\n",
    "# f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### import data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5M1xrrq2NqNx"
   },
   "outputs": [],
   "source": [
    "num_samples = 100000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = open('./data/cleaned_data/equation_train1.txt', 'r')\n",
    "equation = f.read().splitlines()[:num_samples]\n",
    "f = open('./data/cleaned_data/integration_train1.txt', 'r')\n",
    "integration = f.read().splitlines()[:num_samples]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = open('./data/cleaned_data/equation_valid.txt', 'r')\n",
    "equation_valid = f.read().splitlines()\n",
    "f = open('./data/cleaned_data/integration_valid.txt', 'r')\n",
    "integration_valid = f.read().splitlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# f = open('./data/cleaned_data/equation_train2.txt', 'r')\n",
    "# equation2 = f.read().splitlines()[:num_samples]\n",
    "# f = open('./data/cleaned_data/integration_train2.txt', 'r')\n",
    "# integration2 = f.read().splitlines()[:num_samples]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# equation = equation1 + equation2\n",
    "# integration = integration1 + integration2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9263"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(equation_valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8qY1maKJGW_c"
   },
   "outputs": [],
   "source": [
    "# f = open('./equation.txt', 'r')\n",
    "# equation = f.read().splitlines()[:num_samples]\n",
    "# f = open('./integration.txt', 'r')\n",
    "# integration = f.read().splitlines()[:num_samples]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hIKeOMyRGW_e"
   },
   "source": [
    "### Preprocess dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "z0wvBoffGW_e"
   },
   "outputs": [],
   "source": [
    "def tokenize(inp, sequence_length):\n",
    "    \"\"\" word to index \"\"\"\n",
    "    tokenizer = tf.keras.preprocessing.text.Tokenizer(filters='')\n",
    "    tokenizer.fit_on_texts(inp)\n",
    "    sequences = tokenizer.texts_to_sequences(inp)\n",
    "    sequences = tf.keras.preprocessing.sequence.pad_sequences(sequences, padding='post', maxlen=sequence_length, truncating='post')\n",
    "    return  sequences, tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2DgKNuFvGW_g"
   },
   "source": [
    "#### Train test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6KU6MZhYGW_h"
   },
   "outputs": [],
   "source": [
    "#sequence_length = 128\n",
    "sequence_length = 512\n",
    "# Tokenize each word into index and return the tokenized list and tokenizer\n",
    "# X , X_tokenizer = tokenize(equation, sequence_length)\n",
    "# Y,  Y_tokenizer = tokenize(integration, sequence_length)\n",
    "#X_train,  X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2)\n",
    "# X_train = X\n",
    "# Y_train = Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pickle\n",
    "\n",
    "# # saving\n",
    "# with open('X_tokenizer.pickle', 'wb') as handle:\n",
    "#     pickle.dump(X_tokenizer, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "    \n",
    "# with open('Y_tokenizer.pickle', 'wb') as handle:\n",
    "#     pickle.dump(Y_tokenizer, handle, protocol=pickle.HIGHEST_PROTOCOL) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "# load\n",
    "with open('X_tokenizer.pickle', 'rb') as handle:\n",
    "    X_tokenizer = pickle.load(handle)\n",
    "    \n",
    "with open('Y_tokenizer.pickle', 'rb') as handle:\n",
    "    Y_tokenizer = pickle.load(handle) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = X_tokenizer.texts_to_sequences(equation)\n",
    "X_train = tf.keras.preprocessing.sequence.pad_sequences(X_train, padding='post', maxlen=sequence_length, truncating='post')\n",
    "Y_train = Y_tokenizer.texts_to_sequences(integration)\n",
    "Y_train = tf.keras.preprocessing.sequence.pad_sequences(Y_train, padding='post', maxlen=sequence_length, truncating='post')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test = X_tokenizer.texts_to_sequences(equation_valid)\n",
    "X_test = tf.keras.preprocessing.sequence.pad_sequences(X_test, padding='post', maxlen=sequence_length, truncating='post')\n",
    "Y_test = Y_tokenizer.texts_to_sequences(integration_valid)\n",
    "Y_test = tf.keras.preprocessing.sequence.pad_sequences(Y_test, padding='post', maxlen=sequence_length, truncating='post')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "x1H_1pABGW_j"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# tokenize by frequency\n",
    "X_tokenizer.word_index['<end>']   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ct2Dum18GW_l"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_vocab_size :  41\n",
      "output_vocab_size :  39\n"
     ]
    }
   ],
   "source": [
    "# vocabulary size # add 1 for 0 padding \n",
    "input_vocab_size = len(X_tokenizer.word_index) + 1\n",
    "output_vocab_size = len(Y_tokenizer.word_index) + 1\n",
    "\n",
    "print(\"input_vocab_size : \", input_vocab_size)\n",
    "print(\"output_vocab_size : \" ,output_vocab_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zTX35iTzGW_n"
   },
   "source": [
    "### Build transformer \n",
    "- building in ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-Cu3POgOGW_p"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras import models, layers\n",
    "from tensorflow.keras import backend as K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tjfYrDJDGW_q"
   },
   "outputs": [],
   "source": [
    "# ### only for model test\n",
    "# sequence_length = 512\n",
    "# input_vocabulary_size = 1000\n",
    "# output_vocabulary_size = 1000\n",
    "# ###\n",
    "BUFFER_SIZE = len(X_train)\n",
    "batch_size = 32\n",
    "d_model = 512\n",
    "embedding_size = 512\n",
    "num_layers = 6\n",
    "num_heads = 8\n",
    "depth = d_model // num_heads\n",
    "dff = 2048\n",
    "dropout_rate = 0.1\n",
    "learning_rate = 10**(-4)\n",
    "training = True\n",
    "epochs = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hw0UuulgMuqr"
   },
   "outputs": [],
   "source": [
    "# BUFFER_SIZE = len(X_train)\n",
    "# batch_size = 32\n",
    "# d_model = 128\n",
    "# #embedding_size = 128\n",
    "# num_layers = 4\n",
    "# num_heads = 4\n",
    "# depth = d_model // num_heads\n",
    "# dff = 128\n",
    "# dropout_rate = 0.0\n",
    "# learning_rate = 10**(-1)\n",
    "# training = True\n",
    "# epochs = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "H_zD8G3qGW_s"
   },
   "outputs": [],
   "source": [
    "dataset_train = tf.data.Dataset.from_tensor_slices((X_train, Y_train)).shuffle(BUFFER_SIZE).batch(batch_size, drop_remainder=True)\n",
    "dataset_test = tf.data.Dataset.from_tensor_slices((X_test, Y_test)).shuffle(BUFFER_SIZE).batch(batch_size, drop_remainder=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nBQuibYA4n0n"
   },
   "source": [
    "#### Positional encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WhIOZjMNKujn"
   },
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1Rz82wEs5biZ"
   },
   "outputs": [],
   "source": [
    "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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "a_b4ou4TYqUN"
   },
   "source": [
    "#### Masking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "s42Uydjkv0hF"
   },
   "source": [
    "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 `0` at those locations, and a `1` otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "U2i8-e1s8ti9"
   },
   "outputs": [],
   "source": [
    "def create_padding_mask(seq):\n",
    "    seq = tf.cast(1 - tf.cast(tf.math.equal(seq, 0), tf.int32), tf.bool)\n",
    "  \n",
    "    return seq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kmzGPEy64qmA"
   },
   "source": [
    "#### Multi-head attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BSV3PPKsYecw"
   },
   "outputs": [],
   "source": [
    "class MultiHeadAttention(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads, causal=False):\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(self.depth)\n",
    "        self.wk = tf.keras.layers.Dense(self.depth)\n",
    "        self.wv = tf.keras.layers.Dense(self.depth)\n",
    "        self.attention = tf.keras.layers.Attention(use_scale=True, causal=causal)\n",
    "        self.dense = tf.keras.layers.Dense(d_model)\n",
    "\n",
    "    \n",
    "    def call(self, inputs, mask):\n",
    "        batch_size = tf.shape(inputs[0])[0]\n",
    "    \n",
    "        q = self.wq(inputs[0])\n",
    "        k = self.wk(inputs[1])\n",
    "        v = self.wv(inputs[2])\n",
    "        for i in range(num_heads):\n",
    "            self_attention = self.attention(inputs=[q, v, k], mask=[None, mask])\n",
    "            if i == 0:\n",
    "                concat_attention = tf.concat([self_attention], axis=2)\n",
    "            else:\n",
    "                concat_attention = tf.concat([concat_attention, self_attention], axis=2)      \n",
    "        self_attention = self.dense(concat_attention)\n",
    "        return self_attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RdDqGayx67vv"
   },
   "source": [
    "#### Point wise feed forward network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gBqzJXGfHK3X"
   },
   "source": [
    "Point wise feed forward network consists of two fully-connected layers with a ReLU activation in between."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ET7xLt0yCT6Z"
   },
   "outputs": [],
   "source": [
    "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": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7e7hKcxn6-zd"
   },
   "source": [
    "#### Encoder layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ncyS-Ms3i2x_"
   },
   "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",
    "\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6LO_48Owmx_o"
   },
   "source": [
    "#### Decoder layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9SoX0-vd1hue"
   },
   "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, causal=True)\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, en_padding_mask, de_padding_mask):\n",
    "\n",
    "        attn1 = self.mha1([x, x, x], de_padding_mask)  # (batch_size, target_seq_len, d_model)\n",
    "        attn1 = self.dropout1(attn1, training=training)\n",
    "        out1 = self.layernorm1(attn1 + x)\n",
    "\n",
    "        attn2 = self.mha2(\n",
    "          [out1, enc_output, enc_output], en_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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SE1H51Ajm0q1"
   },
   "source": [
    "#### Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jpEox7gJ8FCI"
   },
   "outputs": [],
   "source": [
    "class Encoder(tf.keras.layers.Layer):\n",
    "    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,\n",
    "               maximum_position_encoding, rate=0.1):\n",
    "        super(Encoder, self).__init__()\n",
    "\n",
    "        self.d_model = d_model\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)\n",
    "        self.pos_encoding = positional_encoding(maximum_position_encoding, \n",
    "                                              self.d_model)\n",
    "\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",
    "\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": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "p-uO6ls8m2O5"
   },
   "source": [
    "### Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "d5_d5-PLQXwY"
   },
   "outputs": [],
   "source": [
    "class Decoder(tf.keras.layers.Layer):\n",
    "    def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,\n",
    "               maximum_position_encoding, rate=0.1):\n",
    "        super(Decoder, self).__init__()\n",
    "\n",
    "        self.d_model = d_model\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.embedding = tf.keras.layers.Embedding(target_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, en_padding_mask, de_padding_mask):\n",
    "\n",
    "        seq_len = tf.shape(x)[1]\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 = self.dec_layers[i](x, enc_output, training, en_padding_mask, de_padding_mask)\n",
    "\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "y54xnJnuYgJ7"
   },
   "source": [
    "#### Create the Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uERO1y54cOKq"
   },
   "source": [
    "Transformer consists of the encoder, decoder and a final linear layer. The output of the decoder is the input to the linear layer and its output is returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PED3bIpOYkBu"
   },
   "outputs": [],
   "source": [
    "class Transformer(tf.keras.Model):\n",
    "    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, \n",
    "               target_vocab_size, pe_input, pe_target, rate=0.1):\n",
    "        super(Transformer, self).__init__()\n",
    "\n",
    "        self.encoder = Encoder(num_layers, d_model, num_heads, dff, \n",
    "                            input_vocab_size, pe_input, rate)\n",
    "\n",
    "        self.decoder = Decoder(num_layers, d_model, num_heads, dff, \n",
    "                            target_vocab_size, pe_target, rate)\n",
    "        #self.final_layer = tf.keras.layers.Dense(target_vocab_size)\n",
    "        self.final_layer = tf.keras.layers.Dense(target_vocab_size, activation=\"softmax\")\n",
    "\n",
    "    def call(self, inp, tar, training, en_padding_mask, de_padding_mask):\n",
    "\n",
    "        enc_output = self.encoder(inp, training, en_padding_mask)  # (batch_size, inp_seq_len, d_model)\n",
    "\n",
    "        # dec_output.shape == (batch_size, tar_seq_len, d_model)\n",
    "        dec_output = self.decoder(\n",
    "          tar, enc_output, training, en_padding_mask, de_padding_mask)\n",
    "\n",
    "        final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)\n",
    "\n",
    "        return final_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xYEGhEOtzn5W"
   },
   "source": [
    "###  Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "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": 47,
   "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": 48,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7r4scdulztRx"
   },
   "outputs": [],
   "source": [
    "optimizer = tf.keras.optimizers.Adam(learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YgkDE7hzo8r5"
   },
   "source": [
    "### Loss and metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oxGJtoDuYIHL"
   },
   "source": [
    "Since the target sequences are padded, it is important to apply a padding mask when calculating the loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MlhsJMm0TW_B"
   },
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
    "    from_logits=False, reduction='none')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "67oqVHiT0Eiu"
   },
   "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": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy_function(real, pred):\n",
    "    mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
    "    pred_id = tf.cast(tf.argmax(pred, axis=-1), tf.int32)\n",
    "    return mask, pred_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "phlyxMnm-Tpx"
   },
   "outputs": [],
   "source": [
    "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
    "train_accuracy = tf.keras.metrics.Accuracy(\n",
    "    name='train_accuracy')\n",
    "test_loss = tf.keras.metrics.Mean(name='test_loss')\n",
    "test_accuracy = tf.keras.metrics.Accuracy(\n",
    "   name='test_accuracy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aeHumfr7zmMa"
   },
   "source": [
    "### Training and checkpointing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UiysUa--4tOU"
   },
   "outputs": [],
   "source": [
    "#from model import Transformer\n",
    "transformer = Transformer(num_layers, d_model, num_heads, dff,\n",
    "                          input_vocab_size, output_vocab_size, \n",
    "                          pe_input=sequence_length, \n",
    "                          pe_target=sequence_length,\n",
    "                          rate=dropout_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Fzuf06YZp66w"
   },
   "source": [
    "Create the checkpoint path and the checkpoint manager. This will be used to save checkpoints every `n` epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hNhuYfllndLZ"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Latest checkpoint restored!!\n"
     ]
    }
   ],
   "source": [
    "checkpoint_path = \"./checkpoints/train\"\n",
    "\n",
    "ckpt = tf.train.Checkpoint(model=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": 55,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iJwmp9OE29oj"
   },
   "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, None), dtype=tf.int32),\n",
    "    tf.TensorSpec(shape=(None, None), 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",
    "    en_padding_mask = create_padding_mask(inp)\n",
    "    de_padding_mask = create_padding_mask(tar_inp)\n",
    "    with tf.GradientTape() as tape:\n",
    "        predictions = transformer(inp, tar_inp, \n",
    "                                 True, \n",
    "                                 en_padding_mask,\n",
    "                                 de_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",
    "    mask, predictions = accuracy_function(tar_real, predictions)  \n",
    "    train_loss(loss)\n",
    "    train_accuracy(tar_real, predictions, sample_weight=mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_step_signature = [\n",
    "    tf.TensorSpec(shape=(None, None), dtype=tf.int32),\n",
    "    tf.TensorSpec(shape=(None, None), dtype=tf.int32),\n",
    "]\n",
    "\n",
    "@tf.function(input_signature=test_step_signature)\n",
    "def test_step(inp, tar):\n",
    "    tar_inp = tar[:, :-1]\n",
    "    tar_real = tar[:, 1:]\n",
    "    en_padding_mask = create_padding_mask(inp)\n",
    "    de_padding_mask = create_padding_mask(tar_inp)\n",
    "    predictions = transformer(inp, tar_inp, \n",
    "                                 False, \n",
    "                                 en_padding_mask,\n",
    "                                 de_padding_mask)\n",
    "    loss = loss_function(tar_real, predictions)\n",
    "    \n",
    "    mask, predictions = accuracy_function(tar_real, predictions)  \n",
    "    test_loss(loss)\n",
    "    test_accuracy(tar_real, predictions, sample_weight=mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_time = datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "train_log_dir = 'logs/gradient_tape/' + current_time + '/train'\n",
    "test_log_dir = 'logs/gradient_tape/' + current_time + '/test'\n",
    "train_summary_writer = tf.summary.create_file_writer(train_log_dir)\n",
    "test_summary_writer = tf.summary.create_file_writer(test_log_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_loss =[]\n",
    "t_acc = []\n",
    "v_loss = []\n",
    "v_acc = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bbvmaKNiznHZ"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 Batch 0 Train_Loss 2.5416 Train_Accuracy 0.1926\n",
      "Epoch 1 Batch 200 Train_Loss 2.4979 Train_Accuracy 0.2186\n",
      "Epoch 1 Batch 400 Train_Loss 2.4984 Train_Accuracy 0.2196\n",
      "Epoch 1 Batch 600 Train_Loss 2.5006 Train_Accuracy 0.2194\n",
      "Epoch 1 Batch 800 Train_Loss 2.5009 Train_Accuracy 0.2197\n",
      "Epoch 1 Batch 1000 Train_Loss 2.5011 Train_Accuracy 0.2200\n",
      "Epoch 1 Batch 1200 Train_Loss 2.5010 Train_Accuracy 0.2197\n",
      "Epoch 1 Batch 1400 Train_Loss 2.5009 Train_Accuracy 0.2196\n",
      "Epoch 1 Batch 1600 Train_Loss 2.5010 Train_Accuracy 0.2193\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(epochs):\n",
    "    start = time.time()\n",
    "\n",
    "    train_loss.reset_states()\n",
    "    train_accuracy.reset_states()\n",
    "    \n",
    "    test_loss.reset_states()\n",
    "    test_accuracy.reset_states()\n",
    "\n",
    "    # inp -> equation, tar -> integration\n",
    "    #tf.summary.trace_on(graph=True) \n",
    "    for (batch, (inp, tar)) in enumerate(dataset_train):\n",
    "        train_step(inp, tar)\n",
    "\n",
    "        if batch % 200 == 0:\n",
    "            print ('Epoch {} Batch {} Train_Loss {:.4f} Train_Accuracy {:.4f}'.format(\n",
    "              epoch + 1, batch, train_loss.result(), train_accuracy.result()))\n",
    "            \n",
    "    #if (epoch + 1) % 5 == 0:\n",
    "    ckpt_save_path = ckpt_manager.save()\n",
    "    print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,\n",
    "                                                    ckpt_save_path))            \n",
    "           \n",
    "    with train_summary_writer.as_default():\n",
    "        tf.summary.scalar('train_loss', train_loss.result(), step=epoch)\n",
    "        tf.summary.scalar('train_accuracy', train_accuracy.result(), step=epoch)\n",
    "        #tf.summary.trace_export(name=\"test_model\", step=epoch)\n",
    "        #train_summary_writer.flush()\n",
    "        \n",
    "    for (batch, (inp, tar)) in enumerate(dataset_test):\n",
    "        test_step(inp, tar)\n",
    "\n",
    "        if batch % 50 == 0:\n",
    "            print ('Epoch {} Batch {} Test_Loss {:.4f} Test_Accuracy {:.4f}'.format(\n",
    "              epoch + 1, batch, test_loss.result(), test_accuracy.result()))\n",
    "            \n",
    "    with test_summary_writer.as_default():\n",
    "        tf.summary.scalar('test_loss', test_loss.result(), step=epoch)\n",
    "        tf.summary.scalar('test_accuracy', test_accuracy.result(), step=epoch)  \n",
    "        \n",
    "\n",
    "    t_loss.append(train_loss.result())\n",
    "    t_acc.append(train_accuracy.result())\n",
    "    print ('Epoch {} Train_Loss {:.4f} Train_Accuracy {:.4f}'.format(epoch + 1, \n",
    "                                                train_loss.result(), \n",
    "                                                train_accuracy.result()))\n",
    "    v_loss.append(test_loss.result())\n",
    "    v_acc.append(test_accuracy.result())    \n",
    "    print ('Epoch {} Test_Loss {:.4f} Test_Accuracy {:.4f}'.format(epoch + 1, \n",
    "                                                test_loss.result(), \n",
    "                                                test_accuracy.result()))\n",
    "\n",
    "    print ('Time taken for 1 epoch: {} secs\\n'.format(time.time() - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(inp_sentence):\n",
    "    # add '<start> ' and ' <end>'\n",
    "    input_sentence = preprocess_sentence(inp_sentence)\n",
    "    \n",
    "    # tokenize input_sentence\n",
    "    input_sentence = np.asarray([X_tokenizer.word_index[w] for w in input_sentence.split(' ')], dtype=np.int32)\n",
    "    encoder_input = tf.expand_dims(input_sentence, 0)\n",
    "    encoder_input = tf.keras.preprocessing.sequence.pad_sequences(encoder_input, padding='post', maxlen=sequence_length, truncating='post')\n",
    "\n",
    "    # tokenize decoder_input\n",
    "    decoder_input = np.asarray([Y_tokenizer.word_index['<start>']], dtype=np.int32)\n",
    "    output = tf.expand_dims(decoder_input, 0)\n",
    "    #output = tf.keras.preprocessing.sequence.pad_sequences(output, padding='post', maxlen=sequence_length, truncating='post')\n",
    "\n",
    "    for i in range(sequence_length):\n",
    "        en_padding_mask = create_padding_mask(encoder_input)\n",
    "        de_padding_mask = create_padding_mask(output)\n",
    "\n",
    "        # predictions.shape == (batch_size, seq_len, vocab_size)\n",
    "        predictions = transformer(encoder_input, \n",
    "                                  output, \n",
    "                                  False, \n",
    "                                  en_padding_mask,\n",
    "                                  de_padding_mask)\n",
    "\n",
    "        # select the last word from the seq_len dimension\n",
    "        \n",
    "        predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)\n",
    "        \n",
    "        # greedy decoder\n",
    "        predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)\n",
    "\n",
    "#         # return the result if the predicted_id is equal to the end token\n",
    "        if predicted_id == Y_tokenizer.word_index['<end>']:\n",
    "             return tf.squeeze(output, axis=0)\n",
    "\n",
    "#         # concatentate the predicted_id to the output which is given to the decoder\n",
    "#         # as its input.\n",
    "        output = tf.concat([output, predicted_id], axis=-1)\n",
    "\n",
    "    return tf.squeeze(output, axis=0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def translate(sentence):\n",
    "    result = evaluate(sentence)\n",
    "    #print(result.shape)\n",
    "    predicted_sentence = \" \".join([Y_tokenizer.index_word[w.numpy()] for w in result[1:]]) \n",
    "\n",
    "    #print('Input: {}'.format(sentence))\n",
    "    #print('Predicted translation: {}'.format(predicted_sentence))\n",
    "    return predicted_sentence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### beam search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_beam_search(inp_sentence, beam_width):\n",
    "    # add '<start> ' and ' <end>'\n",
    "    input_sentence = preprocess_sentence(inp_sentence)\n",
    "    \n",
    "    # tokenize input_sentence\n",
    "    input_sentence = np.asarray([X_tokenizer.word_index[w] for w in input_sentence.split(' ')], dtype=np.int32)\n",
    "    encoder_input = tf.expand_dims(input_sentence, 0)\n",
    "    encoder_input = tf.keras.preprocessing.sequence.pad_sequences(encoder_input, padding='post', maxlen=sequence_length, truncating='post')\n",
    "\n",
    "    # tokenize decoder_input\n",
    "    decoder_input = np.asarray([Y_tokenizer.word_index['<start>']], dtype=np.int32)\n",
    "    output = tf.expand_dims(decoder_input, 0)\n",
    "    #output = tf.keras.preprocessing.sequence.pad_sequences(output, padding='post', maxlen=sequence_length, truncating='post')\n",
    "    \n",
    "    k = beam_width\n",
    "    sequences = [[list(), 1.0]]\n",
    "    \n",
    "    for i in range(sequence_length):\n",
    "        en_padding_mask = create_padding_mask(encoder_input)\n",
    "        de_padding_mask = create_padding_mask(output)\n",
    "\n",
    "        # predictions.shape == (batch_size, seq_len, vocab_size)\n",
    "        predictions = transformer(encoder_input, \n",
    "                                  output, \n",
    "                                  False, \n",
    "                                  en_padding_mask,\n",
    "                                  de_padding_mask)\n",
    "\n",
    "        # select the last word from the seq_len dimension\n",
    "        predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)\n",
    "        vocab_size = predictions.shape[2]\n",
    "        \n",
    "        #beam search decoder\n",
    "        all_candidates = list()\n",
    "        for i in range(len(sequences)):\n",
    "            seq, score = sequences[i]\n",
    "            if len(seq)>0 and seq[-1] == Y_tokenizer.word_index['<end>']:\n",
    "                pass\n",
    "            else:    \n",
    "                for j in range(vocab_size):\n",
    "                    prob = tf.cast(predictions[0,-1, j], tf.float32)\n",
    "                    candidate = [seq + [j], score * (-log(prob))]\n",
    "                    all_candidates.append(candidate)\n",
    "        # order all candidates by score\n",
    "        ordered = sorted(all_candidates, key=lambda tup:tup[1])\n",
    "        # select k best\n",
    "        sequences = ordered[:k]\n",
    "    return sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def translate_beam_search(sentence, beam_width):\n",
    "    result = evaluate_beam_search(sentence, beam_width)\n",
    "    predicted_sentences = []\n",
    "    print('Input: {}'.format(sentence))\n",
    "    for i in range(len(result)):\n",
    "        predicted_sentence = \" \".join([Y_tokenizer.index_word[w] for w in result[i][0]]) \n",
    "        print('Predicted translation: {}'.format(predicted_sentence))\n",
    "        predicted_sentences.append([predicted_sentence, result[i][1]])\n",
    "    return predicted_sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#result = translate_beam_search(\"+ x 3\", beam_width=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ssh to localhost:6006"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WxAniRctJCno"
   },
   "outputs": [],
   "source": [
    "%tensorboard --logdir logs/gradient_tape --port=6006"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "s_qNSzzyaCbD"
   ],
   "name": "transformer.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "py_36_env",
   "language": "python",
   "name": "py_36_env"
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}