Neural_Machine_Translation_tf.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Neural-Machine-Translation-tf.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyO/wegNYkR/vjwwS5KBy0Zn"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "#### Mounting Google Drive"
      ],
      "metadata": {
        "id": "0dqNGaFLKKAJ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bVenoImvJZs9"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "cd \"/content/drive/MyDrive/IASNLP\""
      ],
      "metadata": {
        "id": "iNpRDc1-KMj8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Importing Necessary Libraries"
      ],
      "metadata": {
        "id": "OlMuBUdIApMX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install sentencepiece\n",
        "!pip install sacrebleu\n",
        "!pip install tensorflow-gpu # ---> For GPU"
      ],
      "metadata": {
        "id": "Iwi6TDKsB0St"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow.keras import Model\n",
        "from tensorflow.keras import Sequential\n",
        "from tensorflow.keras.layers import Dense, Lambda, Layer, Embedding, LayerNormalization\n",
        "import sentencepiece as spm\n",
        "from sacrebleu.metrics import BLEU, CHRF, TER\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from collections import Counter"
      ],
      "metadata": {
        "id": "DXSYYy2b-jMM"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Loading Data"
      ],
      "metadata": {
        "id": "ToLMSNAHBnvJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We load the `train_data`, `train_dev_data`, `test_val_data` and `test_data` as well as the byte-pair encoder tokenizer for English and Bengali i.e. `sp_en_bpe` and `s_ben_bpe`."
      ],
      "metadata": {
        "id": "zXIf-FL7IG7f"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_data, train_dev_data = pd.read_csv(\"train_data.csv\")[['src', 'tgt']], pd.read_csv(\"train_dev.csv\")[['src', 'tgt']] \n",
        "test_val_data, test_data = pd.read_csv(\"test_val.csv\")[['src', 'tgt']], pd.read_csv(\"test_data.csv\")[['src', 'tgt']]\n",
        "sp_en_bpe, sp_ben_bpe = spm.SentencePieceProcessor(), spm.SentencePieceProcessor()\n",
        "sp_en_bpe.load('eng_bpe.model'); sp_ben_bpe.load('ben_bpe.model');"
      ],
      "metadata": {
        "id": "IXsCg4mBBeoo"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Data Preprocessing"
      ],
      "metadata": {
        "id": "auf4TlPFhNtn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We decide on some formats of the data and batches. We take the maximum tokens we can have per tokenized English and Bengali sentence to be `MAX_TOKENS = 60`. We use the vocabulary size of `32000`(i.e. `ENCODER_VOCAB = DECODER_VOCAB = 32000`). A batch size of `256` was choosen for our Mini-Batch Gradient Descent which is used to train our Model."
      ],
      "metadata": {
        "id": "3z7BUnn_G3WL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "MAX_TOKENS = 60\n",
        "VOCAB_SIZE = 32000\n",
        "BATCH_SIZE = 256\n",
        "BUFFER_SIZE = BATCH_SIZE*4"
      ],
      "metadata": {
        "id": "7jTuQSqOAybO"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following two functions help in tokenizing an detokenizing the data based on our sentencepiece byte-pair encoding model."
      ],
      "metadata": {
        "id": "NlEW7GndbCfw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def tokenize(sentence, sp_model, add_bos = True, add_eos = True):\n",
        "    bos = [sp_model.bos_id()] if add_bos else []\n",
        "    eos = [sp_model.eos_id()] if add_eos else []\n",
        "    inputs = bos + sp_model.encode_as_ids(sentence) + eos\n",
        "    return np.reshape(np.array(inputs), [1, -1])\n",
        "def detokenize(tokenized, sp_model, is_bos = True, is_eos = True):\n",
        "    integers = np.squeeze(tokenized).tolist()\n",
        "    if is_eos:\n",
        "        return sp_model.DecodeIdsWithCheck(integers[int(is_bos):integers.index(sp_model.eos_id())])\n",
        "    else:\n",
        "        if sp_model.pad_id() in tokenized:\n",
        "            return sp_model.DecodeIdsWithCheck(integers[int(is_bos):integers.index(sp_model.pad_id())])\n",
        "        else:\n",
        "            return sp_model.DecodeIdsWithCheck(integers[int(is_bos):])"
      ],
      "metadata": {
        "id": "VC-6hav2BrJJ"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We have tokenized the data."
      ],
      "metadata": {
        "id": "5-76mwBnbuzp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Tokenization\n",
        "src_train_data_enc = [np.squeeze(tokenize(train_data['src'].iloc[i], sp_en_bpe)) for i in range(train_data.shape[0])]\n",
        "tgt_train_data_enc = [np.squeeze(tokenize(train_data['tgt'].iloc[i], sp_ben_bpe)) for i in range(train_data.shape[0])]\n",
        "src_train_dev_data_enc = [np.squeeze(tokenize(train_dev_data['src'].iloc[i], sp_en_bpe)) for i in range(train_dev_data.shape[0])]\n",
        "tgt_train_dev_data_enc = [np.squeeze(tokenize(train_dev_data['tgt'].iloc[i], sp_ben_bpe)) for i in range(train_dev_data.shape[0])]\n",
        "src_test_val_data_enc = [np.squeeze(tokenize(test_val_data['src'].iloc[i], sp_en_bpe)) for i in range(test_val_data.shape[0])]\n",
        "tgt_test_val_data_enc = [np.squeeze(tokenize(test_val_data['tgt'].iloc[i], sp_ben_bpe)) for i in range(test_val_data.shape[0])]"
      ],
      "metadata": {
        "id": "nT0QpepuDoFF"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We have padded the data with post-truncation and post padding upto MAX_TOKENS number of tokens."
      ],
      "metadata": {
        "id": "wdyEjbm5hWXg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Padding\n",
        "# The reason why we add MAX_TOKENS+1 after-pad length for Bengali is that when we form the tf.Dataset later, we will shift the target sentence once and will truncate the last token once\n",
        "# Hence, it will get adjusted\n",
        "train_src = tf.keras.preprocessing.sequence.pad_sequences(src_train_data_enc, maxlen = MAX_TOKENS, padding='post', truncating='post')\n",
        "train_tgt = tf.keras.preprocessing.sequence.pad_sequences(tgt_train_data_enc, maxlen = MAX_TOKENS, padding='post', truncating='post')\n",
        "train_dev_src = tf.keras.preprocessing.sequence.pad_sequences(src_train_dev_data_enc, maxlen = MAX_TOKENS, padding='post', truncating='post')\n",
        "train_dev_tgt = tf.keras.preprocessing.sequence.pad_sequences(tgt_train_dev_data_enc, maxlen = MAX_TOKENS, padding='post', truncating='post')\n",
        "test_val_src = tf.keras.preprocessing.sequence.pad_sequences(src_test_val_data_enc, maxlen = MAX_TOKENS, padding='post', truncating='post')\n",
        "test_val_tgt = tf.keras.preprocessing.sequence.pad_sequences(tgt_test_val_data_enc, maxlen = MAX_TOKENS, padding='post', truncating='post')"
      ],
      "metadata": {
        "id": "NT63vC8PcK-2"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Below, we discard all those pairs where the tokenized source or target sentences are larger than MAX_TOKENS number of tokens."
      ],
      "metadata": {
        "id": "FvxtnfKniMRA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Removing all pairs whose either the target sentence or source setence exceeds MAX_TOKENS number of tokens\n",
        "same_idx = list(set([i for i in range(train_src.shape[0]) if train_src[i][-1] == 1 or train_src[i][-1] == 0]).intersection(set([i for i in range(train_tgt.shape[0]) if train_tgt[i][-1] == 1 or train_tgt[i][-1] == 0])))\n",
        "train_src = train_src[same_idx]\n",
        "train_tgt = train_tgt[same_idx]\n",
        "same_idx = list(set([i for i in range(train_dev_src.shape[0]) if train_dev_src[i][-1] == 1 or train_dev_src[i][-1] == 0]).intersection(set([i for i in range(train_dev_tgt.shape[0]) if train_dev_tgt[i][-1] == 1 or train_dev_tgt[i][-1] == 0])))\n",
        "train_dev_src = train_dev_src[same_idx]\n",
        "train_dev_tgt = train_dev_tgt[same_idx]\n",
        "same_idx = list(set([i for i in range(test_val_src.shape[0]) if test_val_src[i][-1] == 1 or test_val_src[i][-1] == 0]).intersection(set([i for i in range(test_val_tgt.shape[0]) if test_val_tgt[i][-1] == 1 or test_val_tgt[i][-1] == 0])))\n",
        "test_val_src = test_val_src[same_idx]\n",
        "test_val_tgt = test_val_tgt[same_idx]"
      ],
      "metadata": {
        "id": "L6-cV6zpcO5e"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_labels = np.zeros(train_tgt.shape)\n",
        "train_labels[:,0:train_tgt.shape[1] -1] = train_tgt[:,1:]\n",
        "train_dev_labels = np.zeros(train_dev_tgt.shape)\n",
        "train_dev_labels[:,0:train_dev_tgt.shape[1] -1] = train_dev_tgt[:,1:]\n",
        "test_val_labels = np.zeros(test_val_tgt.shape)\n",
        "test_val_labels[:,0:test_val_tgt.shape[1] -1] = test_val_tgt[:,1:]"
      ],
      "metadata": {
        "id": "JGMvPvIxOx-0"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"*** Preprocessed Data ***\")\n",
        "print(\"Train Data Shape(Source): \", train_src.shape)\n",
        "print(\"Train Data Shape(Target): \", train_tgt.shape)\n",
        "print(\"Train Labels Shape(Target): \", train_labels.shape)\n",
        "print(\"Train Dev Data Shape(Source): \", train_dev_src.shape)\n",
        "print(\"Train Dev Data Shape(Target): \", train_dev_tgt.shape)\n",
        "print(\"Train Dev Labels Shape(Target): \", train_dev_labels.shape)\n",
        "print(\"Test Val Data Shape(Source): \", test_val_src.shape)\n",
        "print(\"Test Val Data Shape(Target): \", test_val_tgt.shape)\n",
        "print(\"Test Labels Shape(Target): \", test_val_labels.shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rO5HD33XOGxf",
        "outputId": "1787ff08-b65a-4c72-d9d5-71ac18880a37"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "*** Preprocessed Data ***\n",
            "Train Data Shape(Source):  (154836, 60)\n",
            "Train Data Shape(Target):  (154836, 60)\n",
            "Train Labels Shape(Target):  (154836, 60)\n",
            "Train Dev Data Shape(Source):  (3649, 60)\n",
            "Train Dev Data Shape(Target):  (3649, 60)\n",
            "Train Dev Labels Shape(Target):  (3649, 60)\n",
            "Test Val Data Shape(Source):  (1194, 60)\n",
            "Test Val Data Shape(Target):  (1194, 60)\n",
            "Test Labels Shape(Target):  (1194, 60)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Below we prepare the data in the form ready to feed to the Model."
      ],
      "metadata": {
        "id": "xtsMxbAYzePz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_data_prep = tf.data.Dataset.from_tensor_slices((train_src, train_tgt, train_labels))\n",
        "train_dev_data_prep = tf.data.Dataset.from_tensor_slices((train_dev_src, train_dev_tgt, train_dev_labels))\n",
        "test_val_data_prep = tf.data.Dataset.from_tensor_slices((test_val_src, test_val_tgt, test_val_labels))"
      ],
      "metadata": {
        "id": "_OkR5tbtyBQj"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Below, we can see how our data looks"
      ],
      "metadata": {
        "id": "jSGBKDhc2sc-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for encoder_inputs, decoder_inputs, labels in train_data_prep.take(1):\n",
        "    print(f'encoder_inputs.shape: {encoder_inputs.shape}')\n",
        "    print(f'decoder_inputs.shape: {decoder_inputs.shape}')\n",
        "    print(f\"labels.shape: {labels.shape}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ao1TX6GHzWn7",
        "outputId": "384d55e1-d612-4346-c1dc-d6e6d12606e5"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "encoder_inputs.shape: (60,)\n",
            "decoder_inputs.shape: (60,)\n",
            "labels.shape: (60,)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Model"
      ],
      "metadata": {
        "id": "zjXk6pt-28kL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will use the Transformer Architechture with the following hyperparameters."
      ],
      "metadata": {
        "id": "a9UFHuFO3dHe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Transformer parameters\n",
        "d_model = 256 # 512 in the original paper\n",
        "d_k = 32 # 64 in the original paper\n",
        "d_v = 32 # 64 in the original paper\n",
        "n_heads = 4 # 8 in the original paper\n",
        "n_encoder_layers = 3 # 6 in the original paper\n",
        "n_decoder_layers = 3 # 6 in the original paper\n",
        "\n",
        "max_token_length = MAX_TOKENS # 512 in the original paper"
      ],
      "metadata": {
        "id": "v9ULcpgbJ9WP"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class SingleHeadAttention(Layer):\n",
        "  def __init__(self, input_shape=(3, -1, d_model), dropout=.0, masked=None):\n",
        "    super(SingleHeadAttention, self).__init__()\n",
        "    self.q = Dense(d_k, input_shape=(-1, d_model), kernel_initializer='glorot_uniform', \n",
        "                   bias_initializer='glorot_uniform')\n",
        "    self.normalize_q = Lambda(lambda x: x / np.sqrt(d_k))\n",
        "    self.k = Dense(d_k, input_shape=(-1, d_model), kernel_initializer='glorot_uniform', \n",
        "                   bias_initializer='glorot_uniform')\n",
        "    self.v = Dense(d_v, input_shape=(-1, d_model), kernel_initializer='glorot_uniform', \n",
        "                   bias_initializer='glorot_uniform')\n",
        "    self.dropout = dropout\n",
        "    self.masked = masked\n",
        "  \n",
        "  # Inputs: [query, key, value]\n",
        "  def call(self, inputs, training=None):\n",
        "    assert len(inputs) == 3\n",
        "    # We use a lambda layer to divide vector q by sqrt(d_k) according to the equation\n",
        "    q = self.normalize_q(self.q(inputs[0]))\n",
        "    k = self.k(inputs[1])\n",
        "    # The dimensionality of q is (batch_size, query_length, d_k) and that of k is (batch_size, key_length, d_k)\n",
        "    # So we will do a matrix multication by batch after transposing last 2 dimensions of k\n",
        "    # tf.shape(attn_weights) = (batch_size, query_length, key_length)\n",
        "    attn_weights = tf.matmul(q, tf.transpose(k, perm=[0,2,1]))\n",
        "    if self.masked: # Prevent future attentions in decoding self-attention\n",
        "      # Create a matrix where the strict upper triangle (not including main diagonal) is filled with -inf and 0 elsewhere\n",
        "      length = tf.shape(attn_weights)[-1]\n",
        "      #attn_mask = np.triu(tf.fill((length, length), -np.inf), k=1) # We need to use tensorflow functions instead of numpy\n",
        "      attn_mask = tf.fill((length, length), -np.inf)\n",
        "      attn_mask = tf.linalg.band_part(attn_mask, 0, -1) # Get upper triangle\n",
        "      attn_mask = tf.linalg.set_diag(attn_mask, tf.zeros((length))) # Set diagonal to zeros to avoid operations with infinity\n",
        "      # This matrix is added to the attention weights so all future attention will have -inf logits (0 after softmax)\n",
        "      attn_weights += attn_mask\n",
        "    # Softmax along the last dimension\n",
        "    attn_weights = tf.nn.softmax(attn_weights, axis=-1)\n",
        "    if training: # Attention dropout included in the original paper. This is possibly to encourage multihead diversity.\n",
        "      attn_weights = tf.nn.dropout(attn_weights, rate=self.dropout)\n",
        "    v = self.v(inputs[2])\n",
        "    return tf.matmul(attn_weights, v)\n",
        "\n",
        "class MultiHeadAttention_my(Layer):\n",
        "  def __init__(self, dropout=.0, masked=None):\n",
        "    super(MultiHeadAttention_my, self).__init__()\n",
        "    self.attn_heads = list()\n",
        "    for i in range(n_heads): \n",
        "      self.attn_heads.append(SingleHeadAttention(dropout=dropout, masked=masked))\n",
        "    self.linear = Dense(d_model, input_shape=(-1, n_heads * d_v), kernel_initializer='glorot_uniform', \n",
        "                   bias_initializer='glorot_uniform')\n",
        "    \n",
        "  def call(self, x, training=None):\n",
        "    attentions = [self.attn_heads[i](x, training=training) for i in range(n_heads)]\n",
        "    concatenated_attentions = tf.concat(attentions, axis=-1)\n",
        "    return self.linear(concatenated_attentions)\n",
        "\n",
        "class TransformerEncoder(Layer):\n",
        "  def __init__(self, dropout=.1, attention_dropout=.0, **kwargs):\n",
        "    super(TransformerEncoder, self).__init__(**kwargs)\n",
        "    self.dropout_rate = dropout\n",
        "    self.attention_dropout_rate = attention_dropout\n",
        "  def build(self, input_shape):\n",
        "    self.multihead_attention = MultiHeadAttention_my(dropout=self.attention_dropout_rate)\n",
        "    self.dropout1 = tf.keras.layers.Dropout(self.dropout_rate)\n",
        "    self.layer_normalization1 = LayerNormalization(input_shape=input_shape, epsilon=1e-6)\n",
        "    \n",
        "    self.linear1 = Dense(input_shape[-1] * 4, input_shape=input_shape, activation='relu',\n",
        "                        kernel_initializer='glorot_uniform', bias_initializer='glorot_uniform')\n",
        "    self.linear2 = Dense(input_shape[-1], input_shape=self.linear1.compute_output_shape(input_shape),\n",
        "                        kernel_initializer='glorot_uniform', bias_initializer='glorot_uniform')\n",
        "    self.dropout2 = tf.keras.layers.Dropout(self.dropout_rate)\n",
        "    self.layer_normalization2 = LayerNormalization(input_shape=input_shape, epsilon=1e-6)\n",
        "    super(TransformerEncoder, self).build(input_shape)\n",
        "  def call(self, x, training=None):\n",
        "    sublayer1 = self.multihead_attention((x, x, x), training=training)\n",
        "    sublayer1 = self.dropout1(sublayer1, training=training)\n",
        "    layernorm1 = self.layer_normalization1(x + sublayer1)\n",
        "    \n",
        "    sublayer2 = self.linear2(self.linear1(layernorm1))\n",
        "    sublayer1 = self.dropout2(sublayer2, training=training)\n",
        "    layernorm2 = self.layer_normalization2(layernorm1 + sublayer2)\n",
        "    return layernorm2\n",
        "  def compute_output_shape(self, input_shape):\n",
        "    return input_shape\n",
        "\n",
        "class TransformerDecoder(Layer):\n",
        "  def __init__(self, dropout=.0, attention_dropout=.0, **kwargs):\n",
        "    super(TransformerDecoder, self).__init__(**kwargs)\n",
        "    self.dropout_rate = dropout\n",
        "    self.attention_dropout_rate = attention_dropout\n",
        "  def build(self, input_shape):\n",
        "    self.multihead_self_attention = MultiHeadAttention_my(dropout=self.attention_dropout_rate, masked=True)\n",
        "    self.dropout1 = tf.keras.layers.Dropout(self.dropout_rate)\n",
        "    self.layer_normalization1 = LayerNormalization(input_shape=input_shape, epsilon=1e-6)\n",
        "    \n",
        "    self.multihead_encoder_attention = MultiHeadAttention_my(dropout=self.attention_dropout_rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(self.dropout_rate)\n",
        "    self.layer_normalization2 = LayerNormalization(input_shape=input_shape, epsilon=1e-6)\n",
        "    \n",
        "    self.linear1 = Dense(input_shape[-1] * 4, input_shape=input_shape, activation='relu',\n",
        "                        kernel_initializer='glorot_uniform', bias_initializer='glorot_uniform')\n",
        "    self.linear2 = Dense(input_shape[-1], input_shape=self.linear1.compute_output_shape(input_shape),\n",
        "                        kernel_initializer='glorot_uniform', bias_initializer='glorot_uniform')\n",
        "    self.dropout3 = tf.keras.layers.Dropout(self.dropout_rate)\n",
        "    self.layer_normalization3 = LayerNormalization(input_shape=input_shape, epsilon=1e-6)\n",
        "    super(TransformerDecoder, self).build(input_shape)\n",
        "  def call(self, x, hidden, training=None):\n",
        "    sublayer1 = self.multihead_self_attention((x, x, x))\n",
        "    sublayer1 = self.dropout1(sublayer1, training=training)\n",
        "    layernorm1 = self.layer_normalization1(x + sublayer1)\n",
        "    \n",
        "    sublayer2 = self.multihead_encoder_attention((x, hidden, hidden))\n",
        "    sublayer2 = self.dropout2(sublayer2, training=training)\n",
        "    layernorm2 = self.layer_normalization2(layernorm1 + sublayer2)\n",
        "    \n",
        "    sublayer3 = self.linear2(self.linear1(layernorm1))\n",
        "    sublayer3 = self.dropout3(sublayer3, training=training)\n",
        "    layernorm3 = self.layer_normalization2(layernorm2 + sublayer3)\n",
        "    return layernorm3\n",
        "  def compute_output_shape(self, input_shape):\n",
        "    return input_shape\n",
        "\n",
        "class SinusoidalPositionalEncoding(Layer): # This is a TensorFlow implementation of https://github.com/graykode/nlp-tutorial/blob/master/5-1.Transformer/Transformer_Torch.ipynb\n",
        "  def __init__(self):\n",
        "    super(SinusoidalPositionalEncoding, self).__init__()\n",
        "    self.sinusoidal_encoding = np.array([self.get_positional_angle(pos) for pos in range(max_token_length)], dtype=np.float32)\n",
        "    self.sinusoidal_encoding[:, 0::2] = np.sin(self.sinusoidal_encoding[:, 0::2])\n",
        "    self.sinusoidal_encoding[:, 1::2] = np.cos(self.sinusoidal_encoding[:, 1::2])\n",
        "    self.sinusoidal_encoding = tf.cast(self.sinusoidal_encoding, dtype=tf.float32) # Casting the array to Tensor for slicing\n",
        "  def call(self, x):\n",
        "    return x + self.sinusoidal_encoding[:tf.shape(x)[1]]\n",
        "  def compute_output_shape(self, input_shape):\n",
        "    return input_shape\n",
        "  def get_angle(self, pos, dim):\n",
        "    return pos / np.power(10000, 2 * (dim // 2) / d_model)\n",
        "  def get_positional_angle(self, pos):\n",
        "    return [self.get_angle(pos, dim) for dim in range(d_model)]\n",
        "\n",
        "class Transformer(Model):\n",
        "  def __init__(self, dropout=.1, attention_dropout=.0, **kwargs):\n",
        "    super(Transformer, self).__init__(**kwargs)\n",
        "    self.encoding_embedding = Embedding(VOCAB_SIZE, d_model)\n",
        "    self.decoding_embedding = Embedding(VOCAB_SIZE, d_model)\n",
        "    self.pos_encoding = SinusoidalPositionalEncoding()\n",
        "    self.encoder = [TransformerEncoder(dropout=dropout, attention_dropout=attention_dropout) for i in range(n_encoder_layers)]\n",
        "    self.decoder = [TransformerDecoder(dropout=dropout, attention_dropout=attention_dropout) for i in range(n_decoder_layers)]\n",
        "    self.decoder_final = Dense(VOCAB_SIZE, input_shape=(None, d_model))\n",
        "  def call(self, inputs, training=None): # Source_sentence and decoder_input\n",
        "    source_sentence, decoder_input = inputs\n",
        "    embedded_source = self.encoding_embedding(source_sentence)\n",
        "    encoder_output = self.pos_encoding(embedded_source)\n",
        "    for encoder_unit in self.encoder:\n",
        "      encoder_output = encoder_unit(encoder_output, training=training)\n",
        "    \n",
        "    embedded_target = self.decoding_embedding(decoder_input)\n",
        "    decoder_output = self.pos_encoding(embedded_target)\n",
        "    for decoder_unit in self.decoder:\n",
        "      decoder_output = decoder_unit(decoder_output, encoder_output, training=training)\n",
        "    if training:\n",
        "      decoder_output = self.decoder_final(decoder_output)\n",
        "      decoder_output = tf.nn.softmax(decoder_output, axis=-1)\n",
        "    else:\n",
        "      decoder_output = self.decoder_final(decoder_output[:, -1:, :])\n",
        "      decoder_output = tf.nn.softmax(decoder_output, axis=-1)\n",
        "    return decoder_output"
      ],
      "metadata": {
        "id": "KvsIOzPnzt4F"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transformer = Transformer(dropout=0.3) # Instantiating a new transformer model\n",
        "train = [tf.cast(train_src, dtype=tf.float32), tf.cast(train_tgt, dtype=tf.float32)] # Cast the tuples to tensors\n",
        "validation = [tf.cast(train_dev_src, dtype=tf.float32), tf.cast(train_dev_tgt, dtype=tf.float32)]"
      ],
      "metadata": {
        "id": "eRsGRYTR6XQW"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transformer.build([train_src.shape, train_tgt.shape])\n",
        "transformer.summary()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IHXUpv-0eh-b",
        "outputId": "ccda5dc7-6c57-483c-d976-502815e0ea39"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"transformer\"\n",
            "_________________________________________________________________\n",
            " Layer (type)                Output Shape              Param #   \n",
            "=================================================================\n",
            " embedding (Embedding)       multiple                  8192000   \n",
            "                                                                 \n",
            " embedding_1 (Embedding)     multiple                  8192000   \n",
            "                                                                 \n",
            " sinusoidal_positional_encod  multiple                 0         \n",
            " ing (SinusoidalPositionalEn                                     \n",
            " coding)                                                         \n",
            "                                                                 \n",
            " transformer_encoder (Transf  multiple                 658304    \n",
            " ormerEncoder)                                                   \n",
            "                                                                 \n",
            " transformer_encoder_1 (Tran  multiple                 658304    \n",
            " sformerEncoder)                                                 \n",
            "                                                                 \n",
            " transformer_encoder_2 (Tran  multiple                 658304    \n",
            " sformerEncoder)                                                 \n",
            "                                                                 \n",
            " transformer_decoder (Transf  multiple                 790016    \n",
            " ormerDecoder)                                                   \n",
            "                                                                 \n",
            " transformer_decoder_1 (Tran  multiple                 790016    \n",
            " sformerDecoder)                                                 \n",
            "                                                                 \n",
            " transformer_decoder_2 (Tran  multiple                 790016    \n",
            " sformerDecoder)                                                 \n",
            "                                                                 \n",
            " dense (Dense)               multiple                  8224000   \n",
            "                                                                 \n",
            "=================================================================\n",
            "Total params: 28,952,960\n",
            "Trainable params: 28,952,960\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training"
      ],
      "metadata": {
        "id": "0JkChx5Ndp1s"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "transformer.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])"
      ],
      "metadata": {
        "id": "InXkWf-LawJi"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "EPOCHS = 10\n",
        "checkpoint_filepath = '/Models/checkpoint'\n",
        "model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n",
        "    filepath=checkpoint_filepath,\n",
        "    save_weights_only=True,\n",
        "    monitor='val_accuracy',\n",
        "    mode='max',\n",
        "    save_best_only=True)"
      ],
      "metadata": {
        "id": "LMttctIwaNxG"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "with tf.device('/device:GPU:0'):\n",
        "    history = transformer.fit(train, tf.cast(train_labels, dtype=tf.float32), batch_size=256, epochs=EPOCHS, callbacks=[model_checkpoint_callback],\n",
        "                    validation_data=(validation, tf.cast(train_dev_labels, dtype=tf.float32)), validation_batch_size=128)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bzdsT8kmBEMo",
        "outputId": "5e836bb2-0af4-45a4-a02d-bfe4b2636a07"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/10\n",
            "343/605 [================>.............] - ETA: 2:50 - loss: 2.1044 - accuracy: 0.7890"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "transformer.save(\"./Models/\")"
      ],
      "metadata": {
        "id": "6tmeafiFdZ_U"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "plt.plot(history.history['acc'])\n",
        "plt.plot(history.history['val_acc'])\n",
        "plt.title('model accuracy')\n",
        "plt.ylabel('accuracy')\n",
        "plt.xlabel('epoch')\n",
        "plt.legend(['train', 'val'], loc='upper left')\n",
        "plt.show()\n",
        "\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', 'val'], loc='upper left')\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "JbSCoj0cft6O"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dot_img_file = \"d_mod-\"str(d_model)+\"-d_k-\"str(d_k)+\"-d_v-\"str(d_v)+\"-n_heads-\"str(n_heads)+\"-enc_lays-\"str(n_encoder_layers)+\"-dec_lays-\"str(n_decoder_layers)+'-transformer.png'\n",
        "tf.keras.utils.plot_model(model, to_file=dot_img_file, show_shapes=True)"
      ],
      "metadata": {
        "id": "RSK1eufbeQPL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transformer = keras.models.load_model('./Model-2/')"
      ],
      "metadata": {
        "id": "Oly3kYeU3FnV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Testing"
      ],
      "metadata": {
        "id": "KJ_Gx7RVeNla"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def greedy_translate(model, source_sentence, target_sentence_start=sp_ben_bpe.bos_id()):\n",
        "  # Tokenizing and padding\n",
        "  source_seq = tokenize(source_sentence, sp_en_bpe)\n",
        "  source_seq = tf.keras.preprocessing.sequence.pad_sequences(source_seq, maxlen = MAX_TOKENS, padding='post', truncating='post')\n",
        "  predict_seq = [[target_sentence_start]]\n",
        "  \n",
        "  predict_sentence = [target_sentence_start] # Deep copy here to prevent updates on target_sentence_start\n",
        "  while predict_sentence[-1] != sp_ben_bpe.eos_id() and len(predict_sentence) < MAX_TOKENS:\n",
        "    predict_output = model([np.array(source_seq), np.array(predict_seq)], training=None)\n",
        "    predict_label = tf.argmax(predict_output, axis=-1) # Pick the label with highest softmax score\n",
        "    predict_seq = tf.concat([predict_seq, predict_label], axis=-1) # Updating the prediction sequence\n",
        "    predict_sentence.append(predict_label[0][0].numpy())\n",
        "  \n",
        "  if len(predict_sentence) == MAX_TOKENS:\n",
        "      return detokenize(predict_sentence, sp_ben_bpe, True, False)\n",
        "  return detokenize(predict_sentence, sp_ben_bpe)"
      ],
      "metadata": {
        "id": "zfg1AaKDnPQM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow_probability as tfp"
      ],
      "metadata": {
        "id": "pLrtBO__iPPl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sentence = \"I love you.\"\n",
        "print(\"English Sentence: \", sentence)\n",
        "print(\"Bengali Sentence: \", greedy_translate(transformer, sentence))"
      ],
      "metadata": {
        "id": "KtWeJl1JoapJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sentences = list(test_val_data['src'][:250])\n",
        "refs = [list(test_val_data['tgt'][:250])]\n",
        "sys = [greedy_translate(transformer, sentence) for sentence in sentences]"
      ],
      "metadata": {
        "id": "9XNckV-wBgiN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "bleu = BLEU()\n",
        "bleu.corpus_score(sys, refs)"
      ],
      "metadata": {
        "id": "otRhPJC1VKyj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chrf = CHRF()\n",
        "chrf.corpus_score(sys, refs)"
      ],
      "metadata": {
        "id": "mlhToSrvYNqd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ter = TER()\n",
        "ter.corpus_score(sys, refs)"
      ],
      "metadata": {
        "id": "DtIadLMJYe1w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def sampling_translate(model, source_sentence, temperature=0.6, target_sentence_start=sp_ben_bpe.bos_id()):\n",
        "  # Tokenizing and padding\n",
        "  source_seq = tokenize(source_sentence, sp_en_bpe)\n",
        "  source_seq = tf.keras.preprocessing.sequence.pad_sequences(source_seq, maxlen = MAX_TOKENS, padding='post', truncating='post')\n",
        "  predict_seq = [[target_sentence_start]]\n",
        "  logprob = 0\n",
        "  \n",
        "  predict_sentence = [target_sentence_start] # Deep copy here to prevent updates on target_sentence_start\n",
        "  while predict_sentence[-1] != sp_ben_bpe.eos_id() and len(predict_sentence) < MAX_TOKENS:\n",
        "    predict_output = model([np.array(source_seq), np.array(predict_seq)], training=None)\n",
        "    probs = predict_output[-1].numpy()[0]\n",
        "    log_probs = np.log(probs)\n",
        "    u = np.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)\n",
        "    g = -np.log(-np.log(u))\n",
        "    sample_idx = np.argmax(log_probs + g * temperature, axis=-1)\n",
        "    sample_prob = predict_output[-1].numpy()[0][sample_idx]\n",
        "    logprob += np.log(sample_prob)\n",
        "    predict_label = np.arange(predict_output.shape[-1])[sample_idx] # Sample the label softmax score\n",
        "    predict_seq = tf.concat([predict_seq, tf.cast([[predict_label]], dtype=tf.int32)], axis=-1) # Updating the prediction sequence\n",
        "    predict_sentence.append(predict_label)\n",
        "  \n",
        "  if len(predict_sentence) == MAX_TOKENS:\n",
        "      return detokenize(predict_sentence, sp_ben_bpe, True, False), logprob\n",
        "  return detokenize(predict_sentence, sp_ben_bpe), logprob"
      ],
      "metadata": {
        "id": "2dmVo_eQhE8v"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sentence = \"I love you.\"\n",
        "print(\"English Sentence: \", sentence)\n",
        "print(\"Bengali Sentence: \", sampling_translate(transformer, sentence))"
      ],
      "metadata": {
        "id": "CB5iHHlxh1cC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_samples(sentence, n_samples, model=None, temperature=0.6):\n",
        "    samples, log_probs = [], []\n",
        "    for _ in range(n_samples):\n",
        "        sample, logp = sampling_translate(model, sentence, temperature)\n",
        "        samples.append(sample)\n",
        "        log_probs.append(logp)\n",
        "    return samples, log_probs"
      ],
      "metadata": {
        "id": "xIxJ9HqTtxHY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def rouge1_similarity(system, reference):\n",
        "    sys_counter = Counter(system)\n",
        "    ref_counter = Counter(reference)\n",
        "    overlap = 0\n",
        "    for token in sys_counter:\n",
        "        token_count_sys = sys_counter[token]\n",
        "        token_count_ref = ref_counter[token]\n",
        "        overlap += min(token_count_ref, token_count_sys)\n",
        "    precision = overlap / sum(sys_counter.values())\n",
        "    recall = overlap / sum(ref_counter.values())\n",
        "    if precision + recall != 0:\n",
        "        rouge1_score = 2 * ((precision * recall)/(precision + recall))\n",
        "    else:\n",
        "        rouge1_score = 0     \n",
        "    return rouge1_score"
      ],
      "metadata": {
        "id": "uhIlVoVwt7YC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def weighted_avg_overlap(samples, log_probs):\n",
        "    scores = {}\n",
        "    for index_candidate, candidate in enumerate(samples):    \n",
        "        overlap, weight_sum = 0.0, 0.0\n",
        "        for index_sample, (sample, logp) in enumerate(zip(samples, log_probs)):           \n",
        "            if index_candidate == index_sample:\n",
        "                continue\n",
        "            sample_p = float(np.exp(logp))\n",
        "            weight_sum += sample_p\n",
        "            sample_overlap = rouge1_similarity(candidate, sample)\n",
        "            overlap += sample_p * sample_overlap\n",
        "        score = overlap / weight_sum\n",
        "        scores[index_candidate] = score\n",
        "    return scores"
      ],
      "metadata": {
        "id": "7KxiaRENt-C_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def mbr_decode(sentence, n_samples, model=None, temperature=0.6):\n",
        "    samples, log_probs = generate_samples(sentence, n_samples, model, temperature)\n",
        "    scores = weighted_avg_overlap(samples, log_probs)\n",
        "    max_score_key = max(scores, key=scores.get)\n",
        "    return {sample:score for sample, score in zip(samples, scores.values())}"
      ],
      "metadata": {
        "id": "ny5EYLFWt_5u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sentence = \"I love you.\"\n",
        "translated_sentence = mbr_decode(sentence, 2, transformer, 1)\n",
        "print(\"English: \", sentence)\n",
        "print(\"Bengali: \", translated_sentence)"
      ],
      "metadata": {
        "id": "lzrxKRMAaOgr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "3ACvLvfouZz0"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}