{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.8.17","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":41875,"databundleVersionId":5521661,"sourceType":"competition"},{"sourceId":5807641,"sourceType":"datasetVersion","datasetId":3211581},{"sourceId":6304945,"sourceType":"datasetVersion","datasetId":3627105},{"sourceId":6322991,"sourceType":"datasetVersion","datasetId":3638862},{"sourceId":140286135,"sourceType":"kernelVersion"}],"dockerImageVersionId":30527,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"most_common_gene_ontologies = 5000","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --upgrade pip","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#tpu pip installs\n\n!pip install biopython\n!pip install numpy\n!pip install pip install torchmetrics","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n# for dirname, _, filename s in os.walk('/kaggle/'):\n#     for filename in filenames:\n#         print(os.path.join(dirname, filename))\n\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\n\n\nimport torch\nfrom Bio import SeqIO\n#import plotly.graph_objects as go\nfrom collections import Counter\n\nimport sys\n\n#from transformers import BertModel, BertTokenizer\nimport re\n\nfrom timeit import default_timer as timer\nfrom datetime import timedelta\nimport time \nfrom datetime import datetime\n\nfrom glob import glob\n\nfrom scipy.special import softmax\n\n\nfrom torchmetrics.classification import BinaryF1Score\n\nnp.set_printoptions(precision=5)\n\ntorch.set_printoptions(threshold=5,sci_mode=False)\ntorch.manual_seed(0)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n# !python pytorch-xla-env-setup.py --apt-packages libomp5 libopenblas-dev","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch_xla\nimport torch_xla.core.xla_model as xm\n\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\n\n\ndevice = xm.xla_device()\ntorch.set_default_tensor_type('torch.FloatTensor')\ndevice\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n# device","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class EasyFeedForward(torch.nn.Module):\n    \n        def __init__(self, input_size, hidden_size, output_size):\n            super(EasyFeedForward, self).__init__()\n            self.layer1 = torch.nn.Linear(input_size, hidden_size)\n            self.activiation1 = torch.nn.ReLU()\n            self.dropout1 = torch.nn.Dropout(0.5)\n            self.layer2 = torch.nn.Linear(hidden_size, hidden_size)\n            self.activiation2 = torch.nn.ReLU()\n            self.dropout2 = torch.nn.Dropout(0.5)\n            self.layer3 = torch.nn.Linear(hidden_size, hidden_size)\n            self.activiation3 = torch.nn.ReLU()\n            self.layer4 = torch.nn.Linear(hidden_size, hidden_size)\n            self.activiation4 = torch.nn.ReLU()\n            self.layer5 = torch.nn.Linear(hidden_size, hidden_size)\n            self.activiation5 = torch.nn.ReLU()\n            self.layer6 = torch.nn.Linear(hidden_size, hidden_size)\n            self.activiation6 = torch.nn.ReLU()\n            self.layer7 = torch.nn.Linear(hidden_size, output_size)\n            \n        def forward(self, x):\n            output = self.layer1(x)\n            output = self.activiation1(output)\n            output = self.dropout1(output)\n            output = self.layer2(output)\n            output = self.activiation2(output)\n            output = self.dropout2(output)\n            output = self.layer3(output)\n            output = self.activiation3(output)\n            output = self.layer4(output)\n            output = self.activiation4(output)\n            output = self.layer5(output)\n            output = self.activiation5(output)\n            output = self.layer6(output)\n            output = self.activiation6(output)\n            output = self.layer7(output)\n            return output\n\n#(1) init FeedForwardNet \nnn_input_size  = 2560 #last_hidden \nnn_hidden_size = most_common_gene_ontologies\nnn_output_size = most_common_gene_ontologies #3000     \n        \nff_nn_model = EasyFeedForward(nn_input_size,nn_hidden_size,nn_output_size ).to(device)\n\n# print(\"NeuralNet uses gpu/tpu: \" + str(next(ff_nn_model.parameters()).is_cuda))\n# print(\"\")\nff_nn_model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"target_vecs = np.load(\"/kaggle/input/datagen-esm2-t36-3b-last-layer-embedings/target_vecs_full_5000.npy\",allow_pickle=True)\nlhsts = np.load(\"/kaggle/input/4637427/train_embeds_esm2_t36_3B_UR50D.npy\",allow_pickle=True)\nentry_ids_train = np.load(\"/kaggle/input/4637427/train_ids_esm2_t36_3B_UR50D.npy\",allow_pickle=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(target_vecs),len(lhsts) ,len(entry_ids_train)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"target_vecs[0][0]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(lhsts[0]),len(target_vecs[0][0])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"list_of_target_vecs = []\nfor tv in target_vecs:\n    list_of_target_vecs.append(tv[0].float())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# del target_vecs\n# import gc\n# gc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"lhst = torch.tensor(lhsts)\ntargetVecss = torch.stack(list_of_target_vecs,0)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"lhstVecsGPU = lhst.to(device)\n\ntargetVecsGPU = targetVecss.to(device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#split train test\nsplit_num = 135_000\ntrain_lhstVecsGPU, test_lhstVecsGPU     = lhstVecsGPU[:split_num], lhstVecsGPU[split_num:len(lhstVecsGPU)] \ntrain_targetVecsGPU, test_targetVecsGPU = targetVecsGPU[:split_num], targetVecsGPU[split_num:len(targetVecsGPU)] ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ff_nn_model(train_lhstVecsGPU[0].float())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#torch.cuda.empty_cache()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#from datetime import datetime\nprint(\"start\")\nmetric = BinaryF1Score(multidim_average='global') #,threshold=0.25\nmetric.to(device)\n\ncrossEntropyLoss = torch.nn.BCEWithLogitsLoss(reduction=\"mean\")\n\noptimizer = torch.optim.Adam(ff_nn_model.parameters(), lr=0.0001)\nn_epochs = 1500\n\nbatch_size = 4096\n#print(\"is cuda: \", next(ff_nn_model.parameters()))\n\nprint(\"BEGIN TRAINING...\", datetime.now())\nstart = timer()\nfor epoch in range(n_epochs):\n    \n    \n    for i in range(0,len(train_targetVecsGPU), batch_size):\n        lhst = train_lhstVecsGPU[i:i+batch_size]\n        target = train_targetVecsGPU[i:i+batch_size]\n        optimizer.zero_grad()\n        out = ff_nn_model(lhst.float().to(device))\n        loss = crossEntropyLoss(out,target)\n        loss.backward()\n        #print(loss)\n        xm.optimizer_step(optimizer, barrier=True)\n        #optimizer.step()\n        #xm.master_print(loss)\n    if epoch % 1 == 0:\n        \n        #f1 score\n        out = ff_nn_model(test_lhstVecsGPU.float().to(device))  \n        f1_score = metric(out,test_targetVecsGPU)\n  \n        \n        end = timer()\n        print(epoch, loss.detach().item(),\"f1: \", f1_score.detach().cpu().item(),timedelta(seconds=end-start),datetime.now())\n        start = timer()\n        \n        \n          \n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"out = ff_nn_model(test_lhstVecsGPU.float())  \nf1_score = metric(out,test_targetVecsGPU)\nf1_score","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(ff_nn_model, \"new_esm2_3b_fine_tuning_5000_4096_batch_vec_1\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}