{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This is an inference notebook. Find the training notebook [here](https://www.kaggle.com/code/shlomoron/srrf-transformer-tpu-training/notebook).","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-10-19T14:28:41.513234Z","iopub.execute_input":"2023-10-19T14:28:41.513477Z","iopub.status.idle":"2023-10-19T14:28:49.213021Z","shell.execute_reply.started":"2023-10-19T14:28:41.513455Z","shell.execute_reply":"2023-10-19T14:28:49.212213Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"DEBUG = False","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:28:49.214443Z","iopub.execute_input":"2023-10-19T14:28:49.214892Z","iopub.status.idle":"2023-10-19T14:28:49.218505Z","shell.execute_reply.started":"2023-10-19T14:28:49.214869Z","shell.execute_reply":"2023-10-19T14:28:49.217644Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-10-19T14:28:49.219494Z","iopub.execute_input":"2023-10-19T14:28:49.219730Z","iopub.status.idle":"2023-10-19T14:28:49.244819Z","shell.execute_reply.started":"2023-10-19T14:28:49.219711Z","shell.execute_reply":"2023-10-19T14:28:49.244179Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_max_len = 457\nnum_vocab = 5\n\ndef get_model(hidden_dim = 384, max_len = 206):\n    inp = tf.keras.Input([None])\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    return model","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:28:49.246305Z","iopub.execute_input":"2023-10-19T14:28:49.246897Z","iopub.status.idle":"2023-10-19T14:28:49.258791Z","shell.execute_reply.started":"2023-10-19T14:28:49.246876Z","shell.execute_reply":"2023-10-19T14:28:49.258044Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"test_sequences_df = pd.read_csv('/kaggle/input/stanford-ribonanza-rna-folding/test_sequences.csv')\ntest_sequences_df.head()","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:28:49.259778Z","iopub.execute_input":"2023-10-19T14:28:49.260217Z","iopub.status.idle":"2023-10-19T14:28:56.009566Z","shell.execute_reply.started":"2023-10-19T14:28:49.260196Z","shell.execute_reply":"2023-10-19T14:28:56.008688Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_sequences = test_sequences_df.sequence.to_numpy()\nencoding_dict = {'A':1, 'C': 2, 'G': 3, 'U': 4}\nencoding_dict","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:28:56.010477Z","iopub.execute_input":"2023-10-19T14:28:56.010707Z","iopub.status.idle":"2023-10-19T14:28:56.016184Z","shell.execute_reply.started":"2023-10-19T14:28:56.010688Z","shell.execute_reply":"2023-10-19T14:28:56.015332Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"max_len = 457 \ntest_sequences_encoded = []\nfor seq in test_sequences:\n    test_sequences_encoded.append(\n        np.concatenate([np.asarray([encoding_dict[x] for x in seq]), np.zeros((max_len - len(seq)))]).astype(np.float32))","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:28:56.017304Z","iopub.execute_input":"2023-10-19T14:28:56.017919Z","iopub.status.idle":"2023-10-19T14:29:36.908740Z","shell.execute_reply.started":"2023-10-19T14:28:56.017889Z","shell.execute_reply":"2023-10-19T14:29:36.907971Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_ds = tf.data.Dataset.from_tensor_slices(test_sequences_encoded)\nbatch_size = 256\nif DEBUG:\n    test_ds = test_ds.take(8)\n    batch_size = 2\n#test_ds = test_ds.take(10000)\n\ntest_ds = test_ds.padded_batch(batch_size, padding_values=(0.0), padded_shapes=([max_len]), drop_remainder=False)\ntest_ds = test_ds.prefetch(tf.data.AUTOTUNE)\nbatch = next(iter(test_ds))\nbatch.shape","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:29:36.909711Z","iopub.execute_input":"2023-10-19T14:29:36.909936Z","iopub.status.idle":"2023-10-19T14:30:57.877628Z","shell.execute_reply.started":"2023-10-19T14:29:36.909916Z","shell.execute_reply":"2023-10-19T14:30:57.876626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = get_model(hidden_dim = 192,max_len = max_len)\nmodel.load_weights('/kaggle/input/srrf-transformer-tpu-training/weights/model_epoch_199.h5')\nmodel(batch)\nmodel.summary()","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:30:57.878771Z","iopub.execute_input":"2023-10-19T14:30:57.879090Z","iopub.status.idle":"2023-10-19T14:31:03.076616Z","shell.execute_reply.started":"2023-10-19T14:30:57.879060Z","shell.execute_reply":"2023-10-19T14:31:03.075738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"preds = model.predict(test_ds)","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:31:03.078716Z","iopub.execute_input":"2023-10-19T14:31:03.078966Z","iopub.status.idle":"2023-10-19T14:32:08.702654Z","shell.execute_reply.started":"2023-10-19T14:31:03.078945Z","shell.execute_reply":"2023-10-19T14:32:08.701392Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"preds_processed = []\nfor i, pred in enumerate(preds):\n    preds_processed.append(pred[:len(test_sequences[i])])\nconcat_preds = np.concatenate(preds_processed)","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:32:12.120017Z","iopub.execute_input":"2023-10-19T14:32:12.120679Z","iopub.status.idle":"2023-10-19T14:32:12.147164Z","shell.execute_reply.started":"2023-10-19T14:32:12.120648Z","shell.execute_reply":"2023-10-19T14:32:12.146033Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = pd.DataFrame({'id':np.arange(0, len(concat_preds), 1), 'reactivity_DMS_MaP':concat_preds[:,1], 'reactivity_2A3_MaP':concat_preds[:,0]})\nsubmission.to_csv('submission.csv', index=False)\nsubmission.head()","metadata":{"execution":{"iopub.status.busy":"2023-10-19T14:32:08.705130Z","iopub.status.idle":"2023-10-19T14:32:08.705564Z","shell.execute_reply.started":"2023-10-19T14:32:08.705329Z","shell.execute_reply":"2023-10-19T14:32:08.705348Z"},"trusted":true},"execution_count":null,"outputs":[]}]}