Local Development.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def attention(self, query, key, value):\n",
    "    score = tf.matmul(query, key, transpose_b=True)\n",
    "    dim_key = tf.cast(tf.shape(key)[-1], tf.float32)\n",
    "    scaled_score = score / tf.math.sqrt(dim_key)\n",
    "    weights = tf.nn.softmax(scaled_score, axis=-1)\n",
    "    output = tf.matmul(weights, value)\n",
    "    return output, weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = tf.convert_to_tensor(np.random.rand(8,64,128), dtype=tf.float32)\n",
    "key = tf.convert_to_tensor(np.random.rand(8,64,128), dtype=tf.float32)\n",
    "value = tf.convert_to_tensor(np.random.rand(8,64,128), dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([8, 64, 128]),\n",
       " TensorShape([8, 64, 128]),\n",
       " TensorShape([8, 64, 128]))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query.shape, key.shape, value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 64])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "score = tf.matmul(query, key, transpose_b=True)\n",
    "score.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "128.0"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dim_key = tf.cast(tf.shape(key)[-1], tf.float32)\n",
    "dim_key.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 64])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scaled_score = score / tf.math.sqrt(dim_key)\n",
    "scaled_score.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 64])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights = tf.nn.softmax(scaled_score, axis=-1)\n",
    "weights.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 128])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output = tf.matmul(weights, value)\n",
    "output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([8, 64, 128]), TensorShape([8, 64, 64]))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.shape, weights.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## separate_heads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def separate_heads(self, x, batch_size):\n",
    "        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.projection_dim))\n",
    "        return tf.transpose(x, perm=[0, 2, 1, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.convert_to_tensor(np.random.rand(8,64,128), dtype=tf.float32)\n",
    "batch_size = 8\n",
    "num_heads = 8\n",
    "projection_dim = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 128])"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 8, 16])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = tf.reshape(x, (batch_size, -1, num_heads, projection_dim))\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 8, 64, 16])"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.transpose(x, perm=[0, 2, 1, 3]).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MultiHeadSelfAttention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def separate_heads(x, batch_size):\n",
    "        x = tf.reshape(x, (batch_size, -1, num_heads, projection_dim))\n",
    "        return tf.transpose(x, perm=[0, 2, 1, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "def attention(query, key, value):\n",
    "    score = tf.matmul(query, key, transpose_b=True)\n",
    "    dim_key = tf.cast(tf.shape(key)[-1], tf.float32)\n",
    "    scaled_score = score / tf.math.sqrt(dim_key)\n",
    "    weights = tf.nn.softmax(scaled_score, axis=-1)\n",
    "    output = tf.matmul(weights, value)\n",
    "    return output, weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_dim = 128\n",
    "num_heads = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "16"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "projection_dim = embed_dim // num_heads\n",
    "projection_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_dense = layers.Dense(embed_dim)\n",
    "key_dense = layers.Dense(embed_dim)\n",
    "value_dense = layers.Dense(embed_dim)\n",
    "combine_heads = layers.Dense(embed_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tf.convert_to_tensor(np.random.rand(8,64,128), dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = tf.shape(inputs)[0]\n",
    "query = query_dense(inputs) # (batch_size, seq_len, embed_dim)\n",
    "key = key_dense(inputs)  # (batch_size, seq_len, embed_dim)\n",
    "value = value_dense(inputs) # (batch_size, seq_len, embed_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([8, 64, 128]),\n",
       " TensorShape([8, 64, 128]),\n",
       " TensorShape([8, 64, 128]))"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query.shape, key.shape, value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = separate_heads(query, batch_size)  # (batch_size, num_heads, seq_len, projection_dim)\n",
    "key = separate_heads(key, batch_size)  # (batch_size, num_heads, seq_len, projection_dim)\n",
    "value = separate_heads(value, batch_size)  # (batch_size, num_heads, seq_len, projection_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([8, 8, 64, 16]),\n",
       " TensorShape([8, 8, 64, 16]),\n",
       " TensorShape([8, 8, 64, 16]))"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query.shape, key.shape, value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "attention, weights = attention(query, key, value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([8, 8, 64, 16]), TensorShape([8, 8, 64, 64]))"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attention.shape, weights.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 8, 16])"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attention = tf.transpose(attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len, num_heads, projection_dim)\n",
    "attention.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 128])"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "concat_attention = tf.reshape(attention, (batch_size, -1, embed_dim))  # (batch_size, seq_len, embed_dim)\n",
    "concat_attention.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 128])"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output = combine_heads(concat_attention)  # (batch_size, seq_len, embed_dim)\n",
    "output.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TransformerBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "ff_dim = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "ffn = keras.Sequential([layers.Dense(ff_dim, activation=\"relu\"), layers.Dense(embed_dim),])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "rate=0.1\n",
    "layernorm1 = layers.LayerNormalization(epsilon=1e-6)\n",
    "layernorm2 = layers.LayerNormalization(epsilon=1e-6)\n",
    "dropout1 = layers.Dropout(rate)\n",
    "dropout2 = layers.Dropout(rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([8, 64, 128]), TensorShape([8, 64, 128]))"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attn_output = tf.convert_to_tensor(np.random.rand(8,64,128), dtype=tf.float32)\n",
    "inputs = tf.convert_to_tensor(np.random.rand(8,64,128), dtype=tf.float32)\n",
    "attn_output.shape, inputs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_output = dropout1(attn_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 128])"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out1 = layernorm1(inputs + attn_output)\n",
    "out1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 128])"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ffn_output = ffn(out1)\n",
    "ffn_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "ffn_output = dropout2(ffn_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([8, 64, 128])"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_final = layernorm2(out1 + ffn_output)\n",
    "output_final.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main Function Development"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.insert(0,'./model_X')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers_utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_dim = 32  # Embedding size for each token\n",
    "num_heads = 2  # Number of attention heads\n",
    "ff_dim = 32  # Hidden layer size in feed forward network inside transformer\n",
    "maxlen = 128\n",
    "vocab_size = 1000\n",
    "inputs = tf.keras.layers.Input(shape=(maxlen,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer_block = TransformerBlock(embed_dim, num_heads, ff_dim)(embedding_layer, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor 'layer_normalization_3/Identity:0' shape=(None, 128, 32) dtype=float32>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transformer_block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor 'global_average_pooling1d/Identity:0' shape=(None, 32) dtype=float32>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = tf.keras.layers.GlobalAveragePooling1D()(transformer_block)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.keras.layers.Dropout(0.1)(x)\n",
    "x = tf.keras.layers.Dense(20, activation=\"relu\")(x)\n",
    "x = tf.keras.layers.Dropout(0.1)(x)\n",
    "outputs = tf.keras.layers.Dense(2, activation=\"softmax\")(x)\n",
    "\n",
    "model = tf.keras.Model(inputs=inputs, outputs=outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_2 (InputLayer)            [(None, 128)]        0                                            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Shape (TensorFlowOp [(2,)]               0           input_2[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_strided_slice (Tens [()]                 0           tf_op_layer_Shape[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_range (TensorFlowOp [(None,)]            0           tf_op_layer_strided_slice[0][0]  \n",
      "__________________________________________________________________________________________________\n",
      "embedding (Embedding)           (None, 128, 32)      32000       input_2[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "embedding_1 (Embedding)         (None, 32)           4096        tf_op_layer_range[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_add (TensorFlowOpLa [(None, 128, 32)]    0           embedding[0][0]                  \n",
      "                                                                 embedding_1[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Shape_1 (TensorFlow [(3,)]               0           tf_op_layer_add[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_strided_slice_1 (Te [()]                 0           tf_op_layer_Shape_1[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "dense_7 (Dense)                 (None, 128, 32)      1056        tf_op_layer_add[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_1/shape (Te [(4,)]               0           tf_op_layer_strided_slice_1[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_1 (TensorFl [(None, None, 2, 16) 0           dense_7[0][0]                    \n",
      "                                                                 tf_op_layer_Reshape_1/shape[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_transpose_1 (Tensor [(None, 2, None, 16) 0           tf_op_layer_Reshape_1[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dense_6 (Dense)                 (None, 128, 32)      1056        tf_op_layer_add[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape/shape (Tens [(4,)]               0           tf_op_layer_strided_slice_1[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Shape_2 (TensorFlow [(4,)]               0           tf_op_layer_transpose_1[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape (TensorFlow [(None, None, 2, 16) 0           dense_6[0][0]                    \n",
      "                                                                 tf_op_layer_Reshape/shape[0][0]  \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_strided_slice_2 (Te [()]                 0           tf_op_layer_Shape_2[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_transpose (TensorFl [(None, 2, None, 16) 0           tf_op_layer_Reshape[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Cast (TensorFlowOpL [()]                 0           tf_op_layer_strided_slice_2[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_MatMul (TensorFlowO [(None, 2, None, Non 0           tf_op_layer_transpose[0][0]      \n",
      "                                                                 tf_op_layer_transpose_1[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Sqrt (TensorFlowOpL [()]                 0           tf_op_layer_Cast[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "dense_8 (Dense)                 (None, 128, 32)      1056        tf_op_layer_add[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_2/shape (Te [(4,)]               0           tf_op_layer_strided_slice_1[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_truediv (TensorFlow [(None, 2, None, Non 0           tf_op_layer_MatMul[0][0]         \n",
      "                                                                 tf_op_layer_Sqrt[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_2 (TensorFl [(None, None, 2, 16) 0           dense_8[0][0]                    \n",
      "                                                                 tf_op_layer_Reshape_2/shape[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Softmax (TensorFlow [(None, 2, None, Non 0           tf_op_layer_truediv[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_transpose_2 (Tensor [(None, 2, None, 16) 0           tf_op_layer_Reshape_2[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_MatMul_1 (TensorFlo [(None, 2, None, 16) 0           tf_op_layer_Softmax[0][0]        \n",
      "                                                                 tf_op_layer_transpose_2[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_transpose_3 (Tensor [(None, None, 2, 16) 0           tf_op_layer_MatMul_1[0][0]       \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_3/shape (Te [(3,)]               0           tf_op_layer_strided_slice_1[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_3 (TensorFl [(None, None, 32)]   0           tf_op_layer_transpose_3[0][0]    \n",
      "                                                                 tf_op_layer_Reshape_3/shape[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "dense_9 (Dense)                 (None, None, 32)     1056        tf_op_layer_Reshape_3[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout_2 (Dropout)             (None, None, 32)     0           dense_9[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_add_1 (TensorFlowOp [(None, 128, 32)]    0           tf_op_layer_add[0][0]            \n",
      "                                                                 dropout_2[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "layer_normalization_2 (LayerNor (None, 128, 32)      64          tf_op_layer_add_1[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "sequential_1 (Sequential)       (None, 128, 32)      2112        layer_normalization_2[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout_3 (Dropout)             (None, 128, 32)      0           sequential_1[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_add_2 (TensorFlowOp [(None, 128, 32)]    0           layer_normalization_2[0][0]      \n",
      "                                                                 dropout_3[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "layer_normalization_3 (LayerNor (None, 128, 32)      64          tf_op_layer_add_2[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "global_average_pooling1d (Globa (None, 32)           0           layer_normalization_3[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout_4 (Dropout)             (None, 32)           0           global_average_pooling1d[0][0]   \n",
      "__________________________________________________________________________________________________\n",
      "dense_12 (Dense)                (None, 20)           660         dropout_4[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dropout_5 (Dropout)             (None, 20)           0           dense_12[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "dense_13 (Dense)                (None, 2)            42          dropout_5[0][0]                  \n",
      "==================================================================================================\n",
      "Total params: 43,262\n",
      "Trainable params: 43,262\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tf.convert_to_tensor(np.random.rand(8,128), dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.5419823 , 0.45801768],\n",
       "       [0.5396268 , 0.46037325],\n",
       "       [0.556035  , 0.44396505],\n",
       "       [0.55087   , 0.44912994],\n",
       "       [0.550646  , 0.44935405],\n",
       "       [0.55667895, 0.44332105],\n",
       "       [0.5523735 , 0.4476265 ],\n",
       "       [0.5457177 , 0.4542823 ]], dtype=float32)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.predict(inputs, batch_size=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "config = AlbertConfig.from_json_file('/Users/ankur.kumar/Downloads/base_2/assets/albert_config.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = modeling_tf_albert.TFAlbertModel(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tf.convert_to_tensor(np.random.rand(64,128), dtype=tf.int32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = model(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([64, 128, 768]), TensorShape([64, 768]))"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output[0].shape, output[1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Positional Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def GetPosEncodingMatrix(max_len, d_emb):\n",
    "\tpos_enc = np.array([\n",
    "\t\t[pos / np.power(10000, 2 * (j // 2) / d_emb) for j in range(d_emb)] \n",
    "\t\tif pos != 0 else np.zeros(d_emb) \n",
    "\t\t\tfor pos in range(max_len)\n",
    "\t\t\t])\n",
    "\tpos_enc[1:, 0::2] = np.sin(pos_enc[1:, 0::2]) # dim 2i\n",
    "\tpos_enc[1:, 1::2] = np.cos(pos_enc[1:, 1::2]) # dim 2i+1\n",
    "\treturn pos_enc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_len = 512\n",
    "d_emb = 128\n",
    "weights=GetPosEncodingMatrix(max_len, d_emb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00, ...,\n",
       "         0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "       [ 8.41470985e-01,  5.40302306e-01,  7.61720408e-01, ...,\n",
       "         9.99999991e-01,  1.15478198e-04,  9.99999993e-01],\n",
       "       [ 9.09297427e-01, -4.16146837e-01,  9.87046251e-01, ...,\n",
       "         9.99999964e-01,  2.30956395e-04,  9.99999973e-01],\n",
       "       ...,\n",
       "       [ 6.19504237e-02,  9.98079228e-01,  8.15081054e-01, ...,\n",
       "         9.97697292e-01,  5.87445633e-02,  9.98273047e-01],\n",
       "       [ 8.73326668e-01,  4.87135024e-01,  9.69396188e-01, ...,\n",
       "         9.97688239e-01,  5.88598417e-02,  9.98266257e-01],\n",
       "       [ 8.81770401e-01, -4.71678874e-01,  4.41073911e-01, ...,\n",
       "         9.97679168e-01,  5.89751193e-02,  9.98259453e-01]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "flag = argparse.Namespace(a='Ankur',\n",
    "                         b='Ankur2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Per-training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.2.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0,'model_X/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Local Development.ipynb  \u001b[1m\u001b[36mdist\u001b[m\u001b[m/                    setup.sh\r\n",
      "Local Testing.ipynb      \u001b[1m\u001b[36mmodel_X\u001b[m\u001b[m/                 test_modelX.py\r\n",
      "README.md                \u001b[1m\u001b[36mmodel_X.egg-info\u001b[m\u001b[m/        \u001b[1m\u001b[36mtmp\u001b[m\u001b[m/\r\n",
      "\u001b[1m\u001b[36mbuild\u001b[m\u001b[m/                   setup.py\r\n"
     ]
    }
   ],
   "source": [
    "ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pretraining_layers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers_architectures import *\n",
    "from tensorflow.keras.layers import *\n",
    "from tensorflow.keras.models import Model\n",
    "import argparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = argparse.Namespace(vocab_size=1000,\n",
    "                            embed_dim=512,\n",
    "                            ff_dim=32,\n",
    "                            num_heads=8,\n",
    "                            rate=0.1,\n",
    "                            maxlen=128,\n",
    "                           initializer_range=0.2,\n",
    "                           embedding_size=512,\n",
    "                           max_predictions_per_seq=20,\n",
    "                           hidden_act='relu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tf.keras.layers.Input(shape=(config.maxlen,), dtype=tf.int32)\n",
    "pooled_output,sequence_output = VanillaTransformer(config)(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "masked_lm_positions = tf.keras.layers.Input(\n",
    "      shape=(config.max_predictions_per_seq,),\n",
    "      name='masked_lm_positions',\n",
    "      dtype=tf.int32)\n",
    "masked_lm_weights = tf.keras.layers.Input(\n",
    "      shape=(config.max_predictions_per_seq,),\n",
    "      name='masked_lm_weights',\n",
    "      dtype=tf.int32)\n",
    "masked_lm_ids = tf.keras.layers.Input(\n",
    "      shape=(config.max_predictions_per_seq,), name='masked_lm_ids', dtype=tf.int32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor 'pooler_transform/Identity:0' shape=(None, 512) dtype=float32>,\n",
       " <tf.Tensor 'layer_normalization_1/Identity:0' shape=(None, 128, 512) dtype=float32>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pooled_output, sequence_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer_submodel = tf.keras.Model(\n",
    "      inputs=inputs,\n",
    "      outputs=[pooled_output, sequence_output])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor 'GatherV2:0' shape=(None, 512) dtype=float32>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask_lm_input_tensor = gather_indexes(sequence_output, masked_lm_positions)\n",
    "mask_lm_input_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor 'dense_6/Identity:0' shape=(None, 512) dtype=float32>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm_output = tf.keras.layers.Dense(config.embedding_size)(mask_lm_input_tensor)\n",
    "lm_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor 'MatMul:0' shape=(None, 1000) dtype=float32>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_table = transformer_submodel.get_layer('embedding').embeddings\n",
    "lm_output = tf.matmul(lm_output, embedding_table, transpose_b=True)\n",
    "lm_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor 'LogSoftmax:0' shape=(None, 1000) dtype=float32>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm_output = tf.nn.log_softmax(lm_output, axis=-1)\n",
    "lm_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretraining_model = tf.keras.Model([inputs, masked_lm_positions], lm_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_1\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_1 (InputLayer)            [(None, 128)]        0                                            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Shape (TensorFlowOp [(2,)]               0           input_1[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_strided_slice (Tens [()]                 0           tf_op_layer_Shape[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Range (TensorFlowOp [(None,)]            0           tf_op_layer_strided_slice[0][0]  \n",
      "__________________________________________________________________________________________________\n",
      "embedding (Embedding)           (None, 128, 512)     512000      input_1[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "embedding_1 (Embedding)         (None, 512)          65536       tf_op_layer_Range[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_AddV2 (TensorFlowOp [(None, 128, 512)]   0           embedding[0][0]                  \n",
      "                                                                 embedding_1[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Shape_1 (TensorFlow [(3,)]               0           tf_op_layer_AddV2[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_strided_slice_1 (Te [()]                 0           tf_op_layer_Shape_1[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "dense_1 (Dense)                 (None, 128, 512)     262656      tf_op_layer_AddV2[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_1/shape (Te [(4,)]               0           tf_op_layer_strided_slice_1[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_1 (TensorFl [(None, None, 8, 64) 0           dense_1[0][0]                    \n",
      "                                                                 tf_op_layer_Reshape_1/shape[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Transpose_1 (Tensor [(None, 8, None, 64) 0           tf_op_layer_Reshape_1[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dense (Dense)                   (None, 128, 512)     262656      tf_op_layer_AddV2[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape/shape (Tens [(4,)]               0           tf_op_layer_strided_slice_1[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Shape_2 (TensorFlow [(4,)]               0           tf_op_layer_Transpose_1[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape (TensorFlow [(None, None, 8, 64) 0           dense[0][0]                      \n",
      "                                                                 tf_op_layer_Reshape/shape[0][0]  \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_strided_slice_2 (Te [()]                 0           tf_op_layer_Shape_2[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Transpose (TensorFl [(None, 8, None, 64) 0           tf_op_layer_Reshape[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Cast (TensorFlowOpL [()]                 0           tf_op_layer_strided_slice_2[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_BatchMatMulV2 (Tens [(None, 8, None, Non 0           tf_op_layer_Transpose[0][0]      \n",
      "                                                                 tf_op_layer_Transpose_1[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Sqrt (TensorFlowOpL [()]                 0           tf_op_layer_Cast[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "dense_2 (Dense)                 (None, 128, 512)     262656      tf_op_layer_AddV2[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_2/shape (Te [(4,)]               0           tf_op_layer_strided_slice_1[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_RealDiv (TensorFlow [(None, 8, None, Non 0           tf_op_layer_BatchMatMulV2[0][0]  \n",
      "                                                                 tf_op_layer_Sqrt[0][0]           \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_2 (TensorFl [(None, None, 8, 64) 0           dense_2[0][0]                    \n",
      "                                                                 tf_op_layer_Reshape_2/shape[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Softmax (TensorFlow [(None, 8, None, Non 0           tf_op_layer_RealDiv[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Transpose_2 (Tensor [(None, 8, None, 64) 0           tf_op_layer_Reshape_2[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_BatchMatMulV2_1 (Te [(None, 8, None, 64) 0           tf_op_layer_Softmax[0][0]        \n",
      "                                                                 tf_op_layer_Transpose_2[0][0]    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Transpose_3 (Tensor [(None, None, 8, 64) 0           tf_op_layer_BatchMatMulV2_1[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_3/shape (Te [(3,)]               0           tf_op_layer_strided_slice_1[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_3 (TensorFl [(None, None, 512)]  0           tf_op_layer_Transpose_3[0][0]    \n",
      "                                                                 tf_op_layer_Reshape_3/shape[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "dense_3 (Dense)                 (None, None, 512)    262656      tf_op_layer_Reshape_3[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout (Dropout)               (None, None, 512)    0           dense_3[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_AddV2_1 (TensorFlow [(None, 128, 512)]   0           tf_op_layer_AddV2[0][0]          \n",
      "                                                                 dropout[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "layer_normalization (LayerNorma (None, 128, 512)     1024        tf_op_layer_AddV2_1[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "sequential (Sequential)         (None, 128, 512)     33312       layer_normalization[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "dropout_1 (Dropout)             (None, 128, 512)     0           sequential[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_AddV2_2 (TensorFlow [(None, 128, 512)]   0           layer_normalization[0][0]        \n",
      "                                                                 dropout_1[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "layer_normalization_1 (LayerNor (None, 128, 512)     1024        tf_op_layer_AddV2_2[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Shape_3 (TensorFlow [(3,)]               0           layer_normalization_1[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_strided_slice_4 (Te [()]                 0           tf_op_layer_Shape_3[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Range_1 (TensorFlow [(None,)]            0           tf_op_layer_strided_slice_4[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Mul (TensorFlowOpLa [(None,)]            0           tf_op_layer_Range_1[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Mul_1 (TensorFlowOp [()]                 0           tf_op_layer_strided_slice_4[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "masked_lm_positions (InputLayer [(None, 20)]         0                                            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_4 (TensorFl [(None, 1)]          0           tf_op_layer_Mul[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_6/shape (Te [(2,)]               0           tf_op_layer_Mul_1[0][0]          \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_AddV2_3 (TensorFlow [(None, 20)]         0           masked_lm_positions[0][0]        \n",
      "                                                                 tf_op_layer_Reshape_4[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_6 (TensorFl [(None, 512)]        0           layer_normalization_1[0][0]      \n",
      "                                                                 tf_op_layer_Reshape_6/shape[0][0]\n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Reshape_5 (TensorFl [(None,)]            0           tf_op_layer_AddV2_3[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_GatherV2 (TensorFlo [(None, 512)]        0           tf_op_layer_Reshape_6[0][0]      \n",
      "                                                                 tf_op_layer_Reshape_5[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dense_6 (Dense)                 (None, 512)          262656      tf_op_layer_GatherV2[0][0]       \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_MatMul (TensorFlowO [(None, 1000)]       0           dense_6[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_LogSoftmax (TensorF [(None, 1000)]       0           tf_op_layer_MatMul[0][0]         \n",
      "==================================================================================================\n",
      "Total params: 1,926,176\n",
      "Trainable params: 1,348,640\n",
      "Non-trainable params: 577,536\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "pretraining_model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(lm_label_ids, lm_output):\n",
    "    lm_label_ids = tf.keras.backend.reshape(lm_label_ids, [-1])\n",
    "    lm_label_ids = tf.keras.backend.cast(lm_label_ids, dtype=tf.int32)\n",
    "    lm_label_ids_one_hot = tf.keras.backend.one_hot(lm_label_ids,1000)\n",
    "    lm_per_example_loss = -tf.keras.backend.sum(lm_output * lm_label_ids_one_hot, axis=[-1])\n",
    "    lm_example_loss = tf.reshape(lm_per_example_loss, [-1])\n",
    "    return tf.reduce_mean(lm_example_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretraining_model.compile(optimizer='adam',loss=loss_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = {\n",
    "  'input_1': tf.convert_to_tensor(np.random.randint(1,100,(64,128)), tf.int32),\n",
    "  'masked_lm_positions': tf.convert_to_tensor(np.random.randint(1,100,(64,20)),tf.int32)}\n",
    "y = {\n",
    "  'tf_op_layer_LogSoftmax': tf.convert_to_tensor(np.random.randint(1,1000,(64,20)),tf.int32),\n",
    "  'masked_lm_weights': tf.convert_to_tensor(np.random.randint(1,1000,(64,20)),tf.int32)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 64 samples\n",
      "Epoch 1/3\n",
      "64/64 [==============================] - 0s 7ms/sample - loss: 6.9078\n",
      "Epoch 2/3\n",
      "64/64 [==============================] - 0s 7ms/sample - loss: 6.9078\n",
      "Epoch 3/3\n",
      "64/64 [==============================] - 0s 8ms/sample - loss: 6.9078 0s - loss: 6.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x14839f1d0>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pretraining_model.fit(X,y, batch_size=8, epochs=3, use_multiprocessing=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 2:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(lm_label_ids, lm_label_weights, lm_output):\n",
    "    lm_label_ids = tf.keras.backend.reshape(lm_label_ids, [-1])\n",
    "    lm_label_ids = tf.keras.backend.cast(lm_label_ids, dtype=tf.int32)\n",
    "    lm_label_ids_one_hot = tf.keras.backend.one_hot(lm_label_ids,1000)\n",
    "    lm_per_example_loss = -tf.keras.backend.sum(lm_output * lm_label_ids_one_hot, axis=[-1])\n",
    "    lm_example_loss = tf.reshape(lm_per_example_loss, [-1])\n",
    "    \n",
    "    lm_label_weights = tf.keras.backend.cast(lm_label_weights, tf.float32)\n",
    "    lm_label_weights = tf.keras.backend.reshape(lm_label_weights, [-1])\n",
    "    \n",
    "    return tf.reduce_mean(lm_example_loss * lm_label_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = tf.data.Dataset.from_tensor_slices(((X['input_1'],X['masked_lm_positions']), (y['tf_op_layer_LogSoftmax'],y['masked_lm_weights'])))\n",
    "train_dataset = train_dataset.shuffle(buffer_size=10).batch(8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.keras.optimizers.Adam()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Start of epoch 0\n",
      "tf.Tensor(3524.2505, shape=(), dtype=float32)\n",
      "Training loss (for one batch) at step 0: 3524.2505\n",
      "Seen so far: 8 samples\n",
      "tf.Tensor(3410.2292, shape=(), dtype=float32)\n",
      "tf.Tensor(3547.8665, shape=(), dtype=float32)\n",
      "Training loss (for one batch) at step 2: 3547.8665\n",
      "Seen so far: 24 samples\n",
      "tf.Tensor(3537.9368, shape=(), dtype=float32)\n",
      "tf.Tensor(3668.4062, shape=(), dtype=float32)\n",
      "Training loss (for one batch) at step 4: 3668.4062\n",
      "Seen so far: 40 samples\n",
      "tf.Tensor(3395.939, shape=(), dtype=float32)\n",
      "tf.Tensor(3257.4817, shape=(), dtype=float32)\n",
      "Training loss (for one batch) at step 6: 3257.4817\n",
      "Seen so far: 56 samples\n",
      "tf.Tensor(3323.4507, shape=(), dtype=float32)\n",
      "\n",
      "Start of epoch 1\n",
      "tf.Tensor(3603.69, shape=(), dtype=float32)\n",
      "Training loss (for one batch) at step 0: 3603.6899\n",
      "Seen so far: 8 samples\n",
      "tf.Tensor(3347.3254, shape=(), dtype=float32)\n",
      "tf.Tensor(3446.4949, shape=(), dtype=float32)\n",
      "Training loss (for one batch) at step 2: 3446.4949\n",
      "Seen so far: 24 samples\n",
      "tf.Tensor(3693.102, shape=(), dtype=float32)\n",
      "tf.Tensor(3278.1616, shape=(), dtype=float32)\n",
      "Training loss (for one batch) at step 4: 3278.1616\n",
      "Seen so far: 40 samples\n",
      "tf.Tensor(3290.6821, shape=(), dtype=float32)\n",
      "tf.Tensor(3818.607, shape=(), dtype=float32)\n",
      "Training loss (for one batch) at step 6: 3818.6069\n",
      "Seen so far: 56 samples\n",
      "tf.Tensor(3187.4973, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "epochs = 2\n",
    "for epoch in range(epochs):\n",
    "    print(\"\\nStart of epoch %d\" % (epoch,))\n",
    "\n",
    "    # Iterate over the batches of the dataset.\n",
    "    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):\n",
    "\n",
    "        # Open a GradientTape to record the operations run\n",
    "        # during the forward pass, which enables autodifferentiation.\n",
    "        with tf.GradientTape() as tape:\n",
    "\n",
    "            # Run the forward pass of the layer.\n",
    "            # The operations that the layer applies\n",
    "            # to its inputs are going to be recorded\n",
    "            # on the GradientTape.\n",
    "            logits = pretraining_model(x_batch_train, training=True)  # Logits for this minibatch\n",
    "\n",
    "            # Compute the loss value for this minibatch.\n",
    "            loss_value = loss_fn(y_batch_train[0],y_batch_train[1], logits)\n",
    "\n",
    "        # Use the gradient tape to automatically retrieve\n",
    "        # the gradients of the trainable variables with respect to the loss.\n",
    "        print(loss_value)\n",
    "        grads = tape.gradient(loss_value, pretraining_model.trainable_weights)\n",
    "\n",
    "        # Run one step of gradient descent by updating\n",
    "        # the value of the variables to minimize the loss.\n",
    "        optimizer.apply_gradients(zip(grads, pretraining_model.trainable_weights))\n",
    "\n",
    "        # Log every 200 batches.\n",
    "        if step % 2 == 0:\n",
    "            print(\n",
    "                \"Training loss (for one batch) at step %d: %.4f\"\n",
    "                % (step, float(loss_value))\n",
    "            )\n",
    "            print(\"Seen so far: %s samples\" % ((step + 1) * 8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.5 64-bit ('tf2_env': venv)",
   "language": "python",
   "name": "python37564bittf2envvenv9baad5f63bf544adb456767f347f0830"
  },
  "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.7.5"
  },
  "nteract": {
   "version": "nteract-on-jupyter@2.1.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}