{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.8.17","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","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":"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.  \n200 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":{"execution":{"iopub.status.busy":"2023-11-25T23:54:44.539367Z","iopub.execute_input":"2023-11-25T23:54:44.540372Z","iopub.status.idle":"2023-11-25T23:55:25.272018Z","shell.execute_reply.started":"2023-11-25T23:54:44.540330Z","shell.execute_reply":"2023-11-25T23:55:25.270764Z"},"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-25T23:55:25.274205Z","iopub.execute_input":"2023-11-25T23:55:25.275022Z","iopub.status.idle":"2023-11-25T23:55:33.071366Z","shell.execute_reply.started":"2023-11-25T23:55:25.274982Z","shell.execute_reply":"2023-11-25T23:55:33.070389Z"},"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 = 64\n\nif DEBUG:\n    batch_size = 2\n\nnum_vocab = 5\nhidden_dim = 192","metadata":{"execution":{"iopub.status.busy":"2023-11-25T23:55:33.072522Z","iopub.execute_input":"2023-11-25T23:55:33.072809Z","iopub.status.idle":"2023-11-25T23:55:33.077606Z","shell.execute_reply.started":"2023-11-25T23:55:33.072783Z","shell.execute_reply":"2023-11-25T23:55:33.076813Z"},"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-25T23:55:33.079656Z","iopub.execute_input":"2023-11-25T23:55:33.079931Z","iopub.status.idle":"2023-11-25T23:55:33.091438Z","shell.execute_reply.started":"2023-11-25T23:55:33.079909Z","shell.execute_reply":"2023-11-25T23:55:33.090618Z"},"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-25T23:55:33.092578Z","iopub.execute_input":"2023-11-25T23:55:33.092840Z","iopub.status.idle":"2023-11-25T23:55:33.107076Z","shell.execute_reply.started":"2023-11-25T23:55:33.092819Z","shell.execute_reply":"2023-11-25T23:55:33.106203Z"},"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-25T23:55:33.108123Z","iopub.execute_input":"2023-11-25T23:55:33.108369Z","iopub.status.idle":"2023-11-25T23:55:33.122350Z","shell.execute_reply.started":"2023-11-25T23:55:33.108347Z","shell.execute_reply":"2023-11-25T23:55:33.121522Z"},"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-25T23:55:33.123396Z","iopub.execute_input":"2023-11-25T23:55:33.123678Z","iopub.status.idle":"2023-11-25T23:55:33.134596Z","shell.execute_reply.started":"2023-11-25T23:55:33.123637Z","shell.execute_reply":"2023-11-25T23:55:33.133822Z"},"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):\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 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-25T23:55:33.135481Z","iopub.execute_input":"2023-11-25T23:55:33.135751Z","iopub.status.idle":"2023-11-25T23:55:33.149278Z","shell.execute_reply.started":"2023-11-25T23:55:33.135728Z","shell.execute_reply":"2023-11-25T23:55:33.148522Z"},"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-25T23:55:33.150226Z","iopub.execute_input":"2023-11-25T23:55:33.150480Z","iopub.status.idle":"2023-11-25T23:55:33.176681Z","shell.execute_reply.started":"2023-11-25T23:55:33.150457Z","shell.execute_reply":"2023-11-25T23:55:33.175873Z"},"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\nval_dataset, num_val = get_tfrec_dataset(val_files, shuffle = False, batch_size = 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-25T23:55:33.177779Z","iopub.execute_input":"2023-11-25T23:55:33.178034Z","iopub.status.idle":"2023-11-25T23:56:33.308698Z","shell.execute_reply.started":"2023-11-25T23:55:33.178012Z","shell.execute_reply":"2023-11-25T23:56:33.307586Z"},"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-25T23:56:33.311771Z","iopub.execute_input":"2023-11-25T23:56:33.312084Z","iopub.status.idle":"2023-11-25T23:56:33.364375Z","shell.execute_reply.started":"2023-11-25T23:56:33.312054Z","shell.execute_reply":"2023-11-25T23:56:33.363431Z"},"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-25T23:56:33.365550Z","iopub.execute_input":"2023-11-25T23:56:33.365883Z","iopub.status.idle":"2023-11-25T23:56:33.380576Z","shell.execute_reply.started":"2023-11-25T23:56:33.365855Z","shell.execute_reply":"2023-11-25T23:56:33.379707Z"},"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-25T23:56:33.381738Z","iopub.execute_input":"2023-11-25T23:56:33.382006Z","iopub.status.idle":"2023-11-25T23:56:33.394922Z","shell.execute_reply.started":"2023-11-25T23:56:33.381984Z","shell.execute_reply":"2023-11-25T23:56:33.394160Z"},"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, jit_compile = True)\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-25T23:56:33.395816Z","iopub.execute_input":"2023-11-25T23:56:33.396030Z","iopub.status.idle":"2023-11-25T23:56:41.778162Z","shell.execute_reply.started":"2023-11-25T23:56:33.396011Z","shell.execute_reply":"2023-11-25T23:56:41.777065Z"},"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 = 200\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-25T23:56:41.779321Z","iopub.execute_input":"2023-11-25T23:56:41.779625Z","iopub.status.idle":"2023-11-25T23:56:41.784163Z","shell.execute_reply.started":"2023-11-25T23:56:41.779599Z","shell.execute_reply":"2023-11-25T23:56:41.783333Z"},"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-25T23:56:41.785214Z","iopub.execute_input":"2023-11-25T23:56:41.785448Z","iopub.status.idle":"2023-11-25T23:56:42.907699Z","shell.execute_reply.started":"2023-11-25T23:56:41.785428Z","shell.execute_reply":"2023-11-25T23:56:42.906708Z"},"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-25T23:56:42.908763Z","iopub.execute_input":"2023-11-25T23:56:42.909000Z","iopub.status.idle":"2023-11-25T23:56:42.915242Z","shell.execute_reply.started":"2023-11-25T23:56:42.908976Z","shell.execute_reply":"2023-11-25T23:56:42.914376Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"history = model.fit(\n    train_dataset,\n    validation_data=val_dataset,\n    epochs=N_EPOCHS,\n    verbose = 2,\n    callbacks=[\n        save_model_callback(),\n        lr_callback,\n    ]\n)","metadata":{"execution":{"iopub.status.busy":"2023-11-25T23:56:42.916132Z","iopub.execute_input":"2023-11-25T23:56:42.916350Z"},"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":{"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":{}}]}