{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Same parts are taken/inspired from:\nhttps://www.kaggle.com/code/markwijkhuizen/aslfr-transformer-training-inference\nhttps://www.kaggle.com/code/shlomoron/aslfr-a-simple-transformer","metadata":{}},{"cell_type":"code","source":"import gc\nimport os\nimport json\nimport math\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\n\nimport tensorflow as tf\nimport tensorflow_addons as tfa","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:41:45.065196Z","iopub.execute_input":"2023-07-17T17:41:45.06592Z","iopub.status.idle":"2023-07-17T17:41:45.072428Z","shell.execute_reply.started":"2023-07-17T17:41:45.065884Z","shell.execute_reply":"2023-07-17T17:41:45.071378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open (\"/kaggle/input/asl-fingerspelling/character_to_prediction_index.json\", \"r\") as f:\n    char_to_num = json.load(f)\n\npad_token = '^'\nstart_token = '<'\nend_token = '>'\npad_token_idx = 59\nstart_token_idx = 60\nend_token_idx = 61\n\nchar_to_num[pad_token] = pad_token_idx\nchar_to_num[start_token] = start_token_idx\nchar_to_num[end_token] = end_token_idx\n\nnum_to_char = {j:i for i,j in char_to_num.items()}\n\ndf = pd.read_csv(\"/kaggle/input/asl-fingerspelling/train.csv\")\n\nLIP = [\n    61, 185, 40, 39, 37, 0, 267, 269, 270, 409,\n    291, 146, 91, 181, 84, 17, 314, 405, 321, 375,\n    78, 191, 80, 81, 82, 13, 312, 311, 310, 415,\n    95, 88, 178, 87, 14, 317, 402, 318, 324, 308,\n]\nLPOSE = [13, 15, 17, 19, 21]\nRPOSE = [14, 16, 18, 20, 22]\nPOSE = LPOSE + RPOSE\n\nX = [f'x_right_hand_{i}' for i in range(21)] + [f'x_left_hand_{i}' for i in range(21)] + [f'x_pose_{i}' for i in POSE] + [f'x_face_{i}' for i in LIP]\nY = [f'y_right_hand_{i}' for i in range(21)] + [f'y_left_hand_{i}' for i in range(21)] + [f'y_pose_{i}' for i in POSE] + [f'y_face_{i}' for i in LIP]\nZ = [f'z_right_hand_{i}' for i in range(21)] + [f'z_left_hand_{i}' for i in range(21)] + [f'z_pose_{i}' for i in POSE] + [f'z_face_{i}' for i in LIP]\n\nSEL_COLS = X + Y + Z\nFRAME_LEN = 150\nMAX_PHRASE_LENGTH = 64\n\nLIP_IDX_X   = [i for i, col in enumerate(SEL_COLS)  if  \"face\" in col and \"x\" in col]\nRHAND_IDX_X = [i for i, col in enumerate(SEL_COLS)  if \"right\" in col and \"x\" in col]\nLHAND_IDX_X = [i for i, col in enumerate(SEL_COLS)  if  \"left\" in col and \"x\" in col]\nRPOSE_IDX_X = [i for i, col in enumerate(SEL_COLS)  if  \"pose\" in col and int(col[-2:]) in RPOSE and \"x\" in col]\nLPOSE_IDX_X = [i for i, col in enumerate(SEL_COLS)  if  \"pose\" in col and int(col[-2:]) in LPOSE and \"x\" in col]\n\nLIP_IDX_Y   = [i for i, col in enumerate(SEL_COLS)  if  \"face\" in col and \"y\" in col]\nRHAND_IDX_Y = [i for i, col in enumerate(SEL_COLS)  if \"right\" in col and \"y\" in col]\nLHAND_IDX_Y = [i for i, col in enumerate(SEL_COLS)  if  \"left\" in col and \"y\" in col]\nRPOSE_IDX_Y = [i for i, col in enumerate(SEL_COLS)  if  \"pose\" in col and int(col[-2:]) in RPOSE and \"y\" in col]\nLPOSE_IDX_Y = [i for i, col in enumerate(SEL_COLS)  if  \"pose\" in col and int(col[-2:]) in LPOSE and \"y\" in col]\n\nLIP_IDX_Z   = [i for i, col in enumerate(SEL_COLS)  if  \"face\" in col and \"z\" in col]\nRHAND_IDX_Z = [i for i, col in enumerate(SEL_COLS)  if \"right\" in col and \"z\" in col]\nLHAND_IDX_Z = [i for i, col in enumerate(SEL_COLS)  if  \"left\" in col and \"z\" in col]\nRPOSE_IDX_Z = [i for i, col in enumerate(SEL_COLS)  if  \"pose\" in col and int(col[-2:]) in RPOSE and \"z\" in col]\nLPOSE_IDX_Z = [i for i, col in enumerate(SEL_COLS)  if  \"pose\" in col and int(col[-2:]) in LPOSE and \"z\" in col]\n\nmean_std_dir = \"/kaggle/input/aslfr-dataset-tfrecords/mean_std\"\nRHM = np.load(f\"{mean_std_dir}/rh_mean.npy\")\nLHM = np.load(f\"{mean_std_dir}/lh_mean.npy\")\nRPM = np.load(f\"{mean_std_dir}/rp_mean.npy\")\nLPM = np.load(f\"{mean_std_dir}/lp_mean.npy\")\nLIPM = np.load(f\"{mean_std_dir}/lip_mean.npy\")\n\nRHS = np.load(f\"{mean_std_dir}/rh_std.npy\")\nLHS = np.load(f\"{mean_std_dir}/lh_std.npy\")\nRPS = np.load(f\"{mean_std_dir}/rp_std.npy\")\nLPS = np.load(f\"{mean_std_dir}/lp_std.npy\")\nLIPS = np.load(f\"{mean_std_dir}/lip_std.npy\")","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:41:47.933099Z","iopub.execute_input":"2023-07-17T17:41:47.935572Z","iopub.status.idle":"2023-07-17T17:41:48.166817Z","shell.execute_reply.started":"2023-07-17T17:41:47.935537Z","shell.execute_reply":"2023-07-17T17:41:48.165735Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_relevant_data_subset(pq_path):\n    return pd.read_parquet(pq_path, columns=SEL_COLS)\n\nfile_id = df.file_id.iloc[0]\ninpdir = \"/kaggle/input/asl-fingerspelling/train_landmarks\"\npqfile = f\"{inpdir}/{file_id}.parquet\"\nseq_refs = df.loc[df.file_id == file_id]\nseqs = load_relevant_data_subset(pqfile)\n\nseq_id = seq_refs.sequence_id.iloc[0]\nframes = seqs.iloc[seqs.index == seq_id]\nphrase = str(df.loc[df.sequence_id == seq_id].phrase.iloc[0])","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:41:51.672193Z","iopub.execute_input":"2023-07-17T17:41:51.672903Z","iopub.status.idle":"2023-07-17T17:41:54.641124Z","shell.execute_reply.started":"2023-07-17T17:41:51.67287Z","shell.execute_reply":"2023-07-17T17:41:54.640014Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@tf.function()\ndef resize_pad(x):\n    if tf.shape(x)[0] < FRAME_LEN:\n        x = tf.pad(x, ([[0, FRAME_LEN-tf.shape(x)[0]], [0, 0], [0, 0]]), constant_values=float('NaN'))\n    else:\n        x = tf.image.resize(x, (FRAME_LEN, tf.shape(x)[1]))\n    return x\n\n@tf.function(jit_compile=True)\ndef pre_process0(x):\n    lip_x = tf.gather(x, LIP_IDX_X, axis=1)\n    lip_y = tf.gather(x, LIP_IDX_Y, axis=1)\n    lip_z = tf.gather(x, LIP_IDX_Z, axis=1)\n\n    rhand_x = tf.gather(x, RHAND_IDX_X, axis=1)\n    rhand_y = tf.gather(x, RHAND_IDX_Y, axis=1)\n    rhand_z = tf.gather(x, RHAND_IDX_Z, axis=1)\n    \n    lhand_x = tf.gather(x, LHAND_IDX_X, axis=1)\n    lhand_y = tf.gather(x, LHAND_IDX_Y, axis=1)\n    lhand_z = tf.gather(x, LHAND_IDX_Z, axis=1)\n\n    rpose_x = tf.gather(x, RPOSE_IDX_X, axis=1)\n    rpose_y = tf.gather(x, RPOSE_IDX_Y, axis=1)\n    rpose_z = tf.gather(x, RPOSE_IDX_Z, axis=1)\n    \n    lpose_x = tf.gather(x, LPOSE_IDX_X, axis=1)\n    lpose_y = tf.gather(x, LPOSE_IDX_Y, axis=1)\n    lpose_z = tf.gather(x, LPOSE_IDX_Z, axis=1)\n    \n    lip   = tf.concat([lip_x[..., tf.newaxis], lip_y[..., tf.newaxis], lip_z[..., tf.newaxis]], axis=-1)\n    rhand = tf.concat([rhand_x[..., tf.newaxis], rhand_y[..., tf.newaxis], rhand_z[..., tf.newaxis]], axis=-1)\n    lhand = tf.concat([lhand_x[..., tf.newaxis], lhand_y[..., tf.newaxis], lhand_z[..., tf.newaxis]], axis=-1)\n    rpose = tf.concat([rpose_x[..., tf.newaxis], rpose_y[..., tf.newaxis], rpose_z[..., tf.newaxis]], axis=-1)\n    lpose = tf.concat([lpose_x[..., tf.newaxis], lpose_y[..., tf.newaxis], lpose_z[..., tf.newaxis]], axis=-1)\n    \n    hand = tf.concat([rhand, lhand], axis=1)\n    hand = tf.where(tf.math.is_nan(hand), 0.0, hand)\n    mask = tf.math.not_equal(tf.reduce_sum(hand, axis=[1, 2]), 0.0)\n\n    lip = lip[mask]\n    rhand = rhand[mask]\n    lhand = lhand[mask]\n    rpose = rpose[mask]\n    lpose = lpose[mask]\n\n    return lip, rhand, lhand, rpose, lpose\n\n@tf.function()\ndef pre_process1(lip, rhand, lhand, rpose, lpose):\n    lip   = (resize_pad(lip) - LIPM) / LIPS\n    rhand = (resize_pad(rhand) - RHM) / RHS\n    lhand = (resize_pad(lhand) - LHM) / LHS\n    rpose = (resize_pad(rpose) - RPM) / RPS\n    lpose = (resize_pad(lpose) - LPM) / LPS\n\n    x = tf.concat([lip, rhand, lhand, rpose, lpose], axis=1)\n    s = tf.shape(x)\n    x = tf.reshape(x, (s[0], s[1]*s[2]))\n    x = tf.where(tf.math.is_nan(x), 0.0, x)\n    return x\n\npre0 = pre_process0(frames)\npre1 = pre_process1(*pre0)\nINPUT_SHAPE = list(pre1.shape)\nprint(INPUT_SHAPE)","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:41:57.95687Z","iopub.execute_input":"2023-07-17T17:41:57.957296Z","iopub.status.idle":"2023-07-17T17:42:06.731052Z","shell.execute_reply.started":"2023-07-17T17:41:57.957245Z","shell.execute_reply":"2023-07-17T17:42:06.729125Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode_fn(record_bytes):\n    schema = {\n        \"lip\": tf.io.VarLenFeature(tf.float32),\n        \"rhand\": tf.io.VarLenFeature(tf.float32),\n        \"lhand\": tf.io.VarLenFeature(tf.float32),\n        \"rpose\": tf.io.VarLenFeature(tf.float32),\n        \"lpose\": tf.io.VarLenFeature(tf.float32),\n        \"phrase\": tf.io.VarLenFeature(tf.int64)\n    }\n    x = tf.io.parse_single_example(record_bytes, schema)\n\n    lip = tf.reshape(tf.sparse.to_dense(x[\"lip\"]), (-1, 40, 3))\n    rhand = tf.reshape(tf.sparse.to_dense(x[\"rhand\"]), (-1, 21, 3))\n    lhand = tf.reshape(tf.sparse.to_dense(x[\"lhand\"]), (-1, 21, 3))\n    rpose = tf.reshape(tf.sparse.to_dense(x[\"rpose\"]), (-1, 5, 3))\n    lpose = tf.reshape(tf.sparse.to_dense(x[\"lpose\"]), (-1, 5, 3))\n    phrase = tf.sparse.to_dense(x[\"phrase\"])\n\n    return lip, rhand, lhand, rpose, lpose, phrase\n\n@tf.function()\ndef pre_process_fn(lip, rhand, lhand, rpose, lpose, phrase):\n    phrase = tf.pad(phrase, [[1, 0]], constant_values=start_token_idx)\n    phrase = tf.pad(phrase, [[0, 1]], constant_values=end_token_idx)\n    phrase = tf.pad(phrase, [[0, MAX_PHRASE_LENGTH-tf.shape(phrase)[0]]], constant_values=pad_token_idx)\n    return pre_process1(lip, rhand, lhand, rpose, lpose), phrase\n\ntffiles = [f\"/kaggle/input/aslfr-dataset-tfrecords/tfds/{file_id}.tfrecord\" for file_id in df.file_id.unique()]\nval_len = 1#int(0.05 * len(pqfiles))\ntrain_batch_size = 32\nval_batch_size = 32\n\ntrain_dataset =  tf.data.TFRecordDataset(tffiles[val_len:], num_parallel_reads=tf.data.AUTOTUNE).cache().shuffle(70000).map(decode_fn, num_parallel_calls=tf.data.AUTOTUNE).map(pre_process_fn, num_parallel_calls=tf.data.AUTOTUNE).batch(train_batch_size).prefetch(tf.data.AUTOTUNE)\nval_dataset =  tf.data.TFRecordDataset(tffiles[:val_len], num_parallel_reads=tf.data.AUTOTUNE).prefetch(tf.data.AUTOTUNE).map(decode_fn, num_parallel_calls=tf.data.AUTOTUNE).map(pre_process_fn, num_parallel_calls=tf.data.AUTOTUNE).batch(val_batch_size).prefetch(tf.data.AUTOTUNE).cache()\n\nbatch = next(iter(val_dataset))\nbatch[0].shape, batch[1].shape","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:42:16.165553Z","iopub.execute_input":"2023-07-17T17:42:16.165952Z","iopub.status.idle":"2023-07-17T17:42:17.030242Z","shell.execute_reply.started":"2023-07-17T17:42:16.165921Z","shell.execute_reply":"2023-07-17T17:42:17.029235Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TokenEmbedding(tf.keras.layers.Layer):\n    def __init__(self, num_vocab=1000, maxlen=100, num_hid=64):\n        super().__init__()\n        self.num_hid = num_hid\n        self.emb = tf.keras.layers.Embedding(num_vocab, num_hid)\n        self.pos_emb = self.positional_encoding(maxlen, num_hid)\n\n    def call(self, x):\n        maxlen = tf.shape(x)[-1]\n        x = self.emb(x)\n        x = tf.math.multiply(x, tf.math.sqrt(tf.cast(self.num_hid, tf.float32)))\n        return x + self.pos_emb[:maxlen, :]\n    \n    def positional_encoding(self, maxlen, num_hid):\n        depth = num_hid/2\n        positions = tf.range(maxlen, dtype = tf.float32)[..., tf.newaxis]\n        depths = tf.range(depth, dtype = tf.float32)[np.newaxis, :]/depth\n        angle_rates = tf.math.divide(1, tf.math.pow(tf.cast(10000, tf.float32), depths))\n        angle_rads = tf.linalg.matmul(positions, angle_rates)\n        pos_encoding = tf.concat(\n          [tf.math.sin(angle_rads), tf.math.cos(angle_rads)],\n          axis=-1) \n        return pos_encoding\n\n\nclass LandmarkEmbedding(tf.keras.layers.Layer):\n    def __init__(self, num_hid=64, maxlen=100):\n        super().__init__()\n        self.conv1 = tf.keras.layers.Conv1D(\n            num_hid, 11, padding=\"same\", activation=\"relu\"\n        )\n        self.conv2 = tf.keras.layers.Conv1D(\n            num_hid, 11, padding=\"same\", activation=\"relu\"\n        )\n        self.conv3 = tf.keras.layers.Conv1D(\n            num_hid, 11, padding=\"same\", activation=\"relu\"\n        )\n        self.pos_emb = self.positional_encoding(maxlen, num_hid)\n        self.maxlen = maxlen\n        self.num_hid = num_hid\n\n    def call(self, x):\n        x = self.conv1(x)\n        x = self.conv2(x)\n        x = self.conv3(x)\n        \n        x = tf.math.multiply(x, tf.math.sqrt(tf.cast(self.num_hid, tf.float32)))\n        x = x + self.pos_emb\n        \n        return x\n    \n    def positional_encoding(self, maxlen, num_hid):\n        depth = num_hid/2\n        positions = tf.range(maxlen, dtype = tf.float32)[..., tf.newaxis]\n        depths = tf.range(depth, dtype = tf.float32)[np.newaxis, :]/depth\n        angle_rates = tf.math.divide(1, tf.math.pow(tf.cast(10000, tf.float32), depths))\n        angle_rads = tf.linalg.matmul(positions, angle_rates)\n        pos_encoding = tf.concat(\n          [tf.math.sin(angle_rads), tf.math.cos(angle_rads)],\n          axis=-1) \n        return pos_encoding\n    \nle = LandmarkEmbedding(384, FRAME_LEN)(batch[0])\nte = TokenEmbedding(62, 64, 384)(batch[1])\n\nle.shape, te.shape","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:42:21.770044Z","iopub.execute_input":"2023-07-17T17:42:21.770471Z","iopub.status.idle":"2023-07-17T17:42:24.938217Z","shell.execute_reply.started":"2023-07-17T17:42:21.77044Z","shell.execute_reply":"2023-07-17T17:42:24.937338Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#https://www.kaggle.com/code/markwijkhuizen/aslfr-transformer-training-inference?scriptVersionId=135498607&cellId=46\n\ndef scaled_dot_product(q,k,v, softmax, attention_mask):\n    #calculates Q . K(transpose)\n    qkt = tf.matmul(q,k,transpose_b=True)\n    #caculates scaling factor\n    dk = tf.math.sqrt(tf.cast(q.shape[-1],dtype=tf.float32))\n    scaled_qkt = qkt/dk\n    softmax = softmax(scaled_qkt, mask=attention_mask)\n    z = tf.matmul(softmax,v)\n    #shape: (m,Tx,depth), same shape as q,k,v\n    return z\n\nclass MultiHeadAttention(tf.keras.layers.Layer):\n    def __init__(self, d_model, num_of_heads, d_out=None):\n        super(MultiHeadAttention,self).__init__()\n        self.d_model = d_model\n        self.num_of_heads = num_of_heads\n        self.depth = d_model//num_of_heads\n        self.wq = [tf.keras.layers.Dense(self.depth//2, use_bias=False) for i in range(num_of_heads)]\n        self.wk = [tf.keras.layers.Dense(self.depth//2, use_bias=False) for i in range(num_of_heads)]\n        self.wv = [tf.keras.layers.Dense(self.depth//2, use_bias=False) for i in range(num_of_heads)]\n        self.wo = tf.keras.layers.Dense(d_model if d_out is None else d_out, use_bias=False)\n        self.softmax = tf.keras.layers.Softmax()\n        self.supports_masking = True\n        \n    def call(self, q, k, v, attention_mask=None, training=False):\n        \n        multi_attn = []\n        for i in range(self.num_of_heads):\n            Q = self.wq[i](q)\n            K = self.wk[i](k)\n            V = self.wv[i](v)\n            multi_attn.append(scaled_dot_product(Q,K,V, self.softmax, attention_mask))\n            \n        multi_head = tf.concat(multi_attn, axis=-1)\n        multi_head_attention = self.wo(multi_head)\n        return multi_head_attention","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:42:27.753957Z","iopub.execute_input":"2023-07-17T17:42:27.754324Z","iopub.status.idle":"2023-07-17T17:42:27.767374Z","shell.execute_reply.started":"2023-07-17T17:42:27.754291Z","shell.execute_reply":"2023-07-17T17:42:27.766429Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TransformerEncoder(tf.keras.layers.Layer):\n    def __init__(self, embed_dim, num_heads, feed_forward_dim, rate=0.1):\n        super().__init__()\n        self.att = MultiHeadAttention(num_of_heads=num_heads, d_model=embed_dim)\n        self.ffn = tf.keras.Sequential(\n            [\n                tf.keras.layers.Dense(feed_forward_dim, activation=\"relu\"),\n                tf.keras.layers.Dense(embed_dim),\n            ]\n        )\n        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n        self.dropout1 = tf.keras.layers.Dropout(rate)\n        self.dropout2 = tf.keras.layers.Dropout(rate)\n\n    def call(self, inputs, training):\n        attn_output = self.att(inputs, inputs, inputs)\n        attn_output = self.dropout1(attn_output, training=training)\n        out1 = self.layernorm1(inputs + attn_output)\n        ffn_output = self.ffn(out1)\n        ffn_output = self.dropout2(ffn_output, training=training)\n        return self.layernorm2(out1 + ffn_output)\n    \nenc = TransformerEncoder(384, 4, 400)(le)\nenc.shape","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:42:30.716967Z","iopub.execute_input":"2023-07-17T17:42:30.71773Z","iopub.status.idle":"2023-07-17T17:42:30.901245Z","shell.execute_reply.started":"2023-07-17T17:42:30.717692Z","shell.execute_reply":"2023-07-17T17:42:30.90022Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TransformerDecoder(tf.keras.layers.Layer):\n    def __init__(self, embed_dim, num_heads, feed_forward_dim, dropout_rate=0.1):\n        super().__init__()\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        self.self_att = MultiHeadAttention(num_of_heads=num_heads, d_model=embed_dim)\n        self.enc_att = MultiHeadAttention(num_of_heads=num_heads, d_model=embed_dim)\n        self.self_dropout = tf.keras.layers.Dropout(0.5)\n        self.enc_dropout = tf.keras.layers.Dropout(0.1)\n        self.ffn_dropout = tf.keras.layers.Dropout(0.1)\n        self.ffn = tf.keras.Sequential(\n            [\n                tf.keras.layers.Dense(feed_forward_dim, activation=\"relu\"),\n                tf.keras.layers.Dense(embed_dim),\n            ]\n        )\n\n    def causal_attention_mask(self, batch_size, n_dest, n_src, dtype):\n        \"\"\"Masks the upper half of the dot product matrix in self attention.\n\n        This prevents flow of information from future tokens to current token.\n        1's in the lower triangle, counting from the lower right corner.\n        \"\"\"\n        i = tf.range(n_dest)[:, None]\n        j = tf.range(n_src)\n        m = i >= j - n_src + n_dest\n        mask = tf.cast(m, dtype)\n        mask = tf.reshape(mask, [1, n_dest, n_src])\n        mult = tf.concat(\n            [batch_size[..., tf.newaxis], tf.constant([1, 1], dtype=tf.int32)], 0\n        )\n        return tf.tile(mask, mult)\n\n    def call(self, enc_out, target, training):\n        input_shape = tf.shape(target)\n        batch_size = input_shape[0]\n        seq_len = input_shape[1]\n        causal_mask = self.causal_attention_mask(batch_size, seq_len, seq_len, tf.bool)\n        target_att = self.self_att(target, target, target, attention_mask=causal_mask)\n        target_norm = self.layernorm1(target + self.self_dropout(target_att, training=training))\n        enc_out = self.enc_att(target_norm, enc_out, enc_out)\n        enc_out_norm = self.layernorm2(self.enc_dropout(enc_out, training=training) + target_norm)\n        ffn_out = self.ffn(enc_out_norm)\n        ffn_out_norm = self.layernorm3(enc_out_norm + self.ffn_dropout(ffn_out, training=training))\n        return ffn_out_norm\n    \ndec = TransformerDecoder(384, 4, 400)(enc, te)\ndec.shape","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:42:34.000359Z","iopub.execute_input":"2023-07-17T17:42:34.001351Z","iopub.status.idle":"2023-07-17T17:42:34.235618Z","shell.execute_reply.started":"2023-07-17T17:42:34.001306Z","shell.execute_reply":"2023-07-17T17:42:34.234682Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Transformer(tf.keras.Model):\n    def __init__(\n        self,\n        num_hid_enc=64,\n        num_hid_dec=64,\n        num_head=2,\n        num_feed_forward=128,\n        source_maxlen=100,\n        target_maxlen=100,\n        num_layers_enc=4,\n        num_layers_dec=1,\n        num_classes=10,\n    ):\n        super().__init__()\n        self.loss_metric = tf.keras.metrics.Mean(name=\"loss\")\n        self.num_layers_enc = num_layers_enc\n        self.num_layers_dec = num_layers_dec\n        self.target_maxlen = target_maxlen\n        self.num_classes = num_classes\n\n        self.enc_input = LandmarkEmbedding(num_hid=num_hid_enc, maxlen=source_maxlen)\n        self.dec_input = TokenEmbedding(num_vocab=num_classes, maxlen=target_maxlen, num_hid=num_hid_dec)\n\n        self.encoder = tf.keras.Sequential(\n            [self.enc_input]\n            + [\n                TransformerEncoder(num_hid_enc, num_head, num_feed_forward)\n                for _ in range(num_layers_enc)\n            ]\n        )\n\n        for i in range(num_layers_dec):\n            setattr(\n                self,\n                f\"dec_layer_{i}\",\n                TransformerDecoder(num_hid_dec, num_head, num_feed_forward),\n            )\n\n        self.classifier = tf.keras.Sequential([\n            tf.keras.layers.Dropout(0.2),\n            tf.keras.layers.Dense(num_classes)\n        ])\n\n    def decode(self, enc_out, target, training):\n        y = self.dec_input(target)\n        for i in range(self.num_layers_dec):\n            y = getattr(self, f\"dec_layer_{i}\")(enc_out, y, training=training)\n        return y\n\n    def call(self, inputs, training):\n        source = inputs[0]\n        target = inputs[1]\n        x = self.encoder(source, training=training)\n        y = self.decode(x, target, training)\n        return self.classifier(y, training=training)\n\n    @property\n    def metrics(self):\n        return [self.loss_metric]\n\n    def train_step(self, batch):\n        \"\"\"Processes one batch inside model.fit().\"\"\"\n        source = batch[0]\n        target = batch[1]\n        dec_input = target[:, :-1]\n        dec_target = target[:, 1:]\n        with tf.GradientTape() as tape:\n            preds = self([source, dec_input])\n            one_hot = tf.one_hot(dec_target, depth=self.num_classes)\n            mask = tf.math.logical_not(tf.math.equal(dec_target, pad_token_idx))\n            loss = self.compiled_loss(one_hot, preds, sample_weight=mask)\n        trainable_vars = self.trainable_variables\n        gradients = tape.gradient(loss, trainable_vars)\n        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n        self.loss_metric.update_state(loss)\n        return {\"loss\": self.loss_metric.result()}\n\n    def test_step(self, batch):\n        source = batch[0]\n        target = batch[1]\n        dec_input = target[:, :-1]\n        dec_target = target[:, 1:]\n        preds = self([source, dec_input])\n        one_hot = tf.one_hot(dec_target, depth=self.num_classes)\n        mask = tf.math.logical_not(tf.math.equal(dec_target, pad_token_idx))\n        loss = self.compiled_loss(one_hot, preds, sample_weight=mask)\n        self.loss_metric.update_state(loss)\n        return {\"loss\": self.loss_metric.result()}\n\n    def generate(self, source, target_start_token_idx):\n        \"\"\"Performs inference over one batch of inputs using greedy decoding.\"\"\"\n        bs = tf.shape(source)[0]\n        enc = self.encoder(source, training=False)\n        dec_input = tf.ones((bs, 1), dtype=tf.int32) * target_start_token_idx\n        dec_logits = []\n        for i in range(40):\n            dec_out = self.decode(enc, dec_input, False)\n            logits = self.classifier(dec_out, training=False)\n            logits = tf.argmax(logits, axis=-1, output_type=tf.int32)\n            last_logit = logits[:, -1][..., tf.newaxis]\n            dec_logits.append(last_logit)\n            dec_input = tf.concat([dec_input, last_logit], axis=-1)\n        return dec_input","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:42:37.19772Z","iopub.execute_input":"2023-07-17T17:42:37.198098Z","iopub.status.idle":"2023-07-17T17:42:37.223939Z","shell.execute_reply.started":"2023-07-17T17:42:37.198067Z","shell.execute_reply":"2023-07-17T17:42:37.222792Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = Transformer(\n    num_hid_enc=384,\n    num_hid_dec=256,\n    num_head=5,\n    num_feed_forward=384*2,\n    source_maxlen=FRAME_LEN,\n    target_maxlen=64,\n    num_layers_enc=4,\n    num_layers_dec=2,\n    num_classes=len(char_to_num),\n)\n\ntf.keras.backend.clear_session()\nmodel(batch)\nmodel.summary()","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:42:40.520992Z","iopub.execute_input":"2023-07-17T17:42:40.521408Z","iopub.status.idle":"2023-07-17T17:42:42.653014Z","shell.execute_reply.started":"2023-07-17T17:42:40.521378Z","shell.execute_reply":"2023-07-17T17:42:42.651998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DisplayOutputs(tf.keras.callbacks.Callback):\n    def on_epoch_end(self, epoch, logs=None):\n        model.save_weights(\"model.h5\")\n        source = batch[0]\n        target = batch[1].numpy()\n        bs = tf.shape(source)[0]\n        preds = model.generate(source, start_token_idx)\n        preds = preds.numpy()\n        for i in range(bs):\n            target_text = \"\".join([num_to_char[_] for _ in target[i, :]])\n            prediction = \"\"\n            for idx in preds[i, :]:\n                prediction += num_to_char[idx]\n                if idx == end_token_idx:\n                    break\n            print(f\"target:     {target_text.replace('-','')}\")\n            print(f\"prediction: {prediction}\\n\")\n\ndisplay_cb = DisplayOutputs()\ndisplay_cb.on_epoch_end(0)","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:42:47.050963Z","iopub.execute_input":"2023-07-17T17:42:47.051384Z","iopub.status.idle":"2023-07-17T17:42:52.598355Z","shell.execute_reply.started":"2023-07-17T17:42:47.051352Z","shell.execute_reply":"2023-07-17T17:42:52.597339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"N_EPOCHS = 50\nN_WARMUP_EPOCHS = 10\nN_RESUME_EPOCHS = 0\nLR_MAX = 1e-3\nWD_RATIO = 0.05\nWARMUP_METHOD = \"exp\"\n\ndef lrfn(current_step, num_warmup_steps, lr_max, num_cycles=0.50, num_training_steps=N_EPOCHS):\n    \n    if current_step < num_warmup_steps:\n        if WARMUP_METHOD == 'log':\n            return lr_max * 0.10 ** (num_warmup_steps - current_step)\n        else:\n            return lr_max * 2 ** -(num_warmup_steps - current_step)\n    else:\n        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))\n\n        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) * lr_max\n    \ndef plot_lr_schedule(lr_schedule, epochs):\n    fig = plt.figure(figsize=(20, 10))\n    plt.plot([None] + lr_schedule + [None])\n    # X Labels\n    x = np.arange(1, epochs + 1)\n    x_axis_labels = [i if epochs <= 40 or i % 5 == 0 or i == 1 else None for i in range(1, epochs + 1)]\n    plt.xlim([1, epochs])\n    plt.xticks(x, x_axis_labels) # set tick step to 1 and let x axis start at 1\n    \n    # Increase y-limit for better readability\n    plt.ylim([0, max(lr_schedule) * 1.1])\n    \n    # Title\n    schedule_info = f'start: {lr_schedule[0]:.1E}, max: {max(lr_schedule):.1E}, final: {lr_schedule[-1]:.1E}'\n    plt.title(f'Step Learning Rate Schedule, {schedule_info}', size=18, pad=12)\n    \n    # Plot Learning Rates\n    for x, val in enumerate(lr_schedule):\n        if epochs <= 40 or x % 5 == 0 or x is epochs - 1:\n            if x < len(lr_schedule) - 1:\n                if lr_schedule[x - 1] < val:\n                    ha = 'right'\n                else:\n                    ha = 'left'\n            elif x == 0:\n                ha = 'right'\n            else:\n                ha = 'left'\n            plt.plot(x + 1, val, 'o', color='black');\n            offset_y = (max(lr_schedule) - min(lr_schedule)) * 0.02\n            plt.annotate(f'{val:.1E}', xy=(x + 1, val + offset_y), size=12, ha=ha)\n    \n    plt.xlabel('Epoch', size=16, labelpad=5)\n    plt.ylabel('Learning Rate', size=16, labelpad=5)\n    plt.grid()\n    plt.show()\n\n# Learning rate for encoder\nLR_SCHEDULE = [lrfn(step, num_warmup_steps=N_WARMUP_EPOCHS, lr_max=LR_MAX, num_cycles=0.50) for step in range(N_EPOCHS)][N_RESUME_EPOCHS:]\n# Plot Learning Rate Schedule\nplot_lr_schedule(LR_SCHEDULE, epochs=N_EPOCHS)\n# Learning Rate Callback\nlr_callback = tf.keras.callbacks.LearningRateScheduler(lambda step: LR_SCHEDULE[step], verbose=0)\n\n# Custom callback to update weight decay with learning rate\nclass WeightDecayCallback(tf.keras.callbacks.Callback):\n    def __init__(self, wd_ratio=WD_RATIO):\n        self.step_counter = 0\n        self.wd_ratio = wd_ratio\n    \n    def on_epoch_begin(self, epoch, logs=None):\n        model.optimizer.weight_decay = model.optimizer.learning_rate * self.wd_ratio\n        print(f'learning rate: {model.optimizer.learning_rate.numpy():.2e}, weight decay: {model.optimizer.weight_decay.numpy():.2e}')","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:43:57.202919Z","iopub.execute_input":"2023-07-17T17:43:57.203358Z","iopub.status.idle":"2023-07-17T17:43:57.887229Z","shell.execute_reply.started":"2023-07-17T17:43:57.20332Z","shell.execute_reply":"2023-07-17T17:43:57.886266Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=True, label_smoothing=0.25,)\noptimizer = tfa.optimizers.RectifiedAdam(sma_threshold=4)\noptimizer = tfa.optimizers.Lookahead(optimizer, sync_period=5)\nmodel.compile(optimizer=optimizer, loss=loss_fn)","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:47:09.053611Z","iopub.execute_input":"2023-07-17T17:47:09.054Z","iopub.status.idle":"2023-07-17T17:47:09.07306Z","shell.execute_reply.started":"2023-07-17T17:47:09.053969Z","shell.execute_reply":"2023-07-17T17:47:09.071901Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"history = model.fit(train_dataset, validation_data=val_dataset, callbacks=[display_cb, lr_callback, WeightDecayCallback()], epochs=N_EPOCHS)\n# history = model.fit(train_dataset.take(1), validation_data=val_dataset.take(1), callbacks=[display_cb], epochs=1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TFLiteModel(tf.Module):\n    def __init__(self, model):\n        super(TFLiteModel, self).__init__()\n        self.model = model\n    \n    @tf.function(input_signature=[tf.TensorSpec(shape=[None, len(SEL_COLS)], dtype=tf.float32, name='inputs')])\n    def __call__(self, inputs, training=False):\n        x = tf.cast(inputs, tf.float32)\n        x = x[None]\n        x = tf.cond(tf.shape(x)[1] == 0, lambda: tf.zeros((1, 1, len(SEL_COLS))), lambda: tf.identity(x))\n        x = x[0]\n        x = pre_process0(x)\n        x = pre_process1(*x)\n        x = tf.reshape(x, INPUT_SHAPE)\n        x = x[None]\n        x = self.model.generate(x, start_token_idx)\n        x = x[0]\n        idx = tf.argmax(tf.cast(tf.equal(x, end_token_idx), tf.int32))\n        idx = tf.where(tf.math.less(idx, 1), tf.constant(2, dtype=tf.int64), idx)\n        x = x[1:idx]\n        x = tf.one_hot(x, 59)\n        return {'outputs': x}\n\ntflitemodel_base = TFLiteModel(model)\ntflitemodel_base(frames)[\"outputs\"].shape","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:55:48.690355Z","iopub.execute_input":"2023-07-17T17:55:48.691712Z","iopub.status.idle":"2023-07-17T17:56:35.93175Z","shell.execute_reply.started":"2023-07-17T17:55:48.691667Z","shell.execute_reply":"2023-07-17T17:56:35.930675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"keras_model_converter = tf.lite.TFLiteConverter.from_keras_model(tflitemodel_base)\nkeras_model_converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]#, tf.lite.OpsSet.SELECT_TF_OPS]\ntflite_model = keras_model_converter.convert()\n\nwith open('/kaggle/working/model.tflite', 'wb') as f:\n    f.write(tflite_model)\n    \nwith open('inference_args.json', \"w\") as f:\n    json.dump({\"selected_columns\" : SEL_COLS}, f)\n    \n!zip submission.zip  './model.tflite' './inference_args.json'","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:56:39.419263Z","iopub.execute_input":"2023-07-17T17:56:39.41971Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open (\"inference_args.json\", \"r\") as f:\n    SEL_COLS = json.load(f)[\"selected_columns\"]\n    \ndef load_relevant_data_subset(pq_path):\n    return pd.read_parquet(pq_path, columns=SEL_COLS)\n\ndef create_data_gen(file_ids, y_mul=1):\n    def gen():\n        for file_id in file_ids:\n            pqfile = f\"{inpdir}/{file_id}.parquet\"\n            seq_refs = df.loc[df.file_id == file_id]\n            seqs = load_relevant_data_subset(pqfile)\n\n            for seq_id in seq_refs.sequence_id:\n                x = seqs.iloc[seqs.index == seq_id].to_numpy()\n                y = str(df.loc[df.sequence_id == seq_id].phrase.iloc[0])\n                \n                r_nonan = np.sum(np.sum(np.isnan(x[:, RHAND_IDX_X]), axis = 1) == 0)\n                l_nonan = np.sum(np.sum(np.isnan(x[:, LHAND_IDX_X]), axis = 1) == 0)\n                no_nan = max(r_nonan, l_nonan)\n                \n                if y_mul*len(y)<no_nan:\n                    yield x, y\n    return gen\n\npqfiles = df.file_id.unique()\nval_len = int(0.05 * len(pqfiles))\n\ntest_dataset = tf.data.Dataset.from_generator(create_data_gen(pqfiles, 0),\n    output_signature=(tf.TensorSpec(shape=(None, len(SEL_COLS)), dtype=tf.float32), tf.TensorSpec(shape=(), dtype=tf.string))\n).prefetch(buffer_size=2000)","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:40:09.417694Z","iopub.execute_input":"2023-07-17T17:40:09.418213Z","iopub.status.idle":"2023-07-17T17:40:09.482822Z","shell.execute_reply.started":"2023-07-17T17:40:09.418174Z","shell.execute_reply":"2023-07-17T17:40:09.481642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"interpreter = tf.lite.Interpreter(\"model.tflite\")\n\nREQUIRED_SIGNATURE = \"serving_default\"\nREQUIRED_OUTPUT = \"outputs\"\n\nwith open (\"/kaggle/input/asl-fingerspelling/character_to_prediction_index.json\", \"r\") as f:\n    character_map = json.load(f)\nrev_character_map = {j:i for i,j in character_map.items()}\n\nfound_signatures = list(interpreter.get_signature_list().keys())\n\nif REQUIRED_SIGNATURE not in found_signatures:\n    raise KernelEvalException('Required input signature not found.')\n\nprediction_fn = interpreter.get_signature_runner(\"serving_default\")\n\nfor frame, target in test_dataset.skip(100).take(10):\n    output = prediction_fn(inputs=frame)\n    prediction_str = \"\".join([rev_character_map.get(s, \"\") for s in np.argmax(output[REQUIRED_OUTPUT], axis=1)])\n    target = target.numpy().decode(\"utf-8\")\n    print(\"pred =\", prediction_str, \"; target =\", target)","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:40:12.464576Z","iopub.execute_input":"2023-07-17T17:40:12.465637Z","iopub.status.idle":"2023-07-17T17:40:17.457445Z","shell.execute_reply.started":"2023-07-17T17:40:12.465565Z","shell.execute_reply":"2023-07-17T17:40:17.456509Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%timeit -n 10\noutput = prediction_fn(inputs=frame)","metadata":{"execution":{"iopub.status.busy":"2023-07-17T17:40:20.412377Z","iopub.execute_input":"2023-07-17T17:40:20.413537Z","iopub.status.idle":"2023-07-17T17:40:34.995571Z","shell.execute_reply.started":"2023-07-17T17:40:20.413493Z","shell.execute_reply":"2023-07-17T17:40:34.994633Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from Levenshtein import distance\n\nscores = []\n\nfor frame, target in tqdm(test_dataset.take(1000)):\n    output = prediction_fn(inputs=frame)\n    prediction_str = \"\".join([rev_character_map.get(s, \"\") for s in np.argmax(output[REQUIRED_OUTPUT], axis=1)])\n    target = target.numpy().decode(\"utf-8\")\n    score = (len(target) - distance(prediction_str, target)) / len(target)\n    scores.append(score)\n    \nscores = np.array(scores)\nprint(np.sum(scores) / len(scores))","metadata":{},"execution_count":null,"outputs":[]}]}