Seq2seq_Translation_LSTM_with_Attention.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "Seq2seq Translation LSTM with Attention",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "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.8"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/JosephBless/DL/blob/main/Seq2seq_Translation_LSTM_with_Attention.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MOJYWeJjIH1g"
      },
      "source": [
        "# Machine Translation with LSTM and attention\n",
        "\n",
        "This notebook is to show case the attention layer using seq2seq model trained as translator from English to French. The model is composed of a bidirectional LSTM as encoder and an LSTM as the decoder and of course, the decoder and the encoder are fed to an attention layer. The dataset used is one from Udacity's repository and for text preprocessing, SentencePiece is used to convert the input text into sub-wordings.\n",
        "\n",
        "You will be presented the choice to train everything from scratch yourself or load the models that are trained already in order to just test them. If you choose to train the model youself, based on my experience, it takes more than six hours to acheive 95% accuracy on validation set using a GPU. it will much longer with CPU or a TPU. I do not recommend those options.\n",
        "\n",
        "If you are running this notebook on Google Colab, first remember to go and change the runtime type to GPU. Also there's a good chance that the default 12GB of RAM that is assigned to your notebook is not enough. Unfortunately, it seems there's no way to request for more RAM unless you crash the notebook (due to insufficient memory) which will lead to a popup to appear, offering you more memory (it will appear on the lower left corner). By clicking on that, you'll get 25GB of RAM which should be sufficient for you. But remember, you have to start all over and run the cells again. Based on my experience, step 5 where you'll be mapping the text into a sequence of integers will crash the VM.\n",
        "\n",
        "Finally, the code blocks that are marked as **[TRAINING]** should only be executed only if you want to follow the training path. If you want to load the trained models, omit them."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R3D0A7DqJ_m7"
      },
      "source": [
        "# 1. Clone the attention layer repository\n",
        "\n",
        "This is to add the attention layer to Keras since at this moment it is not part of the project."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-qvsjDtzdxPq"
      },
      "source": [
        "!git clone https://github.com/ziadloo/attention_keras.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2gl8XGEvKHgy"
      },
      "source": [
        "# 2. Download the dataset\n",
        "\n",
        "The dataset is composed of 137860 sentences in both English and French. Each sentence is written in one line and corresponding lines of the two files are the same sentences in different languages."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jD3lUGxcm1KV"
      },
      "source": [
        "!wget -P ./attention_keras/data https://github.com/udacity/deep-learning/raw/master/language-translation/data/small_vocab_en\n",
        "!wget -P ./attention_keras/data https://github.com/udacity/deep-learning/raw/master/language-translation/data/small_vocab_fr"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OO1s_3ryKPHi"
      },
      "source": [
        "# 3. Install the SentencePiece library\n",
        "\n",
        "[SentencePiece](https://github.com/google/sentencepiece/blob/master/python/README.md) is a great library for converting texts into sub-words. Sub-words are the prefered way of tokenizing the text since they are something in between character level tokentization and word level tokenization."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Ra30pEJFiyB"
      },
      "source": [
        "!pip install sentencepiece"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3yNuA2DSKU7V"
      },
      "source": [
        "# 4. Configure the Python's path\n",
        "\n",
        "This is in order to help Python find the relative addressing for the `attention_keras` library we just downloaded."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Dd7rHC9Fy0it"
      },
      "source": [
        "import os\n",
        "import sys\n",
        "base_dir = os.path.join(os.getcwd(), \"attention_keras\")\n",
        "sys.path.insert(0, base_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gg92G_VEKnSV"
      },
      "source": [
        "# 5. Train the sub-word mapping\n",
        "\n",
        "Thanks to SentencePiece, it is so easy to have a sub-word mapping for our dataset. By this process we will have a mapping from English words in our dataset to an integer that we can use in our Machine Learning model (and one separate model for French).\n",
        "\n",
        "Once run, there will be four files generated in the `data` folder which we can feed back to the SentencePiece and map our input sentences to integers.\n",
        "\n",
        "**[TRAINING]**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "frhrfsSFF59k"
      },
      "source": [
        "import sentencepiece as spm\n",
        "\n",
        "target_vocab_size_en = 400\n",
        "target_vocab_size_fr = 600\n",
        "\n",
        "spm.SentencePieceTrainer.Train(\n",
        "    f\" --input={base_dir}/data/small_vocab_en --model_type=unigram --hard_vocab_limit=false\" +\n",
        "    f\" --model_prefix={base_dir}/data/en --vocab_size={target_vocab_size_en}\")\n",
        "spm.SentencePieceTrainer.Train(\n",
        "    f\" --input={base_dir}/data/small_vocab_fr --model_type=unigram --hard_vocab_limit=false\" +\n",
        "    f\" --model_prefix={base_dir}/data/fr --vocab_size={target_vocab_size_fr}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_5oDUF5CxDwx"
      },
      "source": [
        "This block loads the sub-word mapping into memory. Make sure you run whether you want to train the model yourself or not. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RMq8sFDQP_M_"
      },
      "source": [
        "import sentencepiece as spm\n",
        "\n",
        "sp_en = spm.SentencePieceProcessor()\n",
        "sp_en.Load(os.path.join(base_dir, \"data\", 'en.model'))\n",
        "\n",
        "sp_fr = spm.SentencePieceProcessor()\n",
        "sp_fr.Load(os.path.join(base_dir, \"data\", 'fr.model'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DECsp_BUmfre"
      },
      "source": [
        "Now that we have our models loaded into `sp_en` and `sp_fr`, we can read the text files and convert then to sequences of integers. Once we are done with this phase, we won't be needing the actual text.\n",
        "\n",
        "The `pad_sequences` function from `keras` is also used to make all the samples of the same length. Since this is a small dataset, all the samples are made as long as the longest one in the dataset.\n",
        "\n",
        "We will need two extra tokens for input language (in our case English) and three extra tokens for the output language (French). The extra tokens are `<end>`, `<empty>`, and `<start>`. Each sample sequence will have an `<end>` token appended to mark the end of the sequence. For the samples other than the longest one, the empty tokens are filled with `<empty>`. And `<start>` is used in the output samples since we need a `<start>` token to kick off the decoder. Since the output samples will have an extra `<start>` in their beginnings, all of them are padded to a size two tokens longer than the longest one (to accomodate for the `<start>` and the `<end>` tokens while the input samples will only have one token longer than the longest input since we only append the `<end>`).\n",
        "\n",
        "Also, while I named them `<end>`, `<empty>`, and `<start>` but they are never used in these forms instead in their integer forms. One last thing, while the `<end>` and `<empty>` might end up having the same ID in English and French, but that's not necessary true. So I have two versions of each mentioned tokens."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xXzah1LNsSUy"
      },
      "source": [
        "import numpy as np\n",
        "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
        "from tensorflow.python.keras.utils import to_categorical\n",
        "\n",
        "with open(os.path.join(base_dir, 'data', 'small_vocab_en'),\n",
        "          'r', encoding='utf-8') as file:\n",
        "  en_text = file.read().split(\"\\n\")\n",
        "\n",
        "with open(os.path.join(base_dir, 'data', 'small_vocab_fr'),\n",
        "          'r', encoding='utf-8') as file:\n",
        "  fr_text = file.read().split(\"\\n\")\n",
        "\n",
        "train_en_X = []\n",
        "train_fr_X = []\n",
        "train_fr_Y = []\n",
        "\n",
        "en_max_len = 0\n",
        "fr_max_len = 0\n",
        "\n",
        "vocab_size_en = sp_en.GetPieceSize()\n",
        "vocab_size_fr = sp_fr.GetPieceSize()\n",
        "\n",
        "# Assuming three extra tokens: <end>: #vocab_size_en | #vocab_size_fr,\n",
        "# <empty>: #vocab_size_en+1 | #vocab_size_fr+1, and <start>: #vocab_size_fr+2\n",
        "\n",
        "end_token_id_en = vocab_size_en\n",
        "empty_token_id_en = vocab_size_en + 1\n",
        "end_token_id_fr = vocab_size_fr\n",
        "empty_token_id_fr = vocab_size_fr + 1\n",
        "start_token_id_fr = vocab_size_fr + 2\n",
        "\n",
        "# The input text only needs two extra tokens while the output needs 3\n",
        "vocab_size_en = vocab_size_en + 2\n",
        "vocab_size_fr = vocab_size_fr + 3\n",
        "\n",
        "\n",
        "for i in range(len(en_text)):\n",
        "  en_seq = sp_en.EncodeAsIds(en_text[i].strip()) + [end_token_id_en]\n",
        "  en_max_len = max(en_max_len, len(en_seq))\n",
        "  train_en_X.append(en_seq)\n",
        "\n",
        "  fr_seq = sp_fr.EncodeAsIds(fr_text[i].strip()) + [end_token_id_fr]\n",
        "  fr_max_len = max(fr_max_len, len(fr_seq))\n",
        "  train_fr_X.append(fr_seq)\n",
        "\n",
        "# Cleaning up the memory (we don't need them anymore)\n",
        "#en_text = []\n",
        "#fr_text = []\n",
        "\n",
        "# Padding all the samples with <empty> token to make them all of the same length\n",
        "# equal to the longest one\n",
        "train_en_X = pad_sequences(train_en_X, maxlen=en_max_len,\n",
        "                           padding=\"post\", value=empty_token_id_en)\n",
        "# maxlen is fr_max_len+1 since we need to accomodate for <start>\n",
        "train_fr_X = pad_sequences(train_fr_X, maxlen=fr_max_len+1,\n",
        "                           padding=\"post\", value=empty_token_id_fr)\n",
        "\n",
        "# Converting the train_fr_Y to a one-hot vector needed by the training phase as\n",
        "# the output\n",
        "train_fr_Y = to_categorical(train_fr_X, num_classes=vocab_size_fr)\n",
        "\n",
        "# Moving the last <empty> to the first position in each input sample\n",
        "train_fr_X = np.roll(train_fr_X, 1, axis=-1)\n",
        "# Changing the first token in each input sample to <start>\n",
        "train_fr_X[:, 0] = start_token_id_fr\n",
        "\n",
        "fr_max_len = fr_max_len + 1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Lk_zw84T__T"
      },
      "source": [
        "# 6. Cutom metrics\n",
        "\n",
        "These are two custom metrics that I think represent accuracy of a translation model better.\n",
        "\n",
        "First, there's `masked_categorical_accuracy` which acts just like `categorical_accuracy` but with a mask. The reason this is a better measure of the accuracy compared to unmasked version is that, in unmasked version we are getting an accuracy even for learning the `<empty>` tokens at the end of the padded sequences. Of course, it is rather easy to learn them since they are all the same single token and they will be pruned off when mapped back to text form. This accuracy measure excludes learning those from the reported accuracy.\n",
        "\n",
        "Second, we have `exact_matched_accuracy`. In this accuracy we are counting a sample learned only if all the tokens in that sample are learned without a miss. So basically, the reported percentage is actually ratio of the sentences learned completely (not the individual tokens)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BN2Y982ETxgD"
      },
      "source": [
        "import tensorflow.keras.backend as K\n",
        "from tensorflow.python.keras.metrics import MeanMetricWrapper\n",
        "\n",
        "class MaskedCategoricalAccuracy(MeanMetricWrapper):\n",
        "\n",
        "    def __init__(self, mask_id, name='masked_categorical_accuracy', dtype=None):\n",
        "        super(MaskedCategoricalAccuracy, self).__init__(\n",
        "            masked_categorical_accuracy, name, dtype=dtype, mask_id=mask_id)\n",
        "\n",
        "\n",
        "def masked_categorical_accuracy(y_true, y_pred, mask_id):\n",
        "    true_ids = K.argmax(y_true, axis=-1)\n",
        "    pred_ids = K.argmax(y_pred, axis=-1)\n",
        "    maskBool = K.not_equal(true_ids, mask_id)\n",
        "    maskInt64 = K.cast(maskBool, 'int64')\n",
        "    maskFloatX = K.cast(maskBool, K.floatx())\n",
        "\n",
        "    count = K.sum(maskFloatX)\n",
        "    equals = K.equal(true_ids * maskInt64,\n",
        "                     pred_ids * maskInt64)\n",
        "    sum = K.sum(K.cast(equals, K.floatx()) * maskFloatX)\n",
        "    return sum / count\n",
        "\n",
        "\n",
        "class ExactMatchedAccuracy(MeanMetricWrapper):\n",
        "\n",
        "    def __init__(self, mask_id, name='exact_matched_accuracy', dtype=None):\n",
        "        super(ExactMatchedAccuracy, self).__init__(\n",
        "            exact_matched_accuracy, name, dtype=dtype, mask_id=mask_id)\n",
        "\n",
        "\n",
        "def exact_matched_accuracy(y_true, y_pred, mask_id):\n",
        "    true_ids = K.argmax(y_true, axis=-1)\n",
        "    pred_ids = K.argmax(y_pred, axis=-1)\n",
        "\n",
        "    maskBool = K.not_equal(true_ids, mask_id)\n",
        "    maskInt64 = K.cast(maskBool, 'int64')\n",
        "\n",
        "    diff = (true_ids - pred_ids) * maskInt64\n",
        "    matches = K.cast(K.not_equal(diff, K.zeros_like(diff)), 'int64')\n",
        "    matches = K.sum(matches, axis=-1)\n",
        "    matches = K.cast(K.equal(matches, K.zeros_like(matches)), K.floatx())\n",
        "\n",
        "    return K.mean(matches)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XETESwmqLuHh"
      },
      "source": [
        "# 7. Defining the models\n",
        "\n",
        "There are three models to define, the trainging model, the encoder model, and the decoder model. The latter two are used after the training phase for the text generation.\n",
        "\n",
        "If you want to load the models from the disk, you need to remember that these models are all using the same layers and weights. So it's not that straight forward to load them completely. The easiest way to do so is to define the models as if you are doing so for the first time and then load the weights for the training model (load just the weights, not the model). Since the training model holds the weight for all the layers, by doing so you are loading the weights for encoder and decoder as well."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OuWOeFLbLtPF"
      },
      "source": [
        "from tensorflow.keras import Input, layers, models\n",
        "from layers.attention import AttentionLayer\n",
        "\n",
        "hidden_dim = 128\n",
        "\n",
        "# Encoder input (English)\n",
        "input_en = Input(batch_shape=(None, en_max_len), name='input_en')\n",
        "\n",
        "# English embedding layer\n",
        "embedding_en = layers.Embedding(vocab_size_en, hidden_dim, name='embedding_en')\n",
        "embedded_en = embedding_en(input_en)\n",
        "\n",
        "# Encoder RNN (LSTM) layer\n",
        "encoder_lstm = layers.Bidirectional(\n",
        "                  layers.LSTM(hidden_dim,\n",
        "                              return_sequences=True, return_state=True),\n",
        "                  name=\"encoder_lstm\")\n",
        "(encoded_en,\n",
        "  forward_h_en, forward_c_en,\n",
        "  backward_h_en, backward_c_en) = encoder_lstm(embedded_en)\n",
        "\n",
        "# Decoder input (French)\n",
        "input_fr = Input(batch_shape=(None, None), name='input_fr')\n",
        "\n",
        "# English embedding layer\n",
        "embedding_fr = layers.Embedding(vocab_size_fr, hidden_dim, name='embedding_fr')\n",
        "embedded_fr = embedding_fr(input_fr)\n",
        "\n",
        "state_h_en = layers.concatenate([forward_h_en, backward_h_en])\n",
        "state_c_en = layers.concatenate([forward_c_en, backward_c_en])\n",
        "\n",
        "# Decoder RNN (LSTM) layer\n",
        "decoder_lstm = layers.LSTM(hidden_dim * 2, return_sequences=True,\n",
        "                           return_state=True, name=\"decoder_lstm\")\n",
        "(encoded_fr,\n",
        "  forward_h_fr, forward_c_fr) = decoder_lstm(embedded_fr,\n",
        "                 initial_state=[state_h_en, state_c_en])\n",
        "\n",
        "# Attention layer\n",
        "attention_layer = AttentionLayer(name='attention_layer')\n",
        "attention_out, attention_states = attention_layer({\"values\": encoded_en,\n",
        "                                                   \"query\": encoded_fr})\n",
        "\n",
        "# Concatenating the decoder output with attention output\n",
        "rnn_output = layers.concatenate([encoded_fr, attention_out], name=\"rnn_output\")\n",
        "\n",
        "# Dense layer\n",
        "dense_layer0 = layers.Dense(2048, activation='relu', name='dense_0')\n",
        "dl0 = dense_layer0(rnn_output)\n",
        "\n",
        "dense_layer1 = layers.Dense(1024, activation='relu', name='dense_1')\n",
        "dl1 = dense_layer1(dl0)\n",
        "\n",
        "dense_layer2 = layers.Dense(512, activation='relu', name='dense_2')\n",
        "dl2 = dense_layer2(dl1)\n",
        "\n",
        "dl2 = layers.Dropout(0.4)(dl2)\n",
        "\n",
        "dense_layer3 = layers.Dense(vocab_size_fr, activation='softmax', name='dense_3')\n",
        "dense_output = dense_layer3(dl2)\n",
        "\n",
        "training_model = models.Model([input_en, input_fr], dense_output)\n",
        "training_model.summary()\n",
        "\n",
        "training_model.compile(optimizer='adam',\n",
        "                       loss='categorical_crossentropy',\n",
        "                       metrics=[MaskedCategoricalAccuracy(empty_token_id_fr),\n",
        "                                ExactMatchedAccuracy(empty_token_id_fr)])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5bXLJwTdvNxV"
      },
      "source": [
        "Now, the generative models (the encoder and the decoder).\n",
        "\n",
        "It is worth mentioning that `attention_state` is made part of the output for the decoder only to be able to extract the attention scores to plot them. If you do not want to plot the attention scores, you can exclude them from the output."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mPcFj_iOvTG9"
      },
      "source": [
        "# The encoder model that encodes English input into encoded output and states\n",
        "encoder_model = models.Model([input_en],\n",
        "                             [encoded_en,\n",
        "                              state_h_en, state_c_en])\n",
        "encoder_model.summary()\n",
        "\n",
        "\n",
        "# The decoder model, to generate the French tokens (in integer form)\n",
        "input_h = layers.Input(batch_shape=(None, hidden_dim * 2),\n",
        "                       name='input_h')\n",
        "input_c = layers.Input(batch_shape=(None, hidden_dim * 2),\n",
        "                       name='input_c')\n",
        "\n",
        "(decoder_output,\n",
        "  output_h,\n",
        "  output_c) = decoder_lstm(embedded_fr,\n",
        "                           initial_state=[input_h, input_c])\n",
        "\n",
        "input_encoded_en = layers.Input(batch_shape=(None, en_max_len, hidden_dim * 2),\n",
        "                                name='input_encoded_en')\n",
        "\n",
        "attention_out, attention_state = attention_layer({\"values\": input_encoded_en,\n",
        "                                                  \"query\": decoder_output})\n",
        "\n",
        "generative_output = layers.concatenate([decoder_output,\n",
        "                                        attention_out],\n",
        "                                       name=\"generative_output\")\n",
        "\n",
        "g0 = dense_layer0(generative_output)\n",
        "g1 = dense_layer1(g0)\n",
        "g2 = dense_layer2(g1)\n",
        "dense_output = dense_layer3(g2)\n",
        "\n",
        "decoder_model = models.Model([input_encoded_en, input_fr,\n",
        "                              input_h, input_c],\n",
        "                             [dense_output, attention_state,\n",
        "                              output_h, output_c])\n",
        "decoder_model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9gQvJw_8rb_m"
      },
      "source": [
        "# 8. Traning the model / loading the weights\n",
        "\n",
        "If you want you can train your model. But at the same time, to save your time, I've included the trained weights for this model that you can simply load. If you decided to train the model yourself, based on my experience, 170 epochs are enough. Also a disclaimer, each epoch took around 130 seconds to complete on a GPU. It takes a lot longer on a CPU or even a TPU.\n",
        "\n",
        "**[TRAINING]**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vjug7prpM66r"
      },
      "source": [
        "from tensorflow.keras.callbacks import EarlyStopping\n",
        "\n",
        "pocket = EarlyStopping(monitor='val_exact_matched_accuracy', min_delta=0.001,\n",
        "                       patience=10, verbose=1, mode='max',\n",
        "                       restore_best_weights = True)\n",
        "\n",
        "history = training_model.fit(x=[train_en_X, train_fr_X], y=train_fr_Y, batch_size=786,\n",
        "                             epochs=200, verbose=1, validation_split=0.3, shuffle=True,\n",
        "                             workers=3, use_multiprocessing=True, callbacks=[pocket])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MmJ0JzqoxDw1"
      },
      "source": [
        "Saving the model weights to disk.\n",
        "\n",
        "**[TRAINING]**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0HoBbUrGOixv"
      },
      "source": [
        "training_model.save_weights(os.path.join(base_dir, \"data\", \"lstm_weights.h5\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h8ocqvoNxDw2"
      },
      "source": [
        "Downloads the saved model weights. Running this cell too fast will lead to an error. After running the previous cell, give it some time and then run this one. In any case, if it failed, just give it another try.\n",
        "\n",
        "**[TRAINING]**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "obsYci3sxDw2"
      },
      "source": [
        "from google.colab import files\n",
        "\n",
        "# This is code to download the model weights into your computer\n",
        "files.download(os.path.join(base_dir, \"data\", \"lstm_weights.h5\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YsxxiambxDw2"
      },
      "source": [
        "This block will plot the history for the loss value and the two accuracy metrics over the course of training for the trainging set and the validation set. You can run it only if you trained the model yourself.\n",
        "\n",
        "**[TRAINING]**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jMSp1LUZxDw2"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "from matplotlib.pyplot import figure\n",
        "\n",
        "figure(num=None, figsize=(11, 7))\n",
        "\n",
        "# Plot training & validation loss values\n",
        "plt.plot(history.history['loss'])\n",
        "plt.plot(history.history['val_loss'])\n",
        "plt.title('Model loss')\n",
        "plt.ylabel('Loss')\n",
        "plt.xlabel('Epoch')\n",
        "plt.legend(['Train', 'Validation'], loc='upper right')\n",
        "plt.show()\n",
        "\n",
        "figure(num=None, figsize=(11, 7))\n",
        "\n",
        "# Plot training & validation masked_categorical_accuracy values\n",
        "plt.plot(history.history['masked_categorical_accuracy'])\n",
        "plt.plot(history.history['val_masked_categorical_accuracy'])\n",
        "plt.title('Model accuracy')\n",
        "plt.ylabel('Accuracy')\n",
        "plt.xlabel('Epoch')\n",
        "plt.legend(['Train', 'Validation'], loc='lower right')\n",
        "plt.show()\n",
        "\n",
        "figure(num=None, figsize=(11, 7))\n",
        "\n",
        "# Plot training & validation exact_matched_accuracy values\n",
        "plt.plot(history.history['exact_matched_accuracy'])\n",
        "plt.plot(history.history['val_exact_matched_accuracy'])\n",
        "plt.title('Model exact match accuracy')\n",
        "plt.ylabel('Accuracy')\n",
        "plt.xlabel('Epoch')\n",
        "plt.legend(['Train', 'Validation'], loc='lower right')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ekKcaXi3sXQj"
      },
      "source": [
        "This block loads the weights from the repo. **Run it if you decided to load the weights instead of training it yourself**. But in case you've made a mistake and ran it after you've trained your model, don't worry. It will load the weights that you've just saved."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I8uoDTqlsiga"
      },
      "source": [
        "training_model.load_weights(\n",
        "    os.path.join(base_dir, \"data\", \"lstm_weights.h5\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oFc3u8igtDWH"
      },
      "source": [
        "# 9. Evaluate the model using the whole dataset\n",
        "\n",
        "In this block, we are going to evaluate the model with the whole dataset. This is specially good if you decided to load the model and not train it so you can see it's accuracy youself."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ru7Dme2PtggN"
      },
      "source": [
        "results = training_model.evaluate(x=[train_en_X, train_fr_X], y=train_fr_Y,\n",
        "                                  batch_size=786, verbose=1,\n",
        "                                  workers=1, use_multiprocessing=False)\n",
        "\n",
        "print('Test loss:', results[0])\n",
        "print('Test masked categorical accuracy:', results[1])\n",
        "print('Test exact matched accuracy:', results[2])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IQiRndk_thd8"
      },
      "source": [
        "# 10. Testing the model with your input and plotting the alignment matrix\n",
        "\n",
        "In this block, you can type in your string in English to be translated to French. At the end, as a bonus, you'll see how the model's attention layer has mapped the words from English to French (also known as the alignment matrix).\n",
        "\n",
        "Just a disclaimer, this model is trained using 170K samples. Do not expect much from it! The provide English sentence is chosen from the dataset, so it should be translted correctly. But your custom ones might not result in a good translation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wdllQOhjVGK8"
      },
      "source": [
        "import numpy as np\n",
        "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
        "import math\n",
        "import sys\n",
        "\n",
        "# The input English string\n",
        "english_string = \"the united states is never freezing during november , but the united states is sometimes rainy in winter .\"\n",
        "\n",
        "# First, let's tokenize the Eglish string, then pad it\n",
        "english_tokens = sp_en.EncodeAsIds(english_string.strip()) + [end_token_id_en]\n",
        "english_tokens = pad_sequences([english_tokens], maxlen=en_max_len,\n",
        "                               padding=\"post\", value=empty_token_id_en)\n",
        "\n",
        "# The encoder, we only need to use it once per each English string\n",
        "(encoded_en_test,\n",
        "  state_h_en_test, state_c_en_test) = encoder_model.predict(english_tokens)\n",
        "\n",
        "# In order to find a better translation, we are using Beam search\n",
        "beam_search_list = [{\n",
        "  \"decoder_input\": {\n",
        "    \"input_encoded_en\": encoded_en_test,\n",
        "    \"input_fr\": np.array([[start_token_id_fr]]),\n",
        "    \"input_h\": state_h_en_test,\n",
        "    \"input_c\": state_c_en_test\n",
        "  },\n",
        "  \"score\": 0.0,\n",
        "  \"parent_node\": None,\n",
        "  \"depth\": 0,\n",
        "  \"attention_weights\": None,\n",
        "}]\n",
        "ended_branches = []\n",
        "\n",
        "beam_size = 10\n",
        "\n",
        "# We are generating up to fr_max_len tokens\n",
        "for i in range(fr_max_len):\n",
        "  new_beam_candidates = []\n",
        "  # Predict the next token for each member of the list\n",
        "  for beam in beam_search_list:\n",
        "    # Use the decoder to predict the next token using the previously\n",
        "    # predicted token\n",
        "    (output,\n",
        "      attention_out,\n",
        "      state_h_en_test,\n",
        "      state_c_en_test) = decoder_model.predict(beam[\"decoder_input\"])\n",
        "    # Find the top beam_size candidates\n",
        "    top_k = np.argpartition(output[0, 0, :], -beam_size)[-beam_size:]\n",
        "    # For each candidate, put it in the list to predict the next token for it\n",
        "    for k in top_k:\n",
        "      if output[0, 0, k].item() > 0.0:\n",
        "        log_k = math.log(output[0, 0, k].item())\n",
        "      else:\n",
        "        log_k = -sys.float_info.max\n",
        "\n",
        "      if k == end_token_id_fr:\n",
        "        ended_branches.append({\n",
        "          \"decoder_input\": {\n",
        "            \"input_encoded_en\": encoded_en_test,\n",
        "            \"input_fr\": np.array([[k]]),\n",
        "            \"input_h\": state_h_en_test,\n",
        "            \"input_c\": state_c_en_test,\n",
        "          },\n",
        "          \"score\": beam[\"score\"] + log_k,\n",
        "          \"parent_node\": beam,\n",
        "          \"depth\": beam[\"depth\"] + 1,\n",
        "          \"attention_weights\": attention_out,\n",
        "        })\n",
        "      else:\n",
        "        new_beam_candidates.append({\n",
        "          \"decoder_input\": {\n",
        "            \"input_encoded_en\": encoded_en_test,\n",
        "            \"input_fr\": np.array([[k]]),\n",
        "            \"input_h\": state_h_en_test,\n",
        "            \"input_c\": state_c_en_test,\n",
        "          },\n",
        "          \"score\": beam[\"score\"] + log_k,\n",
        "          \"parent_node\": beam,\n",
        "          \"depth\": beam[\"depth\"] + 1,\n",
        "          \"attention_weights\": attention_out,\n",
        "        })\n",
        "\n",
        "  # Keeping only the top beam_size in the list\n",
        "  beam_search_list = sorted(new_beam_candidates,\n",
        "                            key=lambda b: b[\"score\"],\n",
        "                            reverse=True)[0:beam_size]\n",
        "\n",
        "# Now that we are done with our beam search, let's take the best score and\n",
        "# detokenize it\n",
        "beam_node = sorted(beam_search_list + ended_branches,\n",
        "                   key=lambda b: b[\"score\"] / b[\"depth\"],\n",
        "                   reverse=True)[0]\n",
        "\n",
        "# Trace the best beam back to the parent node\n",
        "all_french_tokens = []\n",
        "attention_weights = []\n",
        "while beam_node[\"parent_node\"] is not None:\n",
        "    all_french_tokens.append(\n",
        "        beam_node[\"decoder_input\"][\"input_fr\"][0, 0].item())\n",
        "    attention_weights.append(beam_node[\"attention_weights\"])\n",
        "    beam_node = beam_node[\"parent_node\"]\n",
        "\n",
        "# We traced from tail to head, so we need to reserve the order to have it the right way\n",
        "all_french_tokens.reverse()\n",
        "attention_weights.reverse()\n",
        "\n",
        "# If there's any token out of the vocab, exclude it. This includes `<end>`,\n",
        "# `<empty>`, and <start> tokens\n",
        "french_tokens = [t for t in all_french_tokens if t < sp_fr.get_piece_size()]\n",
        "\n",
        "# Voila!\n",
        "french_string = sp_fr.DecodeIds(french_tokens)\n",
        "\n",
        "print(\"The input English string: \", english_string)\n",
        "print(\"The output French string: \", french_string)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TjpjnwDXMuWu"
      },
      "source": [
        "Plotting the alignment matrix"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P_P946p8G0hK"
      },
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Plotting the alignment matrix\n",
        "\n",
        "attention_mat = []\n",
        "for attn in attention_weights:\n",
        "  attention_mat.append(attn.reshape(-1))\n",
        "\n",
        "# We want to have the English tokens on the left axis, so we need to\n",
        "# trasponse the matrix over the diagonal running from upper right to lower left\n",
        "attention_mat = np.flipud(np.transpose(np.flipud(attention_mat)))\n",
        "\n",
        "fig, ax = plt.subplots(figsize=(16, 16))\n",
        "ax.imshow(attention_mat)\n",
        "\n",
        "ax.set_xticks(np.arange(attention_mat.shape[1]))\n",
        "ax.set_yticks(np.arange(attention_mat.shape[0]))\n",
        "\n",
        "def map_en_special_tokens(t):\n",
        "    switcher = {}\n",
        "    switcher[end_token_id_en] = \"<end>\"\n",
        "    switcher[empty_token_id_en] = \"<empty>\"\n",
        "    return switcher.get(t, \"<unknown>\")\n",
        "\n",
        "def map_fr_special_tokens(t):\n",
        "    switcher = {}\n",
        "    switcher[end_token_id_fr] = \"<end>\"\n",
        "    switcher[empty_token_id_fr] = \"<empty>\"\n",
        "    switcher[start_token_id_fr] = \"<start>\"\n",
        "    return switcher.get(t, \"<unknown>\")\n",
        "\n",
        "ax.set_xticklabels([sp_fr.IdToPiece(t)\n",
        "                    if t < sp_fr.get_piece_size() else map_fr_special_tokens(t)\n",
        "                    for t in all_french_tokens])\n",
        "ax.set_yticklabels([sp_en.IdToPiece(t)\n",
        "                    if t < sp_en.get_piece_size() else map_en_special_tokens(t)\n",
        "                    for t in english_tokens[0].tolist()])\n",
        "\n",
        "ax.tick_params(labelsize=12)\n",
        "ax.tick_params(axis='x', labelrotation=90)\n",
        "\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}