{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This analysis is built upon the foundational work by @sergeifironov from the CAFA 5 competition, who published the acclaimed notebook titled \"T5embeds calculation [only few samples]\" I want to extend my gratitude for his insightful contribution.","metadata":{}},{"cell_type":"markdown","source":"This is the code for calculating embeddings from the t5 dataset for CAFA 6 : https://www.kaggle.com/datasets/ahsuna123/t5-embedding-cafa-6. 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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:49:09.164922Z","iopub.execute_input":"2025-10-16T05:49:09.165400Z","iopub.status.idle":"2025-10-16T05:49:17.072299Z","shell.execute_reply.started":"2025-10-16T05:49:09.165376Z","shell.execute_reply":"2025-10-16T05:49:17.071343Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install Bio","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:49:17.073928Z","iopub.execute_input":"2025-10-16T05:49:17.074236Z","iopub.status.idle":"2025-10-16T05:49:22.096467Z","shell.execute_reply.started":"2025-10-16T05:49:17.074207Z","shell.execute_reply":"2025-10-16T05:49:22.095711Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pip install transformers==4.39.3 --quiet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:49:22.097302Z","iopub.execute_input":"2025-10-16T05:49:22.097506Z","iopub.status.idle":"2025-10-16T05:49:35.110806Z","shell.execute_reply.started":"2025-10-16T05:49:22.097486Z","shell.execute_reply":"2025-10-16T05:49:35.109923Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\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\nimport torch\nfrom transformers import T5Tokenizer, T5EncoderModel\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ntokenizer = T5Tokenizer.from_pretrained(\n    \"Rostlab/prot_t5_xl_half_uniref50-enc\",\n    do_lower_case=False\n)\n\nmodel = T5EncoderModel.from_pretrained(\n    \"Rostlab/prot_t5_xl_half_uniref50-enc\"\n).to(device)\n\nif device.type == \"cuda\":\n    model = model.half()\n\nprint(\"✅ ProtT5 loaded successfully!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:49:35.112648Z","iopub.execute_input":"2025-10-16T05:49:35.112889Z","iopub.status.idle":"2025-10-16T05:49:56.476684Z","shell.execute_reply.started":"2025-10-16T05:49:35.112869Z","shell.execute_reply":"2025-10-16T05:49:56.476025Z"}},"outputs":[],"execution_count":null},{"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)\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input/cafa-6-protein-function-prediction'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:49:56.477406Z","iopub.execute_input":"2025-10-16T05:49:56.477873Z","iopub.status.idle":"2025-10-16T05:49:56.490040Z","shell.execute_reply.started":"2025-10-16T05:49:56.477848Z","shell.execute_reply":"2025-10-16T05:49:56.489496Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from pathlib import Path\npath = Path('/kaggle/input/cafa-6-protein-function-prediction')\n!head {path}/'Test (Targets)/testsuperset-taxon-list.tsv'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:49:56.490717Z","iopub.execute_input":"2025-10-16T05:49:56.491025Z","iopub.status.idle":"2025-10-16T05:50:05.038441Z","shell.execute_reply.started":"2025-10-16T05:49:56.491008Z","shell.execute_reply":"2025-10-16T05:50:05.037456Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!head {path}/'Test (Targets)/testsuperset.fasta'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:50:05.039665Z","iopub.execute_input":"2025-10-16T05:50:05.040267Z","iopub.status.idle":"2025-10-16T05:50:05.774786Z","shell.execute_reply.started":"2025-10-16T05:50:05.040230Z","shell.execute_reply":"2025-10-16T05:50:05.773790Z"}},"outputs":[],"execution_count":null},{"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.batch_encode_plus(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\nget_embeddings('MTMDKSELVQKAKLAEQAERYDDMAAAMKAVTEQGHELSNEERNLLSVAYKNVVGARRSS')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:50:05.776015Z","iopub.execute_input":"2025-10-16T05:50:05.776284Z","iopub.status.idle":"2025-10-16T05:50:06.676423Z","shell.execute_reply.started":"2025-10-16T05:50:05.776261Z","shell.execute_reply":"2025-10-16T05:50:06.675816Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fn = '/kaggle/input/cafa-6-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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:50:06.677149Z","iopub.execute_input":"2025-10-16T05:50:06.677406Z","iopub.status.idle":"2025-10-16T05:50:07.587681Z","shell.execute_reply.started":"2025-10-16T05:50:06.677371Z","shell.execute_reply":"2025-10-16T05:50:07.587002Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import tqdm\nfn = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\n\nsequences = SeqIO.parse(fn, \"fasta\")\n\nids = []\nembeds = np.zeros((num_sequences, 1024))\ni = 0\nfor seq in tqdm.tqdm(sequences):\n    ids.append(seq.id)\n    embeds[i] = get_embeddings(str(seq.seq)).detach().cpu().numpy()\n    i += 1\n    break #remove it for full calculation\n        \nnp.save('train_embeds.npy', embeds)\nnp.save('train_ids.npy', np.array(ids))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:50:07.589866Z","iopub.execute_input":"2025-10-16T05:50:07.590066Z","iopub.status.idle":"2025-10-16T05:50:08.142625Z","shell.execute_reply.started":"2025-10-16T05:50:07.590051Z","shell.execute_reply":"2025-10-16T05:50:08.142041Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fn = '/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta'\n\nsequences = SeqIO.parse(fn, \"fasta\")\nnum_sequences = sum(1 for seq in sequences)\nprint(\"Number of sequences in test:\", num_sequences)\nsequences = SeqIO.parse(fn, \"fasta\")\n\n\nids = []\nembeds = np.zeros((num_sequences, 1024))\ni = 0\nfor seq in tqdm.tqdm(sequences):\n    ids.append(seq.id)\n    embeds[i] = get_embeddings(str(seq.seq)).detach().cpu().numpy()\n    i += 1\n    break #remove it for full calculation\n    \nnp.save('test_embeds.npy', embeds)\nnp.save('test_ids.npy', np.array(ids))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T05:50:08.143475Z","iopub.execute_input":"2025-10-16T05:50:08.143736Z","iopub.status.idle":"2025-10-16T05:50:11.585165Z","shell.execute_reply.started":"2025-10-16T05:50:08.143712Z","shell.execute_reply":"2025-10-16T05:50:11.584293Z"}},"outputs":[],"execution_count":null}]}