{"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":"gpu","dataSources":[{"sourceId":51294,"databundleVersionId":6923401,"sourceType":"competition"},{"sourceId":7049090,"sourceType":"datasetVersion","datasetId":3912807},{"sourceId":7095600,"sourceType":"datasetVersion","datasetId":4060502},{"sourceId":7065049,"sourceType":"datasetVersion","datasetId":3874845,"isSourceIdPinned":true}],"dockerImageVersionId":30559,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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\nimport glob","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:13.599502Z","iopub.execute_input":"2023-12-02T05:42:13.599891Z","iopub.status.idle":"2023-12-02T05:42:21.982771Z","shell.execute_reply.started":"2023-12-02T05:42:13.599857Z","shell.execute_reply":"2023-12-02T05:42:21.981886Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"\n\ndef get_strategy(device='TPU-VM'):\n    if \"TPU\" in device:\n        try:\n            tpu = 'local' if device == 'TPU-VM' else None\n            print(\"Connecting to TPU...\")\n            tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect(tpu=tpu)\n            strategy = tf.distribute.TPUStrategy(tpu)\n            IS_TPU = True\n            print(\"Connected to TPU.\")\n        except:\n            print(\"TPU not available. Falling back to GPU...\")\n            device = \"GPU\"\n\n    if device == \"GPU\":\n        IS_TPU = False\n        ngpu = len(tf.config.experimental.list_physical_devices('GPU'))\n        if ngpu > 1:\n            print(\"Using multiple GPUs.\")\n            strategy = tf.distribute.MirroredStrategy()\n        elif ngpu == 1:\n            print(\"Using single GPU.\")\n            strategy = tf.distribute.OneDeviceStrategy(\"GPU:0\")\n        else:\n            print(\"No GPU available. Falling back to CPU...\")\n            device = \"CPU\"\n\n    if device == \"CPU\":\n        print(\"Using CPU.\")\n        strategy = tf.distribute.OneDeviceStrategy(\"CPU\")\n        IS_TPU = False\n\n    AUTO = tf.data.experimental.AUTOTUNE\n    REPLICAS = strategy.num_replicas_in_sync\n    print(f\"REPLICAS: {REPLICAS}\")\n\n    return strategy, REPLICAS, IS_TPU\n\nSTRATEGY, N_REPLICAS, IS_TPU = get_strategy()\n\n","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:21.984467Z","iopub.execute_input":"2023-12-02T05:42:21.985039Z","iopub.status.idle":"2023-12-02T05:42:24.889477Z","shell.execute_reply.started":"2023-12-02T05:42:21.984985Z","shell.execute_reply":"2023-12-02T05:42:24.888516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEBUG = False","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:24.890760Z","iopub.execute_input":"2023-12-02T05:42:24.891155Z","iopub.status.idle":"2023-12-02T05:42:24.895933Z","shell.execute_reply.started":"2023-12-02T05:42:24.891118Z","shell.execute_reply":"2023-12-02T05:42:24.895056Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def swish(x):\n    return x * tf.nn.sigmoid(x)\n\nclass MultiHeadAttention(tf.keras.layers.Layer):\n    def __init__(self, d_model, n_head, d_k, d_v, dropout=0.1):\n        super(MultiHeadAttention, self).__init__()\n        self.n_head = n_head\n        self.d_k = d_k\n        self.d_v = d_v\n        self.d_model = d_model\n        \n        self.w_qs = tf.keras.layers.Dense(n_head * d_k, use_bias=False)\n        self.w_ks = tf.keras.layers.Dense(n_head * d_k, use_bias=False)\n        self.w_vs = tf.keras.layers.Dense(n_head * d_v, use_bias=False)\n        self.fc = tf.keras.layers.Dense(d_model, use_bias=False)\n        #self.qkv = tf.keras.layers.Dense(3*d_model, use_bias=False)\n        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5, num_heads=n_head)\n\n        self.dropout = tf.keras.layers.Dropout(dropout)\n        # self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n\n    def call(self, q, k, v, mask=None, pad_mask=None):\n        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head\n        bs, len_q, len_k, len_v = q.shape[0], q.shape[1], k.shape[1], v.shape[1]\n\n        residual = q#inputs\n        #qkv = self.qkv(inputs)\n        # Pass through the pre-attention projection: b x lq x (n*dv)\n        # Separate different heads: b x lq x n x dv\n        #print(qkv.shape)\n        #qkv = tf.keras.layers.Permute((2, 1, 3))(tf.keras.layers.Reshape((-1, self.n_head, self.d_model * 3 // self.n_head))(qkv))\n        #q, k, v = tf.split(qkv, [self.d_model // self.n_head] * 3, axis=-1)\n        q = tf.keras.layers.Reshape((-1, n_head, d_k))(self.w_qs(q))\n        k = tf.keras.layers.Reshape((-1, n_head, d_k))(self.w_ks(k))\n        v = tf.keras.layers.Reshape((-1, n_head, d_v))(self.w_vs(v))\n        \n        # Transpose for attention dot product: b x n x lq x dv\n        q, k, v = tf.keras.layers.Permute((2, 1, 3))(q), tf.keras.layers.Permute((2, 1, 3))(k) , tf.keras.layers.Permute((2, 1, 3))(v)\n        if mask is not None:\n            mask = mask  # For head axis broadcasting\n\n        if pad_mask is not None:\n            pad_mask = pad_mask[:, :q.shape[2]]\n            pad_mask = tf.expand_dims(pad_mask, axis=-1)\n            pad_mask = tf.matmul(pad_mask, tf.transpose(pad_mask, perm=(0, 2, 1)))\n            pad_mask = tf.expand_dims(pad_mask, axis=1)\n            #pad_mask = tf.cast(pad_mask)\n            q, attn = self.attention(q, k, v, mask=mask, pad_mask=pad_mask)\n        else:\n            q, attn = self.attention(q, k, v, mask=mask)\n\n        #q = tf.transpose(q, (0, 2, 1, 3))\n        q = tf.keras.layers.Reshape((-1, self.d_model))(tf.keras.layers.Permute((2, 1, 3))(q))\n        #q = tf.reshape(q, (bs, len_q, -1))\n\n        return q, attn\n\nclass ScaledDotProductAttention(tf.keras.layers.Layer):\n    def __init__(self, temperature, num_heads, attn_dropout=0.1):\n        super(ScaledDotProductAttention, self).__init__()\n        self.temperature = temperature\n        self.num_heads = num_heads\n        self.dropout = tf.keras.layers.Dropout(attn_dropout)\n        self.gamma = tf.keras.backend.variable(100.0)\n\n    def call(self, q, k, v, mask=None, pad_mask=None):\n        #attn = tf.matmul(q, tf.transpose(k, perm=(0, 1, 3, 2))) / self.temperature\n        attn = tf.matmul(q, k, transpose_b=True) / self.temperature\n        \n        if mask is not None:\n            attn = attn + mask[:,None,:,:] + self.gamma\n            #mask_with_channel = tf.expand_dims(mask, axis=1)\n            #mask_broadcasted = mask_with_channel * tf.ones_like(attn)\n            #gamma = mask_broadcasted + self.gamma\n            #attn = attn + gamma\n        if pad_mask is not None:\n            #attn = tf.cast(attn, dtype=tf.float32)\n            attn = tf.where(pad_mask == 0, -1e-9, attn)\n            \n        \n        attn = tf.keras.layers.Softmax(axis=-1)(attn)\n        attn = self.dropout(attn)#tf.nn.softmax(attn, axis=-1))\n        output = tf.matmul(attn, v)\n        return output, attn\n","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:24.898995Z","iopub.execute_input":"2023-12-02T05:42:24.899392Z","iopub.status.idle":"2023-12-02T05:42:24.920287Z","shell.execute_reply.started":"2023-12-02T05:42:24.899357Z","shell.execute_reply":"2023-12-02T05:42:24.919407Z"},"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.2):\n        super().__init__()\n        \n        self.att = MultiHeadAttention(dim, num_heads, d_k=dim//num_heads, d_v=dim//num_heads, dropout=rate) #num_heads=num_heads, key_dim=dim//num_heads\n        \n        self.ffn = tf.keras.Sequential(\n            [\n                tf.keras.layers.Dense(feed_forward_dim, activation=swish),\n                tf.keras.layers.Dropout(rate),\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, mask=None,pad_mask=None):\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, inputs, mask, pad_mask)\n        attn_output = self.dropout1(attn_output)\n        out1 = self.layernorm1(inputs + attn_output)\n        ffn_output = self.ffn(out1)\n        ffn_output = self.dropout2(ffn_output)\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\n","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:24.921797Z","iopub.execute_input":"2023-12-02T05:42:24.922390Z","iopub.status.idle":"2023-12-02T05:42:24.937881Z","shell.execute_reply.started":"2023-12-02T05:42:24.922354Z","shell.execute_reply":"2023-12-02T05:42:24.936877Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_max_len = 457\nnum_vocab = 12","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:24.939236Z","iopub.execute_input":"2023-12-02T05:42:24.939600Z","iopub.status.idle":"2023-12-02T05:42:24.951556Z","shell.execute_reply.started":"2023-12-02T05:42:24.939568Z","shell.execute_reply":"2023-12-02T05:42:24.950611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-12-02T05:42:24.952637Z","iopub.execute_input":"2023-12-02T05:42:24.952900Z","iopub.status.idle":"2023-12-02T05:42:24.961593Z","shell.execute_reply.started":"2023-12-02T05:42:24.952877Z","shell.execute_reply":"2023-12-02T05:42:24.960591Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model(hidden_dim = 192, max_len = X_max_len):\n\n    inp = tf.keras.Input((457,459), name='inp')\n    bpp = inp[:,:,2:]#tf.keras.Input((206, 206,))\n    mask = tf.squeeze(inp[:,:,1:2],axis=-1)#tf.keras.Input((206,))\n\n    x = tf.squeeze(inp[:,:,:1],axis=-1)\n\n    x = tf.keras.layers.Embedding(num_vocab, hidden_dim, mask_zero=False)(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, bpp, mask)\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\n    x = transformer_block(hidden_dim, 6, hidden_dim*4)(x, bpp, mask)\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)\n    return model","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:24.962681Z","iopub.execute_input":"2023-12-02T05:42:24.963016Z","iopub.status.idle":"2023-12-02T05:42:24.976474Z","shell.execute_reply.started":"2023-12-02T05:42:24.962971Z","shell.execute_reply":"2023-12-02T05:42:24.975648Z"},"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-12-02T05:42:24.977619Z","iopub.execute_input":"2023-12-02T05:42:24.977919Z","iopub.status.idle":"2023-12-02T05:42:32.773541Z","shell.execute_reply.started":"2023-12-02T05:42:24.977895Z","shell.execute_reply":"2023-12-02T05:42:32.772520Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_sequences_df","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:32.777552Z","iopub.execute_input":"2023-12-02T05:42:32.777867Z","iopub.status.idle":"2023-12-02T05:42:32.797487Z","shell.execute_reply.started":"2023-12-02T05:42:32.777840Z","shell.execute_reply":"2023-12-02T05:42:32.796368Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_sequences = test_sequences_df.sequence.to_numpy()","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:32.798990Z","iopub.execute_input":"2023-12-02T05:42:32.799448Z","iopub.status.idle":"2023-12-02T05:42:32.806819Z","shell.execute_reply.started":"2023-12-02T05:42:32.799408Z","shell.execute_reply":"2023-12-02T05:42:32.805645Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_tf_bpp(sample):\n    # Extract required information from the TFRecord sample\n    bpp_filepath = sample['bpp_path']  # Assuming 'bpp_path' key exists and points to the file\n\n    def read_bpp_file(filepath_tensor):\n        # Read file contents\n        file_contents = tf.cond(tf.equal(filepath_tensor[0], ''), lambda: '', lambda: tf.io.read_file(filepath_tensor[0]))\n        lines = tf.cond(tf.equal(file_contents, ''), lambda: tf.constant(['']), lambda: tf.strings.split(file_contents, '\\n'))\n        lines = lines[:-2]\n        # Process each line of the file\n        def process_line(line):\n            values = tf.strings.split([line], ' ').values  # Split based on spaces\n            values = tf.strings.to_number(values, out_type=tf.float32)\n            return tf.cast(values[0] - 1, tf.int32), tf.cast(values[1] - 1, tf.int32), values[2]\n\n        data = tf.map_fn(process_line, lines, dtype=(tf.int32, tf.int32, tf.float32))\n\n        # Initialize matrix\n        arr = tf.zeros((X_max_len, X_max_len), dtype=tf.float32)\n\n        # Update matrix with line data\n        arr = tf.tensor_scatter_nd_update(arr, tf.stack([data[0], data[1]], axis=1), data[2])\n\n        return arr\n\n    bpp_matrix = read_bpp_file(bpp_filepath)\n    return bpp_matrix\n","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:32.808117Z","iopub.execute_input":"2023-12-02T05:42:32.808471Z","iopub.status.idle":"2023-12-02T05:42:32.819925Z","shell.execute_reply.started":"2023-12-02T05:42:32.808410Z","shell.execute_reply":"2023-12-02T05:42:32.818890Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_bpp(sample):\n    bpp = sample['bpp']  # Assuming 'bpp' key exists and points to the file\n\n    # Create a sparse tensor indices and values\n    indices = tf.cast(tf.stack([tf.subtract(bpp[:, 0], 1), tf.subtract(bpp[:, 1], 1)], axis=1), tf.int64)\n    values = bpp[:, 2]\n\n    # Create a sparse tensor\n    adjacency_matrix = tf.sparse.SparseTensor(indices=indices, values=values, dense_shape=(X_max_len, X_max_len))\n    adjacency_matrix = tf.sparse.to_dense(adjacency_matrix, validate_indices=False)\n\n    return adjacency_matrix\n","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:32.821163Z","iopub.execute_input":"2023-12-02T05:42:32.822288Z","iopub.status.idle":"2023-12-02T05:42:32.834482Z","shell.execute_reply.started":"2023-12-02T05:42:32.822259Z","shell.execute_reply":"2023-12-02T05:42:32.833513Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode_tfrec(example):\n    feature_description = {\n        'seq_id': tf.io.VarLenFeature(tf.string),\n        'seq': tf.io.VarLenFeature(tf.string),\n        'structure': tf.io.VarLenFeature(tf.string),\n        'bpp_path': tf.io.VarLenFeature(tf.string),\n        'length': tf.io.VarLenFeature(tf.int64),\n        'inputs': tf.io.VarLenFeature(tf.float32),\n        'mask' : tf.io.VarLenFeature(tf.int64),\n        'bpp' : tf.io.VarLenFeature(tf.float32)\n    }\n    \n    features = tf.io.parse_single_example(example, feature_description)\n    \n    seq_id = tf.sparse.to_dense(features['seq_id'])\n    seq = tf.sparse.to_dense(features['seq'])\n    bpp = tf.sparse.to_dense(features['bpp'])\n    structure =tf.sparse.to_dense(features['structure'])\n    bpp_path = tf.sparse.to_dense(features['bpp_path'])\n    length = tf.sparse.to_dense(features['length'])\n    inputs = tf.sparse.to_dense(features[\"inputs\"])\n    mask = tf.sparse.to_dense(features[\"mask\"])\n    bpp = tf.reshape(bpp,(-1,3))\n    \n    out = {}\n    out['id'] = seq_id\n    out['bpp'] = bpp\n    out['seq']  = seq\n    out['inputs'] = inputs\n    out['mask'] = mask\n    out['structure'] = structure\n    out['bpp_path'] = bpp_path\n    out['length'] = length\n    \n    return out","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:32.835662Z","iopub.execute_input":"2023-12-02T05:42:32.835973Z","iopub.status.idle":"2023-12-02T05:42:32.847248Z","shell.execute_reply.started":"2023-12-02T05:42:32.835948Z","shell.execute_reply":"2023-12-02T05:42:32.846342Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def concat_target(x):\n\n    inputs = x['inputs']\n    \n    bpp = get_bpp(x)\n    \n    target_length = X_max_len\n\n    pad_size = target_length - tf.shape(inputs)[0]\n    \n    mask_ones = tf.ones(tf.shape(inputs)[0], dtype=tf.float32)\n    mask_zeros = tf.zeros((pad_size,), dtype=tf.float32)\n    mask = tf.concat([mask_ones, mask_zeros], axis=0)\n\n    \n    inputs = tf.concat([inputs, tf.zeros(pad_size, dtype=tf.float32)], axis=0)\n    inp = tf.concat([inputs[...,None],mask[...,None],bpp],axis=1)\n    \n    return inp","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:32.848478Z","iopub.execute_input":"2023-12-02T05:42:32.849059Z","iopub.status.idle":"2023-12-02T05:42:32.862419Z","shell.execute_reply.started":"2023-12-02T05:42:32.849024Z","shell.execute_reply":"2023-12-02T05:42:32.861461Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_extended_sequence(seq, struct):\n    mapping = {\n        ('G', '('): 0,\n        ('G', '.'): 1,\n        ('G', ')'): 2,\n        ('A', '('): 3,\n        ('A', '.'): 4,\n        ('A', ')'): 5,\n        ('C', '('): 6,\n        ('C', '.'): 7,\n        ('C', ')'): 8,\n        ('U', '('): 9,\n        ('U', '.'): 10,\n        ('U', ')'): 11\n    }\n    extended_seq = [mapping[(n, s)] for n, s in zip(seq, struct)]\n\n    return tf.constant(extended_seq, dtype=tf.float32)\n\ndef get_np_bpp(bpp):\n    # Create a sparse tensor indices and values\n    indices = tf.cast(tf.stack([tf.subtract(bpp[:, 0], 1), tf.subtract(bpp[:, 1], 1)], axis=1), tf.int64)\n    values = bpp[:, 2]\n\n    # Create a sparse tensor\n    adjacency_matrix = tf.sparse.SparseTensor(indices=indices, values=values, dense_shape=(X_max_len, X_max_len))\n    adjacency_matrix = tf.sparse.to_dense(adjacency_matrix, validate_indices=False)\n\n    return tf.cast(adjacency_matrix,tf.float32)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:32.863712Z","iopub.execute_input":"2023-12-02T05:42:32.864089Z","iopub.status.idle":"2023-12-02T05:42:32.874632Z","shell.execute_reply.started":"2023-12-02T05:42:32.864062Z","shell.execute_reply":"2023-12-02T05:42:32.873524Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_df = pd.read_parquet('/kaggle/input/extra-dataset/TEST_Data.parquet')","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:32.875815Z","iopub.execute_input":"2023-12-02T05:42:32.876109Z","iopub.status.idle":"2023-12-02T05:42:38.853380Z","shell.execute_reply.started":"2023-12-02T05:42:32.876084Z","shell.execute_reply":"2023-12-02T05:42:38.852527Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"def preprocess_data(seq_id, seq, sct, f):    \n    inputs = create_extended_sequence(seq, sct)\n    \n    if f==0:\n        \n        path = os.path.join('/kaggle/input/rna-dataset/test_dir',f'{seq_id}.npy')\n        bpp = get_np_bpp(np.load(path))\n    \n    else:\n        path = os.path.join('/kaggle/input/extra-dataset/test_tfrecord',f'{seq_id}.npy')\n        bpp = get_np_bpp(np.load(path))\n    \n    target_length = X_max_len\n\n    pad_size = target_length - tf.shape(inputs)[0]\n    \n    mask_ones = tf.ones(tf.shape(inputs)[0], dtype=tf.float32)\n    mask_zeros = tf.zeros((pad_size,), dtype=tf.float32)\n    mask = tf.concat([mask_ones, mask_zeros], axis=0)\n\n    \n    inputs = tf.concat([inputs, tf.zeros(pad_size, dtype=tf.float32)], axis=0)\n    inp = tf.concat([inputs[...,None],mask[...,None],bpp],axis=1) \n    return inp\n\n# Create a function to generate input data for prediction\ndef generate_input_data(df):\n    for _, row in df.iterrows():\n        seq_id, seq, sct = row['sequence_id'], row['sequence'], row['structure']\n        if _ < 1265000:f=0\n        else:f=1\n        preprocessed_data = preprocess_data(seq_id, seq, sct, f)\n        yield tf.cast(preprocessed_data,tf.float32)\"\"\"\n\n#data_generator = generate_input_data(data_df)\n#tensor_dataset = tf.data.Dataset.from_generator(lambda: data_generator,output_shapes=(457,459), output_types=tf.float32)\n#batch_size = 256\n#test_ds = tensor_dataset.batch(batch_size)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:38.854606Z","iopub.execute_input":"2023-12-02T05:42:38.854919Z","iopub.status.idle":"2023-12-02T05:42:38.862749Z","shell.execute_reply.started":"2023-12-02T05:42:38.854892Z","shell.execute_reply":"2023-12-02T05:42:38.861816Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#batch=next(iter(test_ds))","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:38.863915Z","iopub.execute_input":"2023-12-02T05:42:38.864227Z","iopub.status.idle":"2023-12-02T05:42:38.875953Z","shell.execute_reply.started":"2023-12-02T05:42:38.864201Z","shell.execute_reply":"2023-12-02T05:42:38.875059Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tffiles = [f'/kaggle/input/rna-dataset/test_tfrecord/records_00{\"{:02d}\".format(i+1)}.tfrecords' for i in range(27)]\n#tffiles = glob.glob(tffiles_path)\nprint(f'Total files {len(tffiles)}')\n","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:38.877157Z","iopub.execute_input":"2023-12-02T05:42:38.877527Z","iopub.status.idle":"2023-12-02T05:42:38.888026Z","shell.execute_reply.started":"2023-12-02T05:42:38.877491Z","shell.execute_reply":"2023-12-02T05:42:38.887152Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_tfrec_dataset(tffiles, shuffle=False, batch_size=True, cache = False, to_filter = False,\n                      calculate_sample_num = True,paded_batch_size=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    ds = ds.map(concat_target, tf.data.AUTOTUNE)\n    #ds = ds.map(get_data, tf.data.AUTOTUNE)\n    if DEBUG:\n        ds = ds.take(1000)\n\n    if cache:\n        ds = ds.cache()\n\n    samples_num = 0\n    print('Counting num of samples.')\n    if calculate_sample_num:\n        samples_num = ds.reduce(0, lambda x,_: x+1).numpy()\n\n\n    if paded_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.batch(batch_size,drop_remainder=False)\n    ds = ds.prefetch(tf.data.AUTOTUNE)\n    return ds, samples_num","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:38.889263Z","iopub.execute_input":"2023-12-02T05:42:38.890104Z","iopub.status.idle":"2023-12-02T05:42:38.898957Z","shell.execute_reply.started":"2023-12-02T05:42:38.890069Z","shell.execute_reply":"2023-12-02T05:42:38.898051Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 256\ntest_ds, num_val = get_tfrec_dataset(tffiles, shuffle = False, batch_size = batch_size,\n                                                 cache = False, to_filter = 'filter_1', calculate_sample_num = True)\nprint(f\"Total samples counted {num_val}!\")","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:38.900373Z","iopub.execute_input":"2023-12-02T05:42:38.900730Z","iopub.status.idle":"2023-12-02T05:42:39.821326Z","shell.execute_reply.started":"2023-12-02T05:42:38.900698Z","shell.execute_reply":"2023-12-02T05:42:39.820346Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch = next(iter(test_ds))","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:39.822644Z","iopub.execute_input":"2023-12-02T05:42:39.822957Z","iopub.status.idle":"2023-12-02T05:42:40.386390Z","shell.execute_reply.started":"2023-12-02T05:42:39.822929Z","shell.execute_reply":"2023-12-02T05:42:40.385434Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"test_ds = tf.data.TFRecordDataset(\n            tffiles, num_parallel_reads=tf.data.AUTOTUNE).prefetch(tf.data.AUTOTUNE)\ntest_ds = test_ds.map(decode_tfrec, tf.data.AUTOTUNE)\ntest_ds = test_ds.map(concat_target, tf.data.AUTOTUNE)\n\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.batch(batch_size,drop_remainder=False)\ntest_ds = test_ds.prefetch(tf.data.AUTOTUNE)\nbatch = next(iter(test_ds))\nbatch.shape\"\"\"","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:40.387961Z","iopub.execute_input":"2023-12-02T05:42:40.388306Z","iopub.status.idle":"2023-12-02T05:42:40.395053Z","shell.execute_reply.started":"2023-12-02T05:42:40.388279Z","shell.execute_reply":"2023-12-02T05:42:40.394064Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"max_len = 457\ntf.keras.backend.clear_session()\nwith STRATEGY.scope():\n    model = get_model(hidden_dim=192, max_len=max_len)\n    \n    model.load_weights('/kaggle/input/srrf-transformer-model/model_epoch_197.h5')\n    print('Weights loaded!')\n    model(batch)\n    model.summary()","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-12-02T05:42:40.396325Z","iopub.execute_input":"2023-12-02T05:42:40.396687Z","iopub.status.idle":"2023-12-02T05:42:46.704339Z","shell.execute_reply.started":"2023-12-02T05:42:40.396653Z","shell.execute_reply":"2023-12-02T05:42:46.703320Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with STRATEGY.scope():\n    preds = model.predict(test_ds)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T05:42:46.705878Z","iopub.execute_input":"2023-12-02T05:42:46.706635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"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":{"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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import polars as pl\nimport matplotlib.pyplot as plt\n#read your sub here\ndf=pl.read_csv(\"submission.csv\")\n#some parameters\nfont_size=6\nid1=269545321\nid2=269724007\nreshape1=391\nreshape2=457\n#get predictions\npred_DMS=df[id1:id2+1]['reactivity_DMS_MaP'].to_numpy().reshape(reshape1,reshape2)\npred_2A3=df[id1:id2+1]['reactivity_2A3_MaP'].to_numpy().reshape(reshape1,reshape2)\n#plot mutate and map\nfig = plt.figure()\nplt.subplot(121)\nplt.title(f'reactivity_DMS_MaP', fontsize=font_size)\nplt.imshow(pred_DMS,vmin=0,vmax=1, cmap='gray_r')\nplt.subplot(122)\nplt.title(f'reactivity_2A3_MaP', fontsize=font_size)\nplt.imshow(pred_2A3,vmin=0,vmax=1, cmap='gray_r')\nplt.tight_layout()\nplt.savefig(f\"plot.png\",dpi=500)\nplt.clf()\nplt.close()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import cv2","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img = cv2.imread('/kaggle/working/plot.png')\nplt.figure(figsize=(20,10))\nplt.imshow(img)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img = cv2.imread('/kaggle/working/plot.png')\nplt.figure(figsize=(30,20))\nplt.imshow(img)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}