{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":67356,"databundleVersionId":8006601,"sourceType":"competition"},{"sourceId":9003850,"sourceType":"datasetVersion","datasetId":5424192}],"dockerImageVersionId":30746,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install mapply\n!pip install scikit-fingerprints\n!pip install atomInSmiles","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-07-21T14:56:35.175640Z","iopub.execute_input":"2024-07-21T14:56:35.176026Z","iopub.status.idle":"2024-07-21T14:57:16.185662Z","shell.execute_reply.started":"2024-07-21T14:56:35.175995Z","shell.execute_reply":"2024-07-21T14:57:16.184118Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom keras.layers import Dense, LayerNormalization, Add, MultiHeadAttention, Dropout\nfrom keras.layers import Embedding, TextVectorization, GlobalAvgPool1D\nfrom keras import Sequential\nfrom sklearn.model_selection import train_test_split\nimport os\nfrom typing import Union\nimport mapply\nfrom skfp.fingerprints import ECFPFingerprint\nfrom rdkit import Chem\nimport atomInSmiles\nimport dask.dataframe as dd\nimport pyarrow as pa\nimport itertools\nimport einops","metadata":{"execution":{"iopub.status.busy":"2024-07-21T14:57:38.440813Z","iopub.status.idle":"2024-07-21T14:57:38.441273Z","shell.execute_reply.started":"2024-07-21T14:57:38.441061Z","shell.execute_reply":"2024-07-21T14:57:38.441093Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# LOSSES and METRICS >>>\n\n@keras.saving.register_keras_serializable(package='belka', name='MultiLabelLoss')\nclass MultiLabelLoss(keras.losses.Loss):\n    \"\"\"\n    * Macro- or Micro-averaged Weighted Masked Binary Focal loss.\n    * Dynamic mini-batch class weights \"alpha\".\n    * Used for binary multilabel classification.\n    \"\"\"\n\n    def __init__(self, epsilon: float, macro: bool, gamma: float = 2.0, nan_mask: int = 2, name='loss', **kwargs):\n        super(MultiLabelLoss, self).__init__(name=name)\n        self.epsilon = epsilon\n        self.gamma = gamma\n        self.macro = macro\n        self.nan_mask = nan_mask\n\n    def call(self, y_true, y_pred):\n\n        # Cast y_true to tf.int32\n        y_true = tf.cast(y_true, dtype=tf.int32)\n\n        # Compute class weights (\"alpha\"): Inverse of Square Root of Number of Samples\n        # Compute \"alpha\" for each label if \"macro\" = True\n        # Assign zero class weights to missing classes\n        # Normalize: sum of sample weights = sample count per label\n        freq = tf.math.bincount(\n            arr=tf.transpose(y_true, perm=[1,0]) if self.macro else y_true,\n            minlength=2, maxlength=2, dtype=tf.float32, axis=-1 if self.macro else 0)\n        alpha = tf.where(condition=tf.equal(freq, 0.0), x=0.0, y=tf.math.rsqrt(freq))\n        ax = 1 if self.macro else None\n        alpha = alpha * tf.reduce_sum(freq, axis=ax, keepdims=True) / tf.reduce_sum(alpha*freq, axis=ax, keepdims=True)\n        alpha = tf.reduce_sum(alpha * tf.one_hot(y_true, axis=-1, depth=2, dtype=tf.float32), axis=-1)\n\n        # Mask and set to zero missing labels\n        y_true = tf.cast(y_true, tf.float32)\n        mask = tf.cast(tf.not_equal(y_true, tf.constant(self.nan_mask, tf.float32)), dtype=tf.float32)\n        y_true = y_true * mask\n\n        # Compute loss\n        y_pred = tf.clip_by_value(y_pred, clip_value_min=self.epsilon, clip_value_max=1.0 - self.epsilon)\n        pt = tf.add(\n            tf.multiply(y_true, y_pred),\n            tf.multiply(tf.subtract(1.0, y_true), tf.subtract(1.0, y_pred)))\n        loss = - alpha * (1.0 - pt) ** self.gamma * tf.math.log(pt) * mask\n        ax = 1 if self.macro else None\n        loss = tf.divide(tf.reduce_sum(loss, axis=ax), tf.reduce_sum(alpha * mask, axis=ax))\n        loss = tf.reduce_mean(loss)\n        return loss\n\n    def get_config(self):\n        config = super(MultiLabelLoss, self).get_config()\n        config.update({\n            'epsilon': self.epsilon,\n            'gamma': self.gamma,\n            'macro': self.macro,\n            'nan_mask': self.nan_mask})\n        return config\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n\n\n@keras.saving.register_keras_serializable(package='belka', name='CategoricalLoss')\nclass CategoricalLoss(keras.losses.Loss):\n    \"\"\"\n    Masked Categorical Focal loss.\n    Dynamic mini-batch class weights (\"alpha\").\n    Used for MLM training.\n    \"\"\"\n    def __init__(self, epsilon: float, mask: int, vocab_size: int, gamma: float = 2.0, name='loss', **kwargs):\n        super(CategoricalLoss, self).__init__(name=name)\n        self.epsilon = epsilon\n        self.gamma = gamma\n        self.mask = mask\n        self.vocab_size = vocab_size\n\n    def call(self, y_true, y_pred):\n\n        # Unpack y_true to masked (y_true) and unmasked (unmasked) arrays\n        unmasked = y_true[:,:,1]\n        y_true = y_true[:,:,0]\n\n        # Reshape inputs\n        y_true = einops.rearrange(y_true, 'b l -> (b l)')\n        y_pred = einops.rearrange(y_pred, 'b l c -> (b l) c')\n\n        # Drop non-masked from y_true\n        mask = tf.not_equal(y_true, self.mask)\n        y_true = tf.boolean_mask(y_true, mask)\n\n        # Compute class weights (\"alpha\"): Inverse of Square Root of Number of Samples\n        # Assign zero class weights to missing classes\n        # Normalize: sum of sample weights = sample count\n        freq = tf.math.bincount(unmasked, minlength=self.vocab_size, dtype=tf.float32)\n        freq = tf.concat([tf.zeros(shape=(2,)), freq[2:]], axis=0)  # Set frequencies for [PAD], [MASK] = 0\n        alpha = tf.where(condition=tf.equal(freq, 0.0), x=0.0, y=tf.math.rsqrt(freq))\n\n        # Convert y_true to one-hot\n        # Apply mask to y_pred\n        y_true = tf.one_hot(y_true, depth=self.vocab_size, axis=-1, dtype=tf.float32)\n        y_pred = tf.boolean_mask(y_pred, mask, axis=0)\n\n        # Compute loss\n        y_pred = tf.clip_by_value(y_pred, clip_value_min=self.epsilon, clip_value_max=1.0 - self.epsilon)\n        pt = tf.add(\n            tf.multiply(y_true, y_pred),\n            tf.multiply(tf.subtract(1.0, y_true), tf.subtract(1.0, y_pred)))\n        loss = - alpha * ((1.0 - pt) ** self.gamma) * (y_true * tf.math.log(y_pred))\n        loss = tf.divide(tf.reduce_sum(loss), tf.reduce_sum(alpha * y_true))\n        return loss\n\n    def get_config(self) -> dict:\n        config = super(CategoricalLoss, self).get_config()\n        config.update({\n            'epsilon': self.epsilon,\n            'gamma': self.gamma,\n            'mask': self.mask,\n            'vocab_size': self.vocab_size})\n        return config\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n\n\n@keras.saving.register_keras_serializable(package='belka', name='BinaryLoss')\nclass BinaryLoss(keras.losses.Loss):\n    \"\"\"\n    Binary Focal loss.\n    Used for FPs training.\n    \"\"\"\n    def __init__(self, name='loss', **kwargs):\n        super(BinaryLoss, self).__init__(name=name)\n        self.loss = tf.keras.losses.BinaryFocalCrossentropy()\n\n    def call(self, y_true, y_pred):\n        y_true = tf.cast(y_true, dtype=tf.float32)\n        y_true = tf.reshape(y_true, shape=(-1, 1))\n        y_pred = tf.reshape(y_pred, shape=(-1, 1))\n        loss = self.loss(y_true, y_pred)\n        return loss\n\n    def get_config(self) -> dict:\n        config = super(BinaryLoss, self).get_config()\n        return config\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n\n\n@keras.saving.register_keras_serializable(package='belka', name='MaskedAUC')\nclass MaskedAUC(keras.metrics.AUC):\n    def __init__(self, mode: str, mask: int, multi_label: bool, num_labels: Union[int, None], vocab_size: int,\n                 name='auc', **kwargs):\n        super(MaskedAUC, self).__init__(curve='PR', multi_label=multi_label, num_labels=num_labels, name=name)\n        self.mode = mode\n        self.multi_label = multi_label\n        self.mask = mask\n        self.num_labels = num_labels\n        self.vocab_size = vocab_size\n\n    def update_state(self, y_true, y_pred, sample_weight=None):\n\n        if self.mode == 'mlm':\n\n            # Unpack y_true to masked (y_true) and unmasked (unmasked) arrays\n            unmasked = y_true[:, :, 1]\n            y_true = y_true[:, :, 0]\n\n            # Reshape inputs\n            y_true = einops.rearrange(y_true, 'b l -> (b l)')\n            y_pred = einops.rearrange(y_pred, 'b l c -> (b l) c')\n\n            # Drop non-masked tokens from y_true\n            mask = tf.not_equal(y_true, self.mask)\n            y_true = tf.boolean_mask(y_true, mask)\n\n            # Convert y_true to one-hot\n            # Apply mask to y_pred\n            y_true = tf.one_hot(y_true, depth=self.vocab_size, axis=-1, dtype=tf.float32)\n            y_pred = tf.boolean_mask(y_pred, mask, axis=0)\n            mask = None\n\n        elif self.mode == 'clf':\n            mask = tf.cast(tf.not_equal(y_true, self.mask), dtype=tf.float32)\n\n        else:\n            y_true = tf.reshape(y_true, shape=(-1,1))\n            y_pred = tf.reshape(y_pred, shape=(-1,1))\n            mask = tf.ones_like(y_pred, dtype=tf.float32)\n\n        # Compute macro-averaged mAP\n        super().update_state(y_true, y_pred, sample_weight=mask)\n\n    def get_config(self) -> dict:\n        config = super(MaskedAUC, self).get_config()\n        config.update({\n            'mode': self.mode,\n            'multi_label': self.multi_label,\n            'mask': self.mask,\n            'num_labels': self.num_labels,\n            'vocab_size': self.vocab_size})\n        return config\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n\n\n# LAYERS >>>\n\nclass FPGenerator(tf.keras.layers.Layer):\n    def __init__(self, name: str = 'fingerprints', **kwargs):\n        super(FPGenerator, self).__init__(name=name)\n        self.transformer = ECFPFingerprint(include_chirality=True, n_jobs=-1)\n\n    def call(self, inputs, *args, **kwargs):\n        \"\"\"\n        Get fingerprints given SMILES string.\n        \"\"\"\n        x = tf.py_function(\n            func=self.get_fingerprints,\n            inp=[inputs],\n            Tout=tf.int8)\n        return x\n\n    def get_fingerprints(self, inputs):\n        x = inputs.numpy().astype(str)\n        x = self.transformer.transform(x)\n        x = tf.constant(x, dtype=tf.int8)\n        return x\n\n\n@keras.saving.register_keras_serializable(package='belka', name='Encodings')\nclass Encodings(keras.layers.Layer):\n    def __init__(self, depth: int, max_length: int, name: str = 'encodings', **kwargs):\n        super(Encodings, self).__init__(name=name)\n        self.depth = depth\n        self.max_length = max_length\n        self.encodings = self._pos_encodings(depth=depth, max_length=max_length)\n\n    def call(self, inputs, training=False, *args, **kwargs):\n        scale = tf.ones_like(inputs) * tf.math.sqrt(tf.cast(self.depth, tf.float32))\n        x = tf.multiply(inputs, scale)\n        x = tf.add(x, self.encodings[tf.newaxis, :tf.shape(x)[1], :])\n        return x\n\n    @staticmethod\n    def _pos_encodings(depth: int, max_length: int):\n        \"\"\"\n        Get positional encodings of shape [max_length, depth]\n        \"\"\"\n        positions = tf.range(max_length, dtype=tf.float32)[:, tf.newaxis]\n        idx = tf.range(depth)[tf.newaxis, :]\n        power = tf.cast(2 * (idx // 2), tf.float32)\n        power /= tf.cast(depth, tf.float32)\n        angles = 1. / tf.math.pow(10000., power)\n        radians = positions * angles\n        sin = tf.math.sin(radians[:, 0::2])\n        cos = tf.math.cos(radians[:, 1::2])\n        encodings = tf.concat([sin, cos], axis=-1)\n        return encodings\n\n    def get_config(self) -> dict:\n        return {\n            'depth': self.depth,\n            'max_length': self.max_length,\n            'name': self.name}\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n\n\n@keras.saving.register_keras_serializable(package='belka', name='Embeddings')\nclass Embeddings(tf.keras.layers.Layer):\n    def __init__(self, max_length: int, depth: int, input_dim: int, name: str = 'embeddings', **kwargs):\n        super(Embeddings, self).__init__(name=name)\n        self.depth = depth\n        self.max_length = max_length\n        self.input_dim = input_dim\n        self.embeddings = Embedding(input_dim=input_dim, output_dim=depth, mask_zero=True)\n        self.encodings = Encodings(**parameters)\n\n    def build(self, input_shape):\n        self.embeddings.build(input_shape=input_shape)\n        super().build(input_shape)\n\n    def compute_mask(self, *args, **kwargs):\n        return self.embeddings.compute_mask(*args, **kwargs)\n\n    def call(self, inputs, training=False, *args, **kwargs):\n        x = self.embeddings(inputs)\n        x = self.encodings(x)\n        return x\n\n    def get_config(self) -> dict:\n        return {\n            'depth': self.depth,\n            'input_dim': self.input_dim,\n            'max_length': self.max_length,\n            'name': self.name}\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n\n\n@keras.saving.register_keras_serializable(package='belka', name='FeedForward')\nclass FeedForward(tf.keras.layers.Layer):\n    def __init__(self, activation: str, depth: int, dropout_rate: float, epsilon: float, name: str = 'ffn', **kwargs):\n        super(FeedForward, self).__init__(name=name)\n        self.activation = activation\n        self.depth = depth\n        self.dropout_rate = dropout_rate\n        self.epsilon = epsilon\n        self.norm = LayerNormalization(epsilon=epsilon)\n        self.dense1 = Dense(units=int(depth * 2), activation=activation)\n        self.dense2 = Dense(units=depth)\n        self.dropout = Dropout(rate=dropout_rate)\n        self.add = Add()\n\n    def build(self, input_shape):\n        super().build(input_shape)\n\n    def call(self, inputs, training=False, *args, **kwargs):\n        x = self.norm(inputs)\n        x = self.dense1(x)\n        x = self.dense2(x)\n        x = self.dropout(x, training=training)\n        x = self.add([x, inputs])\n        return x\n\n    def get_config(self) -> dict:\n        return {\n            'activation': self.activation,\n            'depth': self.depth,\n            'dropout_rate': self.dropout_rate,\n            'epsilon': self.epsilon,\n            'name': self.name}\n\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n\n\n@keras.saving.register_keras_serializable(package='belka', name='SelfAttention')\nclass SelfAttention(tf.keras.layers.Layer):\n    \"\"\"\n    * Self-Attention block with PRE-layer normalization\n    * LayerNorm -> MHA -> Skip connection\n    \"\"\"\n    def __init__(self, causal: bool, depth: int, dropout_rate: float, epsilon: float, max_length: int, num_heads: int,\n                 name: str = 'self_attention', **kwargs):\n        super(SelfAttention, self).__init__(name=name)\n        self.causal = causal\n        self.depth = depth\n        self.dropout_rate = dropout_rate\n        self.epsilon = epsilon\n        self.max_length = max_length\n        self.num_heads = num_heads\n        self.supports_masking = True\n        self.norm = LayerNormalization(epsilon=epsilon)\n        self.mha = MultiHeadAttention(num_heads=num_heads, key_dim=depth, dropout=dropout_rate)\n        self.add = Add()\n\n    def build(self, input_shape):\n        self.mha.build(input_shape=[input_shape, input_shape])\n        super().build(input_shape)\n\n    def call(self, inputs, training=False, *args, **kwargs):\n\n        # Compute attention mask\n        mask = tf.cast(inputs._keras_mask, dtype=tf.float32)\n        m1 = tf.expand_dims(mask, axis=2)\n        m2 = tf.expand_dims(mask, axis=1)\n        mask = tf.cast(tf.linalg.matmul(m1, m2), dtype=tf.bool)\n\n        # Compute outputs\n        x = self.norm(inputs)\n        x = self.mha(\n            query=x,\n            value=x,\n            use_causal_mask=self.causal,\n            training=training,\n            attention_mask=mask)\n        x = self.add([x, inputs])\n\n        return x\n\n    def get_config(self) -> dict:\n        return {\n            'causal': self.causal,\n            'depth': self.depth,\n            'dropout_rate': self.dropout_rate,\n            'epsilon': self.epsilon,\n            'max_length': self.max_length,\n            'name': self.name,\n            'num_heads': self.num_heads}\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n\n\n@keras.saving.register_keras_serializable(package='belka', name='EncoderLayer')\nclass EncoderLayer(tf.keras.layers.Layer):\n    \"\"\"\n    * Encoder layer with PRE-layer normalization: LayerNorm -> Self-Attention -> LayerNorm -> FeedForward.\n    \"\"\"\n    def __init__(self, activation: str, depth: int, dropout_rate: float, epsilon: float, max_length: int,\n                 num_heads: int, name: str = 'encoder_layer', **kwargs):\n        super(EncoderLayer, self).__init__(name=name)\n        self.activation = activation\n        self.depth = depth\n        self.dropout_rate = dropout_rate\n        self.epsilon = epsilon\n        self.max_length = max_length\n        self.num_heads = num_heads\n        self.supports_masking = True\n        self.self_attention = SelfAttention(causal=False, depth=depth, dropout_rate=dropout_rate, epsilon=epsilon,\n                                            max_length=max_length, num_heads=num_heads)\n        self.ffn = FeedForward(activation=activation, depth=depth, dropout_rate=dropout_rate, epsilon=epsilon)\n\n    def build(self, input_shape):\n        super().build(input_shape)\n\n    def call(self, inputs, training=False, *args, **kwargs):\n        x = self.self_attention(inputs, training=training)\n        x = self.ffn(x, training=training)\n        return x\n\n    def get_config(self) -> dict:\n        return {\n            'activation': self.activation,\n            'depth': self.depth,\n            'dropout_rate': self.dropout_rate,\n            'epsilon': self.epsilon,\n            'max_length': self.max_length,\n            'name': self.name,\n            'num_heads': self.num_heads}\n\n\n    @classmethod\n    def from_config(cls, config):\n        return cls(**config)\n\n\n# MODELS >>>\n\n@keras.saving.register_keras_serializable(package='belka', name='SingleOutput')\nclass Belka(tf.keras.Model):\n    def __init__(self, dropout_rate: float, mode: str, num_layers: int, vocab_size: int, **kwargs):\n        super(Belka, self).__init__()\n\n        # Arguments\n        self.dropout_rate = dropout_rate\n        self.num_layers = num_layers\n        self.vocab_size = vocab_size\n        self.mode = mode\n\n        #  Layers\n        self.embeddings = Embeddings(input_dim=vocab_size, name='smiles_emb', **parameters)\n        self.encoder = [EncoderLayer(name='encoder_{}'.format(i), **parameters) for i in range(num_layers)]\n        if mode == 'mlm':\n            self.head = Dense(units=vocab_size, activation='softmax', name='smiles')\n        else:\n            self.head = Sequential([\n                GlobalAvgPool1D(),\n                Dropout(dropout_rate),\n                Dense(units = 3 if mode == 'clf' else 2048, activation='sigmoid')])\n\n    def call(self, inputs, training=False, *args, **kwargs):\n        x = self.embeddings(inputs, training=training)\n        for encoder in self.encoder:\n            x = encoder(x, training=training)\n        x = self.head(x, training=training)\n        return x\n\n    def get_config(self) -> dict:\n        return {\n            'mode': self.mode,\n            'num_layers': self.num_layers,\n            'vocab_size': self.vocab_size}\n\n    @classmethod\n    def from_config(cls, config, custom_objects=None):\n        return cls(**config)\n\n\n# DATASETS >>>\n\ndef train_val_set(batch_size: int, buffer_size: int, masking_rate: float, max_length: int, mode: str, seed: int,\n                  vocab_size: int,working: str, **kwargs) -> tuple:\n    \"\"\"\n    Make train and validation datasets.\n    \"\"\"\n\n    # Constants\n    auto = tf.data.AUTOTUNE\n    encoder = get_smiles_encoder(**parameters)\n\n    # Helper functions >>>\n\n    def encode_smiles(x, **kwargs) -> dict:\n        \"\"\"\n        Encode SMILES strings.\n        \"\"\"\n        x['smiles'] = tf.io.parse_tensor(x['smiles'], out_type=tf.string)\n        x['smiles'] = tf.cast(encoder(x['smiles']), dtype=tf.int32)\n        return x\n\n    def get_model_inputs(x) -> tuple:\n        \"\"\"\n        MLM mode: mask [mask_rate] of non-zero tokens.\n            * 80% of the time: Replace with the [MSK].\n            * 10% of the time: Replace with a random token.\n            * 10% of the time: Keep the token unchanged.\n        \"\"\"\n\n        if mode == 'mlm':\n\n            # Get paddings mask (0 = [PAD])\n            paddings_mask = tf.cast(x['smiles'] != 0, dtype=tf.float32)\n\n            # Get random mask (1 = [MASK])\n            probs = tf.stack([1.0 - masking_rate, masking_rate * 0.8, masking_rate * 0.1, masking_rate * 0.1], axis=0)\n            probs = tf.expand_dims(probs, axis=0)\n            probs = tf.ones_like(x['smiles'], dtype=tf.float32)[:, :4] * probs\n            probs = tf.math.log(probs)\n            mask = tf.multiply(\n                tf.one_hot(\n                    indices=tf.random.categorical(logits=probs, num_samples=tf.shape(x['smiles'])[-1]),\n                    depth=4,\n                    dtype=tf.float32,\n                    seed=seed),\n                tf.expand_dims(paddings_mask, axis=-1))\n            mask = tf.cast(mask, dtype=tf.int32)\n\n            # Compute masked inputs\n            x['masked'] = tf.multiply(\n                mask,\n                tf.stack(\n                    values=[\n                        x['smiles'],\n                        tf.ones_like(x['smiles']),\n                        tf.random.uniform(shape=tf.shape(x['smiles']), minval=2, maxval=vocab_size + 1, dtype=tf.int32),\n                        x['smiles']],\n                    axis=-1))\n            x['masked'] = tf.reduce_sum(x['masked'], axis=-1)\n            mask = tf.reduce_sum(mask[:, :, 1:], axis=-1)\n\n            # Set non-masked values to -1\n            x['smiles'] = tf.stack(\n                values=[(x['smiles'] * mask) - (1 - mask), x['smiles']],\n                axis=-1)\n\n            return x['masked'], x['smiles']\n\n        elif mode == 'fps':\n            return x['smiles'], x['ecfp']\n\n        else:\n            return x['smiles'], x['binds']\n\n    # Read subsets\n    encoder = get_smiles_encoder(**parameters)\n    padded_shapes = {'smiles':  (max_length,), 'ecfp': (2048,), 'binds': (3,)}\n    train = tf.data.Dataset.load(os.path.join(working, 'belka.tfr'), compression='GZIP')\n    if mode == 'mlm':\n        features = ['smiles']\n        subsets = {\n            'train': train,\n            'none': None}\n    elif mode == 'fps':\n        features = ['smiles', 'ecfp']\n        subsets = {\n            'train': train.filter(lambda x: tf.not_equal(x['subset'], 1)),\n            'val': tf.data.Dataset.load(os.path.join(working, 'belka_val.tfr'), compression='GZIP')}\n    else:\n        features = ['smiles', 'binds']\n        subsets = {\n            'train': train.filter(lambda x: tf.equal(x['subset'], 0)),\n            'val': tf.data.Dataset.load(os.path.join(working, 'belka_val.tfr'), compression='GZIP')}\n\n    # Preprocess subsets:  Cache -> [Repeat -> Shuffle] -> Encode SMILES -> Batch -> Get inputs\n    for key in [key for key in subsets.keys() if key != 'none']:\n        subset = subsets[key].map(lambda x: {key: x[key] for key in features}, num_parallel_calls=auto)\n        subset = subset.cache()\n        if key == 'train':\n            subset = subset.repeat().shuffle(buffer_size=buffer_size)\n        subset = subset.map(lambda x: encode_smiles(x), num_parallel_calls=auto)\n        subset = subset.padded_batch(batch_size=batch_size, padded_shapes={\n            key: padded_shapes[key] for key in features})\n        subsets[key] = subset.map(lambda x: get_model_inputs(x), num_parallel_calls=auto)\n    return subsets['train'], subsets['val']\n\n\n# TRAIN & SUBMISSION >>>\n\ndef train_model(model: Union[str, None], epochs: int, initial_epoch: int, mode: str, model_name: str, patience: int,\n                steps_per_epoch: int, validation_steps: int, working: str, **kwargs):\n    \"\"\"\n    Train the model.\n    \"\"\"\n\n    # Train/val subsets\n    train, val = train_val_set(**parameters)\n\n    # Build the model\n    if model is not None:\n        model = load_model(model)\n    else:\n        model = Belka(**parameters)\n        if mode == 'mlm':\n            loss = CategoricalLoss(mask=-1, **parameters)\n            metrics = MaskedAUC(mask=-1, multi_label=False, num_labels=None, **parameters)\n        elif mode == 'fps':\n            loss = BinaryLoss(**parameters)\n            metrics = MaskedAUC(mask=-1, multi_label=False, num_labels=None, **parameters)\n        else:\n            loss = MultiLabelLoss(macro=True, **parameters)\n            metrics = MaskedAUC(mask=2, multi_label=True, num_labels=3, **parameters)\n        model.compile(optimizer=tf.keras.optimizers.Adam(), loss=loss, metrics=metrics)\n\n\n    # Callbacks\n    suffix = {\n        'mlm': '_{epoch:03d}_{loss:.4f}.model.keras',\n        'fps': '_{epoch:03d}_{auc:.4f}_{val_auc:.4f}.model.keras',\n        'clf': '_{epoch:03d}_{auc:.4f}_{val_auc:.4f}.model.keras'}\n    model_saver = keras.callbacks.ModelCheckpoint(\n        monitor='loss', mode='min',\n        filepath=os.path.join(working, model_name + suffix[mode]),\n        save_best_only=False,\n        save_weights_only=False)\n    early_stopping = keras.callbacks.EarlyStopping(\n        monitor='loss', mode='min', patience=patience, restore_best_weights=True)\n    learning_rate = tf.keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=5, monitor='loss')\n    callbacks = [model_saver, early_stopping, learning_rate]\n\n    # Print model summary\n    x, y_true = iter(train).get_next()\n    y_pred = model(x)\n    print(model.summary())\n\n    # Fit the model\n    validation_steps = None if mode == 'mlm' else validation_steps\n    model.fit(train, epochs=epochs, initial_epoch=initial_epoch, steps_per_epoch=steps_per_epoch,\n              validation_data=val, validation_steps=validation_steps, callbacks=callbacks)\n\n    return model\n\n\ndef make_submission(batch_size: dict, max_length: int, model: str, working: str, **kwargs) -> None:\n    \"\"\"\n    Make submission.\n    \"\"\"\n\n    # Make train dataset >>>\n    df = read_parquet(subset='test', **parameters).iloc[:1000]\n    df['smiles'] = df['smiles'].mapply(lambda x: atomInSmiles.smiles_tokenizer(x))\n    ds = tf.data.Dataset.from_tensor_slices(\n        {'smiles': tf.ragged.constant(df['smiles'].tolist())})\n\n    # Tokenize -> Zero-pad -> Batch -> Cast\n    encoder = get_smiles_encoder(**parameters)\n    ds = ds.map(lambda x: tf.cast(encoder(x['smiles']), dtype=tf.int32))\n    ds = ds.padded_batch(batch_size=batch_size, padded_shapes=(max_length,))\n    ds = ds.prefetch(buffer_size=tf.data.AUTOTUNE)\n\n    # Make predictions >>>\n    model = tf.saved_model.load(model)\n    pred = np.zeros(shape=(0,3), dtype=np.float32)\n    for batch in ds:\n        pred = np.concatenate([pred, model.serve(batch)])\n    print('\\r')\n\n    # Write predictions to csv\n    cols = ['BRD4_pred', 'HSA_pred', 'sEH_pred']\n    df[cols] = pred\n    cols = [['BRD4', 'BRD4_pred'], ['HSA', 'HSA_pred'], ['sEH', 'sEH_pred']]\n    df = np.concatenate([df[col].to_numpy() for col in cols], axis=0)\n    df = pd.DataFrame(data=df, columns=['id', 'binds'])\n    df = df.dropna().sort_values(by='id').reset_index(drop=True)\n    df['id'] = df['id'].astype(int)\n    df.to_csv(os.path.join(working, 'submission.csv'), index=False)\n\n    return None\n\n\n# UTILS >>>\n\ndef read_parquet(subset: str, root: str, **kwargs) -> pd.DataFrame:\n    \"\"\"\n    Read and preprocess train/test parquet files.\n    \"\"\"\n\n    # Read train set\n    df = pd.read_parquet(os.path.join(root, f'{subset}.parquet'))\n    df = df.rename(columns={\n        'buildingblock1_smiles': 'block1',\n        'buildingblock2_smiles': 'block2',\n        'buildingblock3_smiles': 'block3',\n        'molecule_smiles': 'smiles'})\n\n    # Group by molecule -> get multiclass labels\n    cols = ['block1', 'block2', 'block3', 'smiles']\n    values = 'binds' if subset == 'train' else 'id'\n    df = df.pivot(index=cols, columns='protein_name', values=values).reset_index()\n\n    return df\n\n\ndef make_parquet(working: str, seed: int, **kwargs) -> None:\n    \"\"\"\n    Make Dask DataFrame:\n\n    * Read and shuffle dataframe.\n    * Stack labels binding affinity [BRD4, HSA, sEH]. Nan mask = 2.\n    * Validation split (at least one non-shared blocks).\n    * Add train/validation/test indicator (0/1/2).\n    * Replace [Dy] DNA-linker with [H]\n    * Get ECFPs\n    * Write to parquet file.\n\n    Source of extra data: https://chemrxiv.org/engage/chemrxiv/article-details/6438943f08c86922ffeffe57\n    Processed by: @chemdatafarmer @Hengck23\n    \"\"\"\n\n    def validation_split(x, test: set):\n        \"\"\"\n        Get train (0), or validation (1) indicators.\n        Train subset: zero-intersection between blocks and \"test\" blocks.\n        Validation subset: non-zero-intersection between blocks and \"test\" blocks.\n        \"\"\"\n\n        blocks = set(x[col] for col in ['block1', 'block2', 'block3'])\n        i = len(blocks.intersection(test))\n        i = 0 if i == 0 else 1\n        i = np.int8(i)\n\n        return i\n\n    def replace_linker(smiles):\n        \"\"\"\n        Replace [Dy] linker with hydrogen.\n        \"\"\"\n        smiles = smiles.replace('[Dy]', '[H]')\n        smiles = Chem.CanonSmiles(smiles)\n        return smiles\n\n    # Iterate over subsets\n    dataset = []\n    for subset in ['test', 'extra', 'train']:\n\n        # Read parquet/csv\n        if subset in ['train', 'test']:\n            df = read_parquet(subset=subset, **parameters)\n        else:\n            df = pd.read_csv(os.path.join(working, 'DNA_Labeled_Data.csv'), usecols=['new_structure', 'read_count'])\n            df = df.rename(columns={'new_structure': 'smiles', 'read_count': 'binds'})\n\n        # Stack binding affinity labels\n        cols = ['BRD4', 'HSA', 'sEH']\n        if subset == 'train':\n            df['binds'] = np.stack([df[col].to_numpy() for col in cols], axis=-1, dtype=np.int8).tolist()\n        elif subset == 'test':\n            df[\"binds\"] = np.tile(np.array([[2, 2, 2]], dtype=np.int8), reps=(df.shape[0], 1)).tolist()\n        else:\n            df['binds'] = df['binds'].mapply(lambda x: [2, 2, np.clip(x, a_min=0, a_max=1)])\n        for col in cols:\n            df = df.drop(columns=[col]) if col in df.columns else df\n\n        # Validation split\n        if subset == 'train':\n            blocks = list(set(df['block1'].to_list()) | set(df['block2'].tolist()) | set(df['block3'].tolist()))\n            _, val, _, _ = train_test_split(blocks, blocks, test_size=0.03, random_state=seed)\n            df['subset'] = df.mapply(lambda x: validation_split(x, test=val), axis=1)\n        elif subset == 'test':\n            df['subset'] = 2\n        else:\n            df['subset'] = 0    # Use extra data only for training\n        cols = ['block1', 'block2', 'block3']\n        for col in cols:\n            df = df.drop(columns=[col]) if col in df.columns else df\n\n        # Replace [Dy] DNA-linker with [H]\n        df['smiles_no_linker'] = df['smiles'].mapply(lambda x: replace_linker(smiles=x))\n\n        # Append dataframe to list\n        dataset.append(df)\n\n    # Concatenate -> Shuffle -> Convert to Dask Dataframe\n    df = pd.concat(dataset)\n    df = df.sample(frac=1.0, ignore_index=True, random_state=seed)\n    df = dd.from_pandas(data=df)\n    df = df.repartition(npartitions=20)\n\n    # Write to parquet\n    df.to_parquet(os.path.join(working, 'belka.parquet'), schema={\n        'smiles': pa.string(),\n        'binds': pa.list_(pa.int8(), 3),\n        'subset': pa.int8(),\n        'smiles_no_linker': pa.string()})\n\n    return None\n\n\ndef make_dataset(working: str, **kwargs) -> None:\n    \"\"\"\n    Make TFR dataset.\n    \"\"\"\n\n    def generator() -> dict:\n        for row in df.itertuples(index=False, name='Row'):\n            yield {\n                'smiles': row.smiles,\n                'smiles_no_linker': row.smiles_no_linker,\n                'binds': row.binds,\n                'subset': row.subset}\n\n    def serialize_smiles(x) -> dict:\n        \"\"\"\n        Serialize smiles to string.\n        \"\"\"\n\n        x['smiles'] = tf.io.serialize_tensor(x['smiles'])\n        return x\n\n    def get_ecfp(x) -> dict:\n        \"\"\"\n        Compute ECFP form \"smiles_no_linker\".\n        \"\"\"\n\n        x['ecfp'] = transformer(x['smiles_no_linker'])\n        x.pop('smiles_no_linker')\n        return x\n\n    # Read dataset\n    df = dd.read_parquet(os.path.join(working, 'belka.parquet'))\n    df = df.compute()\n\n    # Tokenize SMILES\n    df['smiles'] = df['smiles'].mapply(lambda x: atomInSmiles.smiles_tokenizer(x))\n\n    # Write to TFRecords\n    auto = tf.data.AUTOTUNE\n    transformer = FPGenerator()\n    ds = tf.data.Dataset.from_generator(\n        generator=lambda : generator(),\n        output_signature={\n            'smiles': tf.TensorSpec(shape=(None,), dtype=tf.string),\n            'smiles_no_linker': tf.TensorSpec(shape=(), dtype=tf.string),\n            'binds': tf.TensorSpec(shape=(3,), dtype=tf.int8),\n            'subset': tf.TensorSpec(shape=(), dtype=tf.int8)})\n    ds = ds.map(lambda x: serialize_smiles(x))\n    ds = ds.batch(batch_size=1024, num_parallel_calls=auto)\n    ds = ds.map(lambda x: get_ecfp(x), num_parallel_calls=auto)\n    ds = ds.unbatch()\n    ds.save(os.path.join(working, 'belka.tfr'), compression='GZIP')\n    return None\n\n\ndef get_vocab(working: str, **kwargs) -> None:\n    \"\"\"\n    Get vocabulary for SMILES encoding.\n    \"\"\"\n\n    # Read parquet\n    df = dd.read_parquet(os.path.join(working, 'belka.parquet'))\n    df = df.compute()\n\n    # Tokenize SMILES -> Get list if unique tokens\n    df['smiles'] = df['smiles'].mapply(lambda x: list(set(atomInSmiles.smiles_tokenizer(x))))\n    vocab = np.unique(list(itertools.chain.from_iterable(df['smiles'].tolist()))).tolist()\n    vocab = pd.DataFrame(data=vocab)\n    vocab.to_csv(os.path.join(working, 'vocab.txt'), index=False, header=False)\n    return None\n\n\ndef get_smiles_encoder(vocab: str, **kwargs) -> TextVectorization:\n    \"\"\"\n    Get TextVectorization SMILES encoder.\n    \"\"\"\n\n    tokenizer = TextVectorization(\n        standardize=None,\n        split=None,\n        vocabulary=vocab)\n    return tokenizer\n\n\ndef load_model(model: str, **kwargs) -> tf.keras.Model:\n    \"\"\"\n    Load the model.\n    \"\"\"\n\n    model = tf.keras.models.load_model(model, compile=True, custom_objects={\n        'Encodings': Encodings,\n        'Embeddings': Embeddings,\n        'FeedForward': FeedForward,\n        'SelfAttention': SelfAttention,\n        'EncoderLayer': EncoderLayer,\n        'MultiLabelLoss': MultiLabelLoss,\n        'CategoricalLoss': CategoricalLoss,\n        'BinaryLoss': BinaryLoss,\n        'MaskedAUC': MaskedAUC})\n    return model\n\n\ndef set_parameters(\n        activation: str, batch_size: int, buffer_size: Union[int, float],\n        depth: int,\n        dropout_rate: float, epochs: int, epsilon: float, initial_epoch: int,\n        masking_rate: float, max_length: int, mode: str, model: Union[str, None],\n        model_name: str, num_heads: int, num_layers: int,\n        patience: int, root: str, seed: int, steps_per_epoch: int, validation_steps: int, vocab: str, vocab_size: int,\n        working: str) -> dict:\n    \"\"\"\n    Set uniform parameters for the functions in the scope of the project.\n    :param mode: Choose from ['clf', 'fps', 'mlm'],\n    :param vocab_size: Set to N+2, where N - number of tokens ([PAD] = 0, [MASK] = 1).    :return:\n    \"\"\"\n\n    inputs = {\n        'activation': activation,\n        'batch_size': batch_size,\n        'buffer_size': int(buffer_size),\n        'depth': depth,\n        'dropout_rate': dropout_rate,\n        'epochs': epochs,\n        'epsilon': epsilon,\n        'initial_epoch': initial_epoch,\n        'masking_rate': masking_rate,\n        'max_length': max_length,\n        'mode': mode,\n        'model': model,\n        'model_name': model_name,\n        'num_heads': num_heads,\n        'num_layers': num_layers,\n        'patience': patience,\n        'root': root,\n        'seed': seed,\n        'steps_per_epoch': steps_per_epoch,\n        'validation_steps': validation_steps,\n        'vocab': vocab,\n        'vocab_size': vocab_size,\n        'working': working}\n\n    return inputs\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mapply.init(n_workers=-1, progressbar=True)\nparameters = set_parameters(\n    root='/kaggle/input/leash-BELKA',\n    working='/kaggle/working/',\n    vocab='/kaggle/input/belka2024-finalsolution/vocab.txt',\n    model=None,\n    mode='clf',\n    model_name='belka',\n    masking_rate=0.15,\n    batch_size=2048, buffer_size=1e07,\n    epochs=1000, initial_epoch=0, steps_per_epoch=10000, validation_steps=2000,\n    max_length=128, vocab_size=43,\n    depth=32, dropout_rate=0.1, num_heads=8, num_layers=4, activation='gelu',\n    patience=20, epsilon=1e-07, seed=42)","metadata":{},"execution_count":null,"outputs":[]}]}