{"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"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nsub = pd.read_csv(\"/kaggle/input/cafa-5-protein-function-prediction/sample_submission.tsv\", sep= \"\\t\", header = None)\nsub.columns = [\"The Protein ID\", \"The Gene Ontology term (GO) ID\", \"Predicted link probability that GO appear in Protein\"]\nsub.head(5)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T08:32:38.005078Z","iopub.execute_input":"2023-06-14T08:32:38.005742Z","iopub.status.idle":"2023-06-14T08:32:38.31196Z","shell.execute_reply.started":"2023-06-14T08:32:38.005706Z","shell.execute_reply":"2023-06-14T08:32:38.310909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"- **1 - Collect Embedding vectors from pre-trained protein function prediction models (T5, ProtBERT or EMS2) :**\n\nSources for embeddings vectors : \n- *T5* : https://www.kaggle.com/datasets/sergeifironov/t5embeds\n\n- *ProtBERT* : https://www.kaggle.com/datasets/henriupton/protbert-embeddings-for-cafa5\n\n- *EMS2* : https://www.kaggle.com/datasets/viktorfairuschin/cafa-5-ems-2-embeddings-numpy\n\n- **2 - Generate labels from train_terms file** : by considering the top K most common GO terms in all Proteins set, generate for each protein a sparse vector of length K to indicate the true probabilities that each of the K GO terms are in the Protein (0 or 1). Here we retain K = 600\n\n- **3 - Create Pytorch Dataset class that can handle Protein ID and embeddings**.\n\n- **4 - Create Pytorch Model class for prediction** : can be any architecture of Multilabel classification model that can turn embeddings of shape (E,) to probabilities of shape (K,). Here we explore **MultiLayerPerceptron** and **ConvNN1d** Networks.\n\n- **5 - Make Cross Validation w.r.t the F-1 measure and do Hyperparameter tuning thanks to Weights and Biases package (Wandb)**","metadata":{}},{"cell_type":"markdown","source":"# 1. Imports / Config","metadata":{}},{"cell_type":"code","source":"MAIN_DIR = \"/kaggle/input/cafa-5-protein-function-prediction\"\n\n# UTILITARIES\nimport numpy as np\nfrom tqdm import tqdm\nimport time\nimport matplotlib.pyplot as plt\nplt.style.use('ggplot')\n\n# TORCH MODULES FOR METRICS COMPUTATION :\nimport torch\nfrom torch.utils.data import Dataset\nfrom torch import nn\nfrom torch.utils.data import random_split\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom torchmetrics.classification import MultilabelF1Score\nfrom torchmetrics.classification import MultilabelAccuracy\n\nimport pytorch_lightning as pl\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning.loggers import WandbLogger\n\n# WANDB FOR LIGHTNING :\nimport wandb\n\n# FILES VISUALIZATION\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-14T08:32:38.31793Z","iopub.execute_input":"2023-06-14T08:32:38.318609Z","iopub.status.idle":"2023-06-14T08:32:54.114921Z","shell.execute_reply.started":"2023-06-14T08:32:38.318575Z","shell.execute_reply":"2023-06-14T08:32:54.113998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class config:\n    train_sequences_path = MAIN_DIR  + \"/Train/train_sequences.fasta\"\n    train_labels_path = MAIN_DIR + \"/Train/train_terms.tsv\"\n    test_sequences_path = MAIN_DIR + \"/Test (Targets)/testsuperset.fasta\"\n    \n    num_labels = 500\n    n_epochs = 10\n    batch_size = 32\n    lr = 0.001\n    \n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"execution":{"iopub.status.busy":"2023-06-14T08:41:39.795901Z","iopub.execute_input":"2023-06-14T08:41:39.796358Z","iopub.status.idle":"2023-06-14T08:41:39.807449Z","shell.execute_reply.started":"2023-06-14T08:41:39.796323Z","shell.execute_reply":"2023-06-14T08:41:39.806508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(config.device)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T08:35:34.948211Z","iopub.execute_input":"2023-06-14T08:35:34.948582Z","iopub.status.idle":"2023-06-14T08:35:34.956687Z","shell.execute_reply.started":"2023-06-14T08:35:34.948552Z","shell.execute_reply":"2023-06-14T08:35:34.955763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 2. Baseline","metadata":{}},{"cell_type":"markdown","source":"- Use ProtBERT/T5 embedding vectors for each Sequence ID and its associated Protein Sequence. \n\n- Define a Pytorch Dataset to load all ids/sequences of the train+test sets and their respective ProtBert embeddings\n\n- Define Pytorch Model architecture in order to use these embeddings to proceed classification task : to each ID we associate a probability to be associated to each GO term **For this part, just consider the top_n most common GO terms as labels**\n\n- Return desired probabilities for test set","metadata":{}},{"cell_type":"markdown","source":"# 3. Collect ProtBERT Embedding Vectors","metadata":{}},{"cell_type":"markdown","source":"### Here is a Dataset Card I have created if you want to understand better what are pre-trained model embedding vectors, and if you want to download it directly from it : https://www.kaggle.com/datasets/henriupton/protbert-embeddings-for-cafa5","metadata":{}},{"cell_type":"markdown","source":"##### SCRIPT FOR VECTOR EMBEDDINGS COLLECTING #####\n#### RUN THIS ONLY ONE TIME AND SAVE IT IN LOCAL ####\n\n```python\n\nprint(\"Load ProtBERT Model...\")\n# PROT BERT LOADING :\nfrom transformers import BertModel, BertTokenizer\ntokenizer = BertTokenizer.from_pretrained(\"Rostlab/prot_bert\", do_lower_case=False )\nmodel = BertModel.from_pretrained(\"Rostlab/prot_bert\").to(config.device)\n\ndef get_bert_embedding(\n    sequence : str,\n    len_seq_limit : int\n):\n    '''\n    Function to collect last hidden state embedding vector from pre-trained ProtBERT Model\n    \n    INPUTS:\n    - sequence (str) : protein sequence (ex : AAABBB) from fasta file\n    - len_seq_limit (int) : maximum sequence lenght (i.e nb of letters) for truncation\n    \n    OUTPUTS:\n    - output_hidden : last hidden state embedding vector for input sequence of length 1024\n    '''\n    sequence_w_spaces = ' '.join(list(sequence))\n    encoded_input = tokenizer(\n        sequence_w_spaces,\n        truncation=True,\n        max_length=len_seq_limit,\n        padding='max_length',\n        return_tensors='pt').to(config.device)\n    output = model(**encoded_input)\n    output_hidden = output['last_hidden_state'][:,0][0].detach().cpu().numpy()\n    assert len(output_hidden)==1024\n    return output_hidden\n\n### COLLECTING FOR TRAIN SAMPLES :\nprint(\"Loading train set ProtBERT Embeddings...\")\nfasta_train = SeqIO.parse(config.train_sequences_path, \"fasta\")\nprint(\"Total Nb of Elements : \", len(list(fasta_train)))\nfasta_train = SeqIO.parse(config.train_sequences_path, \"fasta\")\nids_list = []\nembed_vects_list = []\nt0 = time.time()\ncheckpoint = 0\nfor item in tqdm(fasta_train):\n    ids_list.append(item.id)\n    embed_vects_list.append(\n        get_bert_embedding(sequence = item.seq, len_seq_limit = 1200))\n    checkpoint+=1\n    if checkpoint>=100:\n        df_res = pd.DataFrame(data={\"id\" : ids_list, \"embed_vect\" : embed_vects_list})\n        np.save('/kaggle/working/train_ids.npy',np.array(ids_list))\n        np.save('/kaggle/working/train_embeddings.npy',np.array(embed_vects_list))\n        checkpoint=0\n        \nnp.save('/kaggle/working/train_ids.npy',np.array(ids_list))\nnp.save('/kaggle/working/train_embeddings.npy',np.array(embed_vects_list))\nprint('Total Elapsed Time:',time.time()-t0)\n\n### COLLECTING FOR TEST SAMPLES :\nprint(\"Loading test set ProtBERT Embeddings...\")\nfasta_test = SeqIO.parse(config.test_sequences_path, \"fasta\")\nprint(\"Total Nb of Elements : \", len(list(fasta_test)))\nfasta_test = SeqIO.parse(config.test_sequences_path, \"fasta\")\nids_list = []\nembed_vects_list = []\nt0 = time.time()\ncheckpoint=0\nfor item in tqdm(fasta_test):\n    ids_list.append(item.id)\n    embed_vects_list.append(\n        get_bert_embedding(sequence = item.seq, len_seq_limit = 1200))\n    checkpoint+=1\n    if checkpoint>=100:\n        np.save('/kaggle/working/test_ids.npy',np.array(ids_list))\n        np.save('/kaggle/working/test_embeddings.npy',np.array(embed_vects_list))\n        checkpoint=0\n        \nnp.save('/kaggle/working/test_ids.npy',np.array(ids_list))\nnp.save('/kaggle/working/test_embeddings.npy',np.array(embed_vects_list))\nprint('Total Elasped Time:',time.time()-t0)\n```","metadata":{"execution":{"iopub.status.busy":"2023-05-05T08:54:24.395576Z","iopub.execute_input":"2023-05-05T08:54:24.396067Z","iopub.status.idle":"2023-05-05T08:54:35.014197Z","shell.execute_reply.started":"2023-05-05T08:54:24.396028Z","shell.execute_reply":"2023-05-05T08:54:35.012847Z"}}},{"cell_type":"markdown","source":"# 4. Collect labels vectors for train/test","metadata":{}},{"cell_type":"markdown","source":"##### SCRIPT FOR LABELS (TARGETS) VECTORS COLLECTING #####\n#### RUN THIS ONLY ONE TIME AND SAVE IT IN LOCAL ####\n\n```python\nprint(\"GENERATE TARGETS FOR ENTRY IDS (\"+str(config.num_labels)+\" MOST COMMON GO TERMS)\")\nids = np.load(\"/kaggle/input/protbert-embeddings-for-cafa5/train_ids.npy\")\nlabels = pd.read_csv(config.train_labels_path, sep = \"\\t\")\n\ntop_terms = labels.groupby(\"term\")[\"EntryID\"].count().sort_values(ascending=False)\nlabels_names = top_terms[:config.num_labels].index.values\ntrain_labels_sub = labels[(labels.term.isin(labels_names)) & (labels.EntryID.isin(ids))]\nid_labels = train_labels_sub.groupby('EntryID')['term'].apply(list).to_dict()\n\ngo_terms_map = {label: i for i, label in enumerate(labels_names)}\nlabels_matrix = np.empty((len(ids), len(labels_names)))\n\nfor index, id in tqdm(enumerate(ids)):\n    id_gos_list = id_labels[id]\n    temp = [go_terms_map[go] for go in labels_names if go in id_gos_list]\n    labels_matrix[index, temp] = 1\n\nnp.save(\"/kaggle/working/train_targets_top\"+str(config.num_labels)+\".npy\", np.array(labels_matrix))\nprint(\"GENERATION FINISHED!\")\n```","metadata":{"execution":{"iopub.status.busy":"2023-05-11T17:45:22.690593Z","iopub.execute_input":"2023-05-11T17:45:22.691183Z","iopub.status.idle":"2023-05-11T17:45:24.979748Z","shell.execute_reply.started":"2023-05-11T17:45:22.691147Z","shell.execute_reply":"2023-05-11T17:45:24.978075Z"}}},{"cell_type":"markdown","source":"# 5. Pytorch Dataset Architecture","metadata":{}},{"cell_type":"code","source":"# Directories for the different embedding vectors : \nembeds_map = {\n    \"T5\" : \"t5embeds\",\n    \"ProtBERT\" : \"protbert-embeddings-for-cafa5\",\n    \"EMS2\" : \"cafa-5-ems-2-embeddings-numpy\"\n}\n\n# Length of the different embedding vectors :\nembeds_dim = {\n    \"T5\" : 1024,\n    \"ProtBERT\" : 1024,\n    \"EMS2\" : 1280\n}","metadata":{"execution":{"iopub.status.busy":"2023-06-14T08:41:44.280877Z","iopub.execute_input":"2023-06-14T08:41:44.281265Z","iopub.status.idle":"2023-06-14T08:41:44.286671Z","shell.execute_reply.started":"2023-06-14T08:41:44.281226Z","shell.execute_reply":"2023-06-14T08:41:44.285739Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ProteinSequenceDataset(Dataset):\n    \n    def __init__(self, datatype, embeddings_source):\n        super(ProteinSequenceDataset).__init__()\n        self.datatype = datatype\n        \n        if embeddings_source in [\"ProtBERT\", \"EMS2\"]:\n            embeds = np.load(\"/kaggle/input/\"+embeds_map[embeddings_source]+\"/\"+datatype+\"_embeddings.npy\")\n            ids = np.load(\"/kaggle/input/\"+embeds_map[embeddings_source]+\"/\"+datatype+\"_ids.npy\")\n        \n        if embeddings_source == \"T5\":\n            embeds = np.load(\"/kaggle/input/\"+embeds_map[embeddings_source]+\"/\"+datatype+\"_embeds.npy\")\n            ids = np.load(\"/kaggle/input/\"+embeds_map[embeddings_source]+\"/\"+datatype+\"_ids.npy\")\n            \n        embeds_list = []\n        for l in range(embeds.shape[0]):\n            embeds_list.append(embeds[l,:])\n        self.df = pd.DataFrame(data={\"EntryID\": ids, \"embed\" : embeds_list})\n        \n        if datatype==\"train\":\n            np_labels = np.load(\n                \"/kaggle/input/train-targets-top\"+str(config.num_labels)+ \\\n                \"/train_targets_top\"+str(config.num_labels)+\".npy\")\n            df_labels = pd.DataFrame(self.df['EntryID'])\n            df_labels['labels_vect']=[row for row in np_labels]\n            self.df = self.df.merge(df_labels, on=\"EntryID\")\n            \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, index):\n        embed = torch.tensor(self.df.iloc[index][\"embed\"] , dtype = torch.float32)\n        if self.datatype==\"train\":\n            targets = torch.tensor(self.df.iloc[index][\"labels_vect\"], dtype = torch.float32)\n            return embed, targets\n        if self.datatype==\"test\":\n            id = self.df.iloc[index][\"EntryID\"]\n            return embed, id\n        ","metadata":{"execution":{"iopub.status.busy":"2023-06-14T08:41:45.229115Z","iopub.execute_input":"2023-06-14T08:41:45.229481Z","iopub.status.idle":"2023-06-14T08:41:45.242981Z","shell.execute_reply.started":"2023-06-14T08:41:45.229449Z","shell.execute_reply":"2023-06-14T08:41:45.241866Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 6. Pytorch Models Architectures","metadata":{}},{"cell_type":"code","source":"class MultiLayerPerceptron(torch.nn.Module):\n\n    def __init__(self, input_dim, num_classes):\n        super(MultiLayerPerceptron, self).__init__()\n\n        self.linear1 = torch.nn.Linear(input_dim, 1012)\n        self.activation1 = torch.nn.Tanh()\n        self.linear2 = torch.nn.Linear(1012, 864)\n        self.activation2 = torch.nn.Tanh()\n        self.linear3 = torch.nn.Linear(864, 712)\n        self.activation3 = torch.nn.Tanh()\n        self.linear4 = torch.nn.Linear(712, num_classes)\n      \n\n    def forward(self, x):\n        x = self.linear1(x)\n        x = self.activation1(x)\n        x = self.linear2(x)\n        x = self.activation2(x)\n        x = self.linear3(x)\n        x = self.activation3(x)\n        x = self.linear4(x)\n        return x\n    \n    \n    #self.lstm=nn.LSTM(units=64, dropout=0.2, recurrent_dropout=0.2)\n     #x = self.lstm(x)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T08:41:46.783648Z","iopub.execute_input":"2023-06-14T08:41:46.784036Z","iopub.status.idle":"2023-06-14T08:41:46.793243Z","shell.execute_reply.started":"2023-06-14T08:41:46.784002Z","shell.execute_reply":"2023-06-14T08:41:46.792261Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CNN1D(nn.Module):\n    def __init__(self, input_dim, num_classes):\n        super(CNN1D, self).__init__()\n        # (batch_size, channels, embed_size)\n        self.conv1 = nn.Conv1d(in_channels=1, out_channels=3, kernel_size=3, dilation=1, padding=1, stride=1)\n        # (batch_size, 3, embed_size)\n        self.pool1 = nn.MaxPool1d(kernel_size=2, stride=2)\n        # (batch_size, 3, embed_size/2 = 512)\n        self.conv2 = nn.Conv1d(in_channels=3, out_channels=8, kernel_size=3, dilation=1, padding=1, stride=1)\n        # (batch_size, 8, embed_size/2 = 512)\n        self.pool2 = nn.MaxPool1d(kernel_size=2, stride=2)\n        \n        self.lstm = nn.LSTM(input_size=8, hidden_size=128, num_layers=1, batch_first=True)\n        # (batch_size, 8, embed_size/4 = 256)\n        self.fc1 = nn.Linear(in_features=128, out_features=256)\n        \n        self.fc2 = nn.Linear(in_features=256, out_features=num_classes)\n\n    def forward(self, x):\n        x = x.reshape(x.shape[0], 1, x.shape[1])\n        x = self.pool1(nn.functional.tanh(self.conv1(x)))\n        x = self.pool2(nn.functional.tanh(self.conv2(x)))\n        x = x.transpose(1, 2)  # Reshape for LSTM input (batch_size, seq_len, input_size)\n        x, _ = self.lstm(x)\n        x = x[:, -1, :]  # Extract the last hidden state of the LSTM\n        \n        #x = torch.flatten(x, 1)\n        x = nn.functional.tanh(self.fc1(x))\n        x = self.fc2(x)\n        return x\n","metadata":{"execution":{"iopub.status.busy":"2023-06-14T08:52:34.903434Z","iopub.execute_input":"2023-06-14T08:52:34.903807Z","iopub.status.idle":"2023-06-14T08:52:34.913803Z","shell.execute_reply.started":"2023-06-14T08:52:34.903776Z","shell.execute_reply":"2023-06-14T08:52:34.912907Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 7. Train the Model","metadata":{}},{"cell_type":"code","source":"def train_model(embeddings_source, model_type=\"linear\", train_size=0.9):\n    \n    train_dataset = ProteinSequenceDataset(datatype=\"train\", embeddings_source = embeddings_source)\n    \n    train_set, val_set = random_split(train_dataset, lengths = [int(len(train_dataset)*train_size), len(train_dataset)-int(len(train_dataset)*train_size)])\n    train_dataloader = torch.utils.data.DataLoader(train_set, batch_size=config.batch_size, shuffle=True)\n    val_dataloader = torch.utils.data.DataLoader(val_set, batch_size=config.batch_size, shuffle=True)\n\n    if model_type == \"linear\":\n        model = MultiLayerPerceptron(input_dim=embeds_dim[embeddings_source], num_classes=config.num_labels).to(config.device)\n    if model_type == \"convolutional\":\n        model = CNN1D(input_dim=embeds_dim[embeddings_source], num_classes=config.num_labels).to(config.device)\n\n    optimizer = torch.optim.Adam(model.parameters(), lr = config.lr)\n    scheduler = ReduceLROnPlateau(optimizer, factor=0.1, patience=1)\n    CrossEntropy = torch.nn.CrossEntropyLoss()\n    f1_score = MultilabelF1Score(num_labels=config.num_labels).to(config.device)\n    n_epochs = config.n_epochs\n\n    print(\"BEGIN TRAINING...\")\n    train_loss_history=[]\n    val_loss_history=[]\n    \n    train_f1score_history=[]\n    val_f1score_history=[]\n    for epoch in range(n_epochs):\n        print(\"EPOCH \", epoch+1)\n        ## TRAIN PHASE :\n        losses = []\n        scores = []\n        for embed, targets in tqdm(train_dataloader):\n            embed, targets = embed.to(config.device), targets.to(config.device)\n            optimizer.zero_grad()\n            preds = model(embed)\n            loss= CrossEntropy(preds, targets)\n            score=f1_score(preds, targets)\n            losses.append(loss.item()) \n            scores.append(score.item())\n            loss.backward()\n            optimizer.step()\n        avg_loss = np.mean(losses)\n        avg_score = np.mean(scores)\n        print(\"Running Average TRAIN Loss : \", avg_loss)\n        print(\"Running Average TRAIN F1-Score : \", avg_score)\n        train_loss_history.append(avg_loss)\n        train_f1score_history.append(avg_score)\n        \n        ## VALIDATION PHASE : \n        losses = []\n        scores = []\n        for embed, targets in val_dataloader:\n            embed, targets = embed.to(config.device), targets.to(config.device)\n            preds = model(embed)\n            loss= CrossEntropy(preds, targets)\n            score=f1_score(preds, targets)\n            losses.append(loss.item())\n            scores.append(score.item())\n        avg_loss = np.mean(losses)\n        avg_score = np.mean(scores)\n        print(\"Running Average VAL Loss : \", avg_loss)\n        print(\"Running Average VAL F1-Score : \", avg_score)\n        val_loss_history.append(avg_loss)\n        val_f1score_history.append(avg_score)\n        \n        scheduler.step(avg_loss)\n        print(\"\\n\")\n        \n    print(\"TRAINING FINISHED\")\n    print(\"FINAL TRAINING SCORE : \", train_f1score_history[-1])\n    print(\"FINAL VALIDATION SCORE : \", val_f1score_history[-1])\n    \n    losses_history = {\"train\" : train_loss_history, \"val\" : val_loss_history}\n    scores_history = {\"train\" : train_f1score_history, \"val\" : val_f1score_history}\n    \n    return model, losses_history, scores_history","metadata":{"execution":{"iopub.status.busy":"2023-06-14T08:41:52.528545Z","iopub.execute_input":"2023-06-14T08:41:52.528901Z","iopub.status.idle":"2023-06-14T08:41:52.55026Z","shell.execute_reply.started":"2023-06-14T08:41:52.528869Z","shell.execute_reply":"2023-06-14T08:41:52.548908Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 7.1. Training Usecase","metadata":{}},{"cell_type":"code","source":"#ems2_model, ems2_losses, ems2_scores = train_model(embeddings_source=\"EMS2\",model_type=\"convolutional\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"t5_model, t5_losses, t5_scores = train_model(embeddings_source=\"T5\",model_type=\"convolutional\")","metadata":{"execution":{"iopub.status.busy":"2023-06-14T08:52:51.092871Z","iopub.execute_input":"2023-06-14T08:52:51.093349Z","iopub.status.idle":"2023-06-14T09:01:32.55421Z","shell.execute_reply.started":"2023-06-14T08:52:51.093317Z","shell.execute_reply":"2023-06-14T09:01:32.553188Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"protbert_model, protbert_losses, protbert_scores = train_model(embeddings_source=\"T5\",model_type=\"linear\")","metadata":{"execution":{"iopub.status.busy":"2023-06-14T09:02:17.37393Z","iopub.execute_input":"2023-06-14T09:02:17.374898Z","iopub.status.idle":"2023-06-14T09:07:43.075812Z","shell.execute_reply.started":"2023-06-14T09:02:17.374861Z","shell.execute_reply":"2023-06-14T09:07:43.074637Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 7.3. Train/Val Losses ","metadata":{}},{"cell_type":"code","source":"plt.figure(figsize = (10, 4))\n#plt.plot(ems2_losses[\"val\"], label = \"EMS2\")\nplt.plot(t5_losses[\"val\"], label = \"T5\")\nplt.plot(protbert_losses[\"val\"], label = \"ProtBERT\") \nplt.title(\"Validation Losses for # Vector Embeddings\")\nplt.xlabel(\"Epochs\")\nplt.ylabel(\"Average Loss\")\nplt.legend()\nplt.show()\n\nplt.figure(figsize = (10, 4))\n#plt.plot(ems2_scores[\"val\"], label = \"EMS2\")\nplt.plot(t5_scores[\"val\"], label = \"T5\")\nplt.plot(protbert_scores[\"val\"], label = \"ProtBERT\")\nplt.title(\"Validation F1-Scores for # Vector Embeddings\")\nplt.xlabel(\"Epochs\")\nplt.ylabel(\"Average F1-Score\")\nplt.legend()\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 8. Pytorch Lightning module","metadata":{}},{"cell_type":"code","source":"class Linear_Lightning(pl.LightningModule):\n    def __init__(self, input_dim, num_classes, train_size, **hparams):\n        super(Linear_Lightning, self).__init__()\n        self.linear1 = torch.nn.Linear(input_dim, 1012)\n        self.activation1 = torch.nn.Tanh()\n        self.linear2 = torch.nn.Linear(1012, 864)\n        self.activation2 = torch.nn.Tanh()\n        self.linear3 = torch.nn.Linear(864, 712)\n        self.activation3 = torch.nn.Tanh()\n        self.linear4 = torch.nn.Linear(712, num_classes)\n        self.activation4 = torch.nn.Softmax()\n        \n        train_dataset = ProteinSequenceDataset(datatype=\"train\", embeddings_source = embeddings_source)\n        self.train_set, self.val_set = random_split(train_dataset, lengths = [int(len(train_dataset)*train_size), len(train_dataset)-int(len(train_dataset)*train_size)])\n        self.loss_fn = torch.nn.CrossEntropyLoss()\n        \n        self.batch_size = batch_size\n        self.lr = lr\n        \n        self.f1_score = MultilabelF1Score(num_labels=num_classes)\n        self.accuracy = MultilabelAccuracy(num_labels=num_classes)\n\n    def forward(self, x):\n        x = self.linear1(x)\n        x = self.activation1(x)\n        x = self.linear2(x)\n        x = self.activation2(x)\n        x = self.linear3(x)\n        x=self.activation3(x)\n        x= self.linear4(x)\n        x=self.activation4(x)\n        return x\n    \n    def training_step(self, batch, batch_idx):\n        embed, targets = batch\n        preds = self(embed)\n        loss = self.loss_fn(preds, targets)\n        f1_score = self.f1_score(preds, targets)\n        acc_score = self.accuracy(preds, targets)\n        \n        logs = {\"train_loss\" : loss, \"f1_score\" : f1_score, \"accuracy_score\" : acc_score}\n        self.log_dict(\n            logs,\n            on_step=True, on_epoch=True, prog_bar=True, logger=True\n        )\n        return {\"loss\":loss, \"log\" : logs}\n    \n    def validation_step(self, batch, batch_idx):\n        embed, targets = batch\n        preds = self(embed)\n        loss= self.loss_fn(preds, targets)\n        f1_score = self.f1_score(preds, targets)\n        acc_score = self.accuracy(preds, targets)\n        \n        return {\"val_loss\":loss, \"f1_score\" : f1_score, \"accuracy_score\" : acc_score}\n    \n    def validation_end(self, outputs):\n        avg_loss = torch.stack([x[\"val_loss\"] for x in ouputs]).mean()\n        logs = {\"val_loss\" : avg_loss}\n        self.log_dict(\n            logs,\n            on_step=True, on_epoch=True, prog_bar=True, logger=True\n        )\n        return {\"avg_val_loss\" : avg_loss, \"log\" : logs}\n        \n    def val_dataloader(self):\n        val_dataloader = torch.utils.data.DataLoader(self.val_set, batch_size=config.batch_size, shuffle=False,)\n        return val_dataloader\n\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)\n        return optimizer\n    \n    def train_dataloader(self):\n        train_dataloader = torch.utils.data.DataLoader(self.train_set, batch_size=self.batch_size, shuffle=False)\n        return train_dataloader","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## IN PROGRESS\n```python\nparam_grid = {\n    \"batch_size\" : [8, 16, 32, 64],\n    \"lr\" : [1e-4, 1e-3, 1e-2]\n}\n\nembeddings_source = \"T5\"\n\nruns_count=0\nfor batch_size in param_grid[\"batch_size\"]:\n    for lr in param_grid[\"lr\"]:\n        runs_count+=1\n        print(\"NEW SESSION RUN (number \"+str(runs_count)+\")\")\n        print(\"batch size : \", batch_size)\n        print(\"learning rate : \", lr)\n        \n        run_name = \"Adam-\"+str(batch_size)+\"-\"+str(lr)\n        \n        logger = WandbLogger(\n            name = run_name,\n            project=\"mlp_model_cafa5\",\n            job_type=\"train\"\n        )\n\n        trainer = Trainer(\n            max_epochs=1,\n            limit_train_batches=5000,\n            logger=logger\n        )\n        model = Linear_Lightning(\n            input_dim=embeds_dim[embeddings_source],\n            num_classes=config.num_labels,\n            train_size=0.8\n        )\n\n        trainer.fit(model)\n```","metadata":{"execution":{"iopub.status.busy":"2023-05-11T21:06:26.806291Z","iopub.execute_input":"2023-05-11T21:06:26.807432Z","iopub.status.idle":"2023-05-11T21:06:30.034869Z","shell.execute_reply.started":"2023-05-11T21:06:26.807367Z","shell.execute_reply":"2023-05-11T21:06:30.033213Z"}}},{"cell_type":"markdown","source":"```python\n# Define sweep config\nsweep_configuration = {\n    'method': 'random',\n    'name': 'sweep',\n    'metric': {'goal': 'maximize', 'name': 'val_acc'},\n    'parameters': \n    {\n        'batch_size': {'values': [16, 32, 64]},\n        'epochs': {'values': [5, 10, 15]},\n        'lr': {'max': 0.1, 'min': 0.0001}\n     }\n}\n\n# Initialize sweep by passing in config. \n# (Optional) Provide a name of the project.\nsweep_id = wandb.sweep(\n  sweep=sweep_configuration, \n  project='my-first-sweep'\n  )\n\n# Define training function that takes in hyperparameter \n# values from `wandb.config` and uses them to train a \n# model and return metric\ndef train_one_epoch(epoch, lr, bs): \n  acc = 0.25 + ((epoch/30) +  (random.random()/10))\n  loss = 0.2 + (1 - ((epoch-1)/10 +  random.random()/5))\n  return acc, loss\n\ndef evaluate_one_epoch(epoch): \n  acc = 0.1 + ((epoch/20) +  (random.random()/10))\n  loss = 0.25 + (1 - ((epoch-1)/10 +  random.random()/6))\n  return acc, loss\n\ndef main():\n    run = wandb.init()\n\n    lr  =  wandb.config.lr\n    bs = wandb.config.batch_size\n    epochs = wandb.config.epochs\n\n    for epoch in np.arange(1, epochs):\n      train_acc, train_loss = train_one_epoch(epoch, lr, bs)\n      val_acc, val_loss = evaluate_one_epoch(epoch)\n\n      wandb.log({\n        'epoch': epoch, \n        'train_acc': train_acc,\n        'train_loss': train_loss, \n        'val_acc': val_acc, \n        'val_loss': val_loss\n      })\n\n# Start sweep job.\nwandb.agent(sweep_id, function=main, count=4)\n```","metadata":{"_kg_hide-input":true}},{"cell_type":"markdown","source":"# 9. Make Predictions","metadata":{}},{"cell_type":"code","source":"def predict(embeddings_source):\n    \n    test_dataset = ProteinSequenceDataset(datatype=\"test\", embeddings_source = embeddings_source)\n    test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)\n    \n    if embeddings_source == \"T5\":\n        model = t5_model\n    if embeddings_source == \"ProtBERT\":\n        model = protbert_model\n    if embeddings_source == \"EMS2\":\n        model = ems2_model\n        \n    model.eval()\n    \n    labels = pd.read_csv(config.train_labels_path, sep = \"\\t\")\n    top_terms = labels.groupby(\"term\")[\"EntryID\"].count().sort_values(ascending=False)\n    labels_names = top_terms[:config.num_labels].index.values\n    print(\"GENERATE PREDICTION FOR TEST SET...\")\n\n    ids_ = np.empty(shape=(len(test_dataloader)*config.num_labels,), dtype=object)\n    go_terms_ = np.empty(shape=(len(test_dataloader)*config.num_labels,), dtype=object)\n    confs_ = np.empty(shape=(len(test_dataloader)*config.num_labels,), dtype=np.float32)\n\n    for i, (embed, id) in tqdm(enumerate(test_dataloader)):\n        embed = embed.to(config.device)\n        confs_[i*config.num_labels:(i+1)*config.num_labels] = torch.nn.functional.sigmoid(model(embed)).squeeze().detach().cpu().numpy()\n        ids_[i*config.num_labels:(i+1)*config.num_labels] = id[0]\n        go_terms_[i*config.num_labels:(i+1)*config.num_labels] = labels_names\n\n    submission_df = pd.DataFrame(data={\"Id\" : ids_, \"GO term\" : go_terms_, \"Confidence\" : confs_})\n    print(\"PREDICTIONS DONE\")\n    return submission_df","metadata":{"execution":{"iopub.status.busy":"2023-06-14T09:08:36.390864Z","iopub.execute_input":"2023-06-14T09:08:36.391806Z","iopub.status.idle":"2023-06-14T09:08:36.404199Z","shell.execute_reply.started":"2023-06-14T09:08:36.39177Z","shell.execute_reply":"2023-06-14T09:08:36.403193Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df1 = predict(\"T5\")\nsub_df2 = predict(\"ProtBERT\")","metadata":{"execution":{"iopub.status.busy":"2023-06-14T09:10:01.849484Z","iopub.execute_input":"2023-06-14T09:10:01.850034Z","iopub.status.idle":"2023-06-14T09:20:34.634727Z","shell.execute_reply.started":"2023-06-14T09:10:01.849997Z","shell.execute_reply":"2023-06-14T09:20:34.63378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df1.head(5)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T09:26:55.347662Z","iopub.execute_input":"2023-06-14T09:26:55.348039Z","iopub.status.idle":"2023-06-14T09:26:55.360848Z","shell.execute_reply.started":"2023-06-14T09:26:55.348007Z","shell.execute_reply":"2023-06-14T09:26:55.359844Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df2.head(5)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T09:26:58.236288Z","iopub.execute_input":"2023-06-14T09:26:58.236653Z","iopub.status.idle":"2023-06-14T09:26:58.250647Z","shell.execute_reply.started":"2023-06-14T09:26:58.236621Z","shell.execute_reply":"2023-06-14T09:26:58.24957Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df=sub_df1\nsubmission_df['Confidence']=(sub_df1['Confidence']+sub_df2['Confidence'])/2","metadata":{"execution":{"iopub.status.busy":"2023-06-14T09:27:20.985562Z","iopub.execute_input":"2023-06-14T09:27:20.985916Z","iopub.status.idle":"2023-06-14T09:27:21.356595Z","shell.execute_reply.started":"2023-06-14T09:27:20.985885Z","shell.execute_reply":"2023-06-14T09:27:21.35558Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df.head(5)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T09:27:45.962915Z","iopub.execute_input":"2023-06-14T09:27:45.963492Z","iopub.status.idle":"2023-06-14T09:27:45.977413Z","shell.execute_reply.started":"2023-06-14T09:27:45.963449Z","shell.execute_reply":"2023-06-14T09:27:45.976185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df.to_csv('submission.tsv', sep='\\t', index=False)","metadata":{"execution":{"iopub.status.busy":"2023-06-14T09:27:54.798993Z","iopub.execute_input":"2023-06-14T09:27:54.799397Z"},"trusted":true},"execution_count":null,"outputs":[]}]}