{"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":"code","source":"import tensorflow as tf\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import StratifiedGroupKFold, GroupKFold, KFold\nimport gc\nfrom tqdm.auto import tqdm\nimport Levenshtein\nimport time\nimport os","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-08-22T07:29:54.431841Z","iopub.execute_input":"2023-08-22T07:29:54.432513Z","iopub.status.idle":"2023-08-22T07:30:04.572869Z","shell.execute_reply.started":"2023-08-22T07:29:54.432458Z","shell.execute_reply":"2023-08-22T07:30:04.571287Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEBUG = True #False for submission\nN = -1 #-1 for all samples\nMODEL_PATH = [\n '/kaggle/input/aslfr-2nd-place-model-weights/aslfr-fp16-192d-17l-ctcattjoint-seed42-fold0-last.h5',\n# '/kaggle/input/aslfr-2nd-place-model-weights/aslfr-fp16-192d-17l-ctcattjoint-seed42-foldall-last.h5',\n# '/kaggle/input/aslfr-2nd-place-model-weights/aslfr-fp16-192d-17l-ctcattjoint-seed43-foldall-last.h5',\n# '/kaggle/input/aslfr-2nd-place-model-weights/aslfr-fp16-192d-17l-ctcattjoint-seed44-foldall-last.h5',\n]","metadata":{"execution":{"iopub.status.busy":"2023-08-22T07:30:04.575261Z","iopub.execute_input":"2023-08-22T07:30:04.576024Z","iopub.status.idle":"2023-08-22T07:30:04.581341Z","shell.execute_reply.started":"2023-08-22T07:30:04.575982Z","shell.execute_reply":"2023-08-22T07:30:04.580094Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CFG:\n seed = 42\n n_splits = 5","metadata":{"execution":{"iopub.status.busy":"2023-08-22T07:30:04.582826Z","iopub.execute_input":"2023-08-22T07:30:04.583713Z","iopub.status.idle":"2023-08-22T07:30:04.599898Z","shell.execute_reply.started":"2023-08-22T07:30:04.583681Z","shell.execute_reply":"2023-08-22T07:30:04.598754Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SEL_COLS = pd.read_parquet('/kaggle/input/asl-fingerspelling/train_landmarks/1019715464.parquet').columns[1:].tolist()\nprint(len(SEL_COLS))\ndef load_relevant_data_subset(pq_path):\n return pd.read_parquet(pq_path, columns=SEL_COLS)\n\ndf = pd.read_csv('/kaggle/input/asl-fingerspelling/train.csv')\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])\n\ndf['fold'] = -1\n\ndf = df.sample(frac=1,random_state=CFG.seed).reset_index(drop=True)\n\nkfold = GroupKFold(n_splits=CFG.n_splits)\nprint(f'{CFG.n_splits}fold training', len(df), 'samples')\nfor fold_idx, (train_idx, valid_idx) in enumerate(kfold.split(df, groups=df.participant_id)):\n df.loc[valid_idx,'fold'] = fold_idx\n print(f'fold{fold_idx}:', 'train', len(train_idx), 'valid', len(valid_idx))\n\nassert not (df['fold']==-1).sum()\nassert len(np.unique(df['fold']))==CFG.n_splits\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2023-08-22T07:30:04.603676Z","iopub.execute_input":"2023-08-22T07:30:04.604177Z","iopub.status.idle":"2023-08-22T07:30:43.401107Z","shell.execute_reply.started":"2023-08-22T07:30:04.604134Z","shell.execute_reply":"2023-08-22T07:30:43.399835Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import json\nwith open('/kaggle/input/asl-fingerspelling/character_to_prediction_index.json') as json_file:\n CHAR_TO_NUM = json.load(json_file)\nNUM_TO_CHAR = dict([(y+1,x) for x,y in CHAR_TO_NUM.items()] )\nNUM_TO_CHAR[60] = 'S'\nNUM_TO_CHAR[61] = 'E'\nNUM_TO_CHAR[0] = 'P'","metadata":{"execution":{"iopub.status.busy":"2023-08-22T07:30:43.402503Z","iopub.execute_input":"2023-08-22T07:30:43.402944Z","iopub.status.idle":"2023-08-22T07:30:43.413186Z","shell.execute_reply.started":"2023-08-22T07:30:43.402897Z","shell.execute_reply":"2023-08-22T07:30:43.411858Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#for the lip_lr function. LEFT[i] is matching with RIGHT[i](i.e LEFT[i](x) == -RIGHT[i](x)).\n#computed from https://github.com/google/mediapipe/blob/master/mediapipe/modules/face_geometry/data/canonical_face_model.obj\n\nLEFT = [\n 248, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264,\n 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281,\n 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298,\n 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315,\n 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332,\n 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349,\n 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 364, 365, 366,\n 367, 368, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383,\n 384, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394, 395, 396, 397, 398, 399, 400,\n 401, 402, 403, 404, 405, 406, 407, 408, 409, 410, 411, 412, 413, 414, 415, 416, 417,\n 418, 419, 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431, 432, 433, 434,\n 435, 436, 437, 438, 439, 440, 441, 442, 443, 444, 445, 446, 447, 448, 449, 450, 451,\n 452, 453, 454, 455, 456, 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, #LFACE\n 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479, 480, 481, 482, 483, 484, 485, 486, 487, 488, #LHAND\n 493, 494, 495, 497, 499, 501, 503, 505, 507, 509, 511, 513, #LPOSE\n 515, 517, 519, 521, #LLEG\n ]\n\nRIGHT = [\n 3, 7, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38,\n 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59,\n 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80,\n 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 95, 96, 97, 98, 99, 100, 101, 102,\n 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120,\n 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138,\n 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 153, 154, 155, 156, 157, 158,\n 159, 160, 161, 162, 163, 165, 166, 167, 169, 170, 171, 172, 173, 174, 176, 177, 178, 179,\n 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 196, 198, 201,\n 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219,\n 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237,\n 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, #RFACE\n 522, 523, 524, 525, 526, 527, 528, 529, 530, 531, 532, 533, 534, 535, 536, 537, 538, 539, 540, 541, 542, #RHAND\n 490, 491, 492, 496, 498, 500, 502, 504, 506, 508, 510, 512, #RPOSE\n 514, 516, 518, 520, #RLEG\n ]\n\nCENTRE = [\n 0, 1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 94, 151, 152, 164, 168, 175, 195, 197, 199, 200, #FACE\n 489, #POSE\n ]\n\nprint(len(LEFT+RIGHT+CENTRE))","metadata":{"execution":{"iopub.status.busy":"2023-08-22T07:30:43.414705Z","iopub.execute_input":"2023-08-22T07:30:43.415128Z","iopub.status.idle":"2023-08-22T07:30:43.449074Z","shell.execute_reply.started":"2023-08-22T07:30:43.415084Z","shell.execute_reply":"2023-08-22T07:30:43.447755Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ROWS_PER_FRAME = 543\nMAX_LEN = 768\nCROP_LEN = MAX_LEN\nNUM_CLASSES = len(NUM_TO_CHAR.values()) #62\nPAD = -100.\n\nLHAND = np.arange(468, 489).tolist()\nRHAND = np.arange(522, 543).tolist()\nPOINT_LANDMARKS = list(range(543))\n\nNUM_NODES = len(POINT_LANDMARKS)\nCHANNELS = 3*NUM_NODES\n\nprint(NUM_NODES)\nprint(CHANNELS)\n\ndef interp1d_(x, target_len, method='random'):\n length = tf.shape(x)[1]\n target_len = tf.maximum(1,target_len)\n if method == 'random':\n if tf.random.uniform(()) < 0.33:\n x = tf.image.resize(x, (target_len,tf.shape(x)[1]),'bilinear')\n else:\n if tf.random.uniform(()) < 0.5:\n x = tf.image.resize(x, (target_len,tf.shape(x)[1]),'bicubic')\n else:\n x = tf.image.resize(x, (target_len,tf.shape(x)[1]),'nearest')\n else:\n x = tf.image.resize(x, (target_len,tf.shape(x)[1]),method)\n return x\n\ndef tf_nan_mean(x, axis=0, keepdims=False):\n return tf.reduce_sum(tf.where(tf.math.is_nan(x), tf.zeros_like(x), x), axis=axis, keepdims=keepdims) / tf.reduce_sum(tf.where(tf.math.is_nan(x), tf.zeros_like(x), tf.ones_like(x)), axis=axis, keepdims=keepdims)\n\ndef tf_nan_std(x, center=None, axis=0, keepdims=False):\n if center is None:\n center = tf_nan_mean(x, axis=axis, keepdims=True)\n d = x - center\n return tf.math.sqrt(tf_nan_mean(d * d, axis=axis, keepdims=keepdims))\n\ndef filter_nans_tf(x, ref_point=POINT_LANDMARKS):\n mask = tf.math.logical_not(tf.reduce_all(tf.math.is_nan(tf.gather(x,ref_point,axis=1)), axis=[-2,-1]))\n x = tf.boolean_mask(x, mask, axis=0)\n return x\n\ndef is_left_handed(x, left=LHAND, right=RHAND):\n lhand = tf.gather(x, left, axis=1)\n rhand = tf.gather(x, right, axis=1)\n lhand_nans = tf.reduce_sum(tf.cast(tf.math.is_nan(lhand), tf.int32))\n rhand_nans = tf.reduce_sum(tf.cast(tf.math.is_nan(rhand), tf.int32))\n return lhand_nans < rhand_nans\n\ndef flip_lr(x):\n x,y,z = tf.unstack(x, axis=-1)\n x = 1-x\n new_x = tf.stack([x,y,z], -1)\n new_x = tf.transpose(new_x, [1,0,2])\n l_x = tf.gather(new_x, LEFT, axis=0)\n r_x = tf.gather(new_x, RIGHT, axis=0)\n c_x = tf.gather(new_x, CENTRE, axis=0)\n# new_x = tf.tensor_scatter_nd_update(new_x, tf.constant(left)[...,None], r_x) <-weird behavior in tflite!!!:(\n# new_x = tf.tensor_scatter_nd_update(new_x, tf.constant(right)[...,None], l_x)\n new_xr = tf.scatter_nd(tf.constant(LEFT)[...,None], r_x, tf.shape(new_x))\n new_xl = tf.scatter_nd(tf.constant(RIGHT)[...,None], l_x, tf.shape(new_x))\n new_xc = tf.scatter_nd(tf.constant(CENTRE)[...,None], c_x, tf.shape(new_x))\n new_x = new_xr + new_xl + new_xc\n new_x = tf.transpose(new_x, [1,0,2])\n return new_x\n\nclass Preprocess(tf.keras.layers.Layer):\n def __init__(self, max_len=MAX_LEN, point_landmarks=POINT_LANDMARKS, **kwargs):\n super().__init__(**kwargs)\n self.max_len = max_len\n self.point_landmarks = point_landmarks\n\n def call(self, inputs):\n # if tf.rank(inputs) == 3:\n # x = inputs[None,...]\n # else:\n # x = inputs\n x = inputs\n x = filter_nans_tf(x)\n x = tf.cond(is_left_handed(x), lambda:flip_lr(x), lambda:x)\n x = x[None,...]\n\n if self.max_len is not None:\n x = x[:,:self.max_len]\n length = tf.shape(x)[1]\n\n mean = tf_nan_mean(tf.gather(x, self.point_landmarks, axis=2), axis=[1,2], keepdims=True)\n mean = tf.where(tf.math.is_nan(mean), tf.constant([0.5,0.5,0.],x.dtype), mean)\n x = tf.gather(x, self.point_landmarks, axis=2) #N,T,P,C\n std = tf_nan_std(x, center=mean, axis=[1,2], keepdims=True)\n\n x = (x - mean)/std\n\n x = tf.concat([\n tf.reshape(x, (-1,length,3*len(self.point_landmarks))),\n # tf.reshape(dx, (-1,length,3*len(self.point_landmarks))),\n ], axis = -1)\n\n x = tf.where(tf.math.is_nan(x),tf.constant(0.,x.dtype),x)\n\n return x","metadata":{"execution":{"iopub.status.busy":"2023-08-22T10:49:07.179677Z","iopub.execute_input":"2023-08-22T10:49:07.180677Z","iopub.status.idle":"2023-08-22T10:49:07.232863Z","shell.execute_reply.started":"2023-08-22T10:49:07.180631Z","shell.execute_reply":"2023-08-22T10:49:07.231581Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ECA(tf.keras.layers.Layer):\n def __init__(self, kernel_size=5, **kwargs):\n super().__init__(**kwargs)\n self.supports_masking = True\n self.kernel_size = kernel_size\n self.conv = tf.keras.layers.Conv1D(1, kernel_size=kernel_size, strides=1, padding=\"same\", use_bias=False)\n\n def call(self, inputs, mask=None):\n nn = tf.keras.layers.GlobalAveragePooling1D()(inputs, mask=mask)\n nn = tf.expand_dims(nn, -1)\n nn = self.conv(nn)\n nn = tf.squeeze(nn, -1)\n nn = tf.nn.sigmoid(nn)\n nn = nn[:,None,:]\n return inputs * nn\n\nclass LateDropout(tf.keras.layers.Layer):\n def __init__(self, rate, noise_shape=None, start_step=0, **kwargs):\n super().__init__(**kwargs)\n self.supports_masking = True\n self.rate = rate\n self.start_step = start_step\n self.dropout = tf.keras.layers.Dropout(rate, noise_shape=noise_shape)\n\n def build(self, input_shape):\n super().build(input_shape)\n agg = tf.VariableAggregation.ONLY_FIRST_REPLICA\n self._train_counter = tf.Variable(0, dtype=\"int64\", aggregation=agg, trainable=False)\n\n def call(self, inputs, training=False):\n x = tf.cond(self._train_counter < self.start_step, lambda:inputs, lambda:self.dropout(inputs, training=training))\n if training:\n self._train_counter.assign_add(1)\n return x\n\nclass MaskingConv1D(tf.keras.layers.Layer):\n def __init__(self, filters, kernel_size, groups=1, strides=1,\n dilation_rate=1,\n padding='same',\n use_bias=False,\n kernel_initializer='glorot_uniform',**kwargs):\n super().__init__(**kwargs)\n assert padding == 'same'\n self.filters = filter_dataset_eager_fallback\n self.strides = strides\n self.groups = groups\n self.kernel_size = kernel_size\n self.dilation_rate = dilation_rate\n self.use_bias = use_bias\n self.padding = padding\n self.conv = tf.keras.layers.Conv1D(\n filters,\n kernel_size,\n strides=strides,\n groups=groups,\n dilation_rate=dilation_rate,\n padding=padding,\n use_bias=use_bias,\n kernel_initializer=kernel_initializer)\n self.supports_masking = True\n\n def compute_mask(self, inputs, mask=None):\n if mask is not None:\n if self.strides > 1:\n mask = mask[:,::self.strides]\n return mask\n\n def call(self, inputs, mask=None):\n x = inputs\n if mask is not None:\n x = tf.where(mask[...,None], x, tf.constant(0., dtype=x.dtype))\n x = self.conv(x)\n return x\n\nclass MaskingDWConv1D(tf.keras.layers.Layer):\n def __init__(self, kernel_size, strides=1,\n dilation_rate=1,\n padding='same',\n use_bias=False,\n kernel_initializer='glorot_uniform',**kwargs):\n super().__init__(**kwargs)\n assert padding == 'same'\n self.strides = strides\n self.kernel_size = kernel_size\n self.dilation_rate = dilation_rate\n self.use_bias = use_bias\n self.padding = padding\n self.conv = tf.keras.layers.DepthwiseConv1D(\n kernel_size,\n strides=strides,\n dilation_rate=dilation_rate,\n padding=padding,\n use_bias=use_bias,\n kernel_initializer=kernel_initializer)\n self.supports_masking = True\n\n def compute_mask(self, inputs, mask=None):\n if mask is not None:\n if self.strides > 1:\n mask = mask[:,::self.strides]\n return mask\n\n def call(self, inputs, mask=None):\n x = inputs\n if mask is not None:\n x = tf.where(mask[...,None], x, tf.constant(0., dtype=x.dtype))\n x = self.conv(x)\n return x\n\ndef Conv1DBlock(channel_size,\n kernel_size,\n dilation_rate=1,\n strides=1,\n drop_rate=0.0,\n expand_ratio=2,\n se_ratio=0.25,\n activation='swish',\n name=None):\n '''\n efficient conv1d block, @hoyso48\n '''\n if name is None:\n name = str(tf.keras.backend.get_uid(\"mbblock\"))\n # Expansion phase\n def apply(inputs):\n channels_in = tf.keras.backend.int_shape(inputs)[-1]\n channels_expand = channels_in * expand_ratio\n\n skip = inputs\n \n x = tf.keras.layers.BatchNormalization(momentum=0.95, name=name + 'pre_bn')(inputs)\n \n x = tf.keras.layers.Dense(\n channels_expand,\n use_bias=True,\n activation=activation,\n name=name + '_expand_conv')(x)\n\n # Depthwise Convolution\n x = MaskingDWConv1D(kernel_size,\n dilation_rate=dilation_rate,\n strides=strides,\n use_bias=False,\n name=name + '_dwconv')(x)\n # x = MaskingConv1D(channels_expand,kernel_size,\n # dilation_rate=dilation_rate,\n # groups=8,\n # use_bias=False,\n # name=name + '_dwconv')(x)\n\n x = tf.keras.layers.BatchNormalization(momentum=0.95, name=name + 'conv_bn')(x)\n\n x = ECA()(x)\n\n x = tf.keras.layers.Dense(\n channel_size,\n use_bias=True,\n name=name + '_project_conv')(x)\n\n if drop_rate > 0:\n x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1), name=name + '_drop')(x)\n\n if (channels_in == channel_size) and (strides == 1):\n x = tf.keras.layers.add([x, skip], name=name + '_add')\n return x\n\n return apply\n\ndef Conv2DBlock(channel_size,\n kernel_size,\n strides=1,\n dilation_rate=1,\n drop_rate=0.0,\n expand_ratio=2,\n se_ratio=0.25,\n activation='gelu',\n name=None):\n '''\n efficient conv1d block, @hoyso48\n '''\n if name is None:\n name = str(tf.keras.backend.get_uid(\"mbblock\"))\n # Expansion phase\n def apply(inputs):\n channels_in = tf.keras.backend.int_shape(inputs)[-1]\n channels_expand = channels_in * expand_ratio\n\n skip = inputs\n\n x = tf.keras.layers.Dense(\n channels_expand,\n use_bias=True,\n activation=activation,\n name=name + '_expand_conv')(inputs)\n\n # Depthwise Convolution\n # x = MaskingDWConv2D(kernel_size,\n # strides=strides,\n # dilation_rate=dilation_rate,\n # use_bias=False,\n # name=name + '_dwconv')(x)\n x = MaskingConv2D(\n channels_expand,\n kernel_size,\n strides=strides,\n dilation_rate=dilation_rate,\n use_bias=False,\n groups=channels_expand,\n name=name+'_dwconv',\n )(x)\n\n # x = tf.keras.layers.BatchNormalization(momentum=0.95, name=name + '_bn', fused=False)(x)\n # x = MaskingBatchNorm2D()(x)\n # x = ReduceMask()(x)\n x = tf.keras.layers.LayerNormalization()(x)\n\n\n # x = ECA()(x)\n\n x = tf.keras.layers.Dense(\n channel_size,\n use_bias=True,\n name=name + '_project_conv')(x)\n\n if drop_rate > 0:\n x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1,1), name=name + '_drop')(x)\n\n if (channels_in == channel_size) and (strides == 1):\n x = tf.keras.layers.add([x, skip], name=name + '_add')\n return x\n\n return apply","metadata":{"execution":{"iopub.status.busy":"2023-08-22T10:43:48.892026Z","iopub.execute_input":"2023-08-22T10:43:48.893033Z","iopub.status.idle":"2023-08-22T10:43:48.950201Z","shell.execute_reply.started":"2023-08-22T10:43:48.892991Z","shell.execute_reply":"2023-08-22T10:43:48.948476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MultiHeadAttention(tf.keras.layers.Layer):\n def __init__(self, dim=256, num_heads=4, dropout=0, **kwargs):\n super().__init__(**kwargs)\n self.dim = dim\n self.scale = self.dim ** -0.5\n self.num_heads = num_heads\n # self.qkv = tf.keras.layers.Dense(3 * dim, use_bias=False)\n self.q = tf.keras.layers.Dense(dim, use_bias=False)\n self.k = tf.keras.layers.Dense(dim, use_bias=False)\n self.v = tf.keras.layers.Dense(dim, use_bias=False)\n self.drop1 = tf.keras.layers.Dropout(dropout)\n self.proj = tf.keras.layers.Dense(dim, use_bias=False)\n self.supports_masking = True\n\n def get_causal_mask(self, q, k):\n q_len = tf.shape(q)[1]\n k_len = tf.shape(k)[1]\n i = tf.range(q_len)[:, None]\n j = tf.range(k_len)\n mask = i >= j\n mask = tf.reshape(mask, (q_len, k_len))\n return mask\n\n def merge_input_state(self, input, state, layer):\n if input is not None and state is not None:\n return tf.keras.layers.Concatenate(axis=1)([state, layer(input)])\n elif input is not None and state is None:\n return layer(input)\n elif input is None and state is not None:\n return state\n else:\n raise ValueError\n # return out\n\n def call(self, q, k=None, v=None, key_state=None, value_state=None, return_states=False, use_causal_mask=False):\n q = self.q(q)\n k = self.merge_input_state(k, key_state, self.k)\n v = self.merge_input_state(v, value_state, self.v)\n mask = getattr(k, '_keras_mask', None)\n if mask is not None:\n mask = mask[:,None,None,:]\n if use_causal_mask:\n if mask is not None:\n mask = tf.logical_and(mask, self.get_causal_mask(q,k)[None,None,:,:])\n else:\n mask = self.get_causal_mask(q,k)[None,None,:,:]\n q_ = tf.keras.layers.Permute((2, 1, 3))(tf.keras.layers.Reshape((-1, self.num_heads, self.dim // self.num_heads))(q))\n k_ = tf.keras.layers.Permute((2, 1, 3))(tf.keras.layers.Reshape((-1, self.num_heads, self.dim // self.num_heads))(k))\n v_ = tf.keras.layers.Permute((2, 1, 3))(tf.keras.layers.Reshape((-1, self.num_heads, self.dim // self.num_heads))(v))\n attn = tf.matmul(q_, k_, transpose_b=True) * self.scale\n\n attn = tf.keras.layers.Softmax(axis=-1)(attn, mask=mask)\n attn = self.drop1(attn)\n\n x = attn @ v_\n x = tf.keras.layers.Reshape((-1, self.dim))(tf.keras.layers.Permute((2, 1, 3))(x))\n x = self.proj(x)\n if return_states:\n return x, k, v\n else:\n return x\n\nclass PosEmbedding(tf.keras.layers.Layer):\n def __init__(self, dim=64, max_len=64, **kwargs):\n super().__init__(**kwargs)\n self.pos_emb = tf.keras.layers.Embedding(input_dim=max_len, output_dim=dim)\n self.supports_masking = True\n\n def call(self, x, positions=None):\n if positions is None:\n maxlen = tf.shape(x)[1]\n positions = tf.range(start=0, limit=maxlen, delta=1)\n positions = self.pos_emb(positions)\n return x + positions\n \ndef TransformerDecoderBlock(dim=256, num_heads=4, expand=4, attn_dropout=0.2, drop_rate=0.2, activation='swish', name=''):\n def apply(q,k,v):\n x = q\n # key_mask=None\n x = tf.keras.layers.BatchNormalization(momentum=0.95, name=name + '_bn1')(x)\n x = MultiHeadAttention(dim=dim,num_heads=num_heads,dropout=attn_dropout, name=name + '_self_attn')(x,x,x,use_causal_mask=True)\n # print(x.shape, q.shape)\n # x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1))(x)\n x = tf.keras.layers.Add(name=name + '_add1')([q, x])\n attn_out1 = x\n\n x = tf.keras.layers.BatchNormalization(momentum=0.95, name=name + '_bn2')(x)\n x = MultiHeadAttention(dim=dim,num_heads=num_heads,dropout=attn_dropout, name=name + '_cross_attn')(x,k,v)\n # x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1))(x)\n x = tf.keras.layers.Add(name=name + '_add2')([attn_out1, x])\n attn_out2 = x\n\n x = tf.keras.layers.BatchNormalization(momentum=0.95, name=name + '_bn3')(x)\n x = tf.keras.layers.Dense(dim*expand, use_bias=False, activation=activation, name=name + '_fc1')(x)\n x = tf.keras.layers.Dense(dim, use_bias=False, name=name + '_fc2')(x)\n # x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1))(x)\n x = tf.keras.layers.Add(name=name + '_add3')([attn_out2, x])\n return x\n return apply","metadata":{"execution":{"iopub.status.busy":"2023-08-22T10:43:49.080373Z","iopub.execute_input":"2023-08-22T10:43:49.080837Z","iopub.status.idle":"2023-08-22T10:43:49.124472Z","shell.execute_reply.started":"2023-08-22T10:43:49.080802Z","shell.execute_reply":"2023-08-22T10:43:49.123248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MultiHeadSelfAttention(tf.keras.layers.Layer):\n def __init__(self, dim=256, num_heads=4, dropout=0, **kwargs):\n super().__init__(**kwargs)\n self.dim = dim\n self.scale = self.dim ** -0.5\n self.num_heads = num_heads\n self.qkv = tf.keras.layers.Dense(3 * dim, use_bias=False)\n self.drop1 = tf.keras.layers.Dropout(dropout)\n self.proj = tf.keras.layers.Dense(dim, use_bias=False)\n self.supports_masking = True\n\n def call(self, inputs, mask=None):\n qkv = self.qkv(inputs)\n qkv = tf.keras.layers.Permute((2, 1, 3))(tf.keras.layers.Reshape((-1, self.num_heads, self.dim * 3 // self.num_heads))(qkv))\n q, k, v = tf.split(qkv, [self.dim // self.num_heads] * 3, axis=-1)\n\n attn = tf.matmul(q, k, transpose_b=True) * self.scale\n\n if mask is not None:\n mask = mask[:, None, None, :]\n # print('selfattn mask', mask.shape)\n\n attn = tf.keras.layers.Softmax(axis=-1)(attn, mask=mask)\n attn = self.drop1(attn)\n\n x = attn @ v\n x = tf.keras.layers.Reshape((-1, self.dim))(tf.keras.layers.Permute((2, 1, 3))(x))\n x = self.proj(x)\n return x\n\ndef TransformerBlock(dim=256, num_heads=4, expand=4, attn_dropout=0.2, drop_rate=0.2, activation='swish'):\n def apply(inputs):\n x = inputs\n x = tf.keras.layers.BatchNormalization(momentum=0.95)(x)\n x = MultiHeadSelfAttention(dim=dim,num_heads=num_heads,dropout=attn_dropout)(x)\n x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1))(x)\n x = tf.keras.layers.Add()([inputs, x])\n attn_out = x\n\n x = tf.keras.layers.BatchNormalization(momentum=0.95)(x)\n x = tf.keras.layers.Dense(dim*expand, use_bias=False, activation=activation)(x)\n x = tf.keras.layers.Dense(dim, use_bias=False)(x)\n x = tf.keras.layers.Dropout(drop_rate, noise_shape=(None,1,1))(x)\n x = tf.keras.layers.Add()([attn_out, x])\n return x\n return apply","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model(max_len=MAX_LEN, target_len=64, dim=192, dtype='float32'):\n ################# ENCODER #################\n inp1 = tf.keras.Input((max_len,CHANNELS),dtype=dtype)\n# x = tf.keras.layers.Masking(mask_value=PAD,input_shape=(max_len,CHANNELS))(inp1)\n x = inp1\n ksize = 17\n drop_rate = 0.2\n x = tf.keras.layers.Dense(dim,use_bias=False,name='stem_conv')(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = TransformerBlock(dim,expand=2,num_heads=4,drop_rate=drop_rate,attn_dropout=0.2)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = TransformerBlock(dim,expand=2,num_heads=4,drop_rate=drop_rate,attn_dropout=0.2)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=0,strides=2)(x) #drop_rate=0 since we don't want to drop the whole output\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = TransformerBlock(dim,expand=2,num_heads=4,drop_rate=drop_rate,attn_dropout=0.2)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = Conv1DBlock(dim,ksize,expand_ratio=4,drop_rate=drop_rate)(x)\n x = TransformerBlock(dim,expand=2,num_heads=4,drop_rate=drop_rate,attn_dropout=0.2)(x)\n x = tf.keras.layers.BatchNormalization(momentum=0.95)(x)\n\n encoder = tf.keras.Model(inp1,x,name='encoder')\n\n ################# CTC DECDODER #################\n inp3 = tf.keras.Input((x.shape[1],dim),name='ctc_decoder_inp2',dtype=dtype)\n x = inp3\n x = tf.keras.layers.RNN(tf.keras.layers.GRUCell(dim), return_sequences=True)(x)\n x = tf.keras.layers.Dense(dim*2)(x)\n x = tf.keras.layers.Dropout(0.5)(x)\n x = tf.keras.layers.Dense(NUM_CLASSES,name='ctc_classifier')(x) #include sos, eos token\n ctc_decoder = tf.keras.Model(inp3,x,name='ctc_decoder')\n\n ################# ATT DECODER #################\n inp2 = tf.keras.Input((None,),name='att_decoder_inp1',dtype='int32')\n inp3 = tf.keras.Input((x.shape[1],dim),name='att_decoder_inp2',dtype=dtype)\n\n x = inp3\n# y = tf.keras.layers.Masking(mask_value=0,input_shape=(None,),name='att_decoder_input_masking')(inp2)\n y = inp2\n y = tf.keras.layers.Embedding(NUM_CLASSES,dim,name='att_decoder_token_emb')(y) #include sos token\n y = PosEmbedding(dim,max_len=target_len,name='att_decoder_pos_emb')(y)\n y = TransformerDecoderBlock(dim,expand=2,num_heads=4,attn_dropout=0.2,name='att_decoder_block1')(y,x,x)\n y = tf.keras.layers.Dropout(0.5)(y)\n y = tf.keras.layers.Dense(NUM_CLASSES,name='att_decoder_classifier')(y)\n\n decoder = tf.keras.Model([inp2,inp3],y,name='att_decoder')\n\n ################### MODEL #####################\n inp1 = tf.keras.Input((max_len,CHANNELS),dtype=dtype)\n inp2 = tf.keras.Input((None,),dtype='int32')\n\n x = inp1\n enc_out = encoder(x)\n y = inp2\n dec_out = decoder([y, enc_out])\n ctc_out = ctc_decoder(enc_out)\n model = tf.keras.Model([inp1,inp2], [dec_out,ctc_out])\n\n return model\n\nmodel = get_model()\nmodel.summary()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_list = [get_model() for _ in MODEL_PATH]\n \nfor model, path in zip(model_list, MODEL_PATH):\n model.load_weights(path)","metadata":{"execution":{"iopub.status.busy":"2023-08-22T10:43:55.462081Z","iopub.execute_input":"2023-08-22T10:43:55.462615Z","iopub.status.idle":"2023-08-22T10:43:56.380309Z","shell.execute_reply.started":"2023-08-22T10:43:55.462579Z","shell.execute_reply":"2023-08-22T10:43:56.378956Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CTCGreedyDecoder(tf.keras.layers.Layer):\n def __init__(self, model, pad_token_idx=0, **kwargs):\n super().__init__(**kwargs)\n self.encoder = model.get_layer('encoder')\n self.ctc_decoder = model.get_layer('ctc_decoder')\n self.pad_token_idx = pad_token_idx\n \n def decode_phrase(self, pred):\n x = tf.argmax(pred, axis=1, output_type=tf.int32)\n diff = tf.not_equal(x[:-1], x[1:])\n adjacent_indices = tf.where(diff)[:, 0]\n x = tf.gather(x, adjacent_indices)\n mask = x != self.pad_token_idx\n x = tf.boolean_mask(x, mask, axis=0)\n return x\n \n def call(self, batch_x):\n encoder_out = self.encoder(batch_x)\n ctc_probs = self.ctc_decoder(encoder_out)\n return tf.identity([self.decode_phrase(ctc_probs[0])])","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ATTGreedyDecoder(tf.keras.layers.Layer):\n def __init__(self, model, max_output_length=64, input_strides=2, sos_token_idx=60, eos_token_idx=61, pad_token_idx=0, **kwargs):\n super().__init__(**kwargs)\n self.model = model\n self.encoder = self.model.get_layer('encoder')\n self.decoder = self.model.get_layer('att_decoder')\n self.max_output_length = max_output_length\n self.sos_token_idx = sos_token_idx\n self.eos_token_idx = eos_token_idx\n self.pad_token_idx = pad_token_idx\n self.input_strides = input_strides\n\n def att_inference_module(self, query, query_position, key_state, value_state, encoder_key_state, encoder_value_state):\n x = self.decoder.get_layer('att_decoder_inp1')(query)\n x = self.decoder.get_layer('att_decoder_token_emb')(x)\n x = self.decoder.get_layer('att_decoder_pos_emb')(x, positions=query_position)\n\n q = x\n x = self.decoder.get_layer('att_decoder_block1_bn1')(x)\n x, k, v = self.decoder.get_layer('att_decoder_block1_self_attn')(x, x, x, key_state=key_state, value_state=value_state, return_states=True)\n x = self.decoder.get_layer('att_decoder_block1_add1')([q,x])\n attn_out1 = x\n\n x = self.decoder.get_layer('att_decoder_block1_bn2')(x)\n x = self.decoder.get_layer('att_decoder_block1_cross_attn')(x, None, None, key_state=encoder_key_state, value_state=encoder_value_state)\n x = self.decoder.get_layer('att_decoder_block1_add2')([attn_out1,x])\n attn_out2 = x\n\n x = self.decoder.get_layer('att_decoder_block1_bn3')(x)\n x = self.decoder.get_layer('att_decoder_block1_fc1')(x)\n x = self.decoder.get_layer('att_decoder_block1_fc2')(x)\n x = self.decoder.get_layer('att_decoder_block1_add3')([attn_out2,x])\n out = self.decoder.get_layer('att_decoder_classifier')(x)\n return out, k, v\n\n def compute_input_length(self, batch_x):\n input_length = tf.cast(tf.shape(batch_x)[1], tf.float32)\n input_length = tf.math.ceil(input_length/self.input_strides)\n return tf.cast(input_length, tf.int32)\n \n def call(self, batch_x):\n encoder_out = self.encoder(batch_x)\n input_length = self.compute_input_length(batch_x)\n encoder_key_state = self.decoder.get_layer('att_decoder_block1_cross_attn').k(encoder_out)\n encoder_value_state = self.decoder.get_layer('att_decoder_block1_cross_attn').v(encoder_out)\n\n time = tf.constant(0, dtype=tf.int32)\n predictions = tf.ones((tf.shape(batch_x)[0],1), dtype=tf.int32) * self.sos_token_idx\n pad = tf.ones((tf.shape(batch_x)[0],), dtype=tf.int32) * self.pad_token_idx\n init = True\n key_state = tf.zeros((0,0,192))\n value_state = tf.zeros((0,0,192))\n\n def condition(_time, predictions, key_state, value_state, init):\n return tf.logical_and(tf.logical_and(_time < self.max_output_length, tf.logical_not(tf.reduce_all(tf.reduce_any(predictions==self.eos_token_idx, axis=1)))), tf.reduce_any(_time < input_length))\n\n def body(_time, predictions, key_state, value_state, init):\n if init:\n out, key_state, value_state = self.att_inference_module(predictions[:,-1:], _time, None, None, encoder_key_state, encoder_value_state)\n init = False\n else:\n out, key_state, value_state = self.att_inference_module(predictions[:,-1:], _time, key_state, value_state, encoder_key_state, encoder_value_state)\n pred_curr = tf.where(tf.logical_or(tf.reduce_any(predictions==self.eos_token_idx, axis=1), _time >= input_length), pad, tf.argmax(out[:,-1], axis=-1, output_type=tf.int32))\n predictions = tf.concat([predictions, pred_curr[...,None]], axis=1)\n return _time+1, predictions, key_state, value_state, init\n\n _, predictions, _, _, _ = tf.while_loop(condition, body, \n shape_invariants=[tf.TensorShape([]),\n tf.TensorShape([None,None]),\n tf.TensorShape([None,None,192]),\n tf.TensorShape([None,None,192]),\n tf.TensorShape([])], \n loop_vars=[time, predictions, key_state, value_state, init])\n return predictions","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_ctc_initial_states(log_probs, blank_idx=0):\n\n blank_probs = log_probs[...,blank_idx]\n states_n = tf.ones_like(blank_probs, dtype=tf.float32) * tf.float32.min\n states_b = tf.math.cumsum(blank_probs)\n \n return states_n, states_b\n \ndef compute_ctc_prefix_scores(beams, log_probs, states_n, states_b, eos_idx=61, blank_idx=0):\n # beams: (N=hypothesis_length)\n # probs: (L=(padded/strided)input_length,M=num_models,V=vocab_size)\n # states_n: (L,M)\n # states_b: (L,M)\n \n N = tf.shape(beams)[0]\n L = tf.shape(states_n)[0]\n V = tf.shape(log_probs)[-1]\n M = tf.shape(states_n)[1]\n new_states_n = tf.ones((L,M,V), dtype=tf.float32) * tf.float32.min\n new_states_b = tf.ones((L,M,V), dtype=tf.float32) * tf.float32.min\n new_states_n = tf.cond(N==1, lambda:log_probs, lambda:new_states_n)\n \n r_sum = tf.math.reduce_logsumexp([states_n, states_b], axis=0) #(B,N)\n last = beams[-1] #(1,)\n\n repeated_idx = last \n \n log_phi_ = tf.repeat(r_sum[None,...], repeats=V, axis=0) #(V,L,M)\n log_phi = tf.tensor_scatter_nd_update(log_phi_, [[repeated_idx]], [states_b])\n log_phi = tf.transpose(log_phi, (1,2,0)) #(L,M,V)\n \n log_phi = tf.cond(N==1, lambda:tf.transpose(log_phi_, (1,2,0)), lambda:log_phi)\n \n def step_function(prev, inputs):\n prev_r_n, prev_r_b = prev\n current_log_phi, current_prob = inputs\n updated_r_n = tf.math.reduce_logsumexp([prev_r_n, current_log_phi], axis=0) + current_prob\n updated_r_b = tf.math.reduce_logsumexp([prev_r_b, prev_r_n], axis=0) + current_prob[...,blank_idx][...,None]\n return updated_r_n, updated_r_b\n \n start = 1\n log_psi = new_states_n[start-1]\n\n sequence_log_phi = log_phi[start-1:L-1]\n sequence_probs = log_probs[start-1+N:L-1+N]\n sequences = (sequence_log_phi, sequence_probs) #((L-start,M,V), (L-start,M,V))\n\n initial_state = (new_states_n[start-1], new_states_b[start-1]) #((M,V),(M,V),(M,V))\n \n log_psi = tf.math.reduce_logsumexp([tf.math.reduce_logsumexp(sequence_log_phi + sequence_probs, axis=0), log_psi], axis=0)\n\n new_states_n, new_states_b = tf.scan(step_function, sequences, initial_state)\n\n log_psi_eos = r_sum[-1]\n model_idx = tf.range(M)\n eos_idxs = tf.stack([model_idx, tf.fill((M,), eos_idx)], axis=-1)\n blank_idxs = tf.stack([model_idx, tf.fill((M,), blank_idx)], axis=-1)\n log_psi = tf.tensor_scatter_nd_update(log_psi, eos_idxs, log_psi_eos)\n log_psi = tf.tensor_scatter_nd_update(log_psi, blank_idxs, tf.fill((M,), tf.float32.min))\n \n return log_psi, new_states_n, new_states_b #(M,V), (L,M,V), (L,M,V)\n\n\nclass EnsembleCTCAttentionJointGreedyDecoder(tf.keras.layers.Layer):\n def __init__(self, model_list, ctc_weight=0.2, input_strides=2, max_output_length=64, blank_idx=0, pad_frame_idx=-100, sos_token_idx=60, eos_token_idx=61, pad_token_idx=0, from_logits=True, **kwargs):\n super().__init__(**kwargs)\n self.encoder_list = [m.get_layer('encoder') for m in model_list]\n self.decoder_list = [m.get_layer('att_decoder') for m in model_list]\n self.ctc_decoder_list = [m.get_layer('ctc_decoder') for m in model_list]\n self.ctc_weight = ctc_weight\n self.input_strides = input_strides\n self.max_output_length = max_output_length\n self.blank_idx = blank_idx\n self.pad_frame_idx = pad_frame_idx\n self.sos_token_idx = sos_token_idx\n self.eos_token_idx = eos_token_idx\n self.pad_token_idx = pad_token_idx\n self.from_logits = from_logits\n \n def compute_input_length(self, batch_x):\n input_length = tf.cast(tf.shape(batch_x)[1], tf.float32)#tf.reduce_sum(tf.cast(mask, tf.float32), axis=-1)\n input_length = tf.math.ceil(input_length/self.input_strides)\n return tf.cast(input_length, tf.int32)\n\n def att_inference_module(self, query, query_position, key_state_list, value_state_list, encoder_key_state_list, encoder_value_state_list):\n outputs = []\n key_states = []\n value_states = []\n for i in range(len(self.decoder_list)):\n decoder = self.decoder_list[i]\n key_state = key_state_list[i] if key_state_list is not None else None\n value_state = value_state_list[i] if value_state_list is not None else None\n encoder_key_state = encoder_key_state_list[i] if encoder_key_state_list is not None else None\n encoder_value_state = encoder_value_state_list[i] if encoder_value_state_list is not None else None\n x = decoder.get_layer('att_decoder_inp1')(query)\n x = decoder.get_layer('att_decoder_token_emb')(x)\n x = decoder.get_layer('att_decoder_pos_emb')(x, positions=query_position)\n\n q = x\n x = decoder.get_layer('att_decoder_block1_bn1')(x)\n x, k, v = decoder.get_layer('att_decoder_block1_self_attn')(x, x, x, key_state=key_state, value_state=value_state, return_states=True)\n x = decoder.get_layer('att_decoder_block1_add1')([q,x])\n attn_out1 = x\n\n x = decoder.get_layer('att_decoder_block1_bn2')(x)\n x = decoder.get_layer('att_decoder_block1_cross_attn')(x, None, None, key_state=encoder_key_state, value_state=encoder_value_state)\n x = decoder.get_layer('att_decoder_block1_add2')([attn_out1,x])\n attn_out2 = x\n\n x = decoder.get_layer('att_decoder_block1_bn3')(x)\n x = decoder.get_layer('att_decoder_block1_fc1')(x)\n x = decoder.get_layer('att_decoder_block1_fc2')(x)\n x = decoder.get_layer('att_decoder_block1_add3')([attn_out2,x])\n out = decoder.get_layer('att_decoder_classifier')(x)\n outputs.append(out)\n key_states.append(k)\n value_states.append(v)\n return tf.identity(outputs), tf.identity(key_states), tf.identity(value_states)\n\n def get_initial_states(self, batch_x):\n encoder_outputs = [enc(batch_x) for enc in self.encoder_list]\n encoder_key_states = [dec.get_layer('att_decoder_block1_cross_attn').k(x) for dec, x in zip(self.decoder_list, encoder_outputs)]\n encoder_value_states = [dec.get_layer('att_decoder_block1_cross_attn').v(x) for dec, x in zip(self.decoder_list, encoder_outputs)]\n key_states = [tf.zeros((0,0,192)) for _ in self.encoder_list]\n value_states = [tf.zeros((0,0,192)) for _ in self.encoder_list]\n ctc_probs = [dec(x)[0] for dec,x in zip(self.ctc_decoder_list, encoder_outputs)]\n \n encoder_key_states = tf.stack(encoder_key_states)\n encoder_value_states = tf.stack(encoder_value_states)\n key_states = tf.stack(key_states)\n value_states = tf.stack(value_states)\n encoder_outputs = tf.stack(encoder_outputs)\n \n if self.from_logits:\n ctc_probs = [tf.nn.softmax(x, axis=-1) for x in ctc_probs]\n ctc_probs = tf.stack([tf.math.log(x) for x in ctc_probs], axis=1)\n ctc_states_n, ctc_states_b = get_ctc_initial_states(ctc_probs, self.blank_idx)\n return encoder_key_states, encoder_value_states, key_states, value_states, ctc_probs, ctc_states_n, ctc_states_b\n\n def call(self, batch_x):\n\n encoder_key_state, encoder_value_state, key_state, value_state, ctc_log_probs, ctc_states_n, ctc_states_b = self.get_initial_states(batch_x)\n input_length = self.compute_input_length(batch_x)\n\n time = tf.constant(0, dtype=tf.int32)\n predictions = tf.ones((tf.shape(batch_x)[0],1), dtype=tf.int32) * self.sos_token_idx#tf.TensorArray(dtype=tf.int32,size=self.max_output_length)\n pad = tf.ones((tf.shape(batch_x)[0],), dtype=tf.int32) * self.pad_token_idx\n init = True\n\n def condition(_time, predictions, ctc_states_n, ctc_states_b, key_state, value_state, init):\n return tf.logical_and(_time < tf.minimum(self.max_output_length, input_length), tf.logical_not(tf.reduce_all(tf.reduce_any(predictions==self.eos_token_idx, axis=1))))\n\n def body(_time, predictions, ctc_states_n, ctc_states_b, key_state, value_state, init):\n if init:\n out, key_state, value_state = self.att_inference_module(predictions[:,-1:], _time, None, None, encoder_key_state, encoder_value_state)\n init = False\n else:\n out, key_state, value_state = self.att_inference_module(predictions[:,-1:], _time, key_state, value_state, encoder_key_state, encoder_value_state)\n\n log_ctc, new_ctc_states_n, new_ctc_states_b = compute_ctc_prefix_scores(predictions[0], \n ctc_log_probs, \n ctc_states_n, \n ctc_states_b, \n self.eos_token_idx, \n self.blank_idx)\n log_ctc = tf.reduce_mean(log_ctc, axis=0) #log-prob ensemble\n out = out[:,0,0] #(M,V)\n if self.from_logits:\n out = tf.nn.softmax(out, axis=-1)\n out = tf.math.log(out)\n log_att = tf.reduce_mean(out, axis=0) #log-prob ensemble\n \n probs_final = (1-self.ctc_weight) * log_att + self.ctc_weight * log_ctc #tf.expand_dims(log_psi, axis=0)\n next_token = tf.argmax(probs_final, axis=-1, output_type=tf.int32)#[0]\n \n ctc_states_n = new_ctc_states_n[...,next_token]\n ctc_states_b = new_ctc_states_b[...,next_token]\n \n predictions = tf.concat([predictions, [next_token[...,None]]], axis=1)\n return _time+1, predictions, ctc_states_n, ctc_states_b, key_state, value_state, init\n\n _, predictions, _, _, _, _, _ = tf.while_loop(condition, body, \n shape_invariants=[tf.TensorShape([]),\n tf.TensorShape([1,None]),\n tf.TensorShape([None,len(self.encoder_list)]),\n tf.TensorShape([None,len(self.encoder_list)]),\n tf.TensorShape([len(self.encoder_list),None,None,None]),\n tf.TensorShape([len(self.encoder_list),None,None,None]),\n tf.TensorShape([])], \n loop_vars=[time, predictions, ctc_states_n, ctc_states_b, key_state, value_state, init])\n return predictions","metadata":{"execution":{"iopub.status.busy":"2023-08-22T10:43:56.528436Z","iopub.execute_input":"2023-08-22T10:43:56.528881Z","iopub.status.idle":"2023-08-22T10:43:56.612084Z","shell.execute_reply.started":"2023-08-22T10:43:56.528844Z","shell.execute_reply":"2023-08-22T10:43:56.610461Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"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 self.preprocess = Preprocess()\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 # Preprocess Data\n x = tf.transpose(tf.reshape(inputs, (-1,3,543)), (0,2,1))\n x = tf.cond(tf.shape(x)[0] == 0, lambda: tf.zeros((1, 543, 3),dtype=tf.float32), lambda: tf.identity(x))\n x = self.preprocess(x)\n x = self.model(x)[0]\n x = x - 1\n idxs = tf.where((0<=x) & (x<=58))[...,0]\n x = tf.gather(x, idxs)\n x = tf.cond(tf.shape(x)[0] == 0, lambda: tf.zeros(1, tf.int32), lambda: tf.identity(x))\n x = tf.one_hot(x, 59)\n return {'outputs': x}","metadata":{"execution":{"iopub.status.busy":"2023-08-22T10:50:24.459026Z","iopub.execute_input":"2023-08-22T10:50:24.459934Z","iopub.status.idle":"2023-08-22T10:50:24.486111Z","shell.execute_reply.started":"2023-08-22T10:50:24.459849Z","shell.execute_reply":"2023-08-22T10:50:24.482812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf.keras.backend.clear_session()\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2023-08-22T10:50:26.066472Z","iopub.execute_input":"2023-08-22T10:50:26.066941Z","iopub.status.idle":"2023-08-22T10:50:33.176159Z","shell.execute_reply.started":"2023-08-22T10:50:26.066883Z","shell.execute_reply":"2023-08-22T10:50:33.174708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tflitemodel_base = TFLiteModel(CTCGreedyDecoder(model_list[0]))\n# tflitemodel_base = TFLiteModel(ATTGreedyDecoder(model_list[0]))\n# tflitemodel_base = TFLiteModel(EnsembleCTCAttentionJointGreedyDecoder(model_list[:3], ctc_weight=0.3))\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()}\npred = tflitemodel_base(frames)[\"outputs\"].numpy().argmax(-1)\n''.join([rev_character_map[x] for x in pred]), phrase","metadata":{"execution":{"iopub.status.busy":"2023-08-22T10:50:33.196862Z","iopub.execute_input":"2023-08-22T10:50:33.197445Z","iopub.status.idle":"2023-08-22T10:50:42.345466Z","shell.execute_reply.started":"2023-08-22T10:50:33.197404Z","shell.execute_reply":"2023-08-22T10:50:42.343872Z"},"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]\nkeras_model_converter._experimental_default_to_single_batch_in_tensor_list_ops = True\nkeras_model_converter.optimizations = [tf.lite.Optimize.DEFAULT]\nkeras_model_converter.target_spec.supported_types = [tf.float16]\ntflite_model = keras_model_converter.convert()\nwith open('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-08-22T10:50:42.362475Z","iopub.execute_input":"2023-08-22T10:50:42.366134Z","iopub.status.idle":"2023-08-22T10:52:14.198067Z","shell.execute_reply.started":"2023-08-22T10:50:42.36608Z","shell.execute_reply":"2023-08-22T10:52:14.196313Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"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, fold=None):\n def gen():\n for file_id in file_ids:\n pqfile = f\"{inpdir}/{file_id}.parquet\"\n if fold is not None:\n seq_refs = df.loc[(df.file_id == file_id) & (df.fold==fold)]\n else:\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)