{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":51294,"databundleVersionId":6923401,"sourceType":"competition"},{"sourceId":6889425,"sourceType":"datasetVersion","datasetId":3860063}],"dockerImageVersionId":30529,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Update (version 18)\nThis version has a massive speedup, from 138s per epoch to 42s per epoch, with large steps_per_executin and BS 128.  \nYou can get another speedup of ~25% (from 42s to 32s) with a global mixed precision policy of bfloat16 (thank you, Jayoo Hwang, for this pointer!) at the cost of 0.001 accuracy. I prefer the accuracy, but if the 10s speedup is important to you, you can get it with the following:  \n\npolicyConfig = 'mixed_bfloat16'  \npolicy = tf.keras.mixed_precision.Policy(policyConfig)  \ntf.keras.mixed_precision.set_global_policy(policy)  \n\nAnd a minor alteration of the positional encoding layer to make it compatible.  \nYou can get another minor speedup with BS 256 (38s, 27s with bfloat16 global mixed precision policy). Since it is so tiny, BS 128 is probably the ideal one. If you have exciting findings regarding speed optimization, particularly if you can achieve another massive speedup or train with mixed precision without sacrificing accuracy, please share it in the comments for the common good of all of us low-budget folks without a private 4090...  \np.s. I always thought that TPU already makes matrix multiplications in bfloat16, so I don't understand why setting bfloat16 mixed precision policy gives another boost. If anyone here knows the answer, please enlighten me.  \np.p.s. Since I now have a massive speedup, I increased the epochs from 200 to 400 because why not. I don't know if that will improve accuracy; we will find out together. So, for anyone interested in the previous results, this notebook can achieve approximately 0.1300-0.1303 validation score with 200 epochs.","metadata":{}},{"cell_type":"markdown","source":"12-layers small (192 dim) transformer structure, same as [this notebook](https://www.kaggle.com/code/iafoss/rna-starter-0-186-lb)  .        \nAdamW optimizer.  \nBatch size ~64~ 128.  \n~200~ 400 epochs.  \n[Inference notebook](https://www.kaggle.com/code/shlomoron/srrf-transformer-tpu-inference)","metadata":{}},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import tensorflow as tf\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport pickle\nimport shutil\nimport math\nimport pandas as pd\nimport gc\nimport os","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-11-30T12:01:28.753988Z","iopub.execute_input":"2023-11-30T12:01:28.754474Z","iopub.status.idle":"2023-11-30T12:01:39.549034Z","shell.execute_reply.started":"2023-11-30T12:01:28.754435Z","shell.execute_reply":"2023-11-30T12:01:39.547647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# TPU boilerplate code","metadata":{}},{"cell_type":"code","source":"# Configure Strategy. Assume TPU...if not set default for GPU\ntpu = None\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect(tpu=\"local\") # \"local\" for 1VM TPU\n    strategy = tf.distribute.TPUStrategy(tpu)\n    print(\"on TPU\")\n    print(\"REPLICAS: \", strategy.num_replicas_in_sync)\nexcept:\n    strategy = tf.distribute.get_strategy()","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:03:31.903417Z","iopub.execute_input":"2023-11-30T12:03:31.903919Z","iopub.status.idle":"2023-11-30T12:03:32.293409Z","shell.execute_reply.started":"2023-11-30T12:03:31.903884Z","shell.execute_reply":"2023-11-30T12:03:32.292007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configs","metadata":{}},{"cell_type":"code","source":"DEBUG = False\n\nPAD_x = 0.0\nPAD_y = np.nan\nX_max_len = 206\nbatch_size = 128\nval_batch_size = 5512\n\n\nif DEBUG:\n    batch_size = 2\n    val_batch_size = 2\n\nnum_vocab = 5\nhidden_dim = 192","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:03:34.553360Z","iopub.execute_input":"2023-11-30T12:03:34.553855Z","iopub.status.idle":"2023-11-30T12:03:34.560920Z","shell.execute_reply.started":"2023-11-30T12:03:34.553803Z","shell.execute_reply":"2023-11-30T12:03:34.559678Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data API pipeline\nThis section applies filtering, preprocessing, shuffling, paddings, and batchings. I already transformed all the data to TFRecords; you can find the TFRecords dataset [here](https://www.kaggle.com/datasets/shlomoron/srrf-tfrecords-ds). I shuffled the samples before creating the TFRecords.","metadata":{}},{"cell_type":"code","source":"tffiles_path = '/kaggle/input/srrf-tfrecords-ds/tfds'\ntffiles = [f'{tffiles_path}/{x}.tfrecord' for x in range(164)]","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:03:42.042758Z","iopub.execute_input":"2023-11-30T12:03:42.043216Z","iopub.status.idle":"2023-11-30T12:03:42.049703Z","shell.execute_reply.started":"2023-11-30T12:03:42.043179Z","shell.execute_reply":"2023-11-30T12:03:42.048407Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"- tf.io.VarLenFeature\n      - tf.io.VarLenFeature is a feature parsing function used in the context of the tf.io.parse_example or tf.io.parse_single_example functions. These functions are commonly used when working with TensorFlow's input pipelines to parse and preprocess data.\n       - VarLenFeature is specifically designed for parsing variable-length features, which are features that can have varying lengths across examples in your dataset. This is often used when dealing with sequences of data, such as sequences of words in natural language processing or sequences of values in time series analysis.","metadata":{}},{"cell_type":"code","source":"# example using tf.io.VarLenFeature to define a variable-length feature for sequences of integers\nimport tensorflow as tf\n\n# Suppose you have a dataset where each example is a sequence of integers.\n# In this case, the sequence length can vary across examples.\n\n# Create a TFRecord file with a few examples\ndef create_tfrecord_example(sequence):\n    feature = {\n        'sequence_feature': tf.train.Feature(int64_list=tf.train.Int64List(value=sequence))\n    }\n    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))\n    return example_proto.SerializeToString()\n\n# Write TFRecord file\nfilename = 'variable_length_sequence.tfrecord'\nwith tf.io.TFRecordWriter(filename) as writer:\n    writer.write(create_tfrecord_example([1, 2, 3]))\n    writer.write(create_tfrecord_example([4, 5]))\n    writer.write(create_tfrecord_example([6, 7, 8, 9]))\n\n# Define the feature description using VarLenFeature\nsequence_feature = tf.io.VarLenFeature(tf.int64)\n\nfeature_description = {\n    'sequence_feature': sequence_feature\n}\n\n# Parse the TFRecord file\nraw_dataset = tf.data.TFRecordDataset([filename])\n\nfor raw_record in raw_dataset.take(3):\n    # Parse each example\n    parsed_example = tf.io.parse_single_example(raw_record, feature_description)\n\n    # Access the parsed variable-length feature\n    parsed_sequence_feature = parsed_example['sequence_feature']\n\n    # Convert the sparse representation to a dense tensor\n    dense_sequence_feature = tf.sparse.to_dense(parsed_sequence_feature)\n\n    print(\"Parsed Sequence Feature:\", dense_sequence_feature.numpy())\n","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:04:03.853258Z","iopub.execute_input":"2023-11-30T12:04:03.853721Z","iopub.status.idle":"2023-11-30T12:04:04.083227Z","shell.execute_reply.started":"2023-11-30T12:04:03.853687Z","shell.execute_reply":"2023-11-30T12:04:04.081799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Decoding TFRecords","metadata":{}},{"cell_type":"code","source":"def decode_tfrec(record_bytes):\n    schema = {}\n    schema[\"id\"] = tf.io.VarLenFeature(dtype=tf.string)\n    schema[\"seq\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"dataset_name_2A3\"] = tf.io.VarLenFeature(dtype=tf.string)\n    schema[\"dataset_name_DMS\"] = tf.io.VarLenFeature(dtype=tf.string)\n    schema[\"reads_2A3\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"reads_DMS\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"signal_to_noise_2A3\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"signal_to_noise_DMS\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"SN_filter_2A3\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"SN_filter_DMS\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"reactivity_2A3\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"reactivity_DMS\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"reactivity_error_2A3\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    schema[\"reactivity_error_DMS\"] = tf.io.VarLenFeature(dtype=tf.float32)\n    features = tf.io.parse_single_example(record_bytes, schema)\n\n    sample_id = tf.sparse.to_dense(features[\"id\"])\n    seq = tf.sparse.to_dense(features[\"seq\"])\n    dataset_name_2A3 = tf.sparse.to_dense(features[\"dataset_name_2A3\"])\n    dataset_name_DMS = tf.sparse.to_dense(features[\"dataset_name_DMS\"])\n    reads_2A3 = tf.sparse.to_dense(features[\"reads_2A3\"])\n    reads_DMS = tf.sparse.to_dense(features[\"reads_DMS\"])\n    signal_to_noise_2A3 = tf.sparse.to_dense(features[\"signal_to_noise_2A3\"])\n    signal_to_noise_DMS = tf.sparse.to_dense(features[\"signal_to_noise_DMS\"])\n    SN_filter_2A3 = tf.sparse.to_dense(features[\"SN_filter_2A3\"])\n    SN_filter_DMS = tf.sparse.to_dense(features[\"SN_filter_DMS\"])\n    reactivity_2A3 = tf.sparse.to_dense(features[\"reactivity_2A3\"])\n    reactivity_DMS = tf.sparse.to_dense(features[\"reactivity_DMS\"])\n    reactivity_error_2A3 = tf.sparse.to_dense(features[\"reactivity_error_2A3\"])\n    reactivity_error_DMS = tf.sparse.to_dense(features[\"reactivity_error_DMS\"])\n\n    out = {}\n    out['seq']  = seq\n    out['SN_filter_2A3']  = SN_filter_2A3\n    out['SN_filter_DMS']  = SN_filter_DMS\n    out['reads_2A3']  = reads_2A3\n    out['reads_DMS']  = reads_DMS\n    out['signal_to_noise_2A3']  = signal_to_noise_2A3\n    out['signal_to_noise_DMS']  = signal_to_noise_DMS\n    out['reactivity_2A3']  = reactivity_2A3\n    out['reactivity_DMS']  = reactivity_DMS\n    return out","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:04:09.403314Z","iopub.execute_input":"2023-11-30T12:04:09.403762Z","iopub.status.idle":"2023-11-30T12:04:09.422875Z","shell.execute_reply.started":"2023-11-30T12:04:09.403730Z","shell.execute_reply":"2023-11-30T12:04:09.421314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Filtering","metadata":{}},{"cell_type":"code","source":"def f1(): return True\ndef f2(): return False\n\ndef filter_function_1(x):\n    SN_filter_2A3 = x['SN_filter_2A3']\n    SN_filter_DMS = x['SN_filter_DMS']\n    return tf.cond((SN_filter_2A3 == 1) and (SN_filter_DMS == 1) , true_fn=f1, false_fn=f2)\n\ndef filter_function_2(x):\n    reads_2A3 = x['reads_2A3']\n    reads_DMS = x['reads_DMS']\n    signal_to_noise_2A3 = x['signal_to_noise_2A3']\n    signal_to_noise_DMS = x['signal_to_noise_DMS']\n    cond = (reads_2A3>100 and signal_to_noise_2A3>0.75) or (reads_DMS>100 and signal_to_noise_DMS>0.75)\n    return tf.cond(cond, true_fn=f1, false_fn=f2)","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:04:13.204397Z","iopub.execute_input":"2023-11-30T12:04:13.204915Z","iopub.status.idle":"2023-11-30T12:04:13.215877Z","shell.execute_reply.started":"2023-11-30T12:04:13.204879Z","shell.execute_reply":"2023-11-30T12:04:13.214049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Preprocessing","metadata":{}},{"cell_type":"code","source":"def nan_below_filter(x):\n    reads_2A3 = x['reads_2A3']\n    reads_DMS = x['reads_DMS']\n    signal_to_noise_2A3 = x['signal_to_noise_2A3']\n    signal_to_noise_DMS = x['signal_to_noise_DMS']\n    reactivity_2A3 = x['reactivity_2A3']\n    reactivity_DMS = x['reactivity_DMS']\n\n    if reads_2A3<100 or signal_to_noise_2A3<0.75:\n        reactivity_2A3 = np.nan+reactivity_2A3\n    if reads_DMS<100 or signal_to_noise_DMS<0.75:\n        reactivity_DMS = np.nan+reactivity_DMS\n\n    x['reactivity_2A3'] = reactivity_2A3\n    x['reactivity_DMS'] = reactivity_DMS\n    return x\n\ndef concat_target(x):\n    reactivity_2A3 = x['reactivity_2A3']\n    reactivity_DMS = x['reactivity_DMS']\n    target = tf.concat([reactivity_2A3[..., tf.newaxis], reactivity_DMS[..., tf.newaxis]], axis = 1)\n    target = tf.clip_by_value(target, 0, 1)\n    return x['seq'], target","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:04:38.782670Z","iopub.execute_input":"2023-11-30T12:04:38.783144Z","iopub.status.idle":"2023-11-30T12:04:38.793883Z","shell.execute_reply.started":"2023-11-30T12:04:38.783109Z","shell.execute_reply":"2023-11-30T12:04:38.792375Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## get_tfrec_dataset","metadata":{}},{"cell_type":"code","source":"def get_tfrec_dataset(tffiles, shuffle, batch_size, cache = False, to_filter = False,\n                      calculate_sample_num = True, to_repeat = False):\n    ds = tf.data.TFRecordDataset(\n        tffiles, num_parallel_reads=tf.data.AUTOTUNE, compression_type = 'GZIP').prefetch(tf.data.AUTOTUNE)\n\n    ds = ds.map(decode_tfrec, tf.data.AUTOTUNE)\n    if to_filter == 'filter_1':\n        ds = ds.filter(filter_function_1)\n    elif to_filter == 'filter_2':\n        ds = ds.filter(filter_function_2)\n    ds = ds.map(nan_below_filter, tf.data.AUTOTUNE)\n    ds = ds.map(concat_target, tf.data.AUTOTUNE)\n\n    if DEBUG:\n        ds = ds.take(8)\n\n    if cache:\n        ds = ds.cache()\n\n    samples_num = 0\n    if calculate_sample_num:\n        samples_num = ds.reduce(0, lambda x,_: x+1).numpy()        \n\n    if shuffle:\n        if shuffle == -1:\n            ds = ds.shuffle(samples_num, reshuffle_each_iteration = True)\n        else:\n            ds = ds.shuffle(shuffle, reshuffle_each_iteration = True)\n\n    if to_repeat:\n        ds = ds.repeat()\n         \n    if batch_size:\n        ds = ds.padded_batch(\n            batch_size, padding_values=(PAD_x, PAD_y), padded_shapes=([X_max_len],[X_max_len, 2]), drop_remainder=True)\n    ds = ds.prefetch(tf.data.AUTOTUNE)\n    return ds, samples_num","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:07:38.229008Z","iopub.execute_input":"2023-11-30T12:07:38.229500Z","iopub.status.idle":"2023-11-30T12:07:38.243484Z","shell.execute_reply.started":"2023-11-30T12:07:38.229465Z","shell.execute_reply":"2023-11-30T12:07:38.242041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define datasets","metadata":{}},{"cell_type":"code","source":"val_len = 5\nif DEBUG:\n    val_len = 1\n\nval_files = tffiles[:val_len]\n\nif DEBUG:\n    train_files = tffiles[val_len:val_len+1]\nelse:\n    train_files = tffiles[val_len:]","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:07:44.202680Z","iopub.execute_input":"2023-11-30T12:07:44.204071Z","iopub.status.idle":"2023-11-30T12:07:44.212906Z","shell.execute_reply.started":"2023-11-30T12:07:44.204007Z","shell.execute_reply":"2023-11-30T12:07:44.210356Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Get datasets","metadata":{}},{"cell_type":"code","source":"train_dataset, num_train = get_tfrec_dataset(train_files, shuffle = -1, batch_size = batch_size,\n                                                  cache = True, to_filter = 'filter_2', calculate_sample_num = True,\n                                            to_repeat = True)\n\nval_dataset, num_val = get_tfrec_dataset(val_files, shuffle = False, batch_size = val_batch_size,\n                                                  cache = True, to_filter = 'filter_1', calculate_sample_num = True)\nprint(num_train)\nprint(num_val)","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:07:47.003313Z","iopub.execute_input":"2023-11-30T12:07:47.004257Z","iopub.status.idle":"2023-11-30T12:10:02.189426Z","shell.execute_reply.started":"2023-11-30T12:07:47.004207Z","shell.execute_reply":"2023-11-30T12:10:02.187201Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch = next(iter(val_dataset))\nbatch[0].shape, batch[1].shape","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:10:17.373036Z","iopub.execute_input":"2023-11-30T12:10:17.373486Z","iopub.status.idle":"2023-11-30T12:10:17.462081Z","shell.execute_reply.started":"2023-11-30T12:10:17.373448Z","shell.execute_reply":"2023-11-30T12:10:17.460430Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"markdown","source":"## Model layers","metadata":{}},{"cell_type":"code","source":"class transformer_block(tf.keras.layers.Layer):\n    def __init__(self, dim, num_heads, feed_forward_dim, rate=0.1):\n        super().__init__()\n        self.att = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=dim//num_heads)\n        self.ffn = tf.keras.Sequential(\n            [\n                tf.keras.layers.Dense(feed_forward_dim, activation=\"relu\"),\n                tf.keras.layers.Dense(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        self.supports_masking = True\n\n    def call(self, inputs, training, mask):\n        att_mask = tf.expand_dims(mask, axis=-1)\n        att_mask = tf.repeat(att_mask, repeats=tf.shape(att_mask)[1], axis=-1)\n\n        attn_output = self.att(inputs, inputs, attention_mask = att_mask)\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\n\nclass positional_encoding_layer(tf.keras.layers.Layer):\n    def __init__(self, num_vocab=5, maxlen=500, hidden_dim=384):\n        super().__init__()\n        self.hidden_dim = hidden_dim\n        self.pos_emb = self.positional_encoding(maxlen-1, hidden_dim)\n        self.supports_masking = True\n\n    def call(self, x):\n        maxlen = tf.shape(x)[-2]\n        x = tf.math.multiply(x, tf.math.sqrt(tf.cast(self.hidden_dim, tf.float32)))\n        return x + self.pos_emb[:maxlen, :]\n\n    def positional_encoding(self, maxlen, hidden_dim):\n        depth = hidden_dim/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","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:10:23.963069Z","iopub.execute_input":"2023-11-30T12:10:23.963513Z","iopub.status.idle":"2023-11-30T12:10:23.984275Z","shell.execute_reply.started":"2023-11-30T12:10:23.963472Z","shell.execute_reply":"2023-11-30T12:10:23.983181Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loss function","metadata":{}},{"cell_type":"code","source":"def loss_fn(labels, targets):\n    labels_mask = tf.math.is_nan(labels)\n    labels = tf.where(labels_mask, tf.zeros_like(labels), labels)\n    mask_count = tf.math.reduce_sum(tf.where(labels_mask, tf.zeros_like(labels), tf.ones_like(labels)))\n    loss = tf.math.abs(labels - targets)\n    loss = tf.where(labels_mask, tf.zeros_like(loss), loss)\n    loss = tf.math.reduce_sum(loss)/mask_count\n    return loss","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:10:27.093147Z","iopub.execute_input":"2023-11-30T12:10:27.093558Z","iopub.status.idle":"2023-11-30T12:10:27.101576Z","shell.execute_reply.started":"2023-11-30T12:10:27.093528Z","shell.execute_reply":"2023-11-30T12:10:27.100541Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"def get_model(hidden_dim = 384, max_len = 206):\n    with strategy.scope():\n        inp = tf.keras.Input([max_len])\n        x = inp\n\n        x = tf.keras.layers.Embedding(num_vocab, hidden_dim, mask_zero=True)(x)\n        x = positional_encoding_layer(num_vocab=num_vocab, maxlen=500, hidden_dim=hidden_dim)(x)\n\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n        x = transformer_block(hidden_dim, 6, hidden_dim*4)(x)\n\n        x = tf.keras.layers.Dropout(0.5)(x)\n        x = tf.keras.layers.Dense(2)(x)\n\n        model = tf.keras.Model(inp, x)\n        loss = loss_fn\n        optimizer = tf.keras.optimizers.AdamW(learning_rate=0.0005)\n        model.compile(loss=loss, optimizer=optimizer, steps_per_execution = 100)\n        return model\n\ntf.keras.backend.clear_session()\n\nmodel = get_model(hidden_dim = 192,max_len = X_max_len)\nmodel(batch[0])\nmodel.summary()","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:10:30.203558Z","iopub.execute_input":"2023-11-30T12:10:30.204115Z","iopub.status.idle":"2023-11-30T12:20:57.335092Z","shell.execute_reply.started":"2023-11-30T12:10:30.204073Z","shell.execute_reply":"2023-11-30T12:20:57.333616Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Learning rate scheduler\nI copied the scheduler from [here](https://www.kaggle.com/code/irohith/aslfr-ctc-based-on-prev-comp-1st-place).","metadata":{}},{"cell_type":"code","source":"N_EPOCHS = 40\nif DEBUG:\n    N_EPOCHS = 5\nN_WARMUP_EPOCHS = 0\nLR_MAX = 5e-4\nWD_RATIO = 0.05\nWARMUP_METHOD = \"exp\"","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:23:18.923728Z","iopub.execute_input":"2023-11-30T12:23:18.924346Z","iopub.status.idle":"2023-11-30T12:23:18.932794Z","shell.execute_reply.started":"2023-11-30T12:23:18.924297Z","shell.execute_reply":"2023-11-30T12:23:18.930987Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def lrfn(current_step, num_warmup_steps, lr_max, num_cycles=0.50, num_training_steps=N_EPOCHS):\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# 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)","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:23:21.623690Z","iopub.execute_input":"2023-11-30T12:23:21.625231Z","iopub.status.idle":"2023-11-30T12:23:22.585197Z","shell.execute_reply.started":"2023-11-30T12:23:21.625171Z","shell.execute_reply":"2023-11-30T12:23:22.583821Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Saving callback","metadata":{}},{"cell_type":"code","source":"save_folder = '/kaggle/working'\ntry:\n    os.mkdir(f'{save_folder}/weights')\nexcept:\n    pass\n\nclass save_model_callback(tf.keras.callbacks.Callback):\n    def __init__(self):\n        super().__init__()\n    def on_epoch_end(self, epoch: int, logs=None):\n        if epoch == 3 or (epoch+1)%25 == 0:\n            self.model.save_weights(f\"{save_folder}/weights/model_epoch_{epoch}.h5\")","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:23:26.013432Z","iopub.execute_input":"2023-11-30T12:23:26.013951Z","iopub.status.idle":"2023-11-30T12:23:26.023330Z","shell.execute_reply.started":"2023-11-30T12:23:26.013896Z","shell.execute_reply":"2023-11-30T12:23:26.022081Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"steps_per_epoch = num_train//batch_size\nval_steps_per_epoch = num_val//val_batch_size\nprint(steps_per_epoch)\nprint(val_steps_per_epoch)","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:23:28.843618Z","iopub.execute_input":"2023-11-30T12:23:28.844122Z","iopub.status.idle":"2023-11-30T12:23:28.851726Z","shell.execute_reply.started":"2023-11-30T12:23:28.844083Z","shell.execute_reply":"2023-11-30T12:23:28.849964Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"history = model.fit(\n    train_dataset,\n    validation_data=val_dataset,\n    epochs=N_EPOCHS,\n    steps_per_epoch = steps_per_epoch,\n    validation_steps=val_steps_per_epoch,\n    verbose = 2,\n    callbacks=[\n        save_model_callback(),\n        lr_callback,\n    ]\n)","metadata":{"execution":{"iopub.status.busy":"2023-11-30T12:23:31.253223Z","iopub.execute_input":"2023-11-30T12:23:31.254095Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Plotting loss","metadata":{}},{"cell_type":"code","source":"plt.plot(history.history['loss'])\nplt.plot(history.history['val_loss'])","metadata":{"execution":{"iopub.status.busy":"2023-10-22T12:08:11.208829Z","iopub.status.idle":"2023-10-22T12:08:11.209136Z","shell.execute_reply.started":"2023-10-22T12:08:11.208992Z","shell.execute_reply":"2023-10-22T12:08:11.209006Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference\nDoing the inference in a separate notebook was easier, so I split it. Sorry about that. Find the inference notebook [HERE](https://www.kaggle.com/code/shlomoron/srrf-transformer-tpu-inference).","metadata":{}}]}