{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"References\n\n - https://www.kaggle.com/code/henriupton/proteinet-pytorch-ems2-t5-protbert-embeddings","metadata":{}},{"cell_type":"code","source":"!pip install -U transformers huggingface_hub","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-13T17:34:18.808241Z","iopub.execute_input":"2025-11-13T17:34:18.808437Z","iopub.status.idle":"2025-11-13T17:34:34.50676Z","shell.execute_reply.started":"2025-11-13T17:34:18.80842Z","shell.execute_reply":"2025-11-13T17:34:34.506014Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install Bio","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-13T17:34:34.508692Z","iopub.execute_input":"2025-11-13T17:34:34.508894Z","iopub.status.idle":"2025-11-13T17:34:39.706133Z","shell.execute_reply.started":"2025-11-13T17:34:34.508872Z","shell.execute_reply":"2025-11-13T17:34:39.705359Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import Bio\nfrom Bio import SeqIO","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-13T17:34:39.70941Z","iopub.execute_input":"2025-11-13T17:34:39.709699Z","iopub.status.idle":"2025-11-13T17:34:39.782942Z","shell.execute_reply.started":"2025-11-13T17:34:39.709668Z","shell.execute_reply":"2025-11-13T17:34:39.782389Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nsub = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/sample_submission.tsv\", sep= \"\\t\", on_bad_lines='skip', 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":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-13T17:34:39.783451Z","iopub.execute_input":"2025-11-13T17:34:39.783696Z","iopub.status.idle":"2025-11-13T17:34:40.089149Z","shell.execute_reply.started":"2025-11-13T17:34:39.783678Z","shell.execute_reply":"2025-11-13T17:34:40.088526Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"MAIN_DIR = \"/kaggle/input/cafa-6-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\nfrom itertools import islice\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":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-13T17:35:29.654851Z","iopub.execute_input":"2025-11-13T17:35:29.655669Z","iopub.status.idle":"2025-11-13T17:35:44.53777Z","shell.execute_reply.started":"2025-11-13T17:35:29.655631Z","shell.execute_reply":"2025-11-13T17:35:44.536989Z"}},"outputs":[],"execution_count":null},{"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/testsuperset.fasta\"\n    \n    num_labels = 5 # 500\n    n_epochs = 8\n    batch_size = 128\n    lr = 0.001\n    \n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \nprint(config.device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-13T17:35:47.123484Z","iopub.execute_input":"2025-11-13T17:35:47.124152Z","iopub.status.idle":"2025-11-13T17:35:47.165479Z","shell.execute_reply.started":"2025-11-13T17:35:47.12412Z","shell.execute_reply":"2025-11-13T17:35:47.164801Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Load ProtBERT Model...\")\n# PROT BERT LOADING :\nfrom transformers import BertModel, BertTokenizer\ntokenizer = BertTokenizer.from_pretrained(\"Rostlab/prot_bert\", do_lower_case=False )\nmodel = BertModel.from_pretrained(\"Rostlab/prot_bert\").to(config.device)\n\ndef get_bert_embedding(\n    sequence : str,\n    len_seq_limit : int\n):\n    '''\n    Function to collect last hidden state embedding vector from pre-trained ProtBERT Model\n\n    INPUTS:\n    - sequence (str) : protein sequence (ex : AAABBB) from fasta file\n    - len_seq_limit (int) : maximum sequence lenght (i.e nb of letters) for truncation\n\n    OUTPUTS:\n    - output_hidden : last hidden state embedding vector for input sequence of length 1024\n    '''\n    sequence_w_spaces = ' '.join(list(sequence))\n    encoded_input = tokenizer(\n        sequence_w_spaces,\n        truncation=True,\n        max_length=len_seq_limit,\n        padding='max_length',\n        return_tensors='pt').to(config.device)\n    output = model(**encoded_input)\n    output_hidden = output['last_hidden_state'][:,0][0].detach().cpu().numpy()\n    assert len(output_hidden)==1024\n    return output_hidden\n\n### COLLECTING FOR TRAIN SAMPLES :\nprint(\"Loading train set ProtBERT Embeddings...\")\nfasta_train = SeqIO.parse(config.train_sequences_path, \"fasta\")\nprint(\"Total Nb of Elements : \", len(list(fasta_train)))\nfasta_train = SeqIO.parse(config.train_sequences_path, \"fasta\")\nids_list = []\nembed_vects_list = []\nt0 = time.time()\ncheckpoint = 0\nfor item in tqdm(fasta_train):\n    ids_list.append(item.id)\n    embed_vects_list.append(\n        get_bert_embedding(sequence = item.seq, len_seq_limit = 1200))\n    checkpoint+=1\n    if checkpoint>=100:\n        df_res = pd.DataFrame(data={\"id\" : ids_list, \"embed_vect\" : embed_vects_list})\n        np.save('/kaggle/working/train_ids.npy',np.array(ids_list))\n        np.save('/kaggle/working/train_embeddings.npy',np.array(embed_vects_list))\n        checkpoint=0\n\nnp.save('/kaggle/working/train_ids.npy',np.array(ids_list))\nnp.save('/kaggle/working/train_embeddings.npy',np.array(embed_vects_list))\nprint('Total Elapsed Time:',time.time()-t0)\n\n### COLLECTING FOR TEST SAMPLES :\nprint(\"Loading test set ProtBERT Embeddings...\")\nfasta_test = SeqIO.parse(config.test_sequences_path, \"fasta\")\nprint(\"Total Nb of Elements : \", len(list(fasta_test)))\nfasta_test = SeqIO.parse(config.test_sequences_path, \"fasta\")\nids_list = []\nembed_vects_list = []\nt0 = time.time()\ncheckpoint=0\nfor item in tqdm(fasta_test):\n    ids_list.append(item.id)\n    embed_vects_list.append(\n        get_bert_embedding(sequence = item.seq, len_seq_limit = 1200))\n    checkpoint+=1\n    if checkpoint>=100:\n        np.save('/kaggle/working/test_ids.npy',np.array(ids_list))\n        np.save('/kaggle/working/test_embeddings.npy',np.array(embed_vects_list))\n        checkpoint=0\n\nnp.save('/kaggle/working/test_ids.npy',np.array(ids_list))\nnp.save('/kaggle/working/test_embeddings.npy',np.array(embed_vects_list))\nprint('Total Elasped Time:',time.time()-t0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-13T17:36:09.467917Z","iopub.execute_input":"2025-11-13T17:36:09.468672Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# print(\"GENERATE TARGETS FOR ENTRY IDS (\"+str(config.num_labels)+\" MOST COMMON GO TERMS)\")\n# ids = np.load(\"/kaggle/working/train_ids.npy\")\n# labels = pd.read_csv(config.train_labels_path, sep = \"\\t\")\n\n# top_terms = labels.groupby(\"term\")[\"EntryID\"].count().sort_values(ascending=False)\n# labels_names = top_terms[:config.num_labels].index.values\n# train_labels_sub = labels[(labels.term.isin(labels_names)) & (labels.EntryID.isin(ids))]\n# id_labels = train_labels_sub.groupby('EntryID')['term'].apply(list).to_dict()\n\n# go_terms_map = {label: i for i, label in enumerate(labels_names)}\n# labels_matrix = np.empty((len(ids), len(labels_names)))\n\n# for 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\n# labels_list = []\n# for l in range(labels_matrix.shape[0]):\n#     labels_list.append(labels_matrix[l, :])\n\n# labels_df = pd.DataFrame(data={\"EntryID\":ids, \"labels_vect\":labels_list})\n# labels_df.to_pickle(\"/kaggle/working/train_targets_top\"+str(config.num_labels)+\".pkl\")\n# print(\"GENERATION FINISHED!\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"GENERATE TARGETS FOR ENTRY IDS (\"+str(config.num_labels)+\" MOST COMMON GO TERMS)\")\nids = np.load(\"/kaggle/working/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)}\n\n# Initialize with zeros instead of empty\nlabels_matrix = np.zeros((len(ids), len(labels_names)))\n\nfor index, id in tqdm(enumerate(ids)):\n    # Use .get() to safely retrieve the list, default to empty list if not found\n    id_gos_list = id_labels.get(id, [])\n    \n    # Only process if the protein has GO terms\n    if id_gos_list:\n        temp = [go_terms_map[go] for go in id_gos_list if go in go_terms_map]\n        if temp:  # Only set values if there are matching GO terms\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!\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}