NMT_Attention_Marathi.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dataset => http://www.manythings.org/anki/ (mar-eng.zip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script><script type=\"text/javascript\">if (window.MathJax) {MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script><script>requirejs.config({paths: { 'plotly': ['https://cdn.plot.ly/plotly-latest.min']},});if(!window._Plotly) {require(['plotly'],function(plotly) {window._Plotly=plotly;});}</script>"
      ],
      "text/vnd.plotly.v1+html": [
       "<script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script><script type=\"text/javascript\">if (window.MathJax) {MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script><script>requirejs.config({paths: { 'plotly': ['https://cdn.plot.ly/plotly-latest.min']},});if(!window._Plotly) {require(['plotly'],function(plotly) {window._Plotly=plotly;});}</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "tf.enable_eager_execution()\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import unicodedata\n",
    "import re\n",
    "import numpy as np\n",
    "import os\n",
    "import time\n",
    "import string\n",
    "\n",
    "import plotly\n",
    "import plotly.plotly as py\n",
    "from plotly.offline import init_notebook_mode, iplot\n",
    "plotly.offline.init_notebook_mode(connected=True)\n",
    "import plotly.graph_objs as go"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_path = './mar.txt' # please set the path according to your system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['I know my rights.\\tमला माझे अधिकार माहीत आहेत.',\n",
       " 'I know that girl.\\tमी त्या मुलीला ओळखतो.',\n",
       " 'I know that girl.\\tमी त्या मुलीला ओळखते.',\n",
       " 'I know the rules.\\tमला नियम माहीत आहेत.',\n",
       " 'I know your name.\\tमला तुझं नाव माहीत आहे.',\n",
       " 'I know your name.\\tमला तुझं नाव ठाऊक आहे.',\n",
       " 'I know your name.\\tमला तुमचं नाव माहीत आहे.',\n",
       " 'I like green tea.\\tमला हिरवा चहा आवडतो.',\n",
       " 'I like his music.\\tमला त्याचं संगीत आवडतं.',\n",
       " 'I like his music.\\tमला त्यांचं संगीत आवडतं.']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lines = open(file_path, encoding='UTF-8').read().strip().split('\\n')\n",
    "lines[5000:5010]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "33725"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "exclude = set(string.punctuation) # Set of all special characters\n",
    "remove_digits = str.maketrans('', '', string.digits) # Set of all digits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_eng_sentence(sent):\n",
    "    '''Function to preprocess English sentence'''\n",
    "    sent = sent.lower() # lower casing\n",
    "    sent = re.sub(\"'\", '', sent) # remove the quotation marks if any\n",
    "    sent = ''.join(ch for ch in sent if ch not in exclude)\n",
    "    sent = sent.translate(remove_digits) # remove the digits\n",
    "    sent = sent.strip()\n",
    "    sent = re.sub(\" +\", \" \", sent) # remove extra spaces\n",
    "    sent = '<start> ' + sent + ' <end>' # add <start> and <end> tokens\n",
    "    return sent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_mar_sentence(sent):\n",
    "    '''Function to preprocess Marathi sentence'''\n",
    "    sent = re.sub(\"'\", '', sent) # remove the quotation marks if any\n",
    "    sent = ''.join(ch for ch in sent if ch not in exclude)\n",
    "    sent = re.sub(\"[२३०८१५७९४६]\", \"\", sent) # remove the digits\n",
    "    sent = sent.strip()\n",
    "    sent = re.sub(\" +\", \" \", sent) # remove extra spaces\n",
    "    sent = '<start> ' + sent + ' <end>' # add <start> and <end> tokens\n",
    "    return sent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['<start> i know my rights <end>',\n",
       "  '<start> मला माझे अधिकार माहीत आहेत <end>'],\n",
       " ['<start> i know that girl <end>', '<start> मी त्या मुलीला ओळखतो <end>'],\n",
       " ['<start> i know that girl <end>', '<start> मी त्या मुलीला ओळखते <end>'],\n",
       " ['<start> i know the rules <end>', '<start> मला नियम माहीत आहेत <end>'],\n",
       " ['<start> i know your name <end>', '<start> मला तुझं नाव माहीत आहे <end>'],\n",
       " ['<start> i know your name <end>', '<start> मला तुझं नाव ठाऊक आहे <end>'],\n",
       " ['<start> i know your name <end>', '<start> मला तुमचं नाव माहीत आहे <end>'],\n",
       " ['<start> i like green tea <end>', '<start> मला हिरवा चहा आवडतो <end>'],\n",
       " ['<start> i like his music <end>', '<start> मला त्याचं संगीत आवडतं <end>'],\n",
       " ['<start> i like his music <end>', '<start> मला त्यांचं संगीत आवडतं <end>']]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Generate pairs of cleaned English and Marathi sentences\n",
    "sent_pairs = []\n",
    "for line in lines:\n",
    "    sent_pair = []\n",
    "    eng, mar = line.split('\\t')\n",
    "    eng = preprocess_eng_sentence(eng)\n",
    "    sent_pair.append(eng)\n",
    "    mar = preprocess_mar_sentence(mar)\n",
    "    sent_pair.append(mar)\n",
    "    sent_pairs.append(sent_pair)\n",
    "sent_pairs[5000:5010]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This class creates a word -> index mapping (e.g,. \"dad\" -> 5) and vice-versa \n",
    "# (e.g., 5 -> \"dad\") for each language,\n",
    "class LanguageIndex():\n",
    "    def __init__(self, lang):\n",
    "        self.lang = lang\n",
    "        self.word2idx = {}\n",
    "        self.idx2word = {}\n",
    "        self.vocab = set()\n",
    "\n",
    "        self.create_index()\n",
    "\n",
    "    def create_index(self):\n",
    "        for phrase in self.lang:\n",
    "            self.vocab.update(phrase.split(' '))\n",
    "\n",
    "        self.vocab = sorted(self.vocab)\n",
    "\n",
    "        self.word2idx['<pad>'] = 0\n",
    "        for index, word in enumerate(self.vocab):\n",
    "            self.word2idx[word] = index + 1\n",
    "\n",
    "        for word, index in self.word2idx.items():\n",
    "            self.idx2word[index] = word"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def max_length(tensor):\n",
    "    return max(len(t) for t in tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_dataset(pairs, num_examples):\n",
    "    # pairs => already created cleaned input, output pairs\n",
    "\n",
    "    # index language using the class defined above    \n",
    "    inp_lang = LanguageIndex(en for en, ma in pairs)\n",
    "    targ_lang = LanguageIndex(ma for en, ma in pairs)\n",
    "    \n",
    "    # Vectorize the input and target languages\n",
    "    \n",
    "    # English sentences\n",
    "    input_tensor = [[inp_lang.word2idx[s] for s in en.split(' ')] for en, ma in pairs]\n",
    "    \n",
    "    # Marathi sentences\n",
    "    target_tensor = [[targ_lang.word2idx[s] for s in ma.split(' ')] for en, ma in pairs]\n",
    "    \n",
    "    # Calculate max_length of input and output tensor\n",
    "    # Here, we'll set those to the longest sentence in the dataset\n",
    "    max_length_inp, max_length_tar = max_length(input_tensor), max_length(target_tensor)\n",
    "    \n",
    "    # Padding the input and output tensor to the maximum length\n",
    "    input_tensor = tf.keras.preprocessing.sequence.pad_sequences(input_tensor, \n",
    "                                                                 maxlen=max_length_inp,\n",
    "                                                                 padding='post')\n",
    "    \n",
    "    target_tensor = tf.keras.preprocessing.sequence.pad_sequences(target_tensor, \n",
    "                                                                  maxlen=max_length_tar, \n",
    "                                                                  padding='post')\n",
    "    \n",
    "    return input_tensor, target_tensor, inp_lang, targ_lang, max_length_inp, max_length_tar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_tensor, target_tensor, inp_lang, targ_lang, max_length_inp, max_length_targ = load_dataset(sent_pairs, len(lines))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(30352, 30352, 3373, 3373)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Creating training and validation sets using an 80-20 split\n",
    "input_tensor_train, input_tensor_val, target_tensor_train, target_tensor_val = train_test_split(input_tensor, target_tensor, test_size=0.1, random_state = 101)\n",
    "\n",
    "# Show length\n",
    "len(input_tensor_train), len(target_tensor_train), len(input_tensor_val), len(target_tensor_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "BUFFER_SIZE = len(input_tensor_train)\n",
    "BATCH_SIZE = 64\n",
    "N_BATCH = BUFFER_SIZE//BATCH_SIZE\n",
    "embedding_dim = 256\n",
    "units = 1024\n",
    "vocab_inp_size = len(inp_lang.word2idx)\n",
    "vocab_tar_size = len(targ_lang.word2idx)\n",
    "\n",
    "dataset = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train)).shuffle(BUFFER_SIZE)\n",
    "dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gru(units):\n",
    "  # If you have a GPU, we recommend using CuDNNGRU(provides a 3x speedup than GRU)\n",
    "  # the code automatically does that.\n",
    "    if tf.test.is_gpu_available():\n",
    "        return tf.keras.layers.CuDNNGRU(units, \n",
    "                                        return_sequences=True, \n",
    "                                        return_state=True, \n",
    "                                        recurrent_initializer='glorot_uniform')\n",
    "    else:\n",
    "        return tf.keras.layers.GRU(units, \n",
    "                                   return_sequences=True, \n",
    "                                   return_state=True, \n",
    "                                   recurrent_activation='sigmoid', \n",
    "                                   recurrent_initializer='glorot_uniform')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(tf.keras.Model):\n",
    "    def __init__(self, vocab_size, embedding_dim, enc_units, batch_sz):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.batch_sz = batch_sz\n",
    "        self.enc_units = enc_units\n",
    "        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n",
    "        self.gru = gru(self.enc_units)\n",
    "        \n",
    "    def call(self, x, hidden):\n",
    "        x = self.embedding(x)\n",
    "        output, state = self.gru(x, initial_state = hidden)        \n",
    "        return output, state\n",
    "    \n",
    "    def initialize_hidden_state(self):\n",
    "        return tf.zeros((self.batch_sz, self.enc_units))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(tf.keras.Model):\n",
    "    def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.batch_sz = batch_sz\n",
    "        self.dec_units = dec_units\n",
    "        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n",
    "        self.gru = gru(self.dec_units)\n",
    "        self.fc = tf.keras.layers.Dense(vocab_size)\n",
    "        \n",
    "        # used for attention\n",
    "        self.W1 = tf.keras.layers.Dense(self.dec_units)\n",
    "        self.W2 = tf.keras.layers.Dense(self.dec_units)\n",
    "        self.V = tf.keras.layers.Dense(1)\n",
    "        \n",
    "    def call(self, x, hidden, enc_output):\n",
    "        # enc_output shape == (batch_size, max_length, hidden_size)\n",
    "        \n",
    "        # hidden shape == (batch_size, hidden size)\n",
    "        # hidden_with_time_axis shape == (batch_size, 1, hidden size)\n",
    "        # we are doing this to perform addition to calculate the score\n",
    "        hidden_with_time_axis = tf.expand_dims(hidden, 1)\n",
    "        \n",
    "        # score shape == (batch_size, max_length, 1)\n",
    "        # we get 1 at the last axis because we are applying tanh(FC(EO) + FC(H)) to self.V\n",
    "        score = self.V(tf.nn.tanh(self.W1(enc_output) + self.W2(hidden_with_time_axis)))\n",
    "        \n",
    "        # attention_weights shape == (batch_size, max_length, 1)\n",
    "        attention_weights = tf.nn.softmax(score, axis=1)\n",
    "        \n",
    "        # context_vector shape after sum == (batch_size, hidden_size)\n",
    "        context_vector = attention_weights * enc_output\n",
    "        context_vector = tf.reduce_sum(context_vector, axis=1)\n",
    "        \n",
    "        # x shape after passing through embedding == (batch_size, 1, embedding_dim)\n",
    "        x = self.embedding(x)\n",
    "        \n",
    "        # x shape after concatenation == (batch_size, 1, embedding_dim + hidden_size)\n",
    "        x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)\n",
    "        \n",
    "        # passing the concatenated vector to the GRU\n",
    "        output, state = self.gru(x)\n",
    "        \n",
    "        # output shape == (batch_size * 1, hidden_size)\n",
    "        output = tf.reshape(output, (-1, output.shape[2]))\n",
    "        \n",
    "        # output shape == (batch_size * 1, vocab)\n",
    "        x = self.fc(output)\n",
    "        \n",
    "        return x, state, attention_weights\n",
    "        \n",
    "    def initialize_hidden_state(self):\n",
    "        return tf.zeros((self.batch_sz, self.dec_units))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = Encoder(vocab_inp_size, embedding_dim, units, BATCH_SIZE)\n",
    "decoder = Decoder(vocab_tar_size, embedding_dim, units, BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.train.AdamOptimizer()\n",
    "\n",
    "\n",
    "def loss_function(real, pred):\n",
    "    mask = 1 - np.equal(real, 0)\n",
    "    loss_ = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=real, logits=pred) * mask\n",
    "    return tf.reduce_mean(loss_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_dir = './training_checkpoints'\n",
    "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n",
    "checkpoint = tf.train.Checkpoint(optimizer=optimizer,\n",
    "                                 encoder=encoder,\n",
    "                                 decoder=decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 Batch 0 Loss 1.5206\n",
      "Epoch 1 Batch 100 Loss 1.0305\n",
      "Epoch 1 Batch 200 Loss 0.7886\n",
      "Epoch 1 Batch 300 Loss 0.7676\n",
      "Epoch 1 Batch 400 Loss 0.7481\n",
      "Epoch 1 Loss 0.8671\n",
      "Time taken for 1 epoch 352.7763113975525 sec\n",
      "\n",
      "Epoch 2 Batch 0 Loss 0.7193\n",
      "Epoch 2 Batch 100 Loss 0.6355\n",
      "Epoch 2 Batch 200 Loss 0.6897\n",
      "Epoch 2 Batch 300 Loss 0.5928\n",
      "Epoch 2 Batch 400 Loss 0.5776\n",
      "Epoch 2 Loss 0.6535\n",
      "Time taken for 1 epoch 351.2787010669708 sec\n",
      "\n",
      "Epoch 3 Batch 0 Loss 0.5011\n",
      "Epoch 3 Batch 100 Loss 0.4631\n",
      "Epoch 3 Batch 200 Loss 0.5464\n",
      "Epoch 3 Batch 300 Loss 0.5196\n",
      "Epoch 3 Batch 400 Loss 0.4962\n",
      "Epoch 3 Loss 0.4708\n",
      "Time taken for 1 epoch 353.4211723804474 sec\n",
      "\n",
      "Epoch 4 Batch 0 Loss 0.3234\n",
      "Epoch 4 Batch 100 Loss 0.3257\n",
      "Epoch 4 Batch 200 Loss 0.3345\n",
      "Epoch 4 Batch 300 Loss 0.3141\n",
      "Epoch 4 Batch 400 Loss 0.3244\n",
      "Epoch 4 Loss 0.3200\n",
      "Time taken for 1 epoch 353.1946771144867 sec\n",
      "\n",
      "Epoch 5 Batch 0 Loss 0.2030\n",
      "Epoch 5 Batch 100 Loss 0.2257\n",
      "Epoch 5 Batch 200 Loss 0.1945\n",
      "Epoch 5 Batch 300 Loss 0.2250\n",
      "Epoch 5 Batch 400 Loss 0.2372\n",
      "Epoch 5 Loss 0.2143\n",
      "Time taken for 1 epoch 354.23998522758484 sec\n",
      "\n",
      "Epoch 6 Batch 0 Loss 0.1239\n",
      "Epoch 6 Batch 100 Loss 0.1192\n",
      "Epoch 6 Batch 200 Loss 0.1396\n",
      "Epoch 6 Batch 300 Loss 0.1327\n",
      "Epoch 6 Batch 400 Loss 0.1459\n",
      "Epoch 6 Loss 0.1476\n",
      "Time taken for 1 epoch 354.23475980758667 sec\n",
      "\n",
      "Epoch 7 Batch 0 Loss 0.1021\n",
      "Epoch 7 Batch 100 Loss 0.0832\n",
      "Epoch 7 Batch 200 Loss 0.0870\n",
      "Epoch 7 Batch 300 Loss 0.1290\n",
      "Epoch 7 Batch 400 Loss 0.1195\n",
      "Epoch 7 Loss 0.1083\n",
      "Time taken for 1 epoch 355.2288661003113 sec\n",
      "\n",
      "Epoch 8 Batch 0 Loss 0.0767\n",
      "Epoch 8 Batch 100 Loss 0.0926\n",
      "Epoch 8 Batch 200 Loss 0.0877\n",
      "Epoch 8 Batch 300 Loss 0.1159\n",
      "Epoch 8 Batch 400 Loss 0.0882\n",
      "Epoch 8 Loss 0.0862\n",
      "Time taken for 1 epoch 352.97160029411316 sec\n",
      "\n",
      "Epoch 9 Batch 0 Loss 0.0666\n",
      "Epoch 9 Batch 100 Loss 0.0602\n",
      "Epoch 9 Batch 200 Loss 0.0587\n",
      "Epoch 9 Batch 300 Loss 0.0814\n",
      "Epoch 9 Batch 400 Loss 0.0761\n",
      "Epoch 9 Loss 0.0727\n",
      "Time taken for 1 epoch 353.38186383247375 sec\n",
      "\n",
      "Epoch 10 Batch 0 Loss 0.0606\n",
      "Epoch 10 Batch 100 Loss 0.0516\n",
      "Epoch 10 Batch 200 Loss 0.0664\n",
      "Epoch 10 Batch 300 Loss 0.0800\n",
      "Epoch 10 Batch 400 Loss 0.0683\n",
      "Epoch 10 Loss 0.0642\n",
      "Time taken for 1 epoch 353.1245322227478 sec\n",
      "\n"
     ]
    }
   ],
   "source": [
    "EPOCHS = 10\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    start = time.time()\n",
    "    \n",
    "    hidden = encoder.initialize_hidden_state()\n",
    "    total_loss = 0\n",
    "    \n",
    "    for (batch, (inp, targ)) in enumerate(dataset):\n",
    "        loss = 0\n",
    "        \n",
    "        with tf.GradientTape() as tape:\n",
    "            enc_output, enc_hidden = encoder(inp, hidden)\n",
    "            \n",
    "            dec_hidden = enc_hidden\n",
    "            \n",
    "            dec_input = tf.expand_dims([targ_lang.word2idx['<start>']] * BATCH_SIZE, 1)       \n",
    "            \n",
    "            # Teacher forcing - feeding the target as the next input\n",
    "            for t in range(1, targ.shape[1]):\n",
    "                # passing enc_output to the decoder\n",
    "                predictions, dec_hidden, _ = decoder(dec_input, dec_hidden, enc_output)\n",
    "                \n",
    "                loss += loss_function(targ[:, t], predictions)\n",
    "                \n",
    "                # using teacher forcing\n",
    "                dec_input = tf.expand_dims(targ[:, t], 1)\n",
    "        \n",
    "        batch_loss = (loss / int(targ.shape[1]))\n",
    "        \n",
    "        total_loss += batch_loss\n",
    "        \n",
    "        variables = encoder.variables + decoder.variables\n",
    "        \n",
    "        gradients = tape.gradient(loss, variables)\n",
    "        \n",
    "        optimizer.apply_gradients(zip(gradients, variables))\n",
    "        \n",
    "        if batch % 100 == 0:\n",
    "            print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,\n",
    "                                                         batch,\n",
    "                                                         batch_loss.numpy()))\n",
    "    # saving (checkpoint) the model every epoch\n",
    "    checkpoint.save(file_prefix = checkpoint_prefix)\n",
    "    \n",
    "    print('Epoch {} Loss {:.4f}'.format(epoch + 1,\n",
    "                                        total_loss / N_BATCH))\n",
    "    print('Time taken for 1 epoch {} sec\\n'.format(time.time() - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.training.checkpointable.util.CheckpointLoadStatus at 0x7fe3b1759be0>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# restoring the latest checkpoint in checkpoint_dir\n",
    "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Inference Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(inputs, encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ):\n",
    "    \n",
    "    attention_plot = np.zeros((max_length_targ, max_length_inp))\n",
    "    sentence = ''\n",
    "    for i in inputs[0]:\n",
    "        if i == 0:\n",
    "            break\n",
    "        sentence = sentence + inp_lang.idx2word[i] + ' '\n",
    "    sentence = sentence[:-1]\n",
    "    \n",
    "    inputs = tf.convert_to_tensor(inputs)\n",
    "    \n",
    "    result = ''\n",
    "\n",
    "    hidden = [tf.zeros((1, units))]\n",
    "    enc_out, enc_hidden = encoder(inputs, hidden)\n",
    "\n",
    "    dec_hidden = enc_hidden\n",
    "    dec_input = tf.expand_dims([targ_lang.word2idx['<start>']], 0)\n",
    "\n",
    "    for t in range(max_length_targ):\n",
    "        predictions, dec_hidden, attention_weights = decoder(dec_input, dec_hidden, enc_out)\n",
    "        \n",
    "        # storing the attention weights to plot later on\n",
    "        attention_weights = tf.reshape(attention_weights, (-1, ))\n",
    "        attention_plot[t] = attention_weights.numpy()\n",
    "\n",
    "        predicted_id = tf.argmax(predictions[0]).numpy()\n",
    "\n",
    "        result += targ_lang.idx2word[predicted_id] + ' '\n",
    "\n",
    "        if targ_lang.idx2word[predicted_id] == '<end>':\n",
    "            return result, sentence, attention_plot\n",
    "        \n",
    "        # the predicted ID is fed back into the model\n",
    "        dec_input = tf.expand_dims([predicted_id], 0)\n",
    "\n",
    "    return result, sentence, attention_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Function to predict (translate) a randomly selected test point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_random_val_sentence():\n",
    "    actual_sent = ''\n",
    "    k = np.random.randint(len(input_tensor_val))\n",
    "    random_input = input_tensor_val[k]\n",
    "    random_output = target_tensor_val[k]\n",
    "    random_input = np.expand_dims(random_input,0)\n",
    "    result, sentence, attention_plot = evaluate(random_input, encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)\n",
    "    print('Input: {}'.format(sentence[8:-6]))\n",
    "    print('Predicted translation: {}'.format(result[:-6]))\n",
    "    for i in random_output:\n",
    "        if i == 0:\n",
    "            break\n",
    "        actual_sent = actual_sent + targ_lang.idx2word[i] + ' '\n",
    "    actual_sent = actual_sent[8:-7]\n",
    "    print('Actual translation: {}'.format(actual_sent))\n",
    "    attention_plot = attention_plot[:len(result.split(' '))-2, 1:len(sentence.split(' '))-1]\n",
    "    sentence, result = sentence.split(' '), result.split(' ')\n",
    "    sentence = sentence[1:-1]\n",
    "    result = result[:-2]\n",
    "    \n",
    "    # use plotly to generate the heat map\n",
    "    trace = go.Heatmap(z = attention_plot, x = sentence, y = result, colorscale='Reds')\n",
    "    data=[trace]\n",
    "    iplot(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call the function to visualize outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_random_val_sentence()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_random_val_sentence()"
   ]
  },
  {
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}