bert.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Imports and parameters"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:14.714784314Z",
     "start_time": "2023-07-29T17:20:13.489626787Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "import re\n",
    "import gc\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from tqdm.notebook import trange, tqdm_notebook\n",
    "from transformers import  BertForTokenClassification, AutoTokenizer\n",
    "from torch.nn import BCELoss\n",
    "from torch.optim import AdamW\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "test_size = 0.05\n",
    "learning_rate = 2e-5\n",
    "epochs = 20\n",
    "validation_size=100\n",
    "model_name ='Rostlab/prot_bert'"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:14.761223365Z",
     "start_time": "2023-07-29T17:20:14.760737154Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at Rostlab/prot_bert were not used when initializing BertForTokenClassification: ['cls.predictions.decoder.weight', 'cls.predictions.decoder.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.bias', 'cls.seq_relationship.bias', 'cls.predictions.transform.dense.weight']\n",
      "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForTokenClassification were not initialized from the model checkpoint at Rostlab/prot_bert and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "bert_model = BertModel.from_pretrained(model_name)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:18.568260005Z",
     "start_time": "2023-07-29T17:20:14.946982765Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Loading and preparing Data Set"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"./data/binding_affinity_hla_cleaned.csv\", index_col=0)\n",
    "df.dropna(inplace=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:21.723339879Z",
     "start_time": "2023-07-29T17:20:18.560408776Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [],
   "source": [
    "hla_group_count = len(df['HLA_group_idx'].unique())\n",
    "hla_gene_count = len(df['HLA_gene'].unique())"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:21.765211176Z",
     "start_time": "2023-07-29T17:20:21.724214360Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [],
   "source": [
    "df['HLA_group_one_hot_encode'] = pd.get_dummies(df['HLA_group_idx']).values.tolist()\n",
    "df['HLA_gene_one_hot_encode'] =  pd.get_dummies(df['HLA_gene']).values.tolist()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:25.232675335Z",
     "start_time": "2023-07-29T17:20:21.766657085Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "outputs": [
    {
     "data": {
      "text/plain": "                                        MHC_sequence     MHC_type  \\\n0  MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...  HLA-B*27:05   \n1  MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...  HLA-B*27:05   \n2  MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...  HLA-B*27:05   \n3  MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...  HLA-B*27:05   \n4  MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...  HLA-B*27:05   \n\n  peptide_sequence  label HLA_gene HLA_allele  HLA_allele_group  \\\n0        ERLKEVQKR      1        B      27:05                27   \n1    KPRKTAEVAGKTL      1        B      27:05                27   \n2        KEARRIIKK      1        B      27:05                27   \n3       EEKITEAKEL      0        B      27:05                27   \n4     SLPSSRAARVPG      0        B      27:05                27   \n\n   HLA_allele_id  HLA_group_idx  \\\n0              5              0   \n1              5              0   \n2              5              0   \n3              5              0   \n4              5              0   \n\n                            HLA_group_one_hot_encode HLA_gene_one_hot_encode  \n0  [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...               [0, 1, 0]  \n1  [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...               [0, 1, 0]  \n2  [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...               [0, 1, 0]  \n3  [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...               [0, 1, 0]  \n4  [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...               [0, 1, 0]  ",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>MHC_sequence</th>\n      <th>MHC_type</th>\n      <th>peptide_sequence</th>\n      <th>label</th>\n      <th>HLA_gene</th>\n      <th>HLA_allele</th>\n      <th>HLA_allele_group</th>\n      <th>HLA_allele_id</th>\n      <th>HLA_group_idx</th>\n      <th>HLA_group_one_hot_encode</th>\n      <th>HLA_gene_one_hot_encode</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...</td>\n      <td>HLA-B*27:05</td>\n      <td>ERLKEVQKR</td>\n      <td>1</td>\n      <td>B</td>\n      <td>27:05</td>\n      <td>27</td>\n      <td>5</td>\n      <td>0</td>\n      <td>[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...</td>\n      <td>[0, 1, 0]</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...</td>\n      <td>HLA-B*27:05</td>\n      <td>KPRKTAEVAGKTL</td>\n      <td>1</td>\n      <td>B</td>\n      <td>27:05</td>\n      <td>27</td>\n      <td>5</td>\n      <td>0</td>\n      <td>[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...</td>\n      <td>[0, 1, 0]</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...</td>\n      <td>HLA-B*27:05</td>\n      <td>KEARRIIKK</td>\n      <td>1</td>\n      <td>B</td>\n      <td>27:05</td>\n      <td>27</td>\n      <td>5</td>\n      <td>0</td>\n      <td>[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...</td>\n      <td>[0, 1, 0]</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...</td>\n      <td>HLA-B*27:05</td>\n      <td>EEKITEAKEL</td>\n      <td>0</td>\n      <td>B</td>\n      <td>27:05</td>\n      <td>27</td>\n      <td>5</td>\n      <td>0</td>\n      <td>[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...</td>\n      <td>[0, 1, 0]</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>MRVTAPRTLLLLLWGAVALTETWAGSHSMRYFHTSVSRPGRGEPRF...</td>\n      <td>HLA-B*27:05</td>\n      <td>SLPSSRAARVPG</td>\n      <td>0</td>\n      <td>B</td>\n      <td>27:05</td>\n      <td>27</td>\n      <td>5</td>\n      <td>0</td>\n      <td>[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...</td>\n      <td>[0, 1, 0]</td>\n    </tr>\n  </tbody>\n</table>\n</div>"
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:25.241602358Z",
     "start_time": "2023-07-29T17:20:25.240030686Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "outputs": [],
   "source": [
    "def clean_for_tokenizer(s: str):\n",
    "    return \" \".join(list(re.sub(r\"[UZOB]\", \"X\", s)))\n",
    "\n",
    "def extract_for_tokenizer(row):\n",
    "    return [clean_for_tokenizer(row[\"MHC_sequence\"]), clean_for_tokenizer(row[\"peptide_sequence\"])]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:25.257165938Z",
     "start_time": "2023-07-29T17:20:25.242543397Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "outputs": [],
   "source": [
    "extracted_for_tokenizer_data = df.apply(axis=1, func=extract_for_tokenizer)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:40.169416544Z",
     "start_time": "2023-07-29T17:20:25.257328624Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "outputs": [],
   "source": [
    "extracted_for_tokenizer_data = extracted_for_tokenizer_data.tolist()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:40.192561140Z",
     "start_time": "2023-07-29T17:20:40.188937842Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "outputs": [],
   "source": [
    "train_data, test_data, train_label, test_label, train_group_encode, test_group_encode, train_gene_encode, test_gene_encode = train_test_split(extracted_for_tokenizer_data, df['label'].tolist(), df['HLA_group_one_hot_encode'].tolist(), df['HLA_gene_one_hot_encode'].tolist(), test_size=test_size, random_state=42)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:41.400234371Z",
     "start_time": "2023-07-29T17:20:40.190853870Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "outputs": [
    {
     "data": {
      "text/plain": "0"
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "del df\n",
    "gc.collect()\n",
    "tokenized_train_data = tokenizer.batch_encode_plus(train_data, add_special_tokens=True, padding=True, return_tensors='pt')\n",
    "del train_data\n",
    "gc.collect()\n",
    "train_label_tensors = torch.tensor(train_label, dtype=torch.float)\n",
    "del train_label\n",
    "gc.collect()\n",
    "train_group_encode_tensors = torch.tensor(train_group_encode)\n",
    "del train_group_encode\n",
    "gc.collect()\n",
    "train_gene_encode_tensors = torch.tensor(train_gene_encode)\n",
    "del train_gene_encode\n",
    "gc.collect()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T16:59:40.121795781Z",
     "start_time": "2023-07-29T16:58:02.516019075Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "outputs": [],
   "source": [
    "class MHCPeptideDataSet(Dataset):\n",
    "    def __init__(self, data, labels, allele_group_one_hot, hla_gene_one_hot):\n",
    "        self.data = data\n",
    "        self.data['allele_group_encoding'] = allele_group_one_hot\n",
    "        self.data['hla_gene_encoding'] = hla_gene_one_hot\n",
    "        self.labels = labels\n",
    "        self.size = self.labels.shape[0]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.size\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        if index < 0 or index >= self.__len__():\n",
    "            raise IndexError(f\"Index {index} is out of bounds for dataset with length {self.__len__()}\")\n",
    "        item = {key: value[index] for (key, value) in self.data.items()}\n",
    "        item['labels'] = self.labels[index]\n",
    "        return item"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:41.402979633Z",
     "start_time": "2023-07-29T17:20:41.402096349Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "outputs": [],
   "source": [
    "dataset = MHCPeptideDataSet(tokenized_train_data, train_label_tensors, allele_group_one_hot=train_group_encode_tensors, hla_gene_one_hot=train_gene_encode_tensors)\n",
    "dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:05:45.711031430Z",
     "start_time": "2023-07-29T16:59:39.914746092Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Creating Model"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "outputs": [],
   "source": [
    "bert_model.to('cuda')\n",
    "class BERTClassification(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.bert = bert_model.bert\n",
    "        self.head = nn.Sequential(\n",
    "            nn.Dropout(0.1),\n",
    "            nn.Linear(int(self.bert.config.hidden_size + hla_group_count + hla_gene_count), 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 1),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "        self.freeze()\n",
    "        # Unfreeze the last layer of the BERT model\n",
    "\n",
    "    def freeze(self):\n",
    "        for param in self.bert.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "    def unfreeze_layer(self, layer: int):\n",
    "        for param in self.bert.encoder.layer[layer].parameters():\n",
    "            param.requires_grad = True\n",
    "\n",
    "\n",
    "    def get_special_params(self):\n",
    "        return list(self.head.parameters()) + list(self.bert.parameters())\n",
    "\n",
    "    def forward(self, allele_group_encoding, hla_gene_encoding, *args, **kwargs):\n",
    "        outputs = self.bert(*args, **kwargs)\n",
    "        return self.head(torch.cat((outputs[0][:, 0, :], allele_group_encoding, hla_gene_encoding), dim=1))\n",
    "\n",
    "model = BERTClassification().to('cuda')"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:41.986389751Z",
     "start_time": "2023-07-29T17:20:41.592328484Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## WarmUp round"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "outputs": [],
   "source": [
    "optimizer = AdamW(model.get_special_params(), lr=learning_rate)\n",
    "scheduler = ReduceLROnPlateau(optimizer, 'min', patience=1, threshold=1e-2)\n",
    "criterion = BCELoss()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:41.988996685Z",
     "start_time": "2023-07-29T17:20:41.988296744Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "outputs": [],
   "source": [
    "tmp = torch.load('./output/model_state')\n",
    "model.load_state_dict(tmp['model_state_dict'])\n",
    "optimizer.load_state_dict(tmp['optimizer_state_dict'])\n",
    "scheduler.load_state_dict(tmp['scheduler_state_dict'])"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:20:42.824196532Z",
     "start_time": "2023-07-29T17:20:41.990167219Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "outputs": [],
   "source": [
    "model.unfreeze_layer(-1)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:06:01.286946936Z",
     "start_time": "2023-07-29T17:06:01.286410287Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "outputs": [],
   "source": [
    "def set_lr(lr:float):\n",
    "    for g in optimizer.param_groups:\n",
    "        g['lr'] = lr"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-28T22:39:58.002544343Z",
     "start_time": "2023-07-28T22:39:57.925631940Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "outputs": [],
   "source": [
    "set_lr(5e-6)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-28T22:40:40.839393757Z",
     "start_time": "2023-07-28T22:40:40.797216154Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "outputs": [
    {
     "data": {
      "text/plain": "  0%|          | 0/20 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "61012703053448beb37aa86474cb0478"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/19612 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d51e621a356043f7ae8618098c3c0d63"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 Loss: 12.2929\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/19612 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ab31900bc78047329a35c7ae1e5e41d1"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[24], line 17\u001B[0m\n\u001B[1;32m     14\u001B[0m         optimizer\u001B[38;5;241m.\u001B[39mstep()\n\u001B[1;32m     16\u001B[0m     \u001B[38;5;66;03m# Statistics\u001B[39;00m\n\u001B[0;32m---> 17\u001B[0m     running_loss \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m loss\u001B[38;5;241m.\u001B[39mitem() \u001B[38;5;241m*\u001B[39m batch_labels\u001B[38;5;241m.\u001B[39msize(\u001B[38;5;241m0\u001B[39m)\n\u001B[1;32m     18\u001B[0m epoch_loss \u001B[38;5;241m=\u001B[39m running_loss \u001B[38;5;241m/\u001B[39m \u001B[38;5;28mlen\u001B[39m(dataloader)\n\u001B[1;32m     19\u001B[0m \u001B[38;5;28mprint\u001B[39m(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mepoch: \u001B[39m\u001B[38;5;132;01m{}\u001B[39;00m\u001B[38;5;124m Loss: \u001B[39m\u001B[38;5;132;01m{:.4f}\u001B[39;00m\u001B[38;5;124m'\u001B[39m\u001B[38;5;241m.\u001B[39mformat(epoch, epoch_loss))\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "for epoch in trange(epochs):\n",
    "    running_loss = 0.0\n",
    "    model.train()\n",
    "    for batch in  tqdm_notebook(dataloader):\n",
    "        batch_db = {key: value.to('cuda') for key, value in batch.items() if key != 'labels'}\n",
    "        batch_labels = batch['labels'].to('cuda')\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Forward\n",
    "        with torch.set_grad_enabled(True):\n",
    "            output = model(**batch_db)\n",
    "            loss = criterion(output.view(-1), batch_labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        # Statistics\n",
    "        running_loss += loss.item() * batch_labels.size(0)\n",
    "    epoch_loss = running_loss / len(dataloader)\n",
    "    print('epoch: {} Loss: {:.4f}'.format(epoch, epoch_loss))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T16:54:03.755291004Z",
     "start_time": "2023-07-28T22:40:43.621444355Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "outputs": [],
   "source": [
    "torch.save(\n",
    "    {\n",
    "        'model_state_dict': model.state_dict(),\n",
    "        'optimizer_state_dict': optimizer.state_dict(),\n",
    "        'scheduler_state_dict': scheduler.state_dict()\n",
    "    },\n",
    "    './output/model_state'\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T16:54:08.031536406Z",
     "start_time": "2023-07-29T16:54:06.290745899Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Test"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "outputs": [
    {
     "data": {
      "text/plain": "0"
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gc.collect()\n",
    "tokenized_test_data = tokenizer.batch_encode_plus(test_data, add_special_tokens=True, padding=True, return_tensors='pt')\n",
    "del test_data\n",
    "gc.collect()\n",
    "test_label_tensors = torch.tensor(test_label, dtype=torch.float)\n",
    "del test_label\n",
    "gc.collect()\n",
    "test_group_encode_tensors = torch.tensor(test_group_encode)\n",
    "del test_group_encode\n",
    "gc.collect()\n",
    "test_gene_encode_tensors = torch.tensor(test_gene_encode)\n",
    "del test_gene_encode\n",
    "gc.collect()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:22:51.438933900Z",
     "start_time": "2023-07-29T17:22:36.339761183Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:13:17.483733484Z",
     "start_time": "2023-07-29T17:13:17.442924766Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "outputs": [],
   "source": [
    "model.eval()\n",
    "test_dataset = MHCPeptideDataSet(tokenized_test_data, test_label_tensors, allele_group_one_hot=test_group_encode_tensors, hla_gene_one_hot=test_gene_encode_tensors)\n",
    "test_dataloader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:23:43.319169192Z",
     "start_time": "2023-07-29T17:23:43.317208686Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "outputs": [],
   "source": [
    "model_outputs = []\n",
    "actual_labels = []"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T17:23:45.071627491Z",
     "start_time": "2023-07-29T17:23:45.029573926Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "outputs": [
    {
     "data": {
      "text/plain": "  0%|          | 0/1401 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4bb7daa492904ae8836b50070d2f93a7"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model.eval()\n",
    "for batch in  tqdm_notebook(test_dataloader):\n",
    "    batch_db = {key: value.to('cuda') for key, value in batch.items() if key != 'labels'}\n",
    "    batch_labels = batch['labels'].to('cuda')\n",
    "\n",
    "    # Forward\n",
    "    with torch.no_grad():\n",
    "        test_output = model(**batch_db)\n",
    "        model_outputs += [i[0] for i in test_output.tolist()]\n",
    "        actual_labels += batch_labels.tolist()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T18:32:59.097613382Z",
     "start_time": "2023-07-29T17:23:46.110316215Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "fine_tuned_bert = model.bert\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "outputs": [],
   "source": [
    "df_test_model = pd.DataFrame(data= {\n",
    "    \"model_output\": model_outputs,\n",
    "    \"labels\": actual_labels,\n",
    "})"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T18:32:59.112334013Z",
     "start_time": "2023-07-29T18:32:59.110561400Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "outputs": [
    {
     "data": {
      "text/plain": "       model_output  labels\n0          0.106708     1.0\n1          0.008202     0.0\n2          0.281923     0.0\n3          0.048741     0.0\n4          0.014424     0.0\n...             ...     ...\n89649      0.003637     0.0\n89650      0.028282     0.0\n89651      0.045913     0.0\n89652      0.977217     1.0\n89653      0.956132     0.0\n\n[89654 rows x 2 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>model_output</th>\n      <th>labels</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>0.106708</td>\n      <td>1.0</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>0.008202</td>\n      <td>0.0</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>0.281923</td>\n      <td>0.0</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>0.048741</td>\n      <td>0.0</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>0.014424</td>\n      <td>0.0</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>89649</th>\n      <td>0.003637</td>\n      <td>0.0</td>\n    </tr>\n    <tr>\n      <th>89650</th>\n      <td>0.028282</td>\n      <td>0.0</td>\n    </tr>\n    <tr>\n      <th>89651</th>\n      <td>0.045913</td>\n      <td>0.0</td>\n    </tr>\n    <tr>\n      <th>89652</th>\n      <td>0.977217</td>\n      <td>1.0</td>\n    </tr>\n    <tr>\n      <th>89653</th>\n      <td>0.956132</td>\n      <td>0.0</td>\n    </tr>\n  </tbody>\n</table>\n<p>89654 rows × 2 columns</p>\n</div>"
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_test_model"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T18:32:59.124196992Z",
     "start_time": "2023-07-29T18:32:59.116438326Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "outputs": [],
   "source": [
    "df_test_model.to_csv(\"./output/checkpoint_model_bert.csv\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-07-29T18:32:59.240356665Z",
     "start_time": "2023-07-29T18:32:59.119147246Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}