{"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":"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\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport time\nimport torch\nfrom transformers import T5Tokenizer, T5Model\nimport re\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":"","metadata":{"trusted":true},"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 ProtT5 Model...\")\n\n# PROT T5 LOADING:\ntokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_bfd', do_lower_case=False)\nmodel = T5Model.from_pretrained(\"Rostlab/prot_t5_xl_bfd\").to(config.device)\n\ndef get_t5_embedding(\n    sequence: str,\n    len_seq_limit: int\n):\n    '''\n    Function to collect encoder embedding vector from pre-trained ProtT5 Model\n\n    INPUTS:\n    - sequence (str): protein sequence (ex: AAABBB) from fasta file\n    - len_seq_limit (int): maximum sequence length (i.e nb of letters) for truncation\n\n    OUTPUTS:\n    - output_hidden: encoder embedding vector for input sequence (1024 dimensions)\n    '''\n    # Add spaces between amino acids and replace rare amino acids\n    sequence_w_spaces = ' '.join(list(sequence))\n    sequence_w_spaces = re.sub(r\"[UZOB]\", \"X\", sequence_w_spaces)\n    \n    # Tokenize\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'\n    ).to(config.device)\n    \n    # Get embeddings\n    # with torch.no_grad():\n    #     output = model(\n    #         input_ids=encoded_input['input_ids'],\n    #         attention_mask=encoded_input['attention_mask'],\n    #         decoder_input_ids=None\n    #     )\n\n    # Get embeddings - use encoder only\n    with torch.no_grad():\n        # For T5, we can use the encoder directly without decoder\n        encoder_output = model.encoder(\n            input_ids=encoded_input['input_ids'],\n            attention_mask=encoded_input['attention_mask'],\n            return_dict=True\n        )\n    \n    # Use encoder embedding (last hidden state)\n    encoder_embedding = encoder_output.last_hidden_state\n    \n    # Use encoder embedding (last hidden state)\n    # output.last_hidden_state gives the encoder output\n    # Take the mean over sequence length or use [0] token\n    # encoder_embedding = output.encoder_last_hidden_state\n    \n    # Option 1: Use first token (similar to ProtBERT approach)\n    output_hidden = encoder_embedding[:, 0][0].detach().cpu().numpy()\n    \n    # Option 2 (alternative): Use mean pooling over all tokens\n    # mask = encoded_input['attention_mask'].unsqueeze(-1).expand(encoder_embedding.size()).float()\n    # output_hidden = torch.sum(encoder_embedding * mask, dim=1) / torch.clamp(mask.sum(dim=1), min=1e-9)\n    # output_hidden = output_hidden[0].detach().cpu().numpy()\n    \n    assert len(output_hidden) == 1024\n    return output_hidden\n\n### COLLECTING FOR TRAIN SAMPLES:\nprint(\"Loading train set ProtT5 Embeddings...\")\nfasta_train = SeqIO.parse(config.train_sequences_path, \"fasta\")\nprint(\"Total Nb of Elements:\", len(list(fasta_train)))\n\nfasta_train = SeqIO.parse(config.train_sequences_path, \"fasta\")\nids_list = []\nembed_vects_list = []\nt0 = time.time()\ncheckpoint = 0\n\nfor item in tqdm(fasta_train):\n    ids_list.append(item.id)\n    embed_vects_list.append(\n        get_t5_embedding(sequence=str(item.seq), len_seq_limit=1200))\n    checkpoint += 1\n    if checkpoint >= 100:\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 ProtT5 Embeddings...\")\nfasta_test = SeqIO.parse(config.test_sequences_path, \"fasta\")\nprint(\"Total Nb of Elements:\", len(list(fasta_test)))\n\nfasta_test = SeqIO.parse(config.test_sequences_path, \"fasta\")\nids_list = []\nembed_vects_list = []\nt0 = time.time()\ncheckpoint = 0\n\nfor item in tqdm(fasta_test):\n    ids_list.append(item.id)\n    embed_vects_list.append(\n        get_t5_embedding(sequence=str(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 Elapsed 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}]}