{"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":"# What is about ?\n\nAn example to start working with the protein language model - \"prot_bert\" from  Rost lab: https://www.rostlab.org/ https://github.com/Rostlab . See  https://huggingface.co/Rostlab/prot_bert \n\nBeware - long sequences lead to crash by RAM. (I.e. for CPU 3000 symbols crash RAM, for GPU already 2000 symbols )\n\n    1500 symbols is okay for Kaggle GPU, but seems nearby the limit - it uses 14.7 Gig of GPU RAM with 16 is max possible.\n    Sorry with 1500 we can stil get \"OutOfMemoryError\" but something like AFTER  processing hundreds units - see e.g. version 13 \n    clip to 1200 first symbols - works at least for the first 1000 proteins.\n    \nThe inference time for CPU is MUCH(!) slower than GPU:  \n\n    CPU: 1000 symbols - 13.3 secs on CPU (Kaggle), 2000 symbols - 41 secs,  100 symbols - 1 sec. So it is more or less linear , but a bit faster. \n    GPU: everything in 0.05 sec - does not depend on length\n\nPS\n\nThe notebook is based on YUSUKE SATO : https://www.kaggle.com/code/wanko123/novoenzyme-protbert-embedding-eda\n\nPSPS\n\nThere are many other protein language models, but sometimes there are problems to install them on Kagge:\nhttps://www.kaggle.com/alexandervc/seqvec-bio-embeddings-install-fail/\n\n","metadata":{}},{"cell_type":"markdown","source":"Details:\n\n    Length 2000 - GPU P100 RAM crash - see Version 9: https://www.kaggle.com/code/alexandervc/protbert-embedding-starter?scriptVersionId=126797669  \n    GPU P100 OutOfMemoryError: CUDA out of memory. Tried to allocate 246.00 MiB (GPU 0; 15.90 GiB total capacity; 14.31 GiB already allocated; 229.75 MiB free; 14.82 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF\n\n    The same for the GPU T3 x 2 - look version 10. https://www.kaggle.com/code/alexandervc/protbert-embedding-starter?scriptVersionId=126798213\n    It is a bit slower - time 0.066 instead of 0.05 for inference .\n    \n    ","metadata":{}},{"cell_type":"markdown","source":"Versions:\n\n    15 - run on first 1000 proteins from the test with the clip to first 1200 symbols - results saved to csv - about 10 minutes on GPU","metadata":{}},{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# data check\nfrom collections import Counter\n\n# transformer\nimport torch\nfrom transformers import BertModel, BertTokenizer\nimport datasets\n\n# analysis\nfrom sklearn.cluster import KMeans\nfrom sklearn.mixture import GaussianMixture\nfrom sklearn.cluster import AgglomerativeClustering\nfrom umap import UMAP\nfrom sklearn.preprocessing import MinMaxScaler\n\n# visualization\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\n# use biopython to extract protein properties\nfrom Bio.SeqUtils.ProtParam import ProteinAnalysis\n\n# util\nfrom tqdm import tqdm\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\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-04-22T18:50:42.867228Z","iopub.execute_input":"2023-04-22T18:50:42.868053Z","iopub.status.idle":"2023-04-22T18:51:07.369503Z","shell.execute_reply.started":"2023-04-22T18:50:42.868013Z","shell.execute_reply":"2023-04-22T18:51:07.368035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load model","metadata":{}},{"cell_type":"code","source":"%%time\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\ntokenizer = BertTokenizer.from_pretrained(\"Rostlab/prot_bert\", do_lower_case=False )\nmodel = BertModel.from_pretrained(\"Rostlab/prot_bert\").to(device)","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:07.372654Z","iopub.execute_input":"2023-04-22T18:51:07.374488Z","iopub.status.idle":"2023-04-22T18:51:27.684552Z","shell.execute_reply.started":"2023-04-22T18:51:07.374442Z","shell.execute_reply":"2023-04-22T18:51:27.683318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Example: embded a sequence ","metadata":{}},{"cell_type":"markdown","source":"    The tokenizer of prot_bert only accept white space splitted sequence.\n    Before input a sequence, we need to process it. See the code below.","metadata":{}},{"cell_type":"code","source":"%%time\nsequence_example = \"MALLHSARVLSGVASAFHPGLAAAASARASSWWAHVEMGPPDPILGVTEAYKRDTNSKKMNLGVGAYRDDNGKPYVLPSVRKAEAQIAAKGLDKEYLPIGGLAEFCRASAELALGENSEVVKSGRFVTVQTISGTGALRIGASFLQRFFKFSRDVFLPKPSWGNHTPIFRDAGMQLQSYRYYDPKTCGFDFTGALEDISKIPEQSVLLLHACAHNPTGVDPRPEQWKEIATVVKKRNLFAFFDMAYQGFASGDGDKDAWAVRHFIEQGINVCLCQSYAKNMGLYGERVGAFTVICKDADEAKRVESQLKILIRPMYSNPPIHGARIASTILTSPDLRKQWLQEVKGMADRIIGMRTQLVSNLKKEGSTHSWQHITDQIGMFCFTGLKPEQVERLTKEFSIYMTKDGRISVAGVTSGNVGYLAHAIHQVTK\"\nprint( len(sequence_example) )\n\nsequence_example = ' '.join(list(sequence_example)) # The tokenizer of prot_bert only accept white space splitted sequence.\nprint(len(sequence_example))\nprint((sequence_example))\n\nencoded_input = tokenizer(sequence_example, return_tensors='pt').to(device)\noutput = model(**encoded_input)\ntype(output) , output","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:27.686142Z","iopub.execute_input":"2023-04-22T18:51:27.686607Z","iopub.status.idle":"2023-04-22T18:51:29.254402Z","shell.execute_reply.started":"2023-04-22T18:51:27.686568Z","shell.execute_reply":"2023-04-22T18:51:29.253273Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"output['last_hidden_state'].shape","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:29.257396Z","iopub.execute_input":"2023-04-22T18:51:29.257879Z","iopub.status.idle":"2023-04-22T18:51:29.264574Z","shell.execute_reply.started":"2023-04-22T18:51:29.257834Z","shell.execute_reply":"2023-04-22T18:51:29.263386Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Embedding: ","metadata":{}},{"cell_type":"code","source":"output['last_hidden_state'][:,0][0].detach().cpu().numpy()","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:29.266029Z","iopub.execute_input":"2023-04-22T18:51:29.266982Z","iopub.status.idle":"2023-04-22T18:51:29.280453Z","shell.execute_reply.started":"2023-04-22T18:51:29.266943Z","shell.execute_reply":"2023-04-22T18:51:29.27926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(output['last_hidden_state'][:,0][0].detach().cpu().numpy())","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:29.282229Z","iopub.execute_input":"2023-04-22T18:51:29.282585Z","iopub.status.idle":"2023-04-22T18:51:29.290914Z","shell.execute_reply.started":"2023-04-22T18:51:29.282548Z","shell.execute_reply":"2023-04-22T18:51:29.28976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Check - second run gives same output - true\n\nI.e. we check that inference does not contain randomness. \n","metadata":{}},{"cell_type":"code","source":"v1 = output['last_hidden_state'][:,0][0].detach().cpu().numpy()","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:29.292516Z","iopub.execute_input":"2023-04-22T18:51:29.2932Z","iopub.status.idle":"2023-04-22T18:51:29.298678Z","shell.execute_reply.started":"2023-04-22T18:51:29.293164Z","shell.execute_reply":"2023-04-22T18:51:29.297262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nsequence_example = \"MALLHSARVLSGVASAFHPGLAAAASARASSWWAHVEMGPPDPILGVTEAYKRDTNSKKMNLGVGAYRDDNGKPYVLPSVRKAEAQIAAKGLDKEYLPIGGLAEFCRASAELALGENSEVVKSGRFVTVQTISGTGALRIGASFLQRFFKFSRDVFLPKPSWGNHTPIFRDAGMQLQSYRYYDPKTCGFDFTGALEDISKIPEQSVLLLHACAHNPTGVDPRPEQWKEIATVVKKRNLFAFFDMAYQGFASGDGDKDAWAVRHFIEQGINVCLCQSYAKNMGLYGERVGAFTVICKDADEAKRVESQLKILIRPMYSNPPIHGARIASTILTSPDLRKQWLQEVKGMADRIIGMRTQLVSNLKKEGSTHSWQHITDQIGMFCFTGLKPEQVERLTKEFSIYMTKDGRISVAGVTSGNVGYLAHAIHQVTK\"\nprint( len(sequence_example) )\n\nsequence_example = ' '.join(list(sequence_example))\nencoded_input = tokenizer(sequence_example, return_tensors='pt').to(device)\noutput = model(**encoded_input)\ntype(output) , output\n\nv2 = output['last_hidden_state'][:,0][0].detach().cpu().numpy()","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:29.300089Z","iopub.execute_input":"2023-04-22T18:51:29.301041Z","iopub.status.idle":"2023-04-22T18:51:29.385202Z","shell.execute_reply.started":"2023-04-22T18:51:29.301004Z","shell.execute_reply":"2023-04-22T18:51:29.384107Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"np.max( np.abs( v1-v2)) # .abs().max()","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:29.38667Z","iopub.execute_input":"2023-04-22T18:51:29.387293Z","iopub.status.idle":"2023-04-22T18:51:29.394766Z","shell.execute_reply.started":"2023-04-22T18:51:29.387256Z","shell.execute_reply":"2023-04-22T18:51:29.393517Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference time measurments ","metadata":{}},{"cell_type":"code","source":"%%time\nimport time\nimport gc \n\nstr_accelerator = 'GPU' if torch.cuda.is_available() else 'CPU'\nprint(str_accelerator )\n\nsequence_example_1= \"MALLHSARVLSGVASAFHPGLAAAASARASSWWAHVEMGPPDPILGVTEAYKRDTNSKKMNLGVGAYRDDNGKPYVLPSVRKAEAQIAAKGLDKEYLPIGGLAEFCRASAELALGENSEVVKSGRFVTVQTISGTGALRIGASFLQRFFKFSRDVFLPKPSWGNHTPIFRDAGMQLQSYRYYDPKTCGFDFTGALEDISKIPEQSVLLLHACAHNPTGVDPRPEQWKEIATVVKKRNLFAFFDMAYQGFASGDGDKDAWAVRHFIEQGINVCLCQSYAKNMGLYGERVGAFTVICKDADEAKRVESQLKILIRPMYSNPPIHGARIASTILTSPDLRKQWLQEVKGMADRIIGMRTQLVSNLKKEGSTHSWQHITDQIGMFCFTGLKPEQVERLTKEFSIYMTKDGRISVAGVTSGNVGYLAHAIHQVTK\"\nsequence_example_1 = sequence_example_1[:100] \nprint( len(sequence_example_1) )\n\nl = [1,2,4,8,10,20]\nif str_accelerator == 'GPU':\n    l = l[:-1]\n    \nlist_lengths = []\nlist_times = []\nfor k in l:\n    t0 = time.time()\n    sequence_example = sequence_example_1 * int(k)\n    list_lengths.append(len(sequence_example))\n    len_tmp = len(sequence_example)\n    sequence_example = ' '.join(list(sequence_example))\n    encoded_input = tokenizer(sequence_example, return_tensors='pt').to(device)\n    output = model(**encoded_input)\n    t = time.time() -t0\n    list_times.append(t)\n    print('len:', len_tmp, 'seconds passed:',   np.round(t,3) )\n    \n    gc.collect() \n","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:29.399758Z","iopub.execute_input":"2023-04-22T18:51:29.400064Z","iopub.status.idle":"2023-04-22T18:51:31.160327Z","shell.execute_reply.started":"2023-04-22T18:51:29.400038Z","shell.execute_reply":"2023-04-22T18:51:31.15921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nstr_accelerator = 'GPU' if torch.cuda.is_available() else 'CPU'\nprint(str_accelerator )\n\nplt.figure(figsize = (20,4) )\nplt.plot(list_lengths, list_times)\nplt.title('Inference time (seconds) dependence on input length ' + str_accelerator, fontsize = 20)\nplt.xlabel('input length', fontsize = 20)\nplt.ylabel('seconds', fontsize = 20)\nplt.show()\n","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:31.162991Z","iopub.execute_input":"2023-04-22T18:51:31.163368Z","iopub.status.idle":"2023-04-22T18:51:31.418095Z","shell.execute_reply.started":"2023-04-22T18:51:31.163331Z","shell.execute_reply":"2023-04-22T18:51:31.416999Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Start looking on the CAFA5 data","metadata":{}},{"cell_type":"code","source":"from Bio import SeqIO\nfn = '/kaggle/input/cafa-5-protein-function-prediction/Train/train_sequences.fasta'\nprint(\"Sequence example:\\n\\n\", next(iter(SeqIO.parse(fn, \"fasta\"))))\nsequences = SeqIO.parse(fn, \"fasta\")\nl = [len(seq) for seq in sequences ] \nprint(pd.Series(l).describe() )\n\nsequences = SeqIO.parse(fn, \"fasta\")\nseq = next(iter(sequences))\ngb = seq\nprint('\\nLength of Sequence:')\nprint(len(gb.seq))\n\nprint('\\nRecord ID:')\nprint(gb.id)\n\nprint('\\nName:')\nprint(gb.name)\n\nprint('\\nDescription:')\nprint(gb.description)\n\n# Annotations \nprint('\\nNumber of Annotations:')\nprint(len(gb.annotations))\n\n# Features \nprint('\\nNumber of Features:')\nprint(len(gb.features))","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:31.419181Z","iopub.execute_input":"2023-04-22T18:51:31.420407Z","iopub.status.idle":"2023-04-22T18:51:33.774972Z","shell.execute_reply.started":"2023-04-22T18:51:31.420369Z","shell.execute_reply":"2023-04-22T18:51:33.773805Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Emdedding for many proteins","metadata":{}},{"cell_type":"code","source":"\nclip_first_n_symbols = 1200\n\nfrom Bio import SeqIO\nfn = '/kaggle/input/cafa-5-protein-function-prediction/Train/train_sequences.fasta'\nsequences = SeqIO.parse(fn, \"fasta\")\nl = [len(seq) for seq in sequences ] \n\ns = (np.array(l)>clip_first_n_symbols ).sum() \nprint('Count Proteins with higher length: ' , s ,'Percent:', np.round( s/len(l)*100 , 1)  )","metadata":{"execution":{"iopub.status.busy":"2023-04-22T19:04:24.977958Z","iopub.execute_input":"2023-04-22T19:04:24.978906Z","iopub.status.idle":"2023-04-22T19:04:26.479795Z","shell.execute_reply.started":"2023-04-22T19:04:24.978855Z","shell.execute_reply":"2023-04-22T19:04:26.478611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nfn = '/kaggle/input/cafa-5-protein-function-prediction/Train/train_sequences.fasta'\n\nsequences = SeqIO.parse(fn, \"fasta\")\ndf_embed = pd.DataFrame(); IX = 0; \nt00 = time.time()\n\nfor seq in sequences:\n    t0 = time.time()\n    sequence_example = seq[:clip_first_n_symbols]\n    list_lengths.append(len(sequence_example))\n    len_tmp = len(sequence_example)\n    sequence_example = ' '.join(list(sequence_example))\n    encoded_input = tokenizer(sequence_example, return_tensors='pt').to(device)\n    output = model(**encoded_input)\n    t = time.time() -t0\n    t2 = time.time() -t00\n    list_times.append(t)\n    \n    prot_id = seq.id\n    if IX % 500 == 0:\n        print(IX,prot_id, 'len:', len_tmp, 'seconds passed:',   np.round(t,2), 'total',  np.round(t2,2), )\n    \n    df_embed[prot_id+'_'+str(IX)] =  output['last_hidden_state'][:,0][0].detach().cpu().numpy()\n    IX += 1 \n    gc.collect() \n    if IX >= 70000:\n        break\n        \n    if IX % 10_000 == 0:\n        fn = 'embed_protbert_train_clip_'+str(clip_first_n_symbols)+'_first_'+str(df_embed.shape[1])+'_prot.csv'\n        print(fn)\n        df_embed.to_csv(fn)\n    \n    \n","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:51:33.784348Z","iopub.execute_input":"2023-04-22T18:51:33.785001Z","iopub.status.idle":"2023-04-22T18:57:58.869292Z","shell.execute_reply.started":"2023-04-22T18:51:33.784965Z","shell.execute_reply":"2023-04-22T18:57:58.867301Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nfn = 'embed_protbert_train_clip_'+str(clip_first_n_symbols)+'_first_'+str(df_embed.shape[1])+'_prot.csv'\nprint(fn)\ndf_embed.to_csv(fn)","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:57:58.870628Z","iopub.execute_input":"2023-04-22T18:57:58.87102Z","iopub.status.idle":"2023-04-22T18:58:00.080062Z","shell.execute_reply.started":"2023-04-22T18:57:58.870982Z","shell.execute_reply":"2023-04-22T18:58:00.078846Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_embed","metadata":{"execution":{"iopub.status.busy":"2023-04-22T18:58:00.081946Z","iopub.execute_input":"2023-04-22T18:58:00.082343Z","iopub.status.idle":"2023-04-22T18:58:00.123842Z","shell.execute_reply.started":"2023-04-22T18:58:00.082305Z","shell.execute_reply":"2023-04-22T18:58:00.122898Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}