{"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":59094,"databundleVersionId":7010844,"sourceType":"competition"},{"sourceId":7712331,"sourceType":"datasetVersion","datasetId":4441094},{"sourceId":143494741,"sourceType":"kernelVersion"}],"dockerImageVersionId":30588,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')\nimport os\nimport random\nimport numpy as np \nimport pandas as pd\nfrom tqdm import tqdm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn.utils import clip_grad_norm_\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForMaskedLM, AutoTokenizer\nimport json","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-02-27T12:01:49.906922Z","iopub.execute_input":"2024-02-27T12:01:49.907626Z","iopub.status.idle":"2024-02-27T12:01:55.213356Z","shell.execute_reply.started":"2024-02-27T12:01:49.907588Z","shell.execute_reply":"2024-02-27T12:01:55.212414Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Seed Everything","metadata":{}},{"cell_type":"code","source":"def seed_everything():\n    seed = 42\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    os.environ['TOKENIZERS_PARALLELISM'] = 'true'\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed(seed)\n        torch.cuda.manual_seed_all(seed)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n    print('-----Seed Set!-----') ","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:01:55.215181Z","iopub.execute_input":"2024-02-27T12:01:55.215609Z","iopub.status.idle":"2024-02-27T12:01:55.221491Z","shell.execute_reply.started":"2024-02-27T12:01:55.215581Z","shell.execute_reply":"2024-02-27T12:01:55.220648Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed_everything()","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:01:55.222596Z","iopub.execute_input":"2024-02-27T12:01:55.222895Z","iopub.status.idle":"2024-02-27T12:01:55.266669Z","shell.execute_reply.started":"2024-02-27T12:01:55.222870Z","shell.execute_reply":"2024-02-27T12:01:55.265702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Download the pretrained ChemBERTa model","metadata":{}},{"cell_type":"code","source":"chemberta = AutoModelForMaskedLM.from_pretrained(\"DeepChem/ChemBERTa-77M-MTR\")\ntokenizer = AutoTokenizer.from_pretrained(\"DeepChem/ChemBERTa-77M-MTR\")\nchemberta._modules[\"lm_head\"] = nn.Identity()","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:01:55.268891Z","iopub.execute_input":"2024-02-27T12:01:55.269223Z","iopub.status.idle":"2024-02-27T12:01:58.634744Z","shell.execute_reply.started":"2024-02-27T12:01:55.269196Z","shell.execute_reply":"2024-02-27T12:01:58.633894Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Read Competition Datasets","metadata":{}},{"cell_type":"code","source":"de_train = pd.read_parquet('../input/open-problems-single-cell-perturbations/de_train.parquet')\nid_map = pd.read_csv('../input/open-problems-single-cell-perturbations/id_map.csv')\nsample_submission = pd.read_csv('../input/open-problems-single-cell-perturbations/sample_submission.csv', index_col='id')","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:01:58.635741Z","iopub.execute_input":"2024-02-27T12:01:58.636124Z","iopub.status.idle":"2024-02-27T12:02:03.661300Z","shell.execute_reply.started":"2024-02-27T12:01:58.636099Z","shell.execute_reply":"2024-02-27T12:02:03.660490Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"xlist  = ['cell_type','sm_name']\n_ylist = ['cell_type','sm_name','sm_lincs_id','SMILES','control']\n\ny = de_train.drop(columns=_ylist)\ny.shape","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:02:03.662441Z","iopub.execute_input":"2024-02-27T12:02:03.662744Z","iopub.status.idle":"2024-02-27T12:02:03.708740Z","shell.execute_reply.started":"2024-02-27T12:02:03.662718Z","shell.execute_reply":"2024-02-27T12:02:03.707796Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Use Scikit-Learn's One Hot Encoder\nThis helps encode each pair (cell_type, sm_name) as a multi-dimensional binary vector","metadata":{}},{"cell_type":"code","source":"from sklearn.preprocessing import OneHotEncoder\n\nencoder = OneHotEncoder()\nencoder.fit(de_train[xlist])\none_hot_encode_features = encoder.transform(de_train[xlist])\none_hot_test = encoder.transform(id_map[xlist])\n\nX = pd.DataFrame(one_hot_encode_features.toarray().astype(float))\ntest = pd.DataFrame(one_hot_test.toarray().astype(float))","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:02:03.710051Z","iopub.execute_input":"2024-02-27T12:02:03.710407Z","iopub.status.idle":"2024-02-27T12:02:04.216860Z","shell.execute_reply.started":"2024-02-27T12:02:03.710379Z","shell.execute_reply":"2024-02-27T12:02:04.216119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"markdown","source":"## ","metadata":{}},{"cell_type":"markdown","source":"## First Data Augmentation\nCompute the mean and std differential expression for each cell type and each small molecule name. Take this as additional input features to the model to be built.","metadata":{}},{"cell_type":"code","source":"de_cell_type = de_train.iloc[:, [0] + list(range(5, de_train.shape[1]))]\nde_sm_name = de_train.iloc[:, [1] + list(range(5, de_train.shape[1]))]\nmean_cell_type = de_cell_type.groupby('cell_type').mean().reset_index()\nmean_sm_name = de_sm_name.groupby('sm_name').mean().reset_index()\n\nstd_cell_type = de_cell_type.groupby('cell_type').std().reset_index()\nstd_sm_name = de_sm_name.groupby('sm_name').std().reset_index()","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:02:04.217871Z","iopub.execute_input":"2024-02-27T12:02:04.218154Z","iopub.status.idle":"2024-02-27T12:02:04.960791Z","shell.execute_reply.started":"2024-02-27T12:02:04.218129Z","shell.execute_reply":"2024-02-27T12:02:04.959999Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"I also consider the 25%, 50%, and 75% percentiles (see below). Not that I explored different combinations of these features. In particular, when using all additional features (i.e., mean, std, 25%, 50%, and 75% percentiles) I refer to the corresponding models as \"heavy\".","metadata":{}},{"cell_type":"code","source":"desc_cell_type = de_cell_type.groupby('cell_type').describe().reset_index()\ndesc_cell_type.drop(['mean', 'count', 'std', 'min', 'max'], axis=1, level=1, inplace=True)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:02:04.961967Z","iopub.execute_input":"2024-02-27T12:02:04.962319Z","iopub.status.idle":"2024-02-27T12:04:43.928349Z","shell.execute_reply.started":"2024-02-27T12:02:04.962287Z","shell.execute_reply":"2024-02-27T12:04:43.927444Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build a Mapping Between Each SM Name and Its SMILES\nThis helps build Chem BERTa features for molecule SMILES","metadata":{}},{"cell_type":"code","source":"sm_name2smiles = {smname:smiles for smname, smiles in zip(de_train['sm_name'], de_train['SMILES'])}\ntest_smiles = list(map(sm_name2smiles.get, id_map['sm_name'].values))","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:04:43.933138Z","iopub.execute_input":"2024-02-27T12:04:43.933429Z","iopub.status.idle":"2024-02-27T12:04:43.939344Z","shell.execute_reply.started":"2024-02-27T12:04:43.933404Z","shell.execute_reply":"2024-02-27T12:04:43.938403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Below, we define a variable to control whether we want to reproduce the leaderboard or train different, potentially better instances of our proposed models","metadata":{}},{"cell_type":"code","source":"reproduce_leaderboard_score = True # Either reproduce the leaderboard or build new (different) models to predict","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:04:43.940579Z","iopub.execute_input":"2024-02-27T12:04:43.940862Z","iopub.status.idle":"2024-02-27T12:04:43.951352Z","shell.execute_reply.started":"2024-02-27T12:04:43.940837Z","shell.execute_reply":"2024-02-27T12:04:43.950344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Thanks to ALEKSEY TREPETSKY (I upvoted) https://www.kaggle.com/code/alekseytrepetsky/create-chemberta-embed/notebook, I could either build my own ChemBERTa features or use the ones she/he has created and shared publicly.","metadata":{}},{"cell_type":"code","source":"chemberta.eval()\ndef build_ChemBERTa_features(smiles_list):\n    embeddings = torch.zeros(len(smiles_list), 384)\n    embeddings_mean = torch.zeros(len(smiles_list), 384)\n\n    with torch.no_grad():\n        for i, smiles in enumerate(tqdm(smiles_list)):\n            encoded_input = tokenizer(smiles, return_tensors=\"pt\", padding=False, truncation=True)\n            model_output = chemberta(**encoded_input)\n            \n            embedding = model_output[0][::,0,::]\n            embeddings[i] = embedding\n            \n            embedding = torch.mean(model_output[0], 1)\n            embeddings_mean[i] = embedding\n            \n    return embeddings.numpy(), embeddings_mean.numpy()","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:04:43.952541Z","iopub.execute_input":"2024-02-27T12:04:43.952852Z","iopub.status.idle":"2024-02-27T12:04:43.965658Z","shell.execute_reply.started":"2024-02-27T12:04:43.952826Z","shell.execute_reply":"2024-02-27T12:04:43.964727Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if reproduce_leaderboard_score:\n    train_chem_feat = np.load('../input/create-chemberta-embed/train_ChemBERTa_v2_77MTR_cls_pad_True.npy')\n    test_chem_feat = np.load('../input/create-chemberta-embed/test_ChemBERTa_v2_77MTR_cls_pad_True.npy')\n    train_chem_feat_mean = np.load('../input/create-chemberta-embed/train_ChemBERTa_v2_77MTR_mean_pad_True.npy')\n    test_chem_feat_mean = np.load('../input/create-chemberta-embed/test_ChemBERTa_v2_77MTR_mean_pad_True.npy')\nelse:\n    train_chem_feat, train_chem_feat_mean = build_ChemBERTa_features(de_train.SMILES)\n    test_chem_feat, test_chem_feat_mean = build_ChemBERTa_features(test_smiles)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:04:43.966711Z","iopub.execute_input":"2024-02-27T12:04:43.966995Z","iopub.status.idle":"2024-02-27T12:04:44.059130Z","shell.execute_reply.started":"2024-02-27T12:04:43.966971Z","shell.execute_reply":"2024-02-27T12:04:44.058133Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Now Define the Function to Combine the Created Features in Different Ways\nFor each combination, we create and train 3 deep learning architectures (see below)","metadata":{}},{"cell_type":"code","source":"def combine_features(add_dfs, chem_feats, main_df, add_vecs=None, use_description=False):\n    new_vecs = []\n    chem_feat_dim = 600 if reproduce_leaderboard_score else 384\n    if len(add_dfs) > 0:\n        add_len = sum(add_df.shape[1]-1 for add_df in add_dfs)+chem_feat_dim*len(chem_feats)+add_vecs.shape[1] if\\\n        add_vecs is not None else sum(add_df.shape[1]-1 for add_df in add_dfs)+chem_feat_dim*len(chem_feats)\n    else:\n        add_len = chem_feat_dim*len(chem_feats)+add_vecs.shape[1] if\\\n        add_vecs is not None else chem_feat_dim*len(chem_feats)\n    if use_description and reproduce_leaderboard_score:\n        add_len += (desc_cell_type.shape[1]-1)//3\n    for i in range(len(main_df)):\n        if add_vecs is not None:\n            vec_ = (add_vecs.iloc[i,:].values).copy()\n        else:\n            vec_ = np.array([])\n        for df in add_dfs:\n            if 'cell_type' in df.columns:\n                values = df[df['cell_type']==main_df.iloc[i]['cell_type']].values.squeeze()[1:].astype(float)\n                vec_ = np.concatenate([vec_, values])\n            else:\n                assert 'sm_name' in df.columns\n                values = df[df['sm_name']==main_df.iloc[i]['sm_name']].values.squeeze()[1:].astype(float)\n                vec_ = np.concatenate([vec_, values])\n        for chem_feat in chem_feats:\n            vec_ = np.concatenate([vec_, chem_feat[i]])\n        final_vec = np.concatenate([vec_,np.zeros(add_len-vec_.shape[0],)])\n        new_vecs.append(final_vec)\n    return np.stack(new_vecs, axis=0).astype(float).reshape(len(main_df), 1, add_len)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:04:44.060505Z","iopub.execute_input":"2024-02-27T12:04:44.061235Z","iopub.status.idle":"2024-02-27T12:04:44.073971Z","shell.execute_reply.started":"2024-02-27T12:04:44.061199Z","shell.execute_reply":"2024-02-27T12:04:44.072977Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_vec = combine_features([mean_cell_type,std_cell_type,mean_sm_name,std_sm_name],\\\n                [train_chem_feat, train_chem_feat_mean], de_train, X)\ntest_vec = combine_features([mean_cell_type,std_cell_type,mean_sm_name,std_sm_name],\\\n                   [test_chem_feat, test_chem_feat_mean], id_map, test)\n\nX_vec_light = combine_features([mean_cell_type,mean_sm_name],\\\n                [train_chem_feat, train_chem_feat_mean], de_train, X)\ntest_vec_light = combine_features([mean_cell_type,mean_sm_name],\\\n                   [test_chem_feat, test_chem_feat_mean], id_map, test)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:04:44.075389Z","iopub.execute_input":"2024-02-27T12:04:44.075788Z","iopub.status.idle":"2024-02-27T12:05:44.871090Z","shell.execute_reply.started":"2024-02-27T12:04:44.075754Z","shell.execute_reply":"2024-02-27T12:05:44.870304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not reproduce_leaderboard_score:\n    X_vec = np.concatenate([X_vec, np.zeros((X_vec.shape[0], 1, 2))], axis=-1)\n    test_vec = np.concatenate([test_vec, np.zeros((test_vec.shape[0], 1, 2))], axis=-1)\n    X_vec_light = np.concatenate([X_vec_light, np.zeros((X_vec_light.shape[0], 1, 1))], axis=-1)\n    test_vec_light = np.concatenate([test_vec_light, np.zeros((test_vec_light.shape[0], 1, 1))], axis=-1)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:05:44.872259Z","iopub.execute_input":"2024-02-27T12:05:44.872536Z","iopub.status.idle":"2024-02-27T12:05:44.878447Z","shell.execute_reply.started":"2024-02-27T12:05:44.872512Z","shell.execute_reply":"2024-02-27T12:05:44.877662Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_vec_light.shape","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:05:44.879638Z","iopub.execute_input":"2024-02-27T12:05:44.880318Z","iopub.status.idle":"2024-02-27T12:05:44.891749Z","shell.execute_reply.started":"2024-02-27T12:05:44.880293Z","shell.execute_reply":"2024-02-27T12:05:44.890890Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_vec_heavy = combine_features([desc_cell_type,mean_cell_type,mean_sm_name],\\\n                [train_chem_feat,train_chem_feat_mean], de_train, X, use_description=True)\ntest_vec_heavy = combine_features([desc_cell_type,mean_cell_type,mean_sm_name],\\\n                   [test_chem_feat,test_chem_feat], id_map, test, use_description=True)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:05:44.892898Z","iopub.execute_input":"2024-02-27T12:05:44.893175Z","iopub.status.idle":"2024-02-27T12:06:19.564741Z","shell.execute_reply.started":"2024-02-27T12:05:44.893148Z","shell.execute_reply":"2024-02-27T12:06:19.563494Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_vec_heavy.shape","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.566088Z","iopub.execute_input":"2024-02-27T12:06:19.566414Z","iopub.status.idle":"2024-02-27T12:06:19.572193Z","shell.execute_reply.started":"2024-02-27T12:06:19.566384Z","shell.execute_reply":"2024-02-27T12:06:19.571310Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Evaluation Metric Function","metadata":{}},{"cell_type":"code","source":"def mrrmse_np(y_pred, y_true):\n    return np.sqrt(np.square(y_true - y_pred).mean(axis=1)).mean()","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.573431Z","iopub.execute_input":"2024-02-27T12:06:19.573715Z","iopub.status.idle":"2024-02-27T12:06:19.581947Z","shell.execute_reply.started":"2024-02-27T12:06:19.573680Z","shell.execute_reply":"2024-02-27T12:06:19.581117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Additional Loss Functions\nI discovered experimentally that by combining different loss functions (others defined in the models), we can achieve a better performance. In particular, I used the binary cross entropy loss to push each predicted value in the target to a value other than 0 (i.e., push the value to a strictly positive or negative value). This is motivated by the fact that several values in the target are close to zero, and I wanted to make sure models do not learn this naively. The rest of the loss functions are suited for regression tasks and used normally to enforce the predicted value to be close to the target.","metadata":{}},{"cell_type":"code","source":"class LogCoshLoss(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n    def forward(self, y_prime_t, y_t):\n        ey_t = (y_t - y_prime_t)/3 # divide by 3 to avoid numerical overflow in cosh\n        return torch.mean(torch.log(torch.cosh(ey_t + 1e-12)))","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.583201Z","iopub.execute_input":"2024-02-27T12:06:19.583532Z","iopub.status.idle":"2024-02-27T12:06:19.592225Z","shell.execute_reply.started":"2024-02-27T12:06:19.583499Z","shell.execute_reply":"2024-02-27T12:06:19.591406Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"hidden_dims_reproduce_leaderboard = {'conv': {'heavy': 13400, 'light': 4576, 'initial': 8992},\n                                    'rnn': {'linear': {'heavy': 99968, 'light': 24192, 'initial': 29568},\n                                           'input_shape': {'heavy': [779,142], 'light': [187,202], 'initial': [229,324]}\n                                           }}\n\nhidden_dims_new = {'conv': {'heavy': 11144, 'light': 4520, 'initial': 8936},\n                                    'rnn': {'linear': {'heavy': 36480, 'light': 13952, 'initial': 19968},\n                                           'input_shape': {'heavy': [283,325], 'light': [107,349], 'initial': [154,479]}\n                                           }}\n\ndims_dict = hidden_dims_reproduce_leaderboard if reproduce_leaderboard_score else hidden_dims_new","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.593328Z","iopub.execute_input":"2024-02-27T12:06:19.595531Z","iopub.status.idle":"2024-02-27T12:06:19.603588Z","shell.execute_reply.started":"2024-02-27T12:06:19.595498Z","shell.execute_reply":"2024-02-27T12:06:19.602830Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Modeling","metadata":{}},{"cell_type":"code","source":"class Conv(nn.Module):\n    def __init__(self, scheme):\n        super(Conv, self).__init__()\n        self.name = 'Conv'\n        self.conv_block = nn.Sequential(nn.Conv1d(1, 8, 5, stride=1, padding=0),\n                                        nn.Dropout(0.3),\n                                        nn.Conv1d(8, 8, 5, stride=1, padding=0),\n                                        nn.ReLU(),\n                                        nn.Conv1d(8, 16, 5, stride=2, padding=0),\n                                        nn.Dropout(0.3),\n                                        nn.AvgPool1d(11),\n                                        nn.Conv1d(16, 8, 3, stride=3, padding=0),\n                                        nn.Flatten())\n        self.scheme = scheme\n        self.linear = nn.Sequential(\n                nn.Linear(dims_dict['conv'][self.scheme], 1024),\n                nn.Dropout(0.3),\n                nn.ReLU(),\n                nn.Linear(1024, 512),\n                nn.Dropout(0.3),\n                nn.ReLU())\n        self.head1 = nn.Linear(512, 18211)\n        \n        self.loss1 = nn.MSELoss()\n        self.loss2 = LogCoshLoss()\n        self.loss3 = nn.L1Loss()\n        self.loss4 = nn.BCELoss()\n        \n    def forward(self, x, y=None):\n        if y is None:\n            out = self.conv_block(x)\n            out = self.head1(self.linear(out))\n            return out\n        else:\n            out = self.conv_block(x)\n            out = self.head1(self.linear(out))\n            loss1 = 0.4*self.loss1(out, y) + 0.3*self.loss2(out, y) + 0.3*self.loss3(out, y)\n            yhat = torch.sigmoid(out)\n            yy = torch.sigmoid(y)\n            loss2 = self.loss4(yhat, yy)\n            return 0.8*loss1 + 0.2*loss2\n        \n\nclass LSTM(nn.Module):\n    def __init__(self, scheme):\n        super(LSTM, self).__init__()\n        self.name = 'LSTM'\n        self.scheme = scheme\n        self.lstm = nn.LSTM(dims_dict['rnn']['input_shape'][self.scheme][1], 128, num_layers=2, batch_first=True)\n        self.linear = nn.Sequential(\n            nn.Linear(dims_dict['rnn']['linear'][self.scheme], 1024),\n            nn.Dropout(0.3),\n            nn.ReLU(),\n            nn.Linear(1024, 512),\n            nn.Dropout(0.3),\n            nn.ReLU())\n        self.head1 = nn.Linear(512, 18211)\n        \n        self.loss1 = nn.MSELoss()\n        self.loss2 = LogCoshLoss()\n        self.loss3 = nn.L1Loss()\n        self.loss4 = nn.BCELoss()\n        \n    def forward(self, x, y=None):\n        shape1, shape2 = dims_dict['rnn']['input_shape'][self.scheme]\n        x = x.reshape(x.shape[0],shape1,shape2)\n        if y is None:\n            out, (hn, cn) = self.lstm(x)\n            out = out.reshape(out.shape[0],-1)\n            out = torch.cat([out, hn.reshape(hn.shape[1], -1)], dim=1)\n            out = self.head1(self.linear(out))\n            return out\n        else:\n            out, (hn, cn) = self.lstm(x)\n            out = out.reshape(out.shape[0],-1)\n            out = torch.cat([out, hn.reshape(hn.shape[1], -1)], dim=1)\n            out = self.head1(self.linear(out))\n            loss1 = 0.4*self.loss1(out, y) + 0.3*self.loss2(out, y) + 0.3*self.loss3(out, y)\n            yhat = torch.sigmoid(out)\n            yy = torch.sigmoid(y)\n            loss2 = self.loss4(yhat, yy)\n            return 0.8*loss1 + 0.2*loss2\n        \n        \nclass GRU(nn.Module):\n    def __init__(self, scheme):\n        super(GRU, self).__init__()\n        self.name = 'GRU'\n        self.scheme = scheme\n        self.gru = nn.GRU(dims_dict['rnn']['input_shape'][self.scheme][1], 128, num_layers=2, batch_first=True)\n        self.linear = nn.Sequential(\n            nn.Linear(dims_dict['rnn']['linear'][self.scheme], 1024),\n            nn.Dropout(0.3),\n            nn.ReLU(),\n            nn.Linear(1024, 512),\n            nn.Dropout(0.3),\n            nn.ReLU())\n        self.head1 = nn.Linear(512, 18211)\n        \n        self.loss1 = nn.MSELoss()\n        self.loss2 = LogCoshLoss()\n        self.loss3 = nn.L1Loss()\n        self.loss4 = nn.BCELoss()\n        \n    def forward(self, x, y=None):\n        shape1, shape2 = dims_dict['rnn']['input_shape'][self.scheme]\n        x = x.reshape(x.shape[0],shape1,shape2)\n        if y is None:\n            out, hn = self.gru(x)\n            out = out.reshape(out.shape[0],-1)\n            out = torch.cat([out, hn.reshape(hn.shape[1], -1)], dim=1)\n            out = self.head1(self.linear(out))\n            return out\n        else:\n            out, hn = self.gru(x)\n            out = out.reshape(out.shape[0],-1)\n            out = torch.cat([out, hn.reshape(hn.shape[1], -1)], dim=1)\n            out = self.head1(self.linear(out))\n            loss1 = 0.4*self.loss1(out, y) + 0.3*self.loss2(out, y) + 0.3*self.loss3(out, y)\n            yhat = torch.sigmoid(out)\n            yy = torch.sigmoid(y)\n            loss2 = self.loss4(yhat, yy)\n            return 0.8*loss1 + 0.2*loss2","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.604661Z","iopub.execute_input":"2024-02-27T12:06:19.604931Z","iopub.status.idle":"2024-02-27T12:06:19.633892Z","shell.execute_reply.started":"2024-02-27T12:06:19.604908Z","shell.execute_reply":"2024-02-27T12:06:19.633125Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create Dataset Class","metadata":{}},{"cell_type":"code","source":"class Dataset:\n    def __init__(self, data_x, data_y=None):\n        super(Dataset, self).__init__()\n        self.data_x = data_x\n        self.data_y = data_y\n\n    def __len__(self):\n        return len(self.data_x)\n    \n    def __getitem__(self, idx):\n        if self.data_y is not None:\n            return self.data_x[idx], self.data_y[idx]\n        else:\n            return self.data_x[idx]","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.634714Z","iopub.execute_input":"2024-02-27T12:06:19.634987Z","iopub.status.idle":"2024-02-27T12:06:19.647119Z","shell.execute_reply.started":"2024-02-27T12:06:19.634965Z","shell.execute_reply":"2024-02-27T12:06:19.646423Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define Data Augmentation Function\nIn the following function, we augment the training data by randomly dropping 30% of our 1-dimensional input feature vectors' entries. Input features are of shape (batch, 1, d)","metadata":{}},{"cell_type":"code","source":"import random\ndef augment_data(x_, y_):\n    copy_x = x_.copy()\n    new_x = []\n    new_y = y_.copy()\n    dim = x_.shape[2]\n    k = int(0.3*dim)\n    for i in range(x_.shape[0]):\n        idx = random.sample(range(dim), k=k)\n        copy_x[i,:,idx] = 0\n        new_x.append(copy_x[i])\n    return np.stack(new_x, axis=0), new_y","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.648223Z","iopub.execute_input":"2024-02-27T12:06:19.648543Z","iopub.status.idle":"2024-02-27T12:06:19.656104Z","shell.execute_reply.started":"2024-02-27T12:06:19.648513Z","shell.execute_reply":"2024-02-27T12:06:19.655318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define Helper and Main Training Functions\nGRU experienced numerical overflow with a learning rate of 0.001, so I used 0.0003 instead","metadata":{}},{"cell_type":"code","source":"def train_step(dataloader, model, opt, clip_norm):\n    model.train()\n    train_losses = []\n    for x, target in dataloader:\n        if torch.cuda.is_available():\n            model.cuda()\n            x = x.cuda()\n            target = target.cuda()\n        loss = model(x, target)\n        train_losses.append(loss.item())\n        opt.zero_grad()\n        loss.backward()\n        clip_grad_norm_(model.parameters(), clip_norm)\n        opt.step()\n    return np.mean(train_losses)\n\ndef validation_step(dataloader, model):\n    model.eval()\n    val_losses = []\n    val_mrrmse = []\n    for x, target in dataloader:\n        if torch.cuda.is_available():\n            model.cuda()\n            x = x.cuda()\n            target = target.cuda()\n        loss = model(x,target)\n        pred = model(x).detach().cpu().numpy()\n        val_mrrmse.append(mrrmse_np(pred, target.cpu().numpy()))\n        val_losses.append(loss.item())\n    return np.mean(val_losses), np.mean(val_mrrmse)\n\n\ndef train_function(model, x_train, y_train, x_val, y_val, epochs=20, clip_norm=1.0):\n    if model.name in ['GRU'] or not reproduce_leaderboard_score:\n        print('lr', 0.0003)\n        opt = torch.optim.Adam(model.parameters(), lr=0.0003)\n    else:\n        opt = torch.optim.Adam(model.parameters(), lr=0.001)\n    model.cuda()\n    x_train_aug, y_train_aug = augment_data(x_train, y_train)\n    x_train_aug = np.concatenate([x_train, x_train_aug], axis=0)\n    y_train_aug = np.concatenate([y_train, y_train_aug], axis=0)\n    data_x_train = torch.FloatTensor(x_train_aug)\n    data_y_train = torch.FloatTensor(y_train_aug)\n    data_x_val = torch.FloatTensor(x_val)\n    data_y_val = torch.FloatTensor(y_val)\n    train_dataloader = DataLoader(Dataset(data_x_train, data_y_train), num_workers=4, batch_size=16, shuffle=True)\n    val_dataloader = DataLoader(Dataset(data_x_val, data_y_val), num_workers=4, batch_size=32, shuffle=False)\n    best_loss = np.inf\n    best_weights = None\n    train_losses = []\n    val_losses = []\n    for e in range(epochs):\n        loss = train_step(train_dataloader, model, opt, clip_norm)\n        val_losses.append(loss.item())\n        val_loss, val_mrrmse = validation_step(val_dataloader, model)\n        if val_mrrmse < best_loss:\n            best_loss = val_mrrmse\n            best_weights = model.state_dict()\n            print('BEST ----> ')\n        print(f\"{model.name} Epoch {e}, train_loss {round(loss,3)}, val_loss {round(val_loss, 3)}, val_mrrmse {val_mrrmse}\")\n    model.load_state_dict(best_weights)\n    return model","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.657302Z","iopub.execute_input":"2024-02-27T12:06:19.657560Z","iopub.status.idle":"2024-02-27T12:06:19.672647Z","shell.execute_reply.started":"2024-02-27T12:06:19.657538Z","shell.execute_reply":"2024-02-27T12:06:19.671823Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.model_selection import KFold as KF\nsplits = 5\nkf_cv = KF(n_splits=splits, shuffle=True, random_state=42)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.673615Z","iopub.execute_input":"2024-02-27T12:06:19.673853Z","iopub.status.idle":"2024-02-27T12:06:19.760941Z","shell.execute_reply.started":"2024-02-27T12:06:19.673831Z","shell.execute_reply":"2024-02-27T12:06:19.760116Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def cross_validate_models(X, y, epochs=120, scheme='initial', clip_norm=1.0):\n    trained_models = []\n    for i,(train_idx,val_idx) in enumerate(kf_cv.split(X)):\n        print(f\"\\nSplit {i+1}/{splits}...\")\n        x_train, x_val = X[train_idx], X[val_idx]\n        y_train, y_val = y.values[train_idx], y.values[val_idx]\n        for Model in [LSTM, Conv, GRU]:\n            model = Model(scheme)\n            model = train_function(model, x_train, y_train, x_val, y_val, epochs=epochs, clip_norm=clip_norm)\n            model.to('cpu')\n            trained_models.append(model)\n            torch.cuda.empty_cache()\n    return trained_models","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.766890Z","iopub.execute_input":"2024-02-27T12:06:19.767586Z","iopub.status.idle":"2024-02-27T12:06:19.774648Z","shell.execute_reply.started":"2024-02-27T12:06:19.767553Z","shell.execute_reply":"2024-02-27T12:06:19.773913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define Inference Functions","metadata":{}},{"cell_type":"code","source":"def inference_pytorch(model, dataloader):\n    model.eval()\n    preds = []\n    for x in dataloader:\n        if torch.cuda.is_available():\n            model.cuda()\n            x = x.cuda()\n        pred = model(x).detach().cpu().numpy()\n        preds.append(pred)\n    model.to('cpu')\n    torch.cuda.empty_cache()\n    return np.concatenate(preds, axis=0)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.775749Z","iopub.execute_input":"2024-02-27T12:06:19.776112Z","iopub.status.idle":"2024-02-27T12:06:19.784907Z","shell.execute_reply.started":"2024-02-27T12:06:19.776048Z","shell.execute_reply":"2024-02-27T12:06:19.784006Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def average_prediction(X_test, trained_models):\n    all_preds = []\n    test_dataloader = DataLoader(Dataset(torch.FloatTensor(X_test)), num_workers=4, batch_size=64, shuffle=False)\n    for i,model in enumerate(trained_models):\n        current_pred = inference_pytorch(model, test_dataloader)\n        all_preds.append(current_pred)\n    return np.stack(all_preds, axis=1).mean(axis=1)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.785978Z","iopub.execute_input":"2024-02-27T12:06:19.787325Z","iopub.status.idle":"2024-02-27T12:06:19.795002Z","shell.execute_reply.started":"2024-02-27T12:06:19.787292Z","shell.execute_reply":"2024-02-27T12:06:19.794245Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def weighted_average_prediction(X_test, trained_models, model_wise=[0.25, 0.35, 0.40], fold_wise=None):\n    all_preds = []\n    test_dataloader = DataLoader(Dataset(torch.FloatTensor(X_test)), num_workers=4, batch_size=64, shuffle=False)\n    for i,model in enumerate(trained_models):\n        current_pred = inference_pytorch(model, test_dataloader)\n        current_pred = model_wise[i%3]*current_pred\n        if fold_wise:\n            current_pred = fold_wise[i//3]*current_pred\n        all_preds.append(current_pred)\n    return np.stack(all_preds, axis=1).sum(axis=1)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.795892Z","iopub.execute_input":"2024-02-27T12:06:19.796178Z","iopub.status.idle":"2024-02-27T12:06:19.807985Z","shell.execute_reply.started":"2024-02-27T12:06:19.796155Z","shell.execute_reply":"2024-02-27T12:06:19.807120Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train and Reproduce the Leaderboard!","metadata":{}},{"cell_type":"code","source":"def reproduce(epochs=1):\n    trained_models = {'initial': [], 'light': [], 'heavy': []}\n    for scheme, clip_norm, input_features in zip(['initial', 'light', 'heavy'], [5.0, 1.0, 1.0], [X_vec, X_vec_light, X_vec_heavy]):\n        seed_everything()\n        models = cross_validate_models(input_features, y, epochs=epochs, scheme=scheme, clip_norm=clip_norm)\n        trained_models[scheme].extend(models)\n    return trained_models","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.809050Z","iopub.execute_input":"2024-02-27T12:06:19.809336Z","iopub.status.idle":"2024-02-27T12:06:19.818161Z","shell.execute_reply.started":"2024-02-27T12:06:19.809313Z","shell.execute_reply":"2024-02-27T12:06:19.817333Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trained_models = reproduce(epochs=250)","metadata":{"execution":{"iopub.status.busy":"2024-02-27T12:06:19.819223Z","iopub.execute_input":"2024-02-27T12:06:19.821028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred1 = average_prediction(test_vec, trained_models['initial'])\npred2 = weighted_average_prediction(test_vec, trained_models['initial'],\\\n                                    model_wise=[0.29, 0.33, 0.38], fold_wise=[0.25, 0.15, 0.2, 0.15, 0.25])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred3 = average_prediction(test_vec_light, trained_models['light'])\npred4 = weighted_average_prediction(test_vec_light, trained_models['light'],\\\n                                    model_wise=[0.29, 0.33, 0.38], fold_wise=[0.25, 0.15, 0.2, 0.15, 0.25])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred5 = average_prediction(test_vec_heavy, trained_models['heavy'])\npred6 = weighted_average_prediction(test_vec_heavy, trained_models['heavy'],\\\n                                    model_wise=[0.29, 0.33, 0.38], fold_wise=[0.25, 0.15, 0.2, 0.15, 0.25])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Read Submission Sample File","metadata":{}},{"cell_type":"code","source":"col = list(de_train.columns[5:])\nsubmission = sample_submission.copy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission[col] = 0.18*pred1 + 0.15*pred2 + 0.23*pred3 + 0.15*pred4 + 0.15*pred5 + 0.14*pred6\ndf1 = submission.copy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Ensemble Prediction","metadata":{}},{"cell_type":"code","source":"submission[col] =  0.23*pred1 + 0.15*pred2 + 0.13*pred3 + 0.15*pred4 + 0.20*pred5 + 0.14*pred6\ndf2 = submission.copy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission[col] = 0.17*pred1 + 0.16*pred2 + 0.17*pred3 + 0.16*pred4 + 0.18*pred5 + 0.16*pred6\ndf3 = submission.copy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_sub = 0.34*df1 + 0.33*df2 + 0.33*df3","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Save Submission Dataframe","metadata":{}},{"cell_type":"code","source":"df_sub.to_csv('submission.csv')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_sub","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}