{"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":"markdown","source":"# This notebook forked from https://www.kaggle.com/datasets/henriupton/protbert-embeddings-for-cafa5, and ProtT5 is used\n","metadata":{}},{"cell_type":"markdown","source":"# CAFA 5 Competition : Protein Function Prediction","metadata":{}},{"cell_type":"markdown","source":"# Problem Framing","metadata":{}},{"cell_type":"markdown","source":"This Kaggle competition aims to predict the function of proteins using their amino-acid sequences and additional data. Understanding protein function is crucial for comprehending cellular processes and developing new treatments for diseases. With the abundance of genomic sequence data available, assigning accurate biological functions to proteins becomes challenging due to their multifunctionality and interactions with various partners. This competition, hosted by the Function Community of Special Interest (Function-COSI), brings together computational biologists, experimental biologists, and biocurators to improve protein function prediction through data science and machine learning approaches. The goal is to contribute to advancements in medicine, agriculture, and overall human and animal health.","metadata":{}},{"cell_type":"markdown","source":"![image-intro](https://upload.wikimedia.org/wikipedia/commons/thumb/4/4a/Computational_solvent_mapping_of_AMA1_using_FTMAP.TIF/lossy-page1-500px-Computational_solvent_mapping_of_AMA1_using_FTMAP.TIF.jpg)","metadata":{}},{"cell_type":"markdown","source":"# What to submit ?","metadata":{}},{"cell_type":"markdown","source":"This competition evaluates participants' predictions of Gene Ontology (GO) terms for protein sequences. The evaluation is performed on a test set of proteins that initially have no assigned functions but may accumulate experimental annotations after the submission deadline. The test set is divided into three subontologies: Molecular Function (MF), Biological Process (BP), and Cellular Component (CC). Participants are scored separately for each subontology. The final performance measure is the arithmetic mean of the maximum F-measures calculated on the three subontologies. Weighted precision and recall are used, taking into account the hierarchical structure of the GO. The evaluation code is publicly available. The leaderboard displays performance on a subset of proteins not included in the final test set, so generalization performance is crucial. Submission files should contain protein-target pairs with corresponding GO terms and estimated probabilities, within a specific score range. The predictions are propagated to parent terms if not explicitly listed. There is a limit on the number of terms associated with each protein. If a protein is not listed in the submission file, it is assumed that all predictions for that protein are zero.","metadata":{}},{"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-05-29T08:01:39.447248Z","iopub.execute_input":"2023-05-29T08:01:39.447822Z","iopub.status.idle":"2023-05-29T08:01:39.698426Z","shell.execute_reply.started":"2023-05-29T08:01:39.447786Z","shell.execute_reply":"2023-05-29T08:01:39.697513Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The evaluation metric used in this competition is the weighted F-measure, which combines precision and recall, taking into account the hierarchical structure of the Gene Ontology (GO). The formula for calculating the weighted F-measure is as follows:\n\n$Weighted \\ F-measure = \\frac{(1 + β^2) * (weighted \\ precision * weighted recall)}{((β^2 * weighted \\ precision) + weighted \\ recall)}$\n\nwhere:\n- β is a parameter that controls the trade-off between precision and recall. In this competition, β is set to 1, resulting in an equal weighting of precision and recall.\n- Weighted precision is the precision score, considering the weights of the predicted terms. It is calculated as the sum of the products of the predicted term's weight and its true positive count, divided by the sum of the weights of all predicted terms.\n- Weighted recall is the recall score, considering the weights of the true positive terms. It is calculated as the sum of the products of the true positive term's weight and its count, divided by the sum of the weights of all true positive terms.\n\nNote: The specific formulas for weighted precision and weighted recall are provided in the competition materials and utilize additional information such as term weights and true positive counts.","metadata":{}},{"cell_type":"markdown","source":"# General Baseline in this Notebook","metadata":{}},{"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- *ProtT5* : https://www.kaggle.com/datasets/horikitasaku/prott5-embedding-for-cafa5\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":"![baseline-image](https://www.researchgate.net/publication/334642149/figure/fig1/AS:783995214249986@1563930433525/Flow-chart-of-STRING2GO-based-protein-function-prediction-method.png)\n\n## **Flow-chart of STRING2GO-based protein function prediction method**","metadata":{}},{"cell_type":"markdown","source":"# So now lets get started !","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-05-29T08:01:41.717619Z","iopub.execute_input":"2023-05-29T08:01:41.718799Z","iopub.status.idle":"2023-05-29T08:01:41.74511Z","shell.execute_reply.started":"2023-05-29T08:01:41.718763Z","shell.execute_reply":"2023-05-29T08:01:41.744062Z"},"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 = 128\n    lr = 0.001\n    \n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:01:41.889712Z","iopub.execute_input":"2023-05-29T08:01:41.890276Z","iopub.status.idle":"2023-05-29T08:01:41.89724Z","shell.execute_reply.started":"2023-05-29T08:01:41.89024Z","shell.execute_reply":"2023-05-29T08:01:41.895736Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(config.device)","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:01:42.053089Z","iopub.execute_input":"2023-05-29T08:01:42.053419Z","iopub.status.idle":"2023-05-29T08:01:42.058935Z","shell.execute_reply.started":"2023-05-29T08:01:42.053371Z","shell.execute_reply":"2023-05-29T08:01:42.057739Z"},"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":"## I've put both data and code here: https://www.kaggle.com/datasets/horikitasaku/prott5-embedding-for-cafa5","metadata":{}},{"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    \"ProtT5\": \"prott5-embedding-for-cafa5\"\n}\n\n# Length of the different embedding vectors :\nembeds_dim = {\n    \"T5\" : 1024,\n    \"ProtBERT\" : 1024,\n    \"EMS2\" : 1280,\n    \"ProtT5\":1024\n}","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:01:43.978932Z","iopub.execute_input":"2023-05-29T08:01:43.979285Z","iopub.status.idle":"2023-05-29T08:01:43.984278Z","shell.execute_reply.started":"2023-05-29T08:01:43.979256Z","shell.execute_reply":"2023-05-29T08:01:43.98313Z"},"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\",\"ProtT5\"]:\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            df_labels = pd.read_pickle(\n                \"/kaggle/input/train-tragets-top-\"+str(config.num_labels)+ \\\n                \"/train_targets_top\"+str(config.num_labels)+\".pkl\")\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-05-29T08:02:52.731963Z","iopub.execute_input":"2023-05-29T08:02:52.73248Z","iopub.status.idle":"2023-05-29T08:02:52.747481Z","shell.execute_reply.started":"2023-05-29T08:02:52.732437Z","shell.execute_reply":"2023-05-29T08:02:52.746294Z"},"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.ReLU()\n        self.linear2 = torch.nn.Linear(1012, 712)\n        self.activation2 = torch.nn.ReLU()\n        self.linear3 = torch.nn.Linear(712, 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        return x","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:02:54.046916Z","iopub.execute_input":"2023-05-29T08:02:54.047262Z","iopub.status.idle":"2023-05-29T08:02:54.056352Z","shell.execute_reply.started":"2023-05-29T08:02:54.047233Z","shell.execute_reply":"2023-05-29T08:02:54.054455Z"},"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        # (batch_size, 8, embed_size/4 = 256)\n        self.fc1 = nn.Linear(in_features=int(8 * input_dim/4), out_features=128)\n        self.fc2 = nn.Linear(in_features=128, 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.relu(self.conv1(x)))\n        x = self.pool2(nn.functional.relu(self.conv2(x)))\n        x = torch.flatten(x, 1)\n        x = nn.functional.relu(self.fc1(x))\n        x = self.fc2(x)\n        return x\n","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:02:54.072118Z","iopub.execute_input":"2023-05-29T08:02:54.072457Z","iopub.status.idle":"2023-05-29T08:02:54.083106Z","shell.execute_reply.started":"2023-05-29T08:02:54.072431Z","shell.execute_reply":"2023-05-29T08:02:54.082181Z"},"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-05-29T08:02:56.023704Z","iopub.execute_input":"2023-05-29T08:02:56.024122Z","iopub.status.idle":"2023-05-29T08:02:56.04136Z","shell.execute_reply.started":"2023-05-29T08:02:56.024085Z","shell.execute_reply":"2023-05-29T08:02:56.039904Z"},"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=\"linear\")","metadata":{"execution":{"iopub.status.busy":"2023-05-17T14:52:11.100958Z","iopub.execute_input":"2023-05-17T14:52:11.101636Z","iopub.status.idle":"2023-05-17T14:54:40.639291Z","shell.execute_reply.started":"2023-05-17T14:52:11.101593Z","shell.execute_reply":"2023-05-17T14:54:40.638164Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# t5_model, t5_losses, t5_scores = train_model(embeddings_source=\"T5\",model_type=\"linear\")","metadata":{"execution":{"iopub.status.busy":"2023-05-17T14:55:25.520425Z","iopub.execute_input":"2023-05-17T14:55:25.521455Z","iopub.status.idle":"2023-05-17T14:57:48.225494Z","shell.execute_reply.started":"2023-05-17T14:55:25.521402Z","shell.execute_reply":"2023-05-17T14:57:48.224439Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"protbert_model, protbert_losses, protbert_scores = train_model(embeddings_source=\"ProtBERT\",model_type=\"linear\")","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:05:41.008969Z","iopub.execute_input":"2023-05-29T08:05:41.009458Z","iopub.status.idle":"2023-05-29T08:07:58.964565Z","shell.execute_reply.started":"2023-05-29T08:05:41.009373Z","shell.execute_reply":"2023-05-29T08:07:58.96346Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"prott5_model, prott5_losses, prott5_scores = train_model(embeddings_source=\"ProtT5\",model_type=\"linear\")","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:02:57.90207Z","iopub.execute_input":"2023-05-29T08:02:57.90301Z","iopub.status.idle":"2023-05-29T08:05:14.999063Z","shell.execute_reply.started":"2023-05-29T08:02:57.902972Z","shell.execute_reply":"2023-05-29T08:05:14.998068Z"},"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\")\n# plt.plot(t5_losses[\"val\"], label = \"T5\")\nplt.plot(protbert_losses[\"val\"], label = \"ProtBERT\")\nplt.plot(prott5_losses[\"val\"], label = \"ProtT5\")\n\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\")\n# plt.plot(t5_scores[\"val\"], label = \"T5\")\nplt.plot(protbert_scores[\"val\"], label = \"ProtBERT\")\nplt.plot(prott5_scores[\"val\"], label = \"ProtBERT\")\n\nplt.title(\"Validation F1-Scores for # Vector Embeddings\")\nplt.xlabel(\"Epochs\")\nplt.ylabel(\"Average F1-Score\")\nplt.legend()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:08:05.796934Z","iopub.execute_input":"2023-05-29T08:08:05.797306Z","iopub.status.idle":"2023-05-29T08:08:06.417327Z","shell.execute_reply.started":"2023-05-29T08:08:05.797275Z","shell.execute_reply":"2023-05-29T08:08:06.416313Z"},"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.ReLU()\n        self.linear2 = torch.nn.Linear(1012, 712)\n        self.activation2 = torch.nn.ReLU()\n        self.linear3 = torch.nn.Linear(712, num_classes)\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        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":{"execution":{"iopub.status.busy":"2023-05-29T08:08:11.37395Z","iopub.execute_input":"2023-05-29T08:08:11.374323Z","iopub.status.idle":"2023-05-29T08:08:11.394019Z","shell.execute_reply.started":"2023-05-29T08:08:11.374292Z","shell.execute_reply":"2023-05-29T08:08:11.393049Z"},"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    if embeddings_source == \"ProtT5\":\n        model = prott5_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-05-29T08:08:12.967812Z","iopub.execute_input":"2023-05-29T08:08:12.96817Z","iopub.status.idle":"2023-05-29T08:08:12.982458Z","shell.execute_reply.started":"2023-05-29T08:08:12.968141Z","shell.execute_reply":"2023-05-29T08:08:12.981454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df = predict(\"ProtT5\")","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:08:13.976975Z","iopub.execute_input":"2023-05-29T08:08:13.977644Z","iopub.status.idle":"2023-05-29T08:09:51.895479Z","shell.execute_reply.started":"2023-05-29T08:08:13.977602Z","shell.execute_reply":"2023-05-29T08:09:51.894462Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df.head(50)","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:09:51.897774Z","iopub.execute_input":"2023-05-29T08:09:51.898159Z","iopub.status.idle":"2023-05-29T08:09:51.916727Z","shell.execute_reply.started":"2023-05-29T08:09:51.898112Z","shell.execute_reply":"2023-05-29T08:09:51.915901Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(submission_df)","metadata":{"execution":{"iopub.status.busy":"2023-05-29T08:09:51.917884Z","iopub.execute_input":"2023-05-29T08:09:51.918822Z","iopub.status.idle":"2023-05-29T08:09:51.924648Z","shell.execute_reply.started":"2023-05-29T08:09:51.918788Z","shell.execute_reply":"2023-05-29T08:09:51.923762Z"},"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-05-29T08:09:51.926508Z","iopub.execute_input":"2023-05-29T08:09:51.927016Z","iopub.status.idle":"2023-05-29T08:14:02.178135Z","shell.execute_reply.started":"2023-05-29T08:09:51.926985Z","shell.execute_reply":"2023-05-29T08:14:02.177138Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### *I appreciate any feedback or Upvote as support*\n### *Also open for teaming, contact me in comments or mp*","metadata":{}}]}