train_symbol_decoder.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Code is inspired from: https://www.tensorflow.org/text/tutorials/transformer"
      ],
      "metadata": {
        "id": "URdc__Dpl-3b"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XFG0NDRu5mYQ"
      },
      "outputs": [],
      "source": [
        "# Install the most re version of TensorFlow to use the improved\n",
        "# masking support for `tf.keras.layers.MultiHeadAttention`.\n",
        "!apt install --allow-change-held-packages libcudnn8=8.1.0.77-1+cuda11.2\n",
        "!pip uninstall -y -q tensorflow keras tensorflow-estimator tensorflow-text\n",
        "!pip install protobuf~=3.20.3\n",
        "#!pip install -q tensorflow_datasets\n",
        "#!pip install -q -U tensorflow-text==2.14.0 tensorflow==2.14.0\n",
        "!pip install -q -U tensorflow-text==2.15.0 tensorflow==2.15.0\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0GYpLBSjxJmG"
      },
      "source": [
        "Import the necessary modules:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JjJJyJTZYebt"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "import time\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "\n",
        "import tensorflow_text\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "df_train_org=pd.read_csv(\"data/train_data_fixed_sn07.csv\", header=0, usecols=[0,1])\n",
        "df_val_org=pd.read_csv(\"data/validation_data_fixed_sn07.csv\", header=0, usecols=[0,1])\n"
      ],
      "metadata": {
        "id": "lHZHcCIXFP1u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataset_tf = tf.data.Dataset.from_tensor_slices((df_train_org['Sememe'], df_train_org['English']))\n",
        "val_dataset_tf = tf.data.Dataset.from_tensor_slices((df_val_org['Sememe'], df_val_org['English']))"
      ],
      "metadata": {
        "id": "Y6zqKzMwHKxr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QToMl0NanZPr"
      },
      "outputs": [],
      "source": [
        "model_name = 'ted_hrlr_translate_pt_en_converter'\n",
        "tf.keras.utils.get_file(\n",
        "    f'{model_name}.zip',\n",
        "    f'https://storage.googleapis.com/download.tensorflow.org/models/{model_name}.zip',\n",
        "    cache_dir='.', cache_subdir='', extract=True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#model_name = '/content/drive/MyDrive/sememe_en_tokenizer'\n",
        "tokenizers = tf.saved_model.load(model_name)\n",
        "tokenizers.en.get_vocab_size().numpy()"
      ],
      "metadata": {
        "id": "J6tu_0tcXupG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizers_pre = tf.saved_model.load(model_name)"
      ],
      "metadata": {
        "id": "5AWrlgXQYKzS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#tokenizers = tf.saved_model.load(model_name)\n",
        "tokenizers.en.get_vocab_size().numpy()"
      ],
      "metadata": {
        "id": "T51k1S12j8_R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h5dbGnPXnuI1"
      },
      "outputs": [],
      "source": [
        "tokenizers = tf.saved_model.load(model_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s-PCJijfcZ9_"
      },
      "outputs": [],
      "source": [
        "[item for item in dir(tokenizers.en) if not item.startswith('_')]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KRbke-iaaHFI"
      },
      "outputs": [],
      "source": [
        "lengths = []\n",
        "\n",
        "for sememe_examples, en_examples in train_dataset_tf.batch(1024):\n",
        "  sememe_tokens = tokenizers.en.tokenize(sememe_examples)\n",
        "  #sememe_tokens = custom_tokenizer(sememe_examples,sememe_dict)\n",
        "  lengths.append(sememe_tokens.row_lengths())\n",
        "\n",
        "  en_tokens = tokenizers.en.tokenize(en_examples)\n",
        "  #en_tokens = custom_tokenizer(en_examples,sememe_dict)\n",
        "  lengths.append(en_tokens.row_lengths())\n",
        "  print('.', end='', flush=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9ucA1q3GaK_n"
      },
      "outputs": [],
      "source": [
        "all_lengths = np.concatenate(lengths)\n",
        "\n",
        "plt.hist(all_lengths, np.linspace(0, 500, 101))\n",
        "plt.ylim(plt.ylim())\n",
        "max_length = max(all_lengths)\n",
        "plt.plot([max_length, max_length], plt.ylim())\n",
        "plt.title(f'Maximum tokens per example: {max_length}');"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6shgzEck3FiV"
      },
      "outputs": [],
      "source": [
        "MAX_TOKENS=100\n",
        "def prepare_batch(pt, en):\n",
        "    pt = tokenizers.en.tokenize(pt)      # Output is ragged.\n",
        "    #print(pt)\n",
        "    #pt = custom_tokenizer(pt,sememe_dict)\n",
        "    pt = pt[:, :MAX_TOKENS]    # Trim to MAX_TOKENS.\n",
        "    pt = pt.to_tensor()  # Convert to 0-padded dense Tensor\n",
        "    en = tokenizers.en.tokenize(en)\n",
        "    #en = custom_tokenizer(en,sememe_dict)\n",
        "    en = en[:, :(MAX_TOKENS+1)]\n",
        "    en_inputs = en[:, :-1].to_tensor()  # Drop the [END] tokens\n",
        "    en_labels = en[:, 1:].to_tensor()   # Drop the [START] tokens\n",
        "\n",
        "    return (pt, en_inputs), en_labels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bcRp7VcQ5m6g"
      },
      "outputs": [],
      "source": [
        "BUFFER_SIZE = 20000\n",
        "BATCH_SIZE = 64"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BUN_jLBTwNxk"
      },
      "outputs": [],
      "source": [
        "def make_batches(ds):\n",
        "  return (\n",
        "      ds\n",
        "      .shuffle(BUFFER_SIZE)\n",
        "      .batch(BATCH_SIZE)\n",
        "      .map(prepare_batch, tf.data.AUTOTUNE)\n",
        "\n",
        "      #.map(tf.contrib.eager.py_func(prepare_batch), tf.data.AUTOTUNE)\n",
        "      .prefetch(buffer_size=tf.data.AUTOTUNE))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BSswr5TKvoNM"
      },
      "outputs": [],
      "source": [
        "# Create training and validation set batches.\n",
        "train_batches = make_batches(train_dataset_tf)\n",
        "val_batches = make_batches(val_dataset_tf)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CAw2XjRwLFWr"
      },
      "outputs": [],
      "source": [
        "for (pt, en), en_labels in train_batches.take(1):\n",
        "  break\n",
        "\n",
        "print(pt.shape)\n",
        "print(en.shape)\n",
        "print(en_labels.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Rz82wEs5biZ"
      },
      "outputs": [],
      "source": [
        "def positional_encoding(length, depth):\n",
        "  depth = depth/2\n",
        "\n",
        "  positions = np.arange(length)[:, np.newaxis]     # (seq, 1)\n",
        "  depths = np.arange(depth)[np.newaxis, :]/depth   # (1, depth)\n",
        "\n",
        "  angle_rates = 1 / (10000**depths)         # (1, depth)\n",
        "  angle_rads = positions * angle_rates      # (pos, depth)\n",
        "\n",
        "  pos_encoding = np.concatenate(\n",
        "      [np.sin(angle_rads), np.cos(angle_rads)],\n",
        "      axis=-1)\n",
        "\n",
        "  return tf.cast(pos_encoding, dtype=tf.float32)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "838tmM1cm9cB"
      },
      "outputs": [],
      "source": [
        "class PositionalEmbedding(tf.keras.layers.Layer):\n",
        "  def __init__(self, vocab_size, d_model):\n",
        "    super().__init__()\n",
        "    self.d_model = d_model\n",
        "    self.embedding = tf.keras.layers.Embedding(vocab_size, d_model, mask_zero=True)\n",
        "    self.pos_encoding = positional_encoding(length=2048, depth=d_model)\n",
        "\n",
        "  def compute_mask(self, *args, **kwargs):\n",
        "    return self.embedding.compute_mask(*args, **kwargs)\n",
        "\n",
        "  def call(self, x):\n",
        "    length = tf.shape(x)[1]\n",
        "    x = self.embedding(x)\n",
        "    # This factor sets the relative scale of the embedding and positonal_encoding.\n",
        "    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
        "    x = x + self.pos_encoding[tf.newaxis, :length, :]\n",
        "    return x\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tfz-EaCEDfUJ"
      },
      "outputs": [],
      "source": [
        "embed_pt = PositionalEmbedding(vocab_size=tokenizers.en.get_vocab_size(), d_model=512)\n",
        "embed_en = PositionalEmbedding(vocab_size=tokenizers.en.get_vocab_size(), d_model=512)\n",
        "\n",
        "pt_emb = embed_pt(pt)\n",
        "en_emb = embed_en(en)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5VLa5QcdPpv5"
      },
      "outputs": [],
      "source": [
        "class BaseAttention(tf.keras.layers.Layer):\n",
        "  def __init__(self, **kwargs):\n",
        "    super().__init__()\n",
        "    self.mha = tf.keras.layers.MultiHeadAttention(**kwargs)\n",
        "    self.layernorm = tf.keras.layers.LayerNormalization()\n",
        "    self.add = tf.keras.layers.Add()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kfHVbJUWv8qp"
      },
      "outputs": [],
      "source": [
        "class CrossAttention(BaseAttention):\n",
        "  def call(self, x, context):\n",
        "    attn_output, attn_scores = self.mha(\n",
        "        query=x,\n",
        "        key=context,\n",
        "        value=context,\n",
        "        return_attention_scores=True)\n",
        "\n",
        "    # Cache the attention scores for plotting later.\n",
        "    self.last_attn_scores = attn_scores\n",
        "\n",
        "    x = self.add([x, attn_output])\n",
        "    x = self.layernorm(x)\n",
        "\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qw1FJV5qRk79"
      },
      "outputs": [],
      "source": [
        "sample_ca = CrossAttention(num_heads=2, key_dim=512)\n",
        "\n",
        "print(pt_emb.shape)\n",
        "print(en_emb.shape)\n",
        "print(sample_ca(en_emb, pt_emb).shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RNqoTpn1wB3i"
      },
      "outputs": [],
      "source": [
        "class GlobalSelfAttention(BaseAttention):\n",
        "  def call(self, x):\n",
        "    attn_output = self.mha(\n",
        "        query=x,\n",
        "        value=x,\n",
        "        key=x)\n",
        "    x = self.add([x, attn_output])\n",
        "    x = self.layernorm(x)\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jPn2D07-Jcmj"
      },
      "outputs": [],
      "source": [
        "sample_gsa = GlobalSelfAttention(num_heads=2, key_dim=512)\n",
        "\n",
        "print(pt_emb.shape)\n",
        "print(sample_gsa(pt_emb).shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4MMQ-AfKD99_"
      },
      "outputs": [],
      "source": [
        "class CausalSelfAttention(BaseAttention):\n",
        "  def call(self, x):\n",
        "    attn_output = self.mha(\n",
        "        query=x,\n",
        "        value=x,\n",
        "        key=x,\n",
        "        use_causal_mask = True)\n",
        "    x = self.add([x, attn_output])\n",
        "    x = self.layernorm(x)\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x4dQuzvlD99_"
      },
      "outputs": [],
      "source": [
        "sample_csa = CausalSelfAttention(num_heads=2, key_dim=512)\n",
        "\n",
        "print(en_emb.shape)\n",
        "print(sample_csa(en_emb).shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bwKlheQ-WVxl"
      },
      "outputs": [],
      "source": [
        "out1 = sample_csa(embed_en(en[:, :3]))\n",
        "out2 = sample_csa(embed_en(en))[:, :3]\n",
        "\n",
        "tf.reduce_max(abs(out1 - out2)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rAYLeu0uwXYK"
      },
      "outputs": [],
      "source": [
        "class FeedForward(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, dff, dropout_rate=0.1):\n",
        "    super().__init__()\n",
        "    self.seq = tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(dff, activation='relu'),\n",
        "      tf.keras.layers.Dense(d_model),\n",
        "      tf.keras.layers.Dropout(dropout_rate)\n",
        "    ])\n",
        "    self.add = tf.keras.layers.Add()\n",
        "    self.layer_norm = tf.keras.layers.LayerNormalization()\n",
        "\n",
        "  def call(self, x):\n",
        "    x = self.add([x, self.seq(x)])\n",
        "    x = self.layer_norm(x)\n",
        "    return x\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r-Y8Yqi1_hUt"
      },
      "outputs": [],
      "source": [
        "sample_ffn = FeedForward(512, 2048)\n",
        "\n",
        "print(en_emb.shape)\n",
        "print(sample_ffn(en_emb).shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ncyS-Ms3i2x_"
      },
      "outputs": [],
      "source": [
        "class EncoderLayer(tf.keras.layers.Layer):\n",
        "  def __init__(self,*, d_model, num_heads, dff, dropout_rate=0.1):\n",
        "    super().__init__()\n",
        "\n",
        "    self.self_attention = GlobalSelfAttention(\n",
        "        num_heads=num_heads,\n",
        "        key_dim=d_model,\n",
        "        dropout=dropout_rate)\n",
        "\n",
        "    self.ffn = FeedForward(d_model, dff)\n",
        "\n",
        "  def call(self, x):\n",
        "    x = self.self_attention(x)\n",
        "    x = self.ffn(x)\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AzZRXdO0mI48"
      },
      "outputs": [],
      "source": [
        "sample_encoder_layer = EncoderLayer(d_model=512, num_heads=8, dff=2048)\n",
        "\n",
        "print(pt_emb.shape)\n",
        "print(sample_encoder_layer(pt_emb).shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jpEox7gJ8FCI"
      },
      "outputs": [],
      "source": [
        "class Encoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, *, num_layers, d_model, num_heads,\n",
        "               dff, vocab_size, dropout_rate=0.1):\n",
        "    super().__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.num_layers = num_layers\n",
        "\n",
        "    self.pos_embedding = PositionalEmbedding(\n",
        "        vocab_size=vocab_size, d_model=d_model)\n",
        "\n",
        "    self.enc_layers = [\n",
        "        EncoderLayer(d_model=d_model,\n",
        "                     num_heads=num_heads,\n",
        "                     dff=dff,\n",
        "                     dropout_rate=dropout_rate)\n",
        "        for _ in range(num_layers)]\n",
        "    self.dropout = tf.keras.layers.Dropout(dropout_rate)\n",
        "\n",
        "  def call(self, x):\n",
        "    # `x` is token-IDs shape: (batch, seq_len)\n",
        "    x = self.pos_embedding(x)  # Shape `(batch_size, seq_len, d_model)`.\n",
        "\n",
        "    # Add dropout.\n",
        "    x = self.dropout(x)\n",
        "\n",
        "    for i in range(self.num_layers):\n",
        "      x = self.enc_layers[i](x)\n",
        "\n",
        "    return x  # Shape `(batch_size, seq_len, d_model)`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SDPXTvYgJH8s"
      },
      "outputs": [],
      "source": [
        "# Instantiate the encoder.\n",
        "sample_encoder = Encoder(num_layers=4,\n",
        "                         d_model=512,\n",
        "                         num_heads=8,\n",
        "                         dff=2048,\n",
        "                         vocab_size=7808)\n",
        "\n",
        "sample_encoder_output = sample_encoder(pt, training=False)\n",
        "\n",
        "# Print the shape.\n",
        "print(pt.shape)\n",
        "print(sample_encoder_output.shape)  # Shape `(batch_size, input_seq_len, d_model)`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9SoX0-vd1hue"
      },
      "outputs": [],
      "source": [
        "class DecoderLayer(tf.keras.layers.Layer):\n",
        "  def __init__(self,\n",
        "               *,\n",
        "               d_model,\n",
        "               num_heads,\n",
        "               dff,\n",
        "               dropout_rate=0.1):\n",
        "    super(DecoderLayer, self).__init__()\n",
        "\n",
        "    self.causal_self_attention = CausalSelfAttention(\n",
        "        num_heads=num_heads,\n",
        "        key_dim=d_model,\n",
        "        dropout=dropout_rate)\n",
        "\n",
        "    self.cross_attention = CrossAttention(\n",
        "        num_heads=num_heads,\n",
        "        key_dim=d_model,\n",
        "        dropout=dropout_rate)\n",
        "\n",
        "    self.ffn = FeedForward(d_model, dff)\n",
        "\n",
        "  def call(self, x, context):\n",
        "    x = self.causal_self_attention(x=x)\n",
        "    x = self.cross_attention(x=x, context=context)\n",
        "\n",
        "    # Cache the last attention scores for plotting later\n",
        "    self.last_attn_scores = self.cross_attention.last_attn_scores\n",
        "\n",
        "    x = self.ffn(x)  # Shape `(batch_size, seq_len, d_model)`.\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ne2Bqx8k71l0"
      },
      "outputs": [],
      "source": [
        "sample_decoder_layer = DecoderLayer(d_model=512, num_heads=8, dff=2048)\n",
        "\n",
        "sample_decoder_layer_output = sample_decoder_layer(\n",
        "    x=en_emb, context=pt_emb)\n",
        "\n",
        "print(en_emb.shape)\n",
        "print(pt_emb.shape)\n",
        "print(sample_decoder_layer_output.shape)  # `(batch_size, seq_len, d_model)`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d5_d5-PLQXwY"
      },
      "outputs": [],
      "source": [
        "class Decoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, *, num_layers, d_model, num_heads, dff, vocab_size,\n",
        "               dropout_rate=0.1):\n",
        "    super(Decoder, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.num_layers = num_layers\n",
        "\n",
        "    self.pos_embedding = PositionalEmbedding(vocab_size=vocab_size,\n",
        "                                             d_model=d_model)\n",
        "    self.dropout = tf.keras.layers.Dropout(dropout_rate)\n",
        "    self.dec_layers = [\n",
        "        DecoderLayer(d_model=d_model, num_heads=num_heads,\n",
        "                     dff=dff, dropout_rate=dropout_rate)\n",
        "        for _ in range(num_layers)]\n",
        "\n",
        "    self.last_attn_scores = None\n",
        "\n",
        "  def call(self, x, context):\n",
        "    # `x` is token-IDs shape (batch, target_seq_len)\n",
        "    x = self.pos_embedding(x)  # (batch_size, target_seq_len, d_model)\n",
        "\n",
        "    x = self.dropout(x)\n",
        "\n",
        "    for i in range(self.num_layers):\n",
        "      x  = self.dec_layers[i](x, context)\n",
        "\n",
        "    self.last_attn_scores = self.dec_layers[-1].last_attn_scores\n",
        "\n",
        "    # The shape of x is (batch_size, target_seq_len, d_model).\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xyHdG_jWPgKu"
      },
      "outputs": [],
      "source": [
        "# Instantiate the decoder.\n",
        "sample_decoder = Decoder(num_layers=4,\n",
        "                         d_model=512,\n",
        "                         num_heads=8,\n",
        "                         dff=2048,\n",
        "                         vocab_size=7808)\n",
        "\n",
        "output = sample_decoder(\n",
        "    x=en,\n",
        "    context=pt_emb)\n",
        "\n",
        "# Print the shapes.\n",
        "print(en.shape)\n",
        "print(pt_emb.shape)\n",
        "print(output.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ioJ4XJAUAReI"
      },
      "outputs": [],
      "source": [
        "sample_decoder.last_attn_scores.shape  # (batch, heads, target_seq, input_seq)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PED3bIpOYkBu"
      },
      "outputs": [],
      "source": [
        "class Transformer(tf.keras.Model):\n",
        "  def __init__(self, *, num_layers, d_model, num_heads, dff,\n",
        "               input_vocab_size, target_vocab_size, dropout_rate=0.1):\n",
        "    super().__init__()\n",
        "    self.encoder = Encoder(num_layers=num_layers, d_model=d_model,\n",
        "                           num_heads=num_heads, dff=dff,\n",
        "                           vocab_size=input_vocab_size,\n",
        "                           dropout_rate=dropout_rate)\n",
        "\n",
        "    self.decoder = Decoder(num_layers=num_layers, d_model=d_model,\n",
        "                           num_heads=num_heads, dff=dff,\n",
        "                           vocab_size=target_vocab_size,\n",
        "                           dropout_rate=dropout_rate)\n",
        "\n",
        "    self.final_layer = tf.keras.layers.Dense(target_vocab_size)\n",
        "\n",
        "  def call(self, inputs):\n",
        "    # To use a Keras model with `.fit` you must pass all your inputs in the\n",
        "    # first argument.\n",
        "    context, x  = inputs\n",
        "\n",
        "    context = self.encoder(context)  # (batch_size, context_len, d_model)\n",
        "\n",
        "    x = self.decoder(x, context)  # (batch_size, target_len, d_model)\n",
        "\n",
        "    # Final linear layer output.\n",
        "    logits = self.final_layer(x)  # (batch_size, target_len, target_vocab_size)\n",
        "\n",
        "    try:\n",
        "      # Drop the keras mask, so it doesn't scale the losses/metrics.\n",
        "      # b/250038731\n",
        "      del logits._keras_mask\n",
        "    except AttributeError:\n",
        "      pass\n",
        "\n",
        "    # Return the final output and the attention weights.\n",
        "    return logits"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mzyo6KDfVyhl"
      },
      "outputs": [],
      "source": [
        "num_layers = 3\n",
        "d_model = 128\n",
        "dff = 512\n",
        "num_heads = 8\n",
        "dropout_rate = 0.1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UiysUa--4tOU"
      },
      "outputs": [],
      "source": [
        "transformer = Transformer(\n",
        "    num_layers=num_layers,\n",
        "    d_model=d_model,\n",
        "    num_heads=num_heads,\n",
        "    dff=dff,\n",
        "    input_vocab_size=tokenizers.en.get_vocab_size().numpy(),\n",
        "    target_vocab_size=tokenizers.en.get_vocab_size().numpy(),\n",
        "    dropout_rate=dropout_rate)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c8eO85hpFHmE"
      },
      "outputs": [],
      "source": [
        "output = transformer((pt, en))\n",
        "\n",
        "print(en.shape)\n",
        "print(pt.shape)\n",
        "print(output.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "olTLrK8pAcLd"
      },
      "outputs": [],
      "source": [
        "attn_scores = transformer.decoder.dec_layers[-1].last_attn_scores\n",
        "print(attn_scores.shape)  # (batch, heads, target_seq, input_seq)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_jTvJsXquaHW"
      },
      "source": [
        "Print the summary of the model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IsUPhlfEtOjn"
      },
      "outputs": [],
      "source": [
        "transformer.summary()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iYQdOO1axwEI"
      },
      "outputs": [],
      "source": [
        "class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
        "  def __init__(self, d_model, warmup_steps=4000):\n",
        "    super().__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.d_model = tf.cast(self.d_model, tf.float32)\n",
        "\n",
        "    self.warmup_steps = warmup_steps\n",
        "\n",
        "  def __call__(self, step):\n",
        "    step = tf.cast(step, dtype=tf.float32)\n",
        "    arg1 = tf.math.rsqrt(step)\n",
        "    arg2 = step * (self.warmup_steps ** -1.5)\n",
        "\n",
        "    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7r4scdulztRx"
      },
      "outputs": [],
      "source": [
        "#learning_rate = CustomSchedule(d_model)\n",
        "learning_rate = 1e-4\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,\n",
        "                                     epsilon=1e-9)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xij3MwYVRAAS"
      },
      "outputs": [],
      "source": [
        "#plt.plot(learning_rate(tf.range(40000, dtype=tf.float32)))\n",
        "#plt.ylabel('Learning Rate')\n",
        "#plt.xlabel('Train Step')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "67oqVHiT0Eiu"
      },
      "outputs": [],
      "source": [
        "def masked_loss(label, pred):\n",
        "  mask = label != 0\n",
        "  loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True, reduction='none')\n",
        "  loss = loss_object(label, pred)\n",
        "\n",
        "  mask = tf.cast(mask, dtype=loss.dtype)\n",
        "  loss *= mask\n",
        "\n",
        "  loss = tf.reduce_sum(loss)/tf.reduce_sum(mask)\n",
        "  return loss\n",
        "\n",
        "\n",
        "def masked_accuracy(label, pred):\n",
        "  pred = tf.argmax(pred, axis=2)\n",
        "  label = tf.cast(label, pred.dtype)\n",
        "  match = label == pred\n",
        "\n",
        "  mask = label != 0\n",
        "\n",
        "  match = match & mask\n",
        "\n",
        "  match = tf.cast(match, dtype=tf.float32)\n",
        "  mask = tf.cast(mask, dtype=tf.float32)\n",
        "  return tf.reduce_sum(match)/tf.reduce_sum(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Una1v0hDlIsT"
      },
      "outputs": [],
      "source": [
        "transformer.compile(\n",
        "    loss=masked_loss,\n",
        "    optimizer=optimizer,\n",
        "    metrics=[masked_accuracy])"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "checkpoint_path = \"fixed_sn07/cp.ckpt\"\n",
        "checkpoint_dir = os.path.dirname(checkpoint_path)\n",
        "\n",
        "# Create a callback that saves the model's weights\n",
        "cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,\n",
        "                                                 save_weights_only=True,\n",
        "                                                 verbose=1)\n",
        "\n",
        "#transformer.load_weights(checkpoint_path)"
      ],
      "metadata": {
        "id": "FEDE_s7E4Y35"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jg35qKvVlctp"
      },
      "outputs": [],
      "source": [
        "transformer.fit(train_batches,\n",
        "                epochs=60,\n",
        "                validation_data=val_batches,\n",
        "                callbacks=[cp_callback])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eY_uXsOhSmbb"
      },
      "outputs": [],
      "source": [
        "class Translator(tf.Module):\n",
        "  def __init__(self, tokenizers, transformer):\n",
        "    self.tokenizers = tokenizers\n",
        "    self.transformer = transformer\n",
        "\n",
        "  def __call__(self, sentence, max_length=MAX_TOKENS):\n",
        "    # The input sentence is Portuguese, hence adding the `[START]` and `[END]` tokens.\n",
        "    assert isinstance(sentence, tf.Tensor)\n",
        "    if len(sentence.shape) == 0:\n",
        "      sentence = sentence[tf.newaxis]\n",
        "\n",
        "    sentence = self.tokenizers.en.tokenize(sentence).to_tensor()\n",
        "\n",
        "    encoder_input = sentence\n",
        "\n",
        "    # As the output language is English, initialize the output with the\n",
        "    # English `[START]` token.\n",
        "    start_end = self.tokenizers.en.tokenize([''])[0]\n",
        "    start = start_end[0][tf.newaxis]\n",
        "    end = start_end[1][tf.newaxis]\n",
        "\n",
        "    # `tf.TensorArray` is required here (instead of a Python list), so that the\n",
        "    # dynamic-loop can be traced by `tf.function`.\n",
        "    output_array = tf.TensorArray(dtype=tf.int64, size=0, dynamic_size=True)\n",
        "    output_array = output_array.write(0, start)\n",
        "\n",
        "    for i in tf.range(max_length):\n",
        "      output = tf.transpose(output_array.stack())\n",
        "      predictions = self.transformer([encoder_input, output], training=False)\n",
        "\n",
        "      # Select the last token from the `seq_len` dimension.\n",
        "      predictions = predictions[:, -1:, :]  # Shape `(batch_size, 1, vocab_size)`.\n",
        "\n",
        "      predicted_id = tf.argmax(predictions, axis=-1)\n",
        "\n",
        "      # Concatenate the `predicted_id` to the output which is given to the\n",
        "      # decoder as its input.\n",
        "      output_array = output_array.write(i+1, predicted_id[0])\n",
        "\n",
        "      if predicted_id == end:\n",
        "        break\n",
        "\n",
        "    output = tf.transpose(output_array.stack())\n",
        "    # The output shape is `(1, tokens)`.\n",
        "    text = tokenizers.en.detokenize(output)[0]  # Shape: `()`.\n",
        "\n",
        "    tokens = tokenizers.en.lookup(output)[0]\n",
        "\n",
        "    # `tf.function` prevents us from using the attention_weights that were\n",
        "    # calculated on the last iteration of the loop.\n",
        "    # So, recalculate them outside the loop.\n",
        "    self.transformer([encoder_input, output[:,:-1]], training=False)\n",
        "    attention_weights = self.transformer.decoder.last_attn_scores\n",
        "\n",
        "    return text, tokens, attention_weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-NjbvpHUTEia"
      },
      "outputs": [],
      "source": [
        "translator = Translator(tokenizers, transformer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QfHSRdejTFsC"
      },
      "outputs": [],
      "source": [
        "def print_translation(uncorr_in,sentence, tokens, ground_truth):\n",
        "  print(f'{\"Input:\":15s}: {sentence}')\n",
        "  print(f'{\"Uncorrupted Input:\":15s}: {uncorr_in}')\n",
        "  print(f'{\"Prediction\":15s}: {tokens.numpy().decode(\"utf-8\")}')\n",
        "  print(f'{\"Ground truth\":15s}: {ground_truth}')\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tf.constant(df_val_org['Sememe'][150])"
      ],
      "metadata": {
        "id": "XihlXfTAn_1z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gKlxYO0JTXzD"
      },
      "outputs": [],
      "source": [
        "def plot_attention_head(in_tokens, translated_tokens, attention):\n",
        "  # The model didn't generate `<START>` in the output. Skip it.\n",
        "  translated_tokens = translated_tokens[1:]\n",
        "\n",
        "  ax = plt.gca()\n",
        "  ax.matshow(attention)\n",
        "  ax.set_xticks(range(len(in_tokens)))\n",
        "  ax.set_yticks(range(len(translated_tokens)))\n",
        "\n",
        "  labels = [label.decode('utf-8') for label in in_tokens.numpy()]\n",
        "  ax.set_xticklabels(\n",
        "      labels, rotation=90)\n",
        "\n",
        "  labels = [label.decode('utf-8') for label in translated_tokens.numpy()]\n",
        "  ax.set_yticklabels(labels)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yI4YWU2uXDeW"
      },
      "outputs": [],
      "source": [
        "head = 0\n",
        "# Shape: `(batch=1, num_heads, seq_len_q, seq_len_k)`.\n",
        "attention_heads = tf.squeeze(attention_weights, 0)\n",
        "attention = attention_heads[head]\n",
        "attention.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SMEpyioWTmSN"
      },
      "outputs": [],
      "source": [
        "in_tokens = tf.convert_to_tensor([sentence])\n",
        "in_tokens = tokenizers.pt.tokenize(in_tokens).to_tensor()\n",
        "in_tokens = tokenizers.pt.lookup(in_tokens)[0]\n",
        "in_tokens"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GzvIo5uYTnHG"
      },
      "outputs": [],
      "source": [
        "translated_tokens"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lrNh47D1ToBD"
      },
      "outputs": [],
      "source": [
        "plot_attention_head(in_tokens, translated_tokens, attention)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iMZr-rI_TrGh"
      },
      "outputs": [],
      "source": [
        "def plot_attention_weights(sentence, translated_tokens, attention_heads):\n",
        "  in_tokens = tf.convert_to_tensor([sentence])\n",
        "  in_tokens = tokenizers.pt.tokenize(in_tokens).to_tensor()\n",
        "  in_tokens = tokenizers.pt.lookup(in_tokens)[0]\n",
        "\n",
        "  fig = plt.figure(figsize=(16, 8))\n",
        "\n",
        "  for h, head in enumerate(attention_heads):\n",
        "    ax = fig.add_subplot(2, 4, h+1)\n",
        "\n",
        "    plot_attention_head(in_tokens, translated_tokens, head)\n",
        "\n",
        "    ax.set_xlabel(f'Head {h+1}')\n",
        "\n",
        "  plt.tight_layout()\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lBMujUb1Tr4C"
      },
      "outputs": [],
      "source": [
        "plot_attention_weights(sentence,\n",
        "                       translated_tokens,\n",
        "                       attention_weights[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NZhv5h4AT_n5"
      },
      "outputs": [],
      "source": [
        "class ExportTranslator(tf.Module):\n",
        "  def __init__(self, translator):\n",
        "    self.translator = translator\n",
        "\n",
        "  @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)])\n",
        "  def __call__(self, sentence):\n",
        "    (result,\n",
        "     tokens,\n",
        "     attention_weights) = self.translator(sentence, max_length=MAX_TOKENS)\n",
        "\n",
        "    return result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wm1_eRPvUCUm"
      },
      "outputs": [],
      "source": [
        "translator = ExportTranslator(translator)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_v--e1XmUFw3"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(translator, export_dir=('fixed_24percent'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5KJSQEzlUGo-"
      },
      "outputs": [],
      "source": [
        "reloaded = tf.saved_model.load('fixed_24percent')"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}