{"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 is the code for calculating embeddings from the t5 dataset https://www.kaggle.com/datasets/sergeifironov/t5embeds. Unfortunately, it is impossible to run it on Kaggle resources, even with a batch size of 1 you need A100 for evaluation.","metadata":{}},{"cell_type":"code","source":"!pip install obonet\n!pip install pyvis","metadata":{"execution":{"iopub.status.busy":"2023-04-25T10:49:15.285721Z","iopub.execute_input":"2023-04-25T10:49:15.286101Z","iopub.status.idle":"2023-04-25T10:49:37.232064Z","shell.execute_reply.started":"2023-04-25T10:49:15.286069Z","shell.execute_reply":"2023-04-25T10:49:37.230774Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import EsmModel, EsmTokenizer\nimport torch\nimport torch.nn as nn\n\nimport os, gc\nimport json\nfrom typing import Dict\nfrom collections import Counter\n\nimport random\nimport obonet\nimport pandas as pd\nimport numpy as np\nfrom Bio import SeqIO\n\nmodel_names_sizes = [\n    (\"facebook/esm2_t36_3B_UR50D\", 2560),\n    (\"esm2_t33_650M_UR50D\", 1280),\n    (\"esm2_t30_150M_UR50D\", 640),\n    (\"esm2_t12_35M_UR50D\", 480),\n    (\"esm2_t6_8M_UR50D\", 320)\n    ]\ni = -1\nfolder_name = model_names_sizes[i][0]\nmodel_name = f'facebook/{model_names_sizes[i][0]}'\nembed_size = model_names_sizes[i][1]\nprint(f'{model_name}, {embed_size}')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\ntokenizer = EsmTokenizer.from_pretrained(model_name)\nmodel = EsmModel.from_pretrained(model_name, add_cross_attention=False, is_decoder=False).to(device)\nmodel.eval()","metadata":{"execution":{"iopub.status.busy":"2023-04-25T10:55:35.814615Z","iopub.execute_input":"2023-04-25T10:55:35.815339Z","iopub.status.idle":"2023-04-25T10:55:37.210307Z","shell.execute_reply.started":"2023-04-25T10:55:35.815302Z","shell.execute_reply":"2023-04-25T10:55:37.209222Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import re\ndef get_embeddings(seq):\n    sequence_examples = [\" \".join(list(re.sub(r\"[UZOB]\", \"X\", seq)))]\n\n    ids = tokenizer(sequence_examples, add_special_tokens=True, padding=\"longest\")\n\n    input_ids = torch.tensor(ids['input_ids']).to(device)\n    attention_mask = torch.tensor(ids['attention_mask']).to(device)\n\n    # generate embeddings\n    with torch.no_grad():\n        embedding_repr = model(input_ids=input_ids,\n                               attention_mask=attention_mask)\n\n    # extract residue embeddings for the first ([0,:]) sequence in the batch and remove padded & special tokens ([0,:7]) \n    emb_0 = embedding_repr.last_hidden_state[0]\n    emb_0_per_protein = emb_0.mean(dim=0)\n    \n    return emb_0_per_protein\n\ncheck = get_embeddings('MTMDKSELVQKAKLAEQAERYDDMAAAMKAVTEQGHELSNEERNLLSVAYKNVVGARRSS')\ncheck.shape","metadata":{"execution":{"iopub.status.busy":"2023-04-25T10:55:37.212047Z","iopub.execute_input":"2023-04-25T10:55:37.213179Z","iopub.status.idle":"2023-04-25T10:55:37.222607Z","shell.execute_reply.started":"2023-04-25T10:55:37.213139Z","shell.execute_reply":"2023-04-25T10:55:37.220988Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fn = '/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\")\nnum_sequences = sum(1 for seq in sequences)\nprint()\nprint(\"Number of sequences in train:\", num_sequences)","metadata":{"execution":{"iopub.status.busy":"2023-04-25T10:55:37.224788Z","iopub.execute_input":"2023-04-25T10:55:37.225238Z","iopub.status.idle":"2023-04-25T10:55:38.220356Z","shell.execute_reply.started":"2023-04-25T10:55:37.225199Z","shell.execute_reply":"2023-04-25T10:55:38.219124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tqdm\nfn = '/kaggle/input/cafa-5-protein-function-prediction/Train/train_sequences.fasta'\n\nsequences = SeqIO.parse(fn, \"fasta\")\n\nids = []\nembeds = np.zeros((num_sequences, embed_size))\ni = 0\nfor seq in tqdm.tqdm(sequences):\n    ids.append(seq.id)\n    # cut sequence to 1022 residuals because max length for ESM is 1024 = <cls> + 1022 + <eos>  \n    embeds[i] = get_embeddings(str(seq.seq)[:1022]).detach().cpu().numpy()\n    i += 1\n    gc.collect()\n    torch.cuda.empty_cache()  \n    break #remove it for full calculation\n\nif not os.path.exists(f'embeds_{folder_name}'):\n    os.makedirs(f'embeds_{folder_name}')\n\ntrain_embeds_out_path = os.path.join(\n    f'embeds_{folder_name}',\n    f'train_embeds_{folder_name}.npy'\n)\ntrain_ids_out_path = os.path.join(\n    f'embeds_{folder_name}',\n    f'train_ids_{folder_name}.npy'\n)\n\nnp.save(train_embeds_out_path, embeds)\nnp.save(train_ids_out_path, np.array(ids))","metadata":{"execution":{"iopub.status.busy":"2023-04-25T10:55:38.224564Z","iopub.execute_input":"2023-04-25T10:55:38.224889Z","iopub.status.idle":"2023-04-25T10:55:39.196467Z","shell.execute_reply.started":"2023-04-25T10:55:38.224858Z","shell.execute_reply":"2023-04-25T10:55:39.195198Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test set:\nimport tqdm\nfn = '/kaggle/input/cafa-5-protein-function-prediction/Test (Targets)/testsuperset.fasta'\n\nsequences = SeqIO.parse(fn, \"fasta\")\n\nids = []\nembeds = np.zeros((num_sequences, embed_size))\n\nprint('Start calcutation for test set')\ni = 0\nfor seq in tqdm.tqdm(sequences):\n    ids.append(seq.id)\n    embeds[i] = get_embeddings(str(seq.seq)[:1022]).detach().cpu().numpy()\n    i += 1\n    gc.collect()\n    torch.cuda.empty_cache()  \n    break #remove it for full calculation\n\nif not os.path.exists(f'embeds_{folder_name}'):\n    os.makedirs(f'embeds_{folder_name}')\n\ntest_embeds_out_path = os.path.join(\n    f'embeds_{folder_name}',\n    f'test_embeds_{folder_name}.npy'\n)\ntest_ids_out_path = os.path.join(\n    f'embeds_{folder_name}',\n    f'test_ids_{folder_name}.npy'\n)\n\nnp.save(test_embeds_out_path, embeds)\nnp.save(test_ids_out_path, np.array(ids))","metadata":{"execution":{"iopub.status.busy":"2023-04-25T10:55:39.1986Z","iopub.execute_input":"2023-04-25T10:55:39.199076Z","iopub.status.idle":"2023-04-25T10:55:39.222888Z","shell.execute_reply.started":"2023-04-25T10:55:39.199034Z","shell.execute_reply":"2023-04-25T10:55:39.221738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}