Transformer_raw.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 3652,
     "status": "ok",
     "timestamp": 1582716604719,
     "user": {
      "displayName": "Davide Montanari",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mBFwK0wQ4sGlm8L2VjR5w4UcXGE1Ul2vGoHiuLkMQ=s64",
      "userId": "12602812058261277045"
     },
     "user_tz": -60
    },
    "id": "xBpgv4TBQZta",
    "outputId": "8075ee38-98b4-41e5-ca3a-6f1cab2c4ca1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mounted at /content/drive\n",
      "/content/drive/My Drive/Colab Notebooks/Embeddings/clean_data/src_tgt/not_dis\n"
     ]
    }
   ],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive', force_remount='True')\n",
    "%cd drive/My\\ Drive/Colab\\ Notebooks/Embeddings/clean_data/src_tgt/not_dis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "smtGcCdPSERm"
   },
   "source": [
    "#Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 0
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 10871,
     "status": "ok",
     "timestamp": 1582716611947,
     "user": {
      "displayName": "Davide Montanari",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mBFwK0wQ4sGlm8L2VjR5w4UcXGE1Ul2vGoHiuLkMQ=s64",
      "userId": "12602812058261277045"
     },
     "user_tz": -60
    },
    "id": "gr37Hp_BSFwV",
    "outputId": "ee3ff973-bc04-4729-f56a-f0e4ecbe5b86"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow 2.x selected.\n",
      "Num GPUs Available:  1\n"
     ]
    }
   ],
   "source": [
    "from __future__ import absolute_import, division, print_function, unicode_literals\n",
    "\n",
    "try:\n",
    "    %tensorflow_version 2.x\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from gensim.models import KeyedVectors\n",
    "from gensim.test.utils import datapath, get_tmpfile\n",
    "from gensim.scripts.glove2word2vec import glove2word2vec\n",
    "from gensim.models import Word2Vec\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as ticker\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import json, pathlib, os, time, datetime, gensim, sys, pickle, time, unicodedata, io\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import regex as re\n",
    "from tensor2tensor.utils.beam_search import beam_search\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "g0vkur0LKo6g"
   },
   "source": [
    "#Tokenization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HfV9HpabT1SD"
   },
   "source": [
    "### Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uFDI7pTmKodG"
   },
   "outputs": [],
   "source": [
    "def unicode_to_ascii(s):\n",
    "    return ''.join(c for c in unicodedata.normalize('NFD', s)\n",
    "        if unicodedata.category(c) != 'Mn')\n",
    "  \n",
    "def preproc(w):\n",
    "    w = unicode_to_ascii(w.strip())\n",
    "    w = re.sub(r\"([?.!,¿])\", r\" \\1 \", w)\n",
    "    w = re.sub(r'[\" \"]+', \" \", w)\n",
    "    w = re.sub(r\"[^a-zA-Z?.!,¿]+\", \" \", w)\n",
    "    w = w.rstrip().strip()\n",
    "    w = '<sos> ' + w + ' <eos>'\n",
    "    return w\n",
    "\n",
    "def create_dataset(src_path, tgt_path):\n",
    "    src = io.open(src_path, encoding='UTF-8').read().strip().split('\\n')\n",
    "    tgt = io.open(tgt_path, encoding='UTF-8').read().strip().split('\\n')\n",
    "    return [preproc(line) for line in src], [preproc(line) for line in tgt]\n",
    "\n",
    "def tokenize(cnc, snt):\n",
    "    tokenizer.fit_on_texts(cnc+snt)\n",
    "    cnc_tensor = tokenizer.texts_to_sequences(cnc)\n",
    "    cnc_tensor = tf.keras.preprocessing.sequence.pad_sequences(cnc_tensor, padding='post', dtype='int64')\n",
    "    snt_tensor = tokenizer.texts_to_sequences(snt)\n",
    "    snt_tensor = tf.keras.preprocessing.sequence.pad_sequences(snt_tensor, padding='post', dtype='int64')\n",
    "    return cnc_tensor, snt_tensor, tokenizer\n",
    "\n",
    "def load_dataset(src_path, tgt_path):\n",
    "    cnc, snt = create_dataset(src_path, tgt_path)\n",
    "    input_tensor, target_tensor, inp_lang_tokenizer = tokenize(cnc, snt)\n",
    "    return input_tensor, target_tensor, tokenizer\n",
    "\n",
    "def convert(lang, tensor):\n",
    "    for t in tensor:\n",
    "        if t!=0:\n",
    "            print (\"%d ----> %s\" % (t, lang.index_word[t]))\n",
    "\n",
    "def max_length(tensor):\n",
    "  return max(len(t) for t in tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0F4GlrEhT5nr"
   },
   "source": [
    "### Dataset creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 0
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 18998,
     "status": "ok",
     "timestamp": 1582716620091,
     "user": {
      "displayName": "Davide Montanari",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mBFwK0wQ4sGlm8L2VjR5w4UcXGE1Ul2vGoHiuLkMQ=s64",
      "userId": "12602812058261277045"
     },
     "user_tz": -60
    },
    "id": "Xa3ID1crMjxD",
    "outputId": "769a65bc-0c90-456f-8473-6f0e7d10690a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "concepts: <sos> participate force helicopter <eos>\n",
      "sentence: <sos> the armed forces participated with their brand new helicopter . <eos>\n"
     ]
    }
   ],
   "source": [
    "tokenizer = tf.keras.preprocessing.text.Tokenizer(filters='')\n",
    "cnc_tensor, snt_tensor, tokenizer = load_dataset('concepts.txt', 'sentences.txt')\n",
    "cnc_tensor_train, cnc_tensor_val, snt_tensor_train, snt_tensor_val = train_test_split(cnc_tensor, snt_tensor, test_size=0.1)\n",
    "\n",
    "BUFFER_SIZE = len(cnc_tensor_train)\n",
    "BATCH_SIZE = 32\n",
    "\n",
    "vocab_size = len(tokenizer.word_index)+1\n",
    "max_length_targ, max_length_inp = max_length(snt_tensor), max_length(cnc_tensor)\n",
    "\n",
    "dataset = tf.data.Dataset.from_tensor_slices((cnc_tensor_train, snt_tensor_train)).shuffle(BUFFER_SIZE)\n",
    "train_dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)\n",
    "\n",
    "for ex in train_dataset.take(1):\n",
    "    inp = ex[0][0]\n",
    "    tar = ex[1][0]\n",
    "    print('concepts: {}\\nsentence: {}'.format(\n",
    "            ' '.join([tokenizer.index_word[i.numpy()] for i in inp if i!=0]), \n",
    "            ' '.join([tokenizer.index_word[i.numpy()] for i in tar if i!=0])\n",
    "        ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tdCeUDAiT6qM"
   },
   "source": [
    "#Weights and Biases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 27490,
     "status": "ok",
     "timestamp": 1582716628594,
     "user": {
      "displayName": "Davide Montanari",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mBFwK0wQ4sGlm8L2VjR5w4UcXGE1Ul2vGoHiuLkMQ=s64",
      "userId": "12602812058261277045"
     },
     "user_tz": -60
    },
    "id": "ow6cdlBJwd3G",
    "outputId": "2eb7d1fe-7012-4d87-cbef-f0a751bf20b9"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
       "                Project page: <a href=\"https://app.wandb.ai/edivairanatnom/transformer-not-dis\" target=\"_blank\">https://app.wandb.ai/edivairanatnom/transformer-not-dis</a><br/>\n",
       "                Run page: <a href=\"https://app.wandb.ai/edivairanatnom/transformer-not-dis/runs/iq0xgstr\" target=\"_blank\">https://app.wandb.ai/edivairanatnom/transformer-not-dis/runs/iq0xgstr</a><br/>\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "tags": []
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# !pip install wandb -q\n",
    "import wandb\n",
    "wandb.login()\n",
    "\n",
    "wandb.init(entity = \"edivairanatnom\", project = \"transformer-not-dis\", sync_tensorboard=True)\n",
    "\n",
    "config = wandb.config\n",
    "config.batch_size = 32         \n",
    "config.num_steps = 50000\n",
    "config.display_step = 10\n",
    "config.num_layers = 6\n",
    "config.d_model = 2048\n",
    "config.dff = 1024\n",
    "config.num_heads = 8\n",
    "config.vocab_size = vocab_size\n",
    "config.dropout_rate = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6rBf8Ohl0Z4E"
   },
   "source": [
    "#Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "fXa2p1E4oXI-"
   },
   "source": [
    "##Pos Encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QjoucfYcoa7m"
   },
   "outputs": [],
   "source": [
    "### POSITION\n",
    "\n",
    "def get_angles(pos, i, d_model):\n",
    "    angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n",
    "    return pos * angle_rates\n",
    "\n",
    "def positional_encoding(position, d_model):\n",
    "    angle_rads = get_angles(np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :], d_model)\n",
    "    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
    "    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
    "    pos_encoding = angle_rads[np.newaxis, ...]\n",
    "    return tf.cast(pos_encoding, dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pbZ-wGTzodji"
   },
   "source": [
    "##Masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "AFT42Wddokw_"
   },
   "outputs": [],
   "source": [
    "def create_padding_mask(seq):\n",
    "    seq = tf.cast(tf.math.equal(seq, 0), tf.float32)\n",
    "    return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)\n",
    "\n",
    "def create_look_ahead_mask(size):\n",
    "    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n",
    "    return mask  # (seq_len, seq_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-zHvlo5YohBv"
   },
   "source": [
    "##Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fwOPVuJLorVC"
   },
   "outputs": [],
   "source": [
    "def scaled_dot_product_attention(q, k, v, mask):\n",
    "\n",
    "    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)\n",
    "    dk = tf.cast(tf.shape(k)[-1], tf.float32)\n",
    "    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n",
    "\n",
    "    if mask is not None:\n",
    "        scaled_attention_logits += (mask * -1e9)  \n",
    "\n",
    "    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)\n",
    "    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)\n",
    "\n",
    "    return output, attention_weights\n",
    "\n",
    "class MultiHeadAttention(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.d_model = d_model\n",
    "        \n",
    "        assert d_model % self.num_heads == 0\n",
    "        self.depth = d_model // self.num_heads\n",
    "        \n",
    "        self.wq = tf.keras.layers.Dense(d_model)\n",
    "        self.wk = tf.keras.layers.Dense(d_model)\n",
    "        self.wv = tf.keras.layers.Dense(d_model)\n",
    "        \n",
    "        self.dense = tf.keras.layers.Dense(d_model)\n",
    "            \n",
    "    def split_heads(self, x, batch_size):\n",
    "        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
    "        return tf.transpose(x, perm=[0, 2, 1, 3])\n",
    "        \n",
    "    def call(self, v, k, q, mask):\n",
    "        batch_size = tf.shape(q)[0]\n",
    "        \n",
    "        q = self.wq(q)  # (batch_size, seq_len, d_model)\n",
    "        k = self.wk(k)  # (batch_size, seq_len, d_model)\n",
    "        v = self.wv(v)  # (batch_size, seq_len, d_model)\n",
    "        \n",
    "        q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)\n",
    "        k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)\n",
    "        v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)\n",
    "        \n",
    "        scaled_attention, attention_weights = scaled_dot_product_attention(q, k, v, mask)\n",
    "        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)\n",
    "        concat_attention = tf.reshape(scaled_attention, (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)\n",
    "\n",
    "        output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)\n",
    "            \n",
    "        return output, attention_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TAYfS2UToosF"
   },
   "source": [
    "##Feedforward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0qp0_aqPoVMv"
   },
   "outputs": [],
   "source": [
    "def point_wise_feed_forward_network(d_model, dff):\n",
    "    return tf.keras.Sequential([\n",
    "        tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)\n",
    "        tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)\n",
    "    ])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CGZ0mr17ozGw"
   },
   "source": [
    "##Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IOkAx9mco7l3"
   },
   "outputs": [],
   "source": [
    "class EncoderLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "\n",
    "        self.mha = MultiHeadAttention(d_model, num_heads)\n",
    "        self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
    "\n",
    "        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        \n",
    "        self.dropout1 = tf.keras.layers.Dropout(rate)\n",
    "        self.dropout2 = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "    def call(self, x, training, mask):\n",
    "\n",
    "        attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)\n",
    "        attn_output = self.dropout1(attn_output, training=training)\n",
    "        out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)\n",
    "        \n",
    "        ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)\n",
    "        ffn_output = self.dropout2(ffn_output, training=training)\n",
    "        out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)\n",
    "        \n",
    "        return out2\n",
    "\n",
    "class Encoder(tf.keras.layers.Layer):\n",
    "    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, maximum_position_encoding, rate=0.1):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.d_model = d_model\n",
    "        self.num_layers = num_layers\n",
    "        self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)\n",
    "        self.pos_encoding = positional_encoding(maximum_position_encoding, self.d_model)\n",
    "        self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) for _ in range(num_layers)]\n",
    "        self.dropout = tf.keras.layers.Dropout(rate)\n",
    "        \n",
    "    def call(self, x, training, mask):\n",
    "        seq_len = tf.shape(x)[1]\n",
    "        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)\n",
    "        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
    "        x += self.pos_encoding[:, :seq_len, :]\n",
    "        x = self.dropout(x, training=training)    \n",
    "        for i in range(self.num_layers):\n",
    "            x = self.enc_layers[i](x, training, mask)\n",
    "\n",
    "        return x  # (batch_size, input_seq_len, d_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vLsihovQo87c"
   },
   "source": [
    "##Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HJHlGhimo8Kl"
   },
   "outputs": [],
   "source": [
    "class DecoderLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "\n",
    "        self.mha1 = MultiHeadAttention(d_model, num_heads)\n",
    "        self.mha2 = MultiHeadAttention(d_model, num_heads)\n",
    "\n",
    "        self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
    "    \n",
    "        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        \n",
    "        self.dropout1 = tf.keras.layers.Dropout(rate)\n",
    "        self.dropout2 = tf.keras.layers.Dropout(rate)\n",
    "        self.dropout3 = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "    \n",
    "    def call(self, x, enc_output, training, look_ahead_mask, padding_mask):   # enc_output.shape == (batch_size, input_seq_len, d_model)\n",
    "\n",
    "        attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)\n",
    "        attn1 = self.dropout1(attn1, training=training)\n",
    "        out1 = self.layernorm1(attn1 + x)\n",
    "    \n",
    "        attn2, attn_weights_block2 = self.mha2(enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)\n",
    "        attn2 = self.dropout2(attn2, training=training)\n",
    "        out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)\n",
    "        \n",
    "        ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)\n",
    "        ffn_output = self.dropout3(ffn_output, training=training)\n",
    "        out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)\n",
    "        \n",
    "        return out3, attn_weights_block1, attn_weights_block2\n",
    "\n",
    "class Decoder(tf.keras.layers.Layer):\n",
    "    def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size, maximum_position_encoding, rate=0.1):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.d_model = d_model\n",
    "        self.num_layers = num_layers\n",
    "        self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)\n",
    "        self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)\n",
    "        self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) for _ in range(num_layers)]\n",
    "        self.dropout = tf.keras.layers.Dropout(rate)\n",
    "    \n",
    "    def call(self, x, enc_output, training, look_ahead_mask, padding_mask):\n",
    "        seq_len = tf.shape(x)[1]\n",
    "        attention_weights = {}\n",
    "        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)\n",
    "        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
    "        x += self.pos_encoding[:, :seq_len, :]\n",
    "        x = self.dropout(x, training=training)\n",
    "\n",
    "        for i in range(self.num_layers):\n",
    "            x, block1, block2 = self.dec_layers[i](x, enc_output, training, look_ahead_mask, padding_mask)\n",
    "            attention_weights['decoder_layer{}_block1'.format(i+1)] = block1\n",
    "            attention_weights['decoder_layer{}_block2'.format(i+1)] = block2\n",
    "      \n",
    "        return x, attention_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QatUCCIjpAnI"
   },
   "source": [
    "##Seq2Seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OzhcAlS6pH6-"
   },
   "outputs": [],
   "source": [
    "class Transformer(tf.keras.Model):\n",
    "    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, target_vocab_size, pe_input, pe_target, rate=0.1):\n",
    "        super(Transformer, self).__init__()\n",
    "        self.encoder = Encoder(num_layers, d_model, num_heads, dff, input_vocab_size, pe_input, rate)\n",
    "        self.decoder = Decoder(num_layers, d_model, num_heads, dff, target_vocab_size, pe_target, rate)\n",
    "        self.final_layer = tf.keras.layers.Dense(target_vocab_size)\n",
    "\n",
    "    def call(self, inp, tar, training, enc_padding_mask, look_ahead_mask, dec_padding_mask):\n",
    "        enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)\n",
    "        dec_output, attention_weights = self.decoder(tar, enc_output, training, look_ahead_mask, dec_padding_mask)\n",
    "        final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)\n",
    "        return final_output, attention_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Fis_2ldupHgT"
   },
   "source": [
    "#Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MFNeaos5ozYc"
   },
   "outputs": [],
   "source": [
    "num_layers = 6\n",
    "d_model = 2048\n",
    "dff = 1024\n",
    "num_heads = 8\n",
    "input_vocab_size = target_vocab_size = len(tokenizer.word_index)+1\n",
    "dropout_rate = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Vjspglgzur0A"
   },
   "source": [
    "#LR Scheduler, Optimizer and Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "75wLMGnNuZMe"
   },
   "outputs": [],
   "source": [
    "class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
    "    def __init__(self, d_model, warmup_steps=4000):\n",
    "        super(CustomSchedule, self).__init__()\n",
    "\n",
    "        self.d_model = d_model\n",
    "        self.d_model = tf.cast(self.d_model, tf.float32)\n",
    "\n",
    "        self.warmup_steps = warmup_steps\n",
    "\n",
    "    def __call__(self, step):\n",
    "        arg1 = tf.math.rsqrt(step)\n",
    "        arg2 = step * (self.warmup_steps ** -1.5)\n",
    "\n",
    "        return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)\n",
    "\n",
    "learning_rate = CustomSchedule(d_model)\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)\n",
    "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')\n",
    "\n",
    "def loss_function(real, pred):\n",
    "    mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
    "    loss_ = loss_object(real, pred)\n",
    "    mask = tf.cast(mask, dtype=loss_.dtype)\n",
    "    loss_ *= mask\n",
    "    return tf.reduce_mean(loss_)\n",
    "\n",
    "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
    "train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')\n",
    "\n",
    "transformer = Transformer(num_layers, d_model, \n",
    "                          num_heads, dff, \n",
    "                          input_vocab_size, \n",
    "                          target_vocab_size, \n",
    "                          pe_input=input_vocab_size, \n",
    "                          pe_target=target_vocab_size, \n",
    "                          rate=dropout_rate)\n",
    "\n",
    "def create_masks(inp, tar):\n",
    "    enc_padding_mask = create_padding_mask(inp)\n",
    "    dec_padding_mask = create_padding_mask(inp)\n",
    "    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n",
    "    dec_target_padding_mask = create_padding_mask(tar)\n",
    "    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n",
    "    return enc_padding_mask, combined_mask, dec_padding_mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DT_ldlnf0J0O"
   },
   "source": [
    "#Checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fsJO0b8W0KXJ"
   },
   "outputs": [],
   "source": [
    "checkpoint_path = \"./checkpoints/train\"\n",
    "ckpt = tf.train.Checkpoint(transformer=transformer, optimizer=optimizer)\n",
    "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=2)\n",
    "\n",
    "if ckpt_manager.latest_checkpoint:\n",
    "    ckpt.restore(ckpt_manager.latest_checkpoint)\n",
    "    print ('Latest checkpoint restored!!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QhSbO2fl0NHl"
   },
   "source": [
    "#Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "loM8llDp0SZv"
   },
   "outputs": [],
   "source": [
    "def evaluate(inp_sentence):\n",
    "    inp_sentence = [tokenizer.word_index[t] for t in inp_sentence.split(' ') if t!=0]\n",
    "    encoder_input = tf.expand_dims(inp_sentence, 0)\n",
    "\n",
    "    decoder_input = [tokenizer.word_index['<sos>']]\n",
    "    output = tf.expand_dims(decoder_input, 0)\n",
    "        \n",
    "    for i in range(max_length_targ):\n",
    "        enc_padding_mask, combined_mask, dec_padding_mask = create_masks(\n",
    "            encoder_input, output)\n",
    "\n",
    "        predictions, attention_weights = transformer(encoder_input, \n",
    "                                                    output,\n",
    "                                                    False,\n",
    "                                                    enc_padding_mask,\n",
    "                                                    combined_mask,\n",
    "                                                    dec_padding_mask)\n",
    "        \n",
    "        # select the last word from the seq_len dimension\n",
    "        predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)\n",
    "\n",
    "        predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)\n",
    "        \n",
    "        # return the result if the predicted_id is equal to the end token\n",
    "        if predicted_id == tokenizer.word_index['<eos>']:\n",
    "            return tf.squeeze(output, axis=0), attention_weights\n",
    "        \n",
    "        # concatentate the predicted_id to the output which is given to the decoder\n",
    "        # as its input.\n",
    "        output = tf.concat([output, predicted_id], axis=-1)\n",
    "\n",
    "    return tf.squeeze(output, axis=0), attention_weights\n",
    "\n",
    "def plot_attention_weights(attention, sentence, result, layer):\n",
    "    fig = plt.figure(figsize=(16, 8))\n",
    "    sentence = [tokenizer.word_index[t] for t in sentence.split(' ') if t!=0]\n",
    "    attention = tf.squeeze(attention[layer], axis=0)\n",
    "    \n",
    "    for head in range(attention.shape[0]):\n",
    "        ax = fig.add_subplot(2, 4, head+1)\n",
    "        ax.matshow(attention[head][:-1, :], cmap='viridis')\n",
    "        fontdict = {'fontsize': 10}\n",
    "        ax.set_xticks(range(len(sentence)+2))\n",
    "        ax.set_yticks(range(len(result)))\n",
    "        ax.set_ylim(len(result)-1.5, -0.5)\n",
    "        ax.set_xticklabels(\n",
    "            [tokenizer.index_word[i] for i in sentence if i!=0], \n",
    "            fontdict=fontdict, rotation=90)\n",
    "        ax.set_yticklabels([tokenizer.index_word[i.numpy()] for i in result if i < target_vocab_size and i!=0], \n",
    "                        fontdict=fontdict)\n",
    "        ax.set_xlabel('Head {}'.format(head+1))\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "def generate(concepts, plot=''):\n",
    "    result, attention_weights = evaluate(concepts)\n",
    "    predicted_sentence = [tokenizer.index_word[t.numpy()] for t in result if t < target_vocab_size and t!=0]\n",
    "    print('Input: {}'.format(concepts))\n",
    "    print('Predicted sentence: {}'.format(predicted_sentence))\n",
    "    \n",
    "    if plot:\n",
    "        plot_attention_weights(attention_weights, concepts, result, plot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ADOODUbBumKG"
   },
   "source": [
    "#Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5ijlohbh0pWI"
   },
   "source": [
    "##Step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ka5gR9dEufV3"
   },
   "outputs": [],
   "source": [
    "EPOCHS = 40\n",
    "\n",
    "train_step_signature = [\n",
    "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
    "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
    "]\n",
    "\n",
    "@tf.function(input_signature=train_step_signature)\n",
    "def train_step(inp, tar):\n",
    "    tar_inp = tar[:, :-1]\n",
    "    tar_real = tar[:, 1:]\n",
    "    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)\n",
    "    with tf.GradientTape() as tape:\n",
    "        predictions, _ = transformer(inp, tar_inp, \n",
    "                                    True, \n",
    "                                    enc_padding_mask, \n",
    "                                    combined_mask, \n",
    "                                    dec_padding_mask)\n",
    "        loss = loss_function(tar_real, predictions)\n",
    "\n",
    "    gradients = tape.gradient(loss, transformer.trainable_variables)    \n",
    "    optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))\n",
    "    \n",
    "    train_loss(loss)\n",
    "    train_accuracy(tar_real, predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "T4AZl-jyuij5"
   },
   "source": [
    "##Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 500459,
     "status": "error",
     "timestamp": 1582719384093,
     "user": {
      "displayName": "Davide Montanari",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mBFwK0wQ4sGlm8L2VjR5w4UcXGE1Ul2vGoHiuLkMQ=s64",
      "userId": "12602812058261277045"
     },
     "user_tz": -60
    },
    "id": "4cMSq_qruh82",
    "outputId": "e1d44672-7e96-46a8-cf8b-21ded3df5408"
   },
   "source": [
    "train_loss_vals = []\n",
    "train_acc_vals = []\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    start = time.time()\n",
    "    train_loss.reset_states()\n",
    "    train_accuracy.reset_states()\n",
    "\n",
    "    for (batch, (inp, tar)) in enumerate(train_dataset):\n",
    "        train_step(inp, tar)\n",
    "\n",
    "        if batch == 0:\n",
    "            print('---------------------------------------')\n",
    "            print ('Epoch {} Loss {:.2f} Accuracy {:.1f}'.format(epoch + 1, train_loss.result(), train_accuracy.result()*100))\n",
    "            print('---------------------------------------\\n')\n",
    "            train_loss_vals.append(train_loss.result().numpy())\n",
    "            train_acc_vals.append(train_accuracy.result().numpy()*100)\n",
    "            generate(' '.join([tokenizer.index_word[i.numpy()] for i in inp[0] if i!=0]), plot='decoder_layer6_block2')\n",
    "            print (\"Gold sentence: {}\\n\".format(' '.join([tokenizer.index_word[i.numpy()] for i in tar[0] if i!=0])))\n",
    "        \n",
    "        wandb.log({'Train Accuracy': train_accuracy.result()*100, 'Train Loss': train_loss.result()})\n",
    "\n",
    "    epochs = range(1, len(train_loss_vals) + 1)\n",
    "\n",
    "    if (epoch + 1) % 5 == 0:\n",
    "        ckpt_save_path = ckpt_manager.save()\n",
    "        print ('Saving checkpoint for epoch {} at {}'.format(epoch+1, ckpt_save_path))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyON6AADqEORbjlaZnCH9rjj",
   "collapsed_sections": [
    "smtGcCdPSERm",
    "HfV9HpabT1SD",
    "0F4GlrEhT5nr",
    "tdCeUDAiT6qM",
    "6rBf8Ohl0Z4E",
    "Fis_2ldupHgT",
    "Vjspglgzur0A",
    "DT_ldlnf0J0O",
    "QhSbO2fl0NHl",
    "5ijlohbh0pWI"
   ],
   "name": "transformer_rawtext.ipynb",
   "provenance": []
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}