{"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"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":59094,"databundleVersionId":6541963,"sourceType":"competition"}],"dockerImageVersionId":30558,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"raw","source":"# **<span style=\"color:#F7B2B0;\">Introduction</span>**\n\nHuman biology can be complex, in part due to the function and interplay of the body's approximately 37 trillion cells, which are organized into tissues, organs, and systems. However, recent advances in single-cell technologies have provided unparalleled insight into the function of cells and tissues at the level of DNA, RNA, and proteins. Yet leveraging single-cell methods to develop medicines requires mapping causal links between chemical perturbations and the downstream impact on cell state. These experiments are costly and labor intensive, and not all cells and tissues are amenable to high-throughput transcriptomic screening. If data science could help accurately predict chemical perturbations in new cell types, it could accelerate and expand the development of new medicines.\nSeveral methods have been developed for drug perturbation prediction, most of which are variations on the autoencoder architecture (Dr.VAE, scGEN, and ChemCPA). However, these methods lack proper benchmarking datasets with diverse cell types to determine how well they generalize. The largest available training dataset is the NIH-funded Connectivity Map (CMap), which comprises over 1.3M small molecule perturbation measurements. However, the CMap includes observations of only 978 genes, less than 5% of all genes. Furthermore, the CMap data is comprised almost entirely of measurements in cancer cell lines, which may not accurately represent human biology.\n\n# **<span style=\"color:#F7B2B0;\">Goal</span>**\n\nThe goal of this competition is to accurately predict chemical perturbations in new cell types could accelerate the discovery and enable the creation of new medicines to treat or cure disease.\n\n# **<span style=\"color:#F7B2B0;\">Install the Libraries</span>**\n","metadata":{}},{"cell_type":"code","source":"!pip install -q wandb tabpfn","metadata":{"execution":{"iopub.status.busy":"2023-09-17T06:51:19.460588Z","iopub.execute_input":"2023-09-17T06:51:19.460940Z","iopub.status.idle":"2023-09-17T06:51:35.809904Z","shell.execute_reply.started":"2023-09-17T06:51:19.460910Z","shell.execute_reply":"2023-09-17T06:51:35.808580Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **<span style=\"color:#F7B2B0;\">Import the Packages</span>**","metadata":{}},{"cell_type":"code","source":"# Importing necessary libraries\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport pyarrow.parquet as pq\nfrom sklearn.preprocessing import MinMaxScaler\nfrom sklearn.model_selection import train_test_split\n\nfrom kaggle_secrets import UserSecretsClient\nimport wandb\nfrom datetime import datetime","metadata":{"_uuid":"7b575297-32e1-4d21-ab73-817900926e53","_cell_guid":"a3899b3e-7985-483d-b155-33283f5b5084","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-09-17T07:29:56.733583Z","iopub.execute_input":"2023-09-17T07:29:56.734133Z","iopub.status.idle":"2023-09-17T07:29:56.741321Z","shell.execute_reply.started":"2023-09-17T07:29:56.734098Z","shell.execute_reply":"2023-09-17T07:29:56.737875Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<img src=\"https://camo.githubusercontent.com/dd842f7b0be57140e68b2ab9cb007992acd131c48284eaf6b1aca758bfea358b/68747470733a2f2f692e696d6775722e636f6d2f52557469567a482e706e67\">\n\n> I will be integrating W&B for visualizations and logging artifacts!\n> \n> [Open Problems = Single Cell Pertubations Project on W&B Dashboard](https://wandb.ai/usharengaraju/open_problems)\n> \n> - To get the API key, create an account in the [website](https://wandb.ai/site) .\n> - Use secrets to use API Keys more securely ","metadata":{}},{"cell_type":"code","source":"# Setup user secrets for login\nuser_secrets = UserSecretsClient()\nwandb_api = user_secrets.get_secret(\"api_key\") \nwandb.login(key=wandb_api)","metadata":{"execution":{"iopub.status.busy":"2023-09-17T07:29:56.742797Z","iopub.execute_input":"2023-09-17T07:29:56.743134Z","iopub.status.idle":"2023-09-17T07:29:56.901093Z","shell.execute_reply.started":"2023-09-17T07:29:56.743104Z","shell.execute_reply":"2023-09-17T07:29:56.900165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **<span style=\"color:#F7B2B0;\">Data Pipeline</span>**","metadata":{}},{"cell_type":"code","source":"fn = '/kaggle/input/open-problems-single-cell-perturbations/de_train.parquet'\ndf_de_train = pd.read_parquet(fn)# , index_col = 0)\nprint(df_de_train.shape)\ndf_de_train.head()","metadata":{"execution":{"iopub.status.busy":"2023-09-17T07:29:56.902511Z","iopub.execute_input":"2023-09-17T07:29:56.902872Z","iopub.status.idle":"2023-09-17T07:29:59.302885Z","shell.execute_reply.started":"2023-09-17T07:29:56.902840Z","shell.execute_reply":"2023-09-17T07:29:59.301892Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"run = wandb.init(project = 'open_problems',\n                 config = {},\n                 save_code = True,\n                 \n)\ntable = wandb.Table(dataframe=df_de_train)\n\nwandb.log({\"Table\":table})\nrun.finish()","metadata":{"execution":{"iopub.status.busy":"2023-09-15T17:05:34.667520Z","iopub.execute_input":"2023-09-15T17:05:34.668283Z","iopub.status.idle":"2023-09-15T17:18:06.472214Z","shell.execute_reply.started":"2023-09-15T17:05:34.668250Z","shell.execute_reply":"2023-09-15T17:18:06.471304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X = df_de_train.iloc[:,:5]\ny=  df_de_train.iloc[:,5:]\nX.shape, y.shape\n","metadata":{"execution":{"iopub.status.busy":"2023-09-17T07:29:59.304541Z","iopub.execute_input":"2023-09-17T07:29:59.305213Z","iopub.status.idle":"2023-09-17T07:29:59.342456Z","shell.execute_reply.started":"2023-09-17T07:29:59.305179Z","shell.execute_reply":"2023-09-17T07:29:59.341337Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.decomposition import TruncatedSVD\nreducer = TruncatedSVD(n_components=35, n_iter=7, random_state=42)\nYr = reducer.fit_transform(y)","metadata":{"execution":{"iopub.status.busy":"2023-09-17T07:44:37.430068Z","iopub.execute_input":"2023-09-17T07:44:37.430458Z","iopub.status.idle":"2023-09-17T07:44:39.275897Z","shell.execute_reply.started":"2023-09-17T07:44:37.430427Z","shell.execute_reply":"2023-09-17T07:44:39.274632Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Yr.shape","metadata":{"execution":{"iopub.status.busy":"2023-09-17T07:44:39.278129Z","iopub.execute_input":"2023-09-17T07:44:39.279017Z","iopub.status.idle":"2023-09-17T07:44:39.289647Z","shell.execute_reply.started":"2023-09-17T07:44:39.278973Z","shell.execute_reply":"2023-09-17T07:44:39.288209Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.preprocessing import LabelEncoder\n\nencoder = LabelEncoder()\n\n# Encode Categorical Data\nfor col in X.columns:\n  X[col]=encoder.fit_transform(X[col])","metadata":{"execution":{"iopub.status.busy":"2023-09-17T07:30:01.362044Z","iopub.execute_input":"2023-09-17T07:30:01.362688Z","iopub.status.idle":"2023-09-17T07:30:01.376150Z","shell.execute_reply.started":"2023-09-17T07:30:01.362654Z","shell.execute_reply":"2023-09-17T07:30:01.374981Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_train, X_val, y_train, y_val = train_test_split(X.values, Yr, test_size=0.2)","metadata":{"_uuid":"1b778375-9872-4f05-a5a2-e2bd3dcf46b2","_cell_guid":"8691ffdb-b3c1-4499-8e24-f48547ec1193","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-09-17T07:44:39.549788Z","iopub.execute_input":"2023-09-17T07:44:39.550174Z","iopub.status.idle":"2023-09-17T07:44:39.557035Z","shell.execute_reply.started":"2023-09-17T07:44:39.550140Z","shell.execute_reply":"2023-09-17T07:44:39.556097Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from wandb.keras import WandbCallback, WandbMetricsLogger\nrun = wandb.init(project = 'open_problems',\n                 save_code = True,\n                 name='tabtransformer'\n                 \n)","metadata":{"execution":{"iopub.status.busy":"2023-09-15T17:18:06.473707Z","iopub.execute_input":"2023-09-15T17:18:06.473967Z","iopub.status.idle":"2023-09-15T17:18:38.627950Z","shell.execute_reply.started":"2023-09-15T17:18:06.473943Z","shell.execute_reply":"2023-09-15T17:18:38.627030Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **<span style=\"color:#F7B2B0;\">Tab transformer</span>**\n\n[Source](https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/structured_data/ipynb/tabtransformer.ipynb)\n\nThe TabTransformer architecture works as follows:\n\n📌 All the categorical features are encoded as embeddings, using the same embedding_dims. This means that each value in each categorical feature will have its own embedding vector.\n\n📌 A column embedding, one embedding vector for each categorical feature, is added (point-wise) to the categorical feature embedding.\n\n📌 The embedded categorical features are fed into a stack of Transformer blocks. Each Transformer block consists of a multi-head self-attention layer followed by a feed-forward layer.\n\n📌 The outputs of the final Transformer layer, which are the contextual embeddings of the categorical features, are concatenated with the input numerical features, and fed into a final MLP block.\n\n📌 A softmax classifer is applied at the end of the model.\n\nThe [paper](https://arxiv.org/pdf/2012.06678.pdf) discusses both addition and concatenation of the column embedding in the Appendix: Experiment and Model Details section. The architecture of TabTransformer is shown below, as presented in the paper.\n\n![](https://i.imgur.com/kSB0jYw.png)","metadata":{}},{"cell_type":"code","source":"from tensorflow import keras\nimport tensorflow as tf \nfrom tensorflow.keras import layers\n\nclass TransformerBlock(layers.Layer):\n    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):\n        super(TransformerBlock, self).__init__()\n        # parametreleri\n        self.att = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)\n        self.ffn = keras.Sequential(\n            [layers.Dense(ff_dim, activation=\"relu\"), layers.Dense(embed_dim),]\n        )\n        # batch-layer\n        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)\n        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)\n        self.dropout1 = layers.Dropout(rate)\n        self.dropout2 = layers.Dropout(rate)\n\n    def call(self, inputs, training):\n        attn_output = self.att(inputs, inputs)\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\nclass TabTransformer(keras.Model):\n\n    def __init__(self, \n            categories,\n            num_continuous,\n            dim,\n            dim_out,\n            depth,\n            heads,\n            attn_dropout,\n            ff_dropout,\n            mlp_hidden,\n            normalize_continuous = True):\n        \n        super(TabTransformer, self).__init__()\n\n        # --> continuous inputs\n        self.normalize_continuous = normalize_continuous\n        if normalize_continuous:\n            self.continuous_normalization = layers.LayerNormalization()\n\n        # --> categorical inputs\n\n        # embedding\n        self.embedding_layers = []\n        for number_of_classes in categories:\n            self.embedding_layers.append(layers.Embedding(input_dim = number_of_classes, output_dim = dim))\n\n        # concatenation\n        self.embedded_concatenation = layers.Concatenate(axis=1)\n\n        # adding transformers\n        self.transformers = []\n        for _ in range(depth):\n            self.transformers.append(TransformerBlock(dim, heads, dim))\n        self.flatten_transformer_output = layers.Flatten()\n\n        # --> MLP\n        self.pre_mlp_concatenation = layers.Concatenate()\n\n        # mlp layers\n        self.mlp_layers = []\n        for size, activation in mlp_hidden:\n            self.mlp_layers.append(layers.Dense(size, activation=activation))\n\n        self.output_layer = layers.Dense(dim_out)\n\n    def call(self, inputs):\n        categorical_inputs = inputs\n#         print(inputs[:,0])\n        # --> categorical\n        embedding_outputs = []\n        for i in range(5):\n            embedding_outputs.append(tf.expand_dims(self.embedding_layers[i](categorical_inputs[:,i]),axis=1))\n#         print(embedding_outputs[0].shape)\n        categorical_inputs = self.embedded_concatenation(embedding_outputs)\n#         categorical_inputs = tf.expand_dims(categorical_inputs,axis=1)\n        for transformer in self.transformers:\n            categorical_inputs = transformer(categorical_inputs)\n        contextual_embedding = self.flatten_transformer_output(categorical_inputs)\n\n\n        for mlp_layer in self.mlp_layers:\n            mlp_input = mlp_layer(contextual_embedding)\n\n        return self.output_layer(mlp_input)","metadata":{"execution":{"iopub.status.busy":"2023-09-15T17:43:09.630420Z","iopub.execute_input":"2023-09-15T17:43:09.630982Z","iopub.status.idle":"2023-09-15T17:43:19.926879Z","shell.execute_reply.started":"2023-09-15T17:43:09.630942Z","shell.execute_reply":"2023-09-15T17:43:19.925826Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"nu = []\nfor col in X.columns:\n    nu.append(len(X[col].unique()))","metadata":{"execution":{"iopub.status.busy":"2023-09-15T17:43:19.928938Z","iopub.execute_input":"2023-09-15T17:43:19.929709Z","iopub.status.idle":"2023-09-15T17:43:19.938946Z","shell.execute_reply.started":"2023-09-15T17:43:19.929672Z","shell.execute_reply":"2023-09-15T17:43:19.938002Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tensorflow.keras.optimizers import Adam\n","metadata":{"execution":{"iopub.status.busy":"2023-09-15T17:43:19.940295Z","iopub.execute_input":"2023-09-15T17:43:19.940918Z","iopub.status.idle":"2023-09-15T17:43:19.955915Z","shell.execute_reply.started":"2023-09-15T17:43:19.940884Z","shell.execute_reply":"2023-09-15T17:43:19.954885Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tabTransformer = TabTransformer(\n    categories = nu, # number of unique elements in each categorical feature\n    num_continuous = 5,      # number of numerical features\n    dim = 16,                # embedding/transformer dimension\n    dim_out = 35,             # dimension of the model output\n    depth = 6,               # number of transformer layers in the stack\n    heads = 8,               # number of attention heads\n    attn_dropout = 0.1,      # attention layer dropout in transformers\n    ff_dropout = 0.1,        # feed-forward layer dropout in transformers\n    mlp_hidden = [(32, 'relu'), (16, 'relu')] # mlp layer dimensions and activations\n)\ntabTransformer.compile(Adam(0.001),'mae',metrics=['mae'])\ntabTransformer.fit(X_train,y_train,validation_data=(X_val,y_val),batch_size=32,epochs=30,callbacks=[WandbMetricsLogger()])","metadata":{"execution":{"iopub.status.busy":"2023-09-15T17:21:05.732281Z","iopub.execute_input":"2023-09-15T17:21:05.732646Z","iopub.status.idle":"2023-09-15T17:21:53.122512Z","shell.execute_reply.started":"2023-09-15T17:21:05.732617Z","shell.execute_reply":"2023-09-15T17:21:53.121060Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"run.finish()","metadata":{"execution":{"iopub.status.busy":"2023-09-15T17:21:53.125374Z","iopub.execute_input":"2023-09-15T17:21:53.126069Z","iopub.status.idle":"2023-09-15T17:22:00.563121Z","shell.execute_reply.started":"2023-09-15T17:21:53.126033Z","shell.execute_reply":"2023-09-15T17:22:00.562240Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **<span style=\"color:#F7B2B0;\">SAINT</span>**\n\nSAINT, performs attention over both rows and columns, and it includes an enhanced embedding method. A new contrastive self-supervised pre-training method is used when labels are scarce. SAINT consistently improves performance over previous deep learning methods, and it even outperforms gradient boosting methods, including XGBoost, CatBoost, and LightGBM, on average over a variety of benchmark tasks.\n\n![](https://i.imgur.com/WVEY6uy.png)","metadata":{}},{"cell_type":"code","source":"run = wandb.init(project = 'open_problems',\n                 save_code = True,\n                 name='SAINT'\n                 \n)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\ntrainds = tf.data.Dataset.from_tensor_slices((X_train, y_train))\ntrainds = trainds.batch(32, drop_remainder = True)\n\nvalds = tf.data.Dataset.from_tensor_slices((X_val, y_val))\nvalds = valds.batch(32, drop_remainder = True)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\nfrom wandb.keras import WandbCallback, WandbMetricsLogger\nfrom tensorflow.keras.layers import *\nfrom tensorflow.keras.activations import gelu, softmax\nfrom tensorflow.keras.models import Sequential\nclass MHA(Layer):\n    '''\n    Multi-Head Attention Layer\n    '''\n    \n    def __init__(self, num_head, dropout = 0):\n        super(MHA, self).__init__()\n        \n        # Constants\n        self.num_head = num_head\n        self.dropout_rate = dropout\n        \n    def build(self, input_shape):\n        query_shape = input_shape\n        d_model = query_shape[-1]\n        units = d_model // self.num_head\n        \n        # Loop for Generate each Attention\n        self.layer_q = []\n        for _ in range(self.num_head):\n            layer = Dense(units, activation = None, use_bias = False)\n            layer.build(query_shape)\n            self.layer_q.append(layer)\n            \n        self.layer_k = []\n        for _ in range(self.num_head):\n            layer = Dense(units, activation = None, use_bias = False)\n            layer.build(query_shape)\n            self.layer_k.append(layer)\n            \n        self.layer_v = []\n        for _ in range(self.num_head):\n            layer = Dense(units, activation = None, use_bias = False)\n            layer.build(query_shape)\n            self.layer_v.append(layer)\n            \n        self.out = Dense(d_model, activation = None, use_bias = False)\n        self.out.build(query_shape)\n        self.dropout = Dropout(self.dropout_rate)\n        self.dropout.build(query_shape)\n        \n    def call(self, x):\n        d_model = x.shape[-1]\n        scale = d_model ** -0.5\n        \n        attention_values = []\n        for i in range(self.num_head):\n            attention_score = softmax(tf.matmul(self.layer_q[i](x), self.layer_k[i](x), transpose_b=True) * scale)\n            attention_final = tf.matmul(attention_score, self.layer_v[i](x))\n            attention_values.append(attention_final)\n            \n        attention_concat = tf.concat(attention_values, axis = -1)\n        out = self.out(self.dropout(attention_concat))\n        \n        return out\n\nclass IMHA(Layer):\n    '''\n    Intersample Multi Head Attention\n    Attend on row(samples) not column(features)\n    '''\n    \n    def __init__(self, num_head, dropout = 0):\n        super(IMHA, self).__init__()\n        \n        # Constants\n        self.num_head = num_head\n        self.dropout_rate = dropout\n        \n    def build(self, input_shape):\n        b, n, d = input_shape\n        query_shape = input_shape\n        units = (d * n) // self.num_head\n        # Loop for Generate each Attention\n        self.layer_q = []\n        for _ in range(self.num_head):\n            layer = Dense(units, activation = None, use_bias = False)\n            layer.build([1, b, int(n * d)])\n            self.layer_q.append(layer)\n            \n        self.layer_k = []\n        for _ in range(self.num_head):\n            layer = Dense(units, activation = None, use_bias = False)\n            layer.build([1, b, int(n * d)])\n            self.layer_k.append(layer)\n            \n        self.layer_v = []\n        for _ in range(self.num_head):\n            layer = Dense(units, activation = None, use_bias = False)\n            layer.build([1, b, int(n * d)])\n            self.layer_v.append(layer)\n            \n        self.out = Dense(d, activation = None, use_bias = False)\n        self.out.build(query_shape)\n        self.dropout = Dropout(self.dropout_rate)\n        self.dropout.build(query_shape)\n        \n    def call(self, x):\n        b, n, d = x.shape\n        scale = d ** -0.5\n        x = tf.reshape(x, (1, b, int(n * d)))\n        attention_values = []\n        \n        for i in range(self.num_head):\n            attention_score = softmax(tf.matmul(self.layer_q[i](x), self.layer_k[i](x), transpose_b=True) * scale)\n            attention_final = tf.matmul(attention_score, self.layer_v[i](x))\n            attention_final = tf.reshape(attention_final, (b, n, int(d / self.num_head)))\n            attention_values.append(attention_final)\n            \n        attention_concat = tf.concat(attention_values, axis = -1)\n        out = self.out(self.dropout(attention_concat))\n        \n        return out\n\nclass FeedForwardNetwork(Layer):\n    def __init__(self, dim, dropout = 0.0):\n        super(FeedForwardNetwork, self).__init__()\n        self.dense = Dense(dim, activation = 'gelu')\n        self.dropout = Dropout(dropout)\n        \n    def call(self, x):\n        return self.dropout(self.dense(x))\n\nclass CustomEmbedding(Layer):\n    def __init__(self, num_categorical, dim):\n        super(CustomEmbedding, self).__init__()\n        self.num_categorical = num_categorical\n        self.dim = dim\n        \n    def build(self, input_shape):\n        b, n = input_shape\n        self.embedding_categorical = Embedding(self.dim * 2, self.dim)\n        self.embedding_categorical.build([b, self.num_categorical])\n        \n        self.embedding_numerical = Dense(self.dim, activation = 'relu')\n        self.embedding_numerical.build([b, int(n - self.num_categorical), 1])\n        \n    def call(self, x):\n        b, n = x.shape\n        categorical_x = x[:, :self.num_categorical]\n        numerical_x = x[:, self.num_categorical:]\n        numerical_x = tf.reshape(numerical_x, (b, int(n - self.num_categorical), 1))\n        \n        embedded_cat = self.embedding_categorical(categorical_x)\n        embedded_num = self.embedding_numerical(numerical_x)\n    \n        embedded_x = tf.concat([embedded_cat, embedded_num], axis = 1)\n        \n        return embedded_x\n\n\nclass SAINT(Layer):\n    def __init__(self, repeat, num_categorical, EMB_DIM, MHA_HEADS, IMHA_HEADS):\n        super(SAINT, self).__init__()\n        self.repeat = repeat\n        self.layer_mha = []\n        self.layer_imha = []\n        self.layer_ffn = []\n        self.layer_layernorm = []\n        self.embedding = CustomEmbedding(num_categorical, EMB_DIM)\n        \n        for _ in range(repeat):\n            mha = MHA(MHA_HEADS)\n            imha = IMHA(IMHA_HEADS)\n            ffn_1 = FeedForwardNetwork(EMB_DIM)\n            ffn_2 = FeedForwardNetwork(EMB_DIM)\n            layernorm_1 = LayerNormalization()\n            layernorm_2 = LayerNormalization()\n            layernorm_3 = LayerNormalization()\n            layernorm_4 = LayerNormalization()\n            \n            self.layer_mha.append(mha)\n            self.layer_imha.append(imha)\n            self.layer_ffn.append(ffn_1)\n            self.layer_ffn.append(ffn_2)\n            self.layer_layernorm.append(layernorm_1)\n            self.layer_layernorm.append(layernorm_2)\n            self.layer_layernorm.append(layernorm_3)\n            self.layer_layernorm.append(layernorm_4)\n            \n    def call(self, x):\n        x = self.embedding(x)\n        # Depth of SAINT Layer\n        for i in range(self.repeat):\n            # Multi-Head part\n            x = self.layer_layernorm[i](self.layer_mha[i](x)) + x\n            x = self.layer_layernorm[i+1](self.layer_ffn[i](x)) + x\n            \n            # Intersample Multi-Head part\n            x = self.layer_layernorm[i+2](self.layer_imha[i](x)) + x\n            x = self.layer_layernorm[i+3](self.layer_ffn[i+1](x)) + x\n       \n        # only using cls token for final output\n        out = x[:, 0] # CLS Token\n        \n        return out\nnu = []\nfor col in X.columns:\n    nu.append(len(X[col].unique()))\n    \nmodel = Sequential([\n            Input(shape = (5),batch_size=32),\n            SAINT(1, 4, 64, 8, 8),\n            Dense(35, activation = 'linear')\n        ])\n\nmodel.compile(tf.keras.optimizers.Adam(0.001), 'mae',metrics=['mae'])\nmodel.summary()\nmodel.fit(trainds,epochs=20,batch_size=32,validation_data=valds,callbacks=[WandbMetricsLogger()])","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"run.finish()","metadata":{},"execution_count":null,"outputs":[]}]}