{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"colab":{"provenance":[{"file_id":"1RASWrcVgKSjY-tPaK3lv0MbniVBfPuSb","timestamp":1769976313974},{"file_id":"1bZlJSEzIZmZinO0D0q6bDmVTf2hQzhVR","timestamp":1769311165235}],"authorship_tag":"ABX9TyMJ9fJvLT9B7E9ZseFatWV/"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":5499219,"sourceType":"datasetVersion","datasetId":3167603},{"sourceId":14207560,"sourceType":"datasetVersion","datasetId":9061916},{"sourceId":8576758,"sourceType":"datasetVersion","datasetId":4804930}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **CAFA6-MODEL-3-v8**","metadata":{"id":"wjfr_b2SLKUY"}},{"cell_type":"markdown","source":"# **Implementation of the CAFA6 Winning Architecture:**\n\n\n\n\n","metadata":{"id":"3whX6Pc5DVjC"}},{"cell_type":"markdown","source":"The primary differences between an advanced architecture and common development approaches used to resolve CAFA6 Challenge, lies in how the biological hierarchy is handled, the inclusion of structural data, and the specific optimization for the competition's scoring metric.\n\n### 1. Multi-Modal vs. Sequence-Only Features\n\nMost baseline projects utilize only protein language model (pLM) embeddings from **ESM2** or **Prot-T5**. A high level solution incorporates a **structural channel** using **Foldseek 3Di embeddings**, which treats 3D residue interactions as a separate \"language\", which is critical because protein structure is often more conserved than sequence across evolution, providing a stronger signal for function in remote homologs.\n\n### 2. Hierarchical Consistency (True Path Rule)\n\nInitially it is common to think the solution as a standard multi-label problem, predicting 1,500+ terms independently. However, the Gene Ontology (GO) is a Directed Acyclic Graph where the **True Path Rule** must hold: if a protein has a specific function, it must also possess all parent functions. The recommended architecture uses the `.obo` file to propagate scores leaf-to-root, ensuring that parent probabilities are always greater than or equal to their children, this will prevent biologically contradictory predictions.\n\n### 3. Loss Function and Evaluation Alignment\n\nIs common the use of **Binary Cross-Entropy (BCE)** loss. While effective for general classification, BCE does not prioritize the ** metric** or account for **Information Accretion (IA)** weights. The recommended solution uses a **Hybrid Loss** (combining **Focal Loss** to handle rare terms and **Soft-F1 Loss** to directly optimize the evaluation metric), which forces the model to focus on specific, high-value functional annotations rather than just common root terms.\n\n### 4. Attention-Based Fusion\n\nMany standard approaches are going to simply concatenate embeddings or use simple **Dense Neural Networks (DNNs)**. A better architecture employs a **Fusion Transformer** with **Multi-Head Self-Attention**. This allows the model to learn complex interdependencies between sequence features and structural motifs, effectively \"weighting\" which biological signals are most relevant for a specific subontology.\n\n### Summary Comparison Table\n\n| Feature | Standard Approach | Possible Winning Approach |\n| --- | --- | --- |\n| **Input Data** | Sequence Embeddings only | Sequence + Foldseek 3Di Structure |\n| **Model** | Logistic Regression or MLP | Fusion Transformer + Residual Blocks |\n| **Hierarchy** | Independent predictions | Mandatory True Path Rule enforcement |\n| **Optimization** | Binary Cross-Entropy (BCE) | Focal + Soft-F1 Hybrid Loss |\n| **Scalability** | Standard memory loading | HDF5 Lazy Loading for multi-TB features |\n\n\n","metadata":{"id":"79iQggZTJWMb"}},{"cell_type":"markdown","source":"# **Complete Python Solution for CAFA6**\n\nThis implementation includes:\n\n* HDF5 lazy-loading pipeline\n* Multi-Tower Fusion Model\n* IA-Weighted Hybrid Loss","metadata":{"id":"73poMt5LI38G"}},{"cell_type":"code","source":"# --------------------------------------------\n# Install whats not available\n# --------------------------------------------\n!pip install obonet\n\n\n","metadata":{"executionInfo":{"elapsed":42592,"status":"ok","timestamp":1769982904026,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"},"user_tz":480},"id":"yy_i0s_InnZj","outputId":"715914f5-731c-4e64-eae9-dc5169453467","trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:51.280861Z","iopub.execute_input":"2026-02-02T19:48:51.281155Z","iopub.status.idle":"2026-02-02T19:48:54.568582Z","shell.execute_reply.started":"2026-02-02T19:48:51.281128Z","shell.execute_reply":"2026-02-02T19:48:54.567793Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport h5py\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.metrics import f1_score\nfrom tqdm import tqdm\nimport networkx as nx\nimport joblib\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.model_selection import train_test_split\nimport obonet\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n#import tensorflow as tf\nimport glob\n\n# TODO: To comment when submit to challenge\n# GOOGLE COLLAB\n# This command will prompt you to authorize Colab\n# to access your Google Drive.\n#from google.colab import drive\n#drive.mount('/content/drive')\n\n# =============================================================================\n# 1. CONFIGURATION AND HYPERPARAMETERS\n#\n# KAGGLE: CAFA 6 Challenge\n# VERSION: 4\n# =============================================================================\n\n# Interface KEYS references\nPROTEIN = 'protein'\nTERM = 'term'\nONTOLOGY = 'ontology'\nTAXON = 'taxon'\n\n# Upper limit for label Quantity\nMAXLABEL_Q = 1500\n\nclass CAFA6Config:\n    # Model dimensions\n\n    # INDIVIDUAL EMBEDDINGS:\n    # ESM2_DIM(Embedding Dim or Hidden Size): A size of 2560 for ESM2_DIM\n    # 2560 size is ESM-2 3B (3B parameter). Compared to smaller variants\n    # like 320 or 1280) allow the model to learn and represent more complex and\n    # rich features of the protein sequence. This leads to better performance on\n    # downstream tasks, such as predicting protein structure or function.\n    ESM2_3B_DIM = 2560  # Dimension for ESM2-3B embeddings\n    ESM2_15B_DIM = 5120\n    PROTT5_DIM = 1024 # Dimension for Prot-T5 embeddings\n    ANKH_DIM = 1536\n    FUSION_DIM = 512\n\n    NUM_GO_TERMS = 1500 # Top frequent terms per subontology\n    LATENT_DIM = 512 # Shared space size for fusion\n    NUM_GO_TERMS = 1500 # Top frequent terms per subontology\n    TOTAL_LABELS = NUM_GO_TERMS * 3\n\n    # Training params\n    BATCH_SIZE = 64\n    LEARNING_RATE = 1e-4\n\n    # TODO: change to 50 again\n    EPOCHS = 50\n    DROPOUT = 0.3\n\n    # CUDA:\n    print(f\"CUDA available: {torch.cuda.is_available()}\")\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n    # Paths (Update for actual environment)\n    # ALERT!: Change on commit!\n    #CAFA6_PATH = \"/kaggle/input/cafa-6-protein-function-prediction\"\n    #CAFA6_TRAIN_PATH = f\"{CAFA6_PATH}/kaggle/input\"\n    #CAFA6_TEST_PATH = f\"{CAFA6_PATH}/kaggle/input\"\n    #EMBEDDING_PATH = f\"{CAFA6_PATH}/kaggle/input\"\n    #H5_PATH = f\"{CAFA6_PATH}/kaggle/input\"\n\n    #CAFA6_PATH = \"/content/drive/MyDrive/Colab Notebooks/project-research-kaggle/data\"\n    CAFA6_PATH = \"/kaggle/input/cafa-6-protein-function-prediction\"\n    #/kaggle/input/cafa-6-protein-function-prediction\n    CAFA6_TRAIN_PATH = f\"{CAFA6_PATH}/Train\"\n    CAFA6_TEST_PATH = f\"{CAFA6_PATH}/Test\"\n\n    #EMBEDDING_PATH = f\"{CAFA6_PATH}/T5Embeds\"\n    EMBEDDING_PATH = \"/kaggle/input/t5embeds\"\n\n    EMBEDDING_PATH2 = \"/kaggle/input/cafa-6-embeddings\"\n    \n    KAGGLE_PATH_OUT= \"/kaggle/working\"\n    \n    EMBEDS_H5_CONSOLIDATED = \"/consolidated_embeds.h5\"\n    EMBEDS_H5_TEST_CONSOLI = \"/consoli_test_embeds.h5\"\n    #H5_PATH = f\"{CAFA6_PATH}/h5\"\n    H5_PATH = KAGGLE_PATH_OUT\n\n    LABELS_PATH = \"train_terms.tsv\"  # TODO\n    OBO_PATH = \"go-basic.obo\"        # TODO\n    #OBO_PATH = os.path.join(BASE_PATH, \"Train/go-basic.obo\") #\n\n    # MULTI EMBEDDING STRATEGY:\n    # =========================\n    # This allows you to add/remove embeddings (e.g., adding Ankh or Foldseek)\n    # without breaking the model architecture.\n    #----------------------------------------------------------------------\n    # ATTENTION: ALWAYS CHECK THIS ORDER1. The Intent: Dimension Matching\n    # In PyTorch, a nn.Linear(input_dim, output_dim) layer is a fixed mathematical\n    # operation. If you define a tower for Prot-T5 with 1024 dimensions but\n    # accidentally pass it an ESM2-3B embedding (2560 dimensions), the matrix\n    # multiplication will fail immediately with a RuntimeError\n    # regarding mismatched sizes.\n    # Tower Definitions: {dataset_key: input_dimension}\n    # KAGGLE's CAFA6 \"Pro Move\": ESM-2 15B + ProtT5 + Ankh concatenated together.\n    # 5120 (ESM) + 1024 (T5) + 1536 (Ankh) = 7680 dim.\n    # This \"Ensemble Vector\" is the same size as one 15B vector, but it often\n    # performs better because it combines three different \"views\" of the\n    # protein (structure view + sequence view), rather than just one very deep view.\n    EMBEDDING_MAP = {\n         \"prott5\": 1024,\n           \"esm2\": 5120,\n           \"ankh\": 1536\n    }\n\n    # --- CONFIGURATION: MAP EMBED FILES HERE ---\n    # Map each model to its specific embedding file and ID file.\n    # i.e.: 'train_embeds.npy' has shape (142246, 1024), which matches 'prott5'\n    # and the index file 'train_ids.npy'wich has also 142246 rows\n    FILE_CONFIG = {\n        \"prott5\": {\n            \"embeds\": EMBEDDING_PATH + \"/train_embeds.npy\",\n            \"ids\":    EMBEDDING_PATH + \"/train_ids.npy\"\n            \n        }\n    ,\n       \"esm2\": { \n            \"embeds\": EMBEDDING_PATH2 + \"/esm2_15B/train_sequences_emb.npy\",\n            \"ids\":    EMBEDDING_PATH2 + \"/esm2_15B/train_sequences_ids.npy\"\n        },\n          \"ankh\": {\n            \"embeds\": EMBEDDING_PATH2 + \"/ankh_large/train_sequences_emb.npy\",\n               \"ids\": EMBEDDING_PATH2 + \"/ankh_large/train_sequences_ids.npy\"\n        }\n        # Add more, etc. here \n    }\n\n    # ESM-2 and Ankh are based on data from UniRef50:\n    # ESM-2 embeddings represent \"How does this amino acid fit into the\n    # evolutionary history of this family?\" (Great for Structure/Folding).\n    # -------------------------------------------------------------------------\n    # Ankh embeddings represent \"What is the functional meaning of this sequence\n    #  span?\" (Great for Function Prediction/Generation).\n    FILE_TEST_CONFIG = {\n        \"prott5\": {\n            \"embeds\": EMBEDDING_PATH + \"/test_embeds.npy\",\n            \"ids\":    EMBEDDING_PATH + \"/test_ids.npy\"\n        }\n        ,\n       \"esm2\": { \n            \"embeds\": EMBEDDING_PATH2 + \"/esm2_15B/testsuperset_emb.npy\",\n            \"ids\":    EMBEDDING_PATH2 + \"/esm2_15B/testsuperset_ids.npy\"\n        },\n          \"ankh\": { \n            \"embeds\": EMBEDDING_PATH2 + \"/ankh_large/testsuperset_emb.npy\",\n               \"ids\": EMBEDDING_PATH2 + \"/ankh_large/testsuperset_ids.npy\"\n        }   \n        # Add more, etc. here\n    }\n\n    # Miscellaneous:\n    SEPARATOR = \"====================================================\"","metadata":{"executionInfo":{"elapsed":1615,"status":"ok","timestamp":1769983376803,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"},"user_tz":480},"id":"9h5IePaoEKY5","outputId":"b521f7d0-8084-464d-83c7-caa7d7935874","trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:54.570617Z","iopub.execute_input":"2026-02-02T19:48:54.570988Z","iopub.status.idle":"2026-02-02T19:48:57.727581Z","shell.execute_reply.started":"2026-02-02T19:48:54.570945Z","shell.execute_reply":"2026-02-02T19:48:57.726930Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Data Loader:**  \n\n* Training set ~60-80%\n* Validation set for finetunning ~10-20% \t(not available)\n* Testing set  \t~10-20%","metadata":{"id":"kUxmYSe1RwMV"}},{"cell_type":"code","source":"# In this class is delegated the load of the data necessary for the ML\nclass PipelineDataLoader:\n    def __init__(self, input_path):\n        self.input_path = input_path\n\n    def loadData(self):\n\n        # File Format\" pandas.read_csv, which expects a plain text file encoded in UTF-8.\n        print(\"LOADING file: testsuperset-taxon-list.tsv\")\n        self.test = pd.read_csv(f\"{self.input_path}/Test/testsuperset-taxon-list.tsv\",sep=\"\\t\",\n                                 names=[PROTEIN,TAXON] )\n        print(\"Test terms testsuperset-taxon-list.tsv dim:\", self.test.shape)\n        print(self.test.head(3))\n        print(CAFA6Config.SEPARATOR)\n\n        # T5_embeds: ML training its only possible using numerical models, therefore\n        # aminoacid names sequences in the file train_sequences.fasta must be converted to\n        # numerical and vector format.\n        # EMBEDDINGS OF PROTEIN SEQUENCES: will be used in this notebook, embeddings are similar\n        # to word embeddings used to train NLP models. These P.Embeddings are ML methods to capture\n        # protein STRUCTURAL and FUNCTIONAL characteristics from it's sequence.\n        # There are publicly available pre-trained protein embedding models, given that PROTEINS\n        # are represented through its aminoacids,sequences,so the language is well known, also\n        # is possible to train a model to learn those protein embeddings, but has no sense redoing\n        # the wheel in something that is quite standard.\n\n        # T5_embeds\n        # File Format: numpy binary data\n        print(\"LOADING file: train_embeds.npy\")\n        self.trainEmbeddings = np.load(f\"{CAFA6Config.EMBEDDING_PATH}/train_embeds.npy\")\n        print(\"train_embeds.npy T5Embeds dim:\", self.trainEmbeddings.shape)\n        print(self.trainEmbeddings[:6])\n        print(CAFA6Config.SEPARATOR)\n\n        # Convert self.trainEmbeddings (T5_embeds) into a pandas dataframe:\n        print(\"CONVERTING object: T5Embeds as DF\")\n        colNum = self.trainEmbeddings.shape[1]\n        self.trainEmbeddingsDF = pd.DataFrame(self.trainEmbeddings,\n                                         columns=[\"Column_\"+str(i) for i in range(1,colNum+1)])\n        print(\"T5Embeds as DF dim:\", self.trainEmbeddingsDF.shape)\n        print(self.trainEmbeddingsDF.head(3))\n        print(CAFA6Config.SEPARATOR)\n\n        # TEST T5_embeds:\n        # File Format: numpy binary data\n        print(\"LOADING file: test_embeds.npy\")\n        self.testEmbeddings = np.load(f\"{CAFA6Config.EMBEDDING_PATH}/test_embeds.npy\")\n        print(\"test_embeds.npy T5Embeds  dim:\", self.testEmbeddings.shape)\n        print(self.testEmbeddings[:6])\n        print(CAFA6Config.SEPARATOR)\n\n        # File Format: numpy binary data\n        print(\"LOADING file: test_ids.npy\")\n        self.testProtIds = np.load(f\"{CAFA6Config.EMBEDDING_PATH}/test_ids.npy\")\n        print(\"test_ids.npy T5Embeds Protein IDs dim:\", self.testProtIds.shape)\n        print(self.testProtIds[:6])\n        print(CAFA6Config.SEPARATOR)\n\n\n    #----------------------------------------------------------------------------------------#\n\n    # //TODO: creation of the h5 subdirectory dinamically when the process is executed\n    def exportNPYEmbeddsToH5(self):\n        # H5 Conversions:\n        convert_npy_to_h5(f\"{self.input_path}/T5Embeds/train_embeds.npy\",\n                          f\"{self.input_path}/h5/train_embeds.h5\", \"prott5\")\n        convert_npy_to_h5(f\"{self.input_path}/T5Embeds/test_embeds.npy\",\n                          f\"{self.input_path}/h5/test_embeds.h5\", \"prott5\")\n        #convert_npy_to_h5(f\"{self.input_path}/T5Embeds/test_ids.npy\",\n        #                  f\"{self.input_path}/h5/test_ids.h5\", \"prott5\")\n        #convert_npy_to_h5(f\"{self.input_path}/T5Embeds/train_ids.npy\",\n        #                  f\"{self.input_path}/h5/train_ids.h5\", \"prott5\")\n\n\n    #----------------------------------------------------------------------------------------#\n\n    def load_GO_Onto(self):\n        # Ontology graph\n        # File Format:\n        print(\"LOADING file: go-basic.obo\")\n        self.go_graph_, self.term_to_aspect_ = load_go_ontology(f\"{self.input_path}/Train/go-basic.obo\")\n        print(\"go-basic.obo file successfully Loaded\")\n        print(CAFA6Config.SEPARATOR)\n\n        return self.go_graph_, self.term_to_aspect_\n\n    #----------------------------------------------------------------------------------------#\n\n    def load_train_terms_tsv(self):\n        # train_terms.tsv : contains a list of annotated terms(functions) for the proteins.\n        # Labels aka GO term ID are extracted to create a label df for the protein embeddings.\n        # File Format\" pandas.read_csv, which expects a plain text file encoded in UTF-8.\n        print(\"LOADING file: train_terms.tsv \")\n        self.train = pd.read_csv(f\"{self.input_path}/Train/train_terms.tsv\",sep=\"\\t\",\n                                 names=[PROTEIN,TERM, ONTOLOGY] )\n        print(\"Train terms train_terms.tsv processed dim:\", self.train.shape)\n        print(self.train.head(3))\n        print(CAFA6Config.SEPARATOR)\n\n        return self.train\n\n    #----------------------------------------------------------------------------------------#\n\n    def load_train_ids_npy(self):\n        # TRAIN T5_embeds:\n        # File Format: numpy binary data\n        print(\"LOADING file: train_ids.npy\")\n        self.trainProtIds = np.load(f\"{CAFA6Config.EMBEDDING_PATH}/train_ids.npy\")\n        print(\"Train Protein train_ids.npy IDs dim:\", self.trainProtIds.shape)\n        print(self.trainProtIds[:6])\n        print(CAFA6Config.SEPARATOR)\n\n        return self.trainProtIds\n\n    #----------------------------------------------------------------------------------------#\n\n        # IA.tsv - Information Accretion** (IA).\n        # **Accretion is the gradual, natural increase, growth, or accumulation of something by external addition.\n        # Information Accretion and IA-Weighted Strategy\n        # The CAFA evaluation specifically weights precision and recall by Information Accretion (IA).\n        # This implies that a model that is exceptionally good at predicting general root terms but\n        # poor at predicting specific leaf terms will score lower than a model that achieves modest\n        # performance across the entire hierarchy.\n\n        # Evaluation Metric: The leaderboard uses weighted scores if an IA file is available,\n        # so an ultimate goal with these losses is to achieve high precision on terms\n        # with high Information Accretion.\n    def load_IA_tsv(self):\n        print(\"LOADING file: IA.tsv \")\n        self.ia = pd.read_csv(f\"{self.input_path}/IA.tsv\",sep=\"\\t\")\n        self.ia.columns = ['term_id', 'ia_score'] # Assign correct column names\n        print(\"IA.tsv processed dim:\", self.ia.shape)\n        print(self.ia.head(3))\n        print(CAFA6Config.SEPARATOR)\n\n        return self.ia\n\n\n","metadata":{"id":"DhnuZqzQCvYZ","executionInfo":{"status":"ok","timestamp":1769982941568,"user_tz":480,"elapsed":55,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.728536Z","iopub.execute_input":"2026-02-02T19:48:57.729015Z","iopub.status.idle":"2026-02-02T19:48:57.742824Z","shell.execute_reply.started":"2026-02-02T19:48:57.728980Z","shell.execute_reply":"2026-02-02T19:48:57.742016Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## **.h5 File Conversions**\n\nThe H5 file is structured like a Dictionary.\nHow the H5 is built: file[\"prott5\"] is a Group. Inside it, the data is stored by Name (Protein ID), like file[\"prott5\"][\"P12345\"].","metadata":{"id":"KkSMWii6CIhB"}},{"cell_type":"code","source":"# =============================================================================\n# .NPY embedding file to HDF5 with compression.\n#\n# Example usage:\n# Run the function\n# convert_parallel_npy_to_h5(FILE_CONFIG, OUTPUT_H5)\n# =============================================================================\ndef convert_parallel_npy_to_h5(file_config, h5_path):\n\n    with h5py.File(h5_path, \"w\") as f:\n\n        for model_name, paths in file_config.items():\n            print(f\"--- Processing {model_name} ---\")\n\n            embed_path = paths[\"embeds\"]\n            id_path = paths[\"ids\"]\n\n            if not os.path.exists(embed_path) or not os.path.exists(id_path):\n                print(f\"Skipping {model_name}: One of the files is missing.\")\n                continue\n\n            # 1. Load the Arrays\n            # mmap_mode='r' is great for huge files; it reads from disk without loading everything to RAM at once\n            embeds = np.load(embed_path, mmap_mode='r')\n            ids = np.load(id_path)\n\n            # 2. Validation\n            if embeds.shape[0] != ids.shape[0]:\n                print(f\"CRITICAL ERROR for {model_name}: Row counts do not match!\")\n                print(f\"Embeds: {embeds.shape[0]}, IDs: {ids.shape[0]}\")\n                continue\n\n            # 3. Create H5 Group\n            group = f.create_group(model_name)\n\n            print(f\"Saving {embeds.shape[0]} proteins...\")\n\n            # 4. Iterate and Save\n            for i, pid in enumerate(ids):\n                # Clean the ID: Numpy often stores strings as bytes (b'P12345')\n                # We need to decode them to normal strings for H5 names\n                if isinstance(pid, bytes):\n                    pid_str = pid.decode('utf-8')\n                else:\n                    pid_str = str(pid)\n\n                # Get the vector (row i)\n                vector = embeds[i]\n\n                # Save to H5\n                # Note: H5 requires unique keys per group.\n                group.create_dataset(pid_str, data=vector)\n\n                # Optional: Progress tracker\n                if i % 10000 == 0:\n                    print(f\"  Processed {i} / {len(ids)}\")\n\n    print(\"\\nDone! Verification:\")\n    with h5py.File(h5_path, \"r\") as f:\n        for k in f.keys():\n            print(f\" - Group: {k}, Count: {len(f[k])}\")\n\n#-------------------------------------------------------------------\n# In Kaggle, your output files are usually in the current working\n# directory or \"/kaggle/working/\" \".\" represents the current directory\n# target_dir = \"/kaggle/working/\" # Use this if you want to be explicit\n#-------------------------------------------------------------------\ndef build_test_h5(file_config, h5_path):\n    print(f\"BUILDING TEST H5: {h5_path} ...\")\n    \n    # Always start fresh\n    if os.path.exists(h5_path):\n        try:\n            os.remove(h5_path)\n        except:\n            pass\n\n    with h5py.File(h5_path, \"w\") as f:\n        for model_name, paths in file_config.items():\n            if not os.path.exists(paths[\"embeds\"]):\n                print(f\"  [SKIP] Missing file: {paths['embeds']}\")\n                continue\n                \n            print(f\"  + Processing {model_name}...\")\n            embeds = np.load(paths[\"embeds\"], mmap_mode='r')\n            ids = np.load(paths[\"ids\"])\n            \n            group = f.create_group(model_name)\n            \n            # --- FIX: Track seen IDs to handle duplicates ---\n            seen_ids = set()\n            \n            for i, pid in enumerate(tqdm(ids, desc=f\"Writing {model_name}\")):\n                if isinstance(pid, bytes): pid_str = pid.decode('utf-8')\n                else: pid_str = str(pid)\n                \n                # CRITICAL FIX: Skip if we already wrote this protein\n                if pid_str in seen_ids:\n                    continue\n                \n                group.create_dataset(pid_str, data=embeds[i])\n                seen_ids.add(pid_str)\n                \n    print(\"SUCCESS: Test H5 Created.\\n\")\ndef build_test_h5_2(file_config, h5_path):\n    \"\"\" Packs .npy files into a structured .h5 file for fast random access. \"\"\"\n    print(f\"BUILDING TEST H5: {h5_path} ...\")\n\n    # If file exists, we can skip or delete.\n    # Given time constraints, if it exists, assume it's good to save time?\n    # No, better to be safe. Delete and rebuild to ensure it matches config.\n    if os.path.exists(h5_path):\n        try:\n            os.remove(h5_path)\n        except:\n            pass\n\n    with h5py.File(h5_path, \"w\") as f:\n        for model_name, paths in file_config.items():\n            embed_path = paths[\"embeds\"]\n            id_path = paths[\"ids\"]\n\n            # Skip if file missing\n            if not os.path.exists(embed_path):\n                print(f\"  [SKIP] Missing file: {embed_path}\")\n                continue\n\n            print(f\"  + Processing {model_name}...\")\n\n            # Load with mmap to save RAM\n            embeds = np.load(embed_path, mmap_mode='r')\n            ids = np.load(id_path)\n\n            group = f.create_group(model_name)\n\n            # Save entries\n            for i, pid in enumerate(tqdm(ids, desc=f\"Writing {model_name}\")):\n                if isinstance(pid, bytes): pid_str = pid.decode('utf-8')\n                else: pid_str = str(pid)\n\n                group.create_dataset(pid_str, data=embeds[i])\n    print(\"SUCCESS: Test H5 Created.\\n\")\n\n#-------------------------------------------------------------------\n# In Kaggle, your output files are usually in the current working\n# directory or \"/kaggle/working/\" \".\" represents the current directory\n# target_dir = \"/kaggle/working/\" # Use this if you want to be explicit\n#-------------------------------------------------------------------\ndef delete_h5_dirs_files(h5FilesPath):\n    target_dir = \".\"\n\n    # 1. Find all .h5 files\n    files_to_delete = glob.glob(os.path.join(target_dir, \"*.h5\"))\n\n    if not files_to_delete:\n        print(\"No .h5 files found to delete.\")\n    else:\n        print(f\"Found {len(files_to_delete)} files: {files_to_delete}\")\n\n        # 2. Loop and delete\n        for file_path in files_to_delete:\n            try:\n                os.remove(file_path)\n                print(f\"Deleted: {file_path}\")\n            except OSError as e:\n                print(f\"Error deleting {file_path}: {e}\")\n\n        print(\"Cleanup complete.\")\n\n\n# =============================================================================\n# .NPY embedding file to HDF5 with compression.\n#\n# Example usage:\n# convert_npy_to_h5(\"train_embeds.npy\", \"embeddings.h5\", \"prott5\")\n# =============================================================================\ndef convert_npy_to_h5(npy_path, h5_path, dataset_name): #=\"prott5\"\n    \"\"\"\n    Converts a massive.npy embedding file to HDF5 with compression.\n    \"\"\"\n    print(f\"Converting {npy_path} to {h5_path}...\")\n    data = np.load(npy_path)\n    with h5py.File(h5_path, 'w') as f:\n        # chunked=True is critical for fast random access in deep learning\n        f.create_dataset(dataset_name, data=data, compression=\"gzip\", chunks=True)\n    print(\"Conversion complete.\")\n\n\n# TODO: check this path out_dir = '/kaggle/working'\n# np.save(os.path.join(out_dir, 'protein_embeddings_train.npy'), train_embeddings)\n\n","metadata":{"id":"hVkLD4ytCJtw","executionInfo":{"status":"ok","timestamp":1769982950288,"user_tz":480,"elapsed":31,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.743776Z","iopub.execute_input":"2026-02-02T19:48:57.743996Z","iopub.status.idle":"2026-02-02T19:48:57.762838Z","shell.execute_reply.started":"2026-02-02T19:48:57.743975Z","shell.execute_reply":"2026-02-02T19:48:57.762064Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## **OBO Processing**","metadata":{"id":"65PE8Sjh9nh9"}},{"cell_type":"code","source":"# =============================================================================\n#  ONTOLOGY PROCESSING &.OBO LOADING\n# =============================================================================\ndef load_go_ontology(obo_path):\n    \"\"\"\n    Parses the.obo file and returns a NetworkX graph with subontology mapping.\n    \"\"\"\n    print(f\"Loading GO ontology from {obo_path}...\")\n    # Load into a directed acyclic graph\n    graph = obonet.read_obo(obo_path)\n    print(f\"✓ Loaded {len(graph)} nodes and {len(graph.edges)} edges.\")\n\n    # Extract subontology (aspect) for each term\n    term_to_aspect = {}\n    for node, data in graph.nodes(data=True):\n        ns = data.get('namespace', '')\n        if ns == 'biological_process':\n            term_to_aspect[node] = 'BPO'\n        elif ns == 'molecular_function':\n            term_to_aspect[node] = 'MFO'\n        elif ns == 'cellular_component':\n            term_to_aspect[node] = 'CCO'\n\n    return graph, term_to_aspect\n\ndef get_ancestors_of_term(graph, term_id):\n    \"\"\"\n    Retrieves all unique ancestors of a GO term in the DAG.\n    \"\"\"\n    if term_id not in graph:\n        return set()\n    return nx.ancestors(graph, term_id)\n","metadata":{"id":"8atSqyhb9uCb","executionInfo":{"status":"ok","timestamp":1769982963239,"user_tz":480,"elapsed":69,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.764810Z","iopub.execute_input":"2026-02-02T19:48:57.765190Z","iopub.status.idle":"2026-02-02T19:48:57.779093Z","shell.execute_reply.started":"2026-02-02T19:48:57.765161Z","shell.execute_reply":"2026-02-02T19:48:57.778391Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Custom Dataset**","metadata":{"id":"Y-Y7uxuBemmb"}},{"cell_type":"code","source":"# =============================================================================\n# MULTI-MODAL HDF5 DATA PIPELINE\n# =============================================================================\nclass MultiProteinDataset(Dataset):\n    def __init__(self, h5_path, protein_ids, keys, labels):\n        self.h5_path = h5_path\n        self.ids = protein_ids\n        self.keys = keys\n        self.labels = labels\n        self.h5_file = None\n        self.shapes = {} # Cache shapes for fallback\n\n    def __getitem__(self, idx):\n        if self.h5_file is None:\n            self.h5_file = h5py.File(self.h5_path, 'r')\n\n        # Ensure PID is a string\n        pid = self.ids[idx]\n        if isinstance(pid, bytes):\n            pid = pid.decode('utf-8')\n        else:\n            pid = str(pid)\n\n        embeddings = []\n        for k in self.keys:\n            try:\n                # TRY to get the data\n                data = self.h5_file[k][pid][()]\n                embeddings.append(torch.from_numpy(data).float())\n                \n                # Cache the shape for safety later\n                if k not in self.shapes:\n                    self.shapes[k] = data.shape\n                    \n            except KeyError:\n                # FALLBACK: If protein is missing, return ZERO VECTOR\n                # Use cached shape if available, otherwise find it\n                if k not in self.shapes:\n                    # Grab valid shape from the first key in the group\n                    first_key = next(iter(self.h5_file[k].keys()))\n                    self.shapes[k] = self.h5_file[k][first_key].shape\n                \n                # Return zeros\n                embeddings.append(torch.zeros(self.shapes[k]).float())\n\n        label = torch.tensor(self.labels[idx], dtype=torch.float32)\n        return embeddings, label\n\n    def __len__(self):\n        return len(self.ids)\n\nclass MultiProteinDataset2(Dataset):\n    \"\"\"\n    Lazy-loads multiple embeddings for a single protein from an HDF5 store.\n    \"\"\"\n    def __init__(self, h5_path, ids, embedding_keys, labels=None):\n        self.h5_file = h5py.File(h5_path, 'r')\n        self.ids = ids\n        self.keys = embedding_keys\n        self.labels = labels\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, idx):\n        # 1. Get the actual Protein ID string (e.g., \"P12345\") using the index\n        pid = self.ids[idx]\n\n        # 2. Use the STRING (pid), not the INT (idx), to access the H5 file\n        embeddings = []\n        for k in self.keys:\n            # We access the group self.h5_file[k], and then the dataset [pid]\n            # NOTICE: We use [pid], not [idx]\n            data_numpy = self.h5_file[k][pid][()]\n            embeddings.append(torch.from_numpy(data_numpy).float())\n\n        label = torch.tensor(self.labels[idx], dtype=torch.float32)\n\n        return embeddings, label\n\n\n","metadata":{"id":"QdUj7vMdsLkh","executionInfo":{"status":"ok","timestamp":1769982970291,"user_tz":480,"elapsed":34,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.780016Z","iopub.execute_input":"2026-02-02T19:48:57.780314Z","iopub.status.idle":"2026-02-02T19:48:57.793086Z","shell.execute_reply.started":"2026-02-02T19:48:57.780280Z","shell.execute_reply":"2026-02-02T19:48:57.792448Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **ARCHITECTURE: Fusion Transformer with residual blocks**","metadata":{"id":"FTIwzxrrdm4e"}},{"cell_type":"code","source":"# =============================================================================\n# 3. ARCHITECTURE: FUSION TRANSFORMER WITH RESIDUAL BLOCKS\n# =============================================================================\n# The ResidualDenseBlock is a building block commonly used in deep neural\n# networks, particularly designed to facilitate the training of very deep\n# models. Its core idea is to allow information to bypass some layers, which\n# helps in mitigating issues like vanishing gradients and enables the network\n# to learn more complex representations.\n#\n# import torch.nn as nn\nclass ResidualDenseBlock(nn.Module):\n    \"\"\"\n    Residual connection block for deep feature extraction.\n    Constructor, where Neural Network layers and operations\n    within the block are defined.\n    \"\"\"\n    def __init__(self, dim, dropout=0.3):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(dim, dim * 2),\n            nn.BatchNorm1d(dim * 2),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(dim * 2, dim)\n        )\n        self.norm = nn.LayerNorm(dim)\n\n    \"\"\"\n        This method defines how data (x) flows through\n        the block during a forward pass.\n    \"\"\"\n    def forward(self, x):\n        return self.norm(x + self.net(x))\n\n\n# =============================================================================\n# DEPRECATED //  class CAFA6FusionModel(nn.Module):\n# =============================================================================\nclass CAFA6FusionModel(nn.Module):\n    \"\"\"\n    A Winning model architecture for CAFA6,\n    must use attention to fuse different embedding modalities.\n    \"\"\"\n    def __init__(self, config):\n        super().__init__()\n        # In this implementation, we assume single input for brevity\n        input_dim = config.ESM2_DIM + config.PROTT5_DIM\n\n        # Shared feature extractor\n        self.input_layer = nn.Linear(input_dim, config.FUSION_DIM)\n\n        # Multi-head self-attention to capture latent biological signals\n        self.attention = nn.MultiheadAttention(config.FUSION_DIM, num_heads=8, batch_first=True)\n\n        # Deep residual stack\n        self.residual_stack = nn.Sequential(\n            ResidualDenseBlock(config.FUSION_DIM, config.DROPOUT),\n            ResidualDenseBlock(config.FUSION_DIM, config.DROPOUT),\n            ResidualDenseBlock(config.FUSION_DIM, config.DROPOUT)\n        )\n\n        # Subontology-specific classification heads\n        self.mf_head = nn.Linear(config.FUSION_DIM, config.NUM_GO_TERMS)\n        self.bp_head = nn.Linear(config.FUSION_DIM, config.NUM_GO_TERMS)\n        self.cc_head = nn.Linear(config.FUSION_DIM, config.NUM_GO_TERMS)\n        # NOTE: nn.Sigmoid() removed here to facilitate use of\n        # binary_cross_entropy_with_logits check class HybridLoss.forward()\n        #self.sigmoid = nn.Sigmoid()\n\n    def forward(self, x):\n        # x:\n        x = F.gelu(self.input_layer(x))\n\n        # Attention block (treat features as a sequence of descriptors)\n        x_att = x.unsqueeze(1)\n        x_att, _ = self.attention(x_att, x_att, x_att)\n        x = x_att.squeeze(1)\n\n        # Backbone processing\n        features = self.residual_stack(x)\n\n        # Parallel prediction for subontologies\n        #mf_out = self.sigmoid(self.mf_head(features))\n        #bp_out = self.sigmoid(self.bp_head(features))\n        #cc_out = self.sigmoid(self.cc_head(features))\n\n        #return torch.cat([mf_out, bp_out, cc_out], dim=1)\n\n        # Returns raw logits for numerical stability during training\n        return torch.cat([\n            self.mf_head(features),\n            self.bp_head(features),\n            self.cc_head(features)\n        ], dim=1)\n","metadata":{"id":"f-Zakr7EcHXC","executionInfo":{"status":"ok","timestamp":1769982976262,"user_tz":480,"elapsed":6,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.793900Z","iopub.execute_input":"2026-02-02T19:48:57.794107Z","iopub.status.idle":"2026-02-02T19:48:57.808057Z","shell.execute_reply.started":"2026-02-02T19:48:57.794089Z","shell.execute_reply":"2026-02-02T19:48:57.807369Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **ARCHITECTURE: Multitower Model Approach**","metadata":{"id":"aaVW7lkpjJDP"}},{"cell_type":"code","source":"class CAFA6MultiTowerModel(nn.Module):\n    \"\"\"\n    Modular Winning Architecture:\n    1. Independent Projection Towers for each embedding type.\n    2. Multi-Head Attention to learn cross-modal interactions.\n    3. Parallel Subontology Heads.[11]\n    \"\"\"\n    def __init__(self, config):\n        super().__init__()\n        # Maps varied input sizes (1024, 2560, etc.) to a fixed shared size.\n        self.keys = list(config.EMBEDDING_MAP.keys())\n\n        # 1. Individual Projection Towers:\n        # Automatic Initialization: By using for k, dim in\n        # config.EMBEDDING_MAP.items(), your model automatically builds a\n        # tower that matches whatever size you specify in the config.\n        self.towers = nn.ModuleDict({\n            k: nn.Sequential(\n                nn.Linear(dim, config.LATENT_DIM),\n                nn.LayerNorm(config.LATENT_DIM),\n                nn.GELU()\n            ) for k, dim in config.EMBEDDING_MAP.items()\n        })\n\n        # 2. Attention Fusion Layer\n        self.attention = nn.MultiheadAttention(\n            config.LATENT_DIM, num_heads=8, batch_first=True\n        )\n\n        # 3. Residual Backbone\n        self.backbone = nn.Sequential(\n            ResidualDenseBlock(config.LATENT_DIM),\n            ResidualDenseBlock(config.LATENT_DIM)\n        )\n\n        # 4. Multi-Label Classifier\n        self.mf_head = nn.Linear(config.LATENT_DIM, config.NUM_GO_TERMS)\n        self.bp_head = nn.Linear(config.LATENT_DIM, config.NUM_GO_TERMS)\n        self.cc_head = nn.Linear(config.LATENT_DIM, config.NUM_GO_TERMS)\n\n    # Forward\n    def forward(self, embedding_list):\n        # Apply projection towers\n        projected = []\n        for i, emb in enumerate(embedding_list):\n            projected.append(self.towers[self.keys[i]](emb))\n\n        # Stack towers for attention (Batch, Num_Towers, Latent_Dim)\n        x = torch.stack(projected, dim=1)\n\n        # Cross-modal fusion\n        attn_out, _ = self.attention(x, x, x)\n\n        # Global pooling (Average tower contributions)\n        x = attn_out.mean(dim=1)\n\n        features = self.backbone(x)\n\n        return torch.cat([\n            self.mf_head(features),\n            self.bp_head(features),\n            self.cc_head(features)\n        ], dim=1)","metadata":{"id":"gqrH9yjyjdor","executionInfo":{"status":"ok","timestamp":1769982984948,"user_tz":480,"elapsed":28,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.809115Z","iopub.execute_input":"2026-02-02T19:48:57.809793Z","iopub.status.idle":"2026-02-02T19:48:57.823304Z","shell.execute_reply.started":"2026-02-02T19:48:57.809758Z","shell.execute_reply":"2026-02-02T19:48:57.822765Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Implementation of IA-WEIGHTED HYBRID LOSS (STABLE LOGITS) instead of the simpler but not INFORMATION ACCRETION ready \"HybridLoss\"**","metadata":{"id":"iXakkNURPP1g"}},{"cell_type":"code","source":"# =============================================================================\n#  IA-WEIGHTED HYBRID LOSS (STABLE LOGITS)\n# =============================================================================\n\nclass IAWeightedHybridLoss(nn.Module):\n    \"\"\"\n    Winning Loss: Focal Loss (Imbalance) + IA-Weighted Soft-F1.\n    Addresses class imbalance and optimizes the leaderboard metric simultaneously.\n    \"\"\"\n    def __init__(self, ia_weights, gamma=2.0, alpha=0.25):\n        super().__init__()\n        # Ensure weights are on the correct device\n        # self.ia_weights = ia_weights.to(CAFA6Config.DEVICE)\n        # register_buffer so weights follow the model to any device\n        self.register_buffer('ia_weights', ia_weights)\n        self.gamma = gamma\n        self.alpha = alpha\n\n    def forward(self, logits, targets):\n        # 1. Stable Focal Loss\n        bce = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')\n        p_t = torch.exp(-bce)\n        focal = (self.alpha * (1 - p_t)**self.gamma * bce).mean()\n\n        # 2. IA-Weighted Soft-F1 Loss\n        preds = torch.sigmoid(logits)\n        tp = (preds * targets).sum(dim=0)\n        fp = (preds * (1 - targets)).sum(dim=0)\n        fn = ((1 - preds) * targets).sum(dim=0)\n\n        # Calculate F1 per term\n        f1_per_term = (2 * tp + 1e-15) / (2 * tp + fp + fn + 1e-15)\n\n        # Weight the average by Information Accretion\n        weighted_soft_f1 = (f1_per_term * self.ia_weights).sum() / self.ia_weights.sum()\n        f1_loss = 1 - weighted_soft_f1\n\n        return 0.6 * focal + 0.4 * f1_loss","metadata":{"id":"p2byL20aPRQd","executionInfo":{"status":"ok","timestamp":1769982991493,"user_tz":480,"elapsed":6,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.824128Z","iopub.execute_input":"2026-02-02T19:48:57.824373Z","iopub.status.idle":"2026-02-02T19:48:57.838266Z","shell.execute_reply.started":"2026-02-02T19:48:57.824353Z","shell.execute_reply":"2026-02-02T19:48:57.837538Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **LOSS and HIERARCHICAL TRUE PATH RULE:**\n\n### **Loss:** **Numerical Stability: Logits vs. Probabilities:**\n\nUsing **torch.nn.functional as F** allows HybridLoss to implement:\n\nA. **F.binary_cross_entropy**\n\nB.  **F.binary_cross_entropy_with_logits**\n\n**A.  binary_cross_entropy:** The model must output\nprobabilities (via a Sigmoid layer) before passing them to the **criterion**.\n\n**B.  If using with_logits:** Your model's final layers should not have a Sigmoid activation. This is the recommended approach for numerical stability to prevent underflow when calculating Focal Loss.\n","metadata":{"id":"-T8wMwgLeOPu"}},{"cell_type":"code","source":"# =============================================================================\n# 5. POST-PROCESSING: HIERARCHICAL TRUE PATH RULE\n#    LOSS FUNCTIONS & POST-PROCESSING\n# =============================================================================\n\ndef apply_true_path_rule(preds, go_graph, go_terms):\n    \"\"\"\n    Propagates scores leaf-to-root via the.obo graph\n    to ensure P(parent) >= P(child).\n    \"\"\"\n    term_to_idx = {term: i for i, term in enumerate(go_terms)}\n    consistent_preds = preds.copy()\n\n    # // TODO : CHECK THIS, about arbitrary root\n    # Sort terms by distance from root (deepest first)\n    # Using an arbitrary root like 'GO:0003674' for Molecular Function\n    sorted_terms = sorted(\n        [t for t in go_terms if t in go_graph],\n        key=lambda x: nx.shortest_path_length(go_graph, source=x,\n                                              target='GO:0003674')\n        if 'GO:0003674' in nx.ancestors(go_graph, x) else 0,\n        reverse=True\n    )\n\n    for term in sorted_terms:\n        idx = term_to_idx[term]\n        parents = list(go_graph.successors(term)) # Edges child -> parent in obonet\n        for p in parents:\n            if p in term_to_idx:\n                p_idx = term_to_idx[p]\n                consistent_preds[:, p_idx] = np.maximum(consistent_preds[:, p_idx],\n                                                        consistent_preds[:, idx])\n\n    return consistent_preds\n\n\n\n","metadata":{"id":"8cCzt3CSeO75","executionInfo":{"status":"ok","timestamp":1769982998655,"user_tz":480,"elapsed":10,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.839167Z","iopub.execute_input":"2026-02-02T19:48:57.839473Z","iopub.status.idle":"2026-02-02T19:48:57.853720Z","shell.execute_reply.started":"2026-02-02T19:48:57.839431Z","shell.execute_reply":"2026-02-02T19:48:57.852989Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## **Inference and Submission Builder:**","metadata":{"id":"-Iq0BUQ7w3KU"}},{"cell_type":"code","source":"def run_inference_h5(model, h5_path, test_ids, target_terms, model_keys, go_graph):\n    \"\"\"\n    Runs inference reading inputs from H5.\n    model_keys: The list of keys the MODEL expects (e.g. ['prott5'])\n    Applies True Path Rule (Consistency Enforcer).\n    \"\"\"\n    print(f\"INFERENCE: Starting prediction on {len(test_ids)} proteins...\")\n    print(f\"INFERENCE: Using modalities -> {model_keys}\")\n\n    model.eval()\n    h5_file = h5py.File(h5_path, 'r')\n\n    predictions = []\n    batch_size = 1000\n\n    with torch.no_grad():\n        for i in tqdm(range(0, len(test_ids), batch_size)):\n            batch_ids = test_ids[i : i+batch_size]\n\n            # 1. Collect inputs for this batch\n            batch_inputs = []\n\n            # For every modality the model was trained on...\n            for k in model_keys:\n                vectors = []\n                for pid in batch_ids:\n                    if isinstance(pid, bytes): pid = pid.decode('utf-8')\n                    else: pid = str(pid)\n\n                    # Robustness: If key missing in H5 (e.g. protein not in ESM2 list), use zeros?\n                    # For now, assume data integrity.\n                    try:\n                        vectors.append(h5_file[k][pid][()])\n                    except KeyError:\n                        # Fallback for missing ID in specific modality (Rare but possible)\n                        # Create zero vector of correct shape\n                        dim = h5_file[k][list(h5_file[k].keys())[0]].shape[0]\n                        vectors.append(np.zeros((dim,), dtype=np.float32))\n\n                tensor_data = torch.tensor(np.array(vectors)).float().to(conf.DEVICE)\n                batch_inputs.append(tensor_data)\n\n            # 2. Forward Pass\n            # If the model is MultiTower, it expects a list.\n            # If it's SingleTower, it might expect a Tensor.\n            # Your CAFA6MultiTowerModel expects a LIST.\n            logits = model(batch_inputs)\n            probs = torch.sigmoid(logits).cpu().numpy()\n            predictions.append(probs)\n\n    h5_file.close()\n\n    # 3. Save\n    print(\"INFERENCE: Formatting submission.tsv...\")\n    all_probs = np.vstack(predictions)\n\n    # 2. APPLY TRUE PATH RULE (The Fix!)\n    print(\"POST-PROCESSING: Enforcing True Path Rule (Leaf-to-Root Propagation)...\")\n    print(\"This enforces: Score(Parent) >= Score(Child). This creates biological consistency.\")\n\n    # APPLY_TRUE_PATH_RULE:\n    all_probs = apply_true_path_rule(all_probs, go_graph, target_terms)\n\n    # 3. SAVE WITH CUTOFF\n    print(\"INFERENCE: Writing clean submission file...\")\n    cutoff = 0.01\n    \n    # Force it to the correct Kaggle output directory\n    with open(\"/kaggle/working/submission.tsv\", \"w\") as f:    \n        for i, pid in enumerate(tqdm(test_ids, desc=\"Writing TSV\")):\n            if isinstance(pid, bytes): pid = pid.decode('utf-8')\n            else: pid = str(pid)\n\n            row_probs = all_probs[i]\n            # Optimization: Cutoff < 0.001\n            indices = np.where(row_probs >= cutoff)[0]\n\n            for idx in indices:\n                f.write(f\"{pid}\\t{target_terms[idx]}\\t{row_probs[idx]:.3f}\\n\")\n\n    print(\"DONE! '/kaggle/working/submission.tsv' is ready.\")","metadata":{"id":"MRqUj4QKw-eY","executionInfo":{"status":"ok","timestamp":1769983006170,"user_tz":480,"elapsed":36,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.854543Z","iopub.execute_input":"2026-02-02T19:48:57.854897Z","iopub.status.idle":"2026-02-02T19:48:57.870171Z","shell.execute_reply.started":"2026-02-02T19:48:57.854859Z","shell.execute_reply":"2026-02-02T19:48:57.869565Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def run_ensemble_inference(h5_path, test_ids, target_terms, model_keys, go_graph, taxonomy_map, constraints_dict):\n    \"\"\"\n    1. Loads top 3 saved models.\n    2. Predicts in batches (reading from H5).\n    3. Averages predictions + True Path + Taxon Filter (ALL IN MEMORY).\n    4. Writes ONE final file with cutoff 0.01.\n    \"\"\"\n    \n    # --- 1. FIND CHECKPOINTS ---\n    # We look for the models you saved during the night shift\n    checkpoints = sorted(glob.glob(\"cafa6_model_epoch_*.pth\"))\n    if not checkpoints:\n        # Fallback for Kaggle paths\n        checkpoints = sorted(glob.glob(\"/kaggle/working/cafa6_model_epoch_*.pth\"))\n    \n    # Pick the last 3 (Best Ensembling Strategy)\n    selected_checkpoints = checkpoints[-3:]\n    if not selected_checkpoints:\n        print(\"WARNING: No checkpoints found! Falling back to 'cafa6_final_model.pth'\")\n        if os.path.exists(\"cafa6_final_model.pth\"):\n             selected_checkpoints = [\"cafa6_final_model.pth\"]\n        else:\n             raise FileNotFoundError(\"CRITICAL: No models found. Did training finish?\")\n\n    print(f\"ENSEMBLE: Using models -> {selected_checkpoints}\")\n\n    # Setup\n    config = CAFA6Config()\n    h5_file = h5py.File(h5_path, 'r') # Open H5 once for reading inputs\n    batch_size = 1000\n    \n    # We will accumulate probabilities here\n    # Shape: (Num_Proteins, Num_Terms)\n    sum_probs = None\n    \n    # --- 2. LOOP THROUGH MODELS ---\n    for ckpt_path in selected_checkpoints:\n        print(f\"--- Processing Model: {os.path.basename(ckpt_path)} ---\")\n        \n        # Load Model Structure & Weights\n        model = CAFA6MultiTowerModel(config).to(config.DEVICE)\n        state_dict = torch.load(ckpt_path, map_location=config.DEVICE)\n        model.load_state_dict(state_dict)\n        model.eval()\n        \n        epoch_preds = []\n        \n        # Prediction Loop\n        with torch.no_grad():\n            for i in tqdm(range(0, len(test_ids), batch_size), desc=\"Predicting\"):\n                batch_ids = test_ids[i : i+batch_size]\n                batch_inputs = [] \n                \n                # Fetch Embeddings from H5\n                for k in model_keys:\n                    vectors = []\n                    for pid in batch_ids:\n                        if isinstance(pid, bytes): pid = pid.decode('utf-8')\n                        else: pid = str(pid)\n                        try:\n                            vectors.append(h5_file[k][pid][()])\n                        except KeyError:\n                            # Zero-vector fallback (Robustness)\n                            dim = h5_file[k][list(h5_file[k].keys())[0]].shape[0]\n                            vectors.append(np.zeros((dim,), dtype=np.float32))\n                    \n                    batch_inputs.append(torch.tensor(np.array(vectors)).float().to(config.DEVICE))\n                \n                # Forward Pass\n                logits = model(batch_inputs)\n                epoch_preds.append(torch.sigmoid(logits).cpu().numpy())\n        \n        # Stack this epoch's predictions\n        current_epoch_probs = np.vstack(epoch_preds)\n        \n        # Add to total sum\n        if sum_probs is None:\n            sum_probs = current_epoch_probs\n        else:\n            sum_probs += current_epoch_probs\n            \n        # Clear RAM\n        del model\n        del current_epoch_probs\n        torch.cuda.empty_cache()\n\n    h5_file.close() # Done with inputs\n\n    # --- 3. IN-MEMORY POST-PROCESSING ---\n    print(\"POST-PROCESSING: Calculating Average...\")\n    avg_probs = sum_probs / len(selected_checkpoints)\n    \n    print(\"POST-PROCESSING: Applying True Path Rule (Consistency)...\")\n    # This modifies avg_probs in place or returns new array\n    final_probs = apply_true_path_rule(avg_probs, go_graph, target_terms)\n    \n    if taxonomy_map and constraints_dict:\n        print(\"POST-PROCESSING: Applying Taxonomy Filter (Biology Check)...\")\n        final_probs = apply_taxon_constraints(final_probs, test_ids, taxonomy_map, constraints_dict)\n\n    # --- 4. FINAL WRITE (Cutoff 0.01) ---\n    cutoff = 0.01 \n    output_path = \"submission.tsv\" # Save to current directory\n    \n    print(f\"WRITING: Saving to {output_path} (Cutoff: {cutoff})...\")\n    \n    with open(output_path, \"w\") as f:\n        # Header (optional for CAFA, usually not needed but check rules)\n        # f.write(\"Protein Id\\tGO Term Id\\tPrediction\\n\") \n        \n        for i, pid in enumerate(tqdm(test_ids, desc=\"Writing File\")):\n            if isinstance(pid, bytes): pid = pid.decode('utf-8')\n            else: pid = str(pid)\n            \n            row_probs = final_probs[i]\n            \n            # FAST FILTERING\n            indices = np.where(row_probs >= cutoff)[0]\n            \n            for idx in indices:\n                # Write directly\n                f.write(f\"{pid}\\t{target_terms[idx]}\\t{row_probs[idx]:.3f}\\n\")\n                \n    print(f\"DONE! Optimized submission saved to {output_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.871253Z","iopub.execute_input":"2026-02-02T19:48:57.871513Z","iopub.status.idle":"2026-02-02T19:48:57.887065Z","shell.execute_reply.started":"2026-02-02T19:48:57.871492Z","shell.execute_reply":"2026-02-02T19:48:57.886422Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def apply_taxon_constraints(submission_probs, test_ids, taxonomy_map, constraints_dict):\n    \"\"\"\n    Zeros out probabilities for terms that are impossible for a given species.\n    taxonomy_map: dict {protein_id: taxon_id}\n    constraints_dict: dict {taxon_id: [allowed_terms]} OR {taxon_id: [forbidden_terms]}\n    \"\"\"\n    print(\"POST-PROCESSING: Applying Taxonomy Constraints...\")\n    \n    modified_probs = submission_probs.copy()\n    changes = 0\n    \n    # Pre-calculate indices for speed if possible, but loop is fine for post-processing\n    for i, pid in enumerate(tqdm(test_ids, desc=\"Taxon Filter\")):\n        pid_str = pid.decode('utf-8') if isinstance(pid, bytes) else str(pid)\n        \n        # Get species of this protein\n        taxon = taxonomy_map.get(pid_str)\n        \n        if taxon and taxon in constraints_dict:\n            # Get the list of forbidden terms (indices) for this taxon\n            forbidden_indices = constraints_dict[taxon]\n            \n            # Zero them out efficiently\n            # We use a mask for speed instead of a loop\n            if len(forbidden_indices) > 0:\n                # Check if we are actually removing anything (non-zero)\n                if np.any(modified_probs[i, forbidden_indices] > 0):\n                    modified_probs[i, forbidden_indices] = 0.0\n                    changes += 1\n                    \n    print(f\"TAXONOMY: Removed impossible predictions for {changes} proteins.\")\n    return modified_probs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.888112Z","iopub.execute_input":"2026-02-02T19:48:57.888440Z","iopub.status.idle":"2026-02-02T19:48:57.902868Z","shell.execute_reply.started":"2026-02-02T19:48:57.888407Z","shell.execute_reply":"2026-02-02T19:48:57.902125Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## **Training and Inference:**","metadata":{"id":"wDX8ZPv475SM"}},{"cell_type":"code","source":"# =============================================================================\n# TRAINING AND INFERENCE SCRIPT\n# =============================================================================\nclass ProteinFunctionPipeline:\n\n    def __init__(self, inputPath):\n        self.pDataLoader = PipelineDataLoader(inputPath)\n\n\n    def run_cafa6_training(self):\n        config = CAFA6Config()\n        ##################\n        # 1. LOAD DATA ###\n        ##################\n        # IA=Information Accretion\n        # IA.tsv processed dim: (40121, 2)\n        #    GO:0000001       0.0\n        # 0  GO:0000002  2.849666\n        # 1  GO:0000011  0.137504 ...etc\n        iaTsvDF = self.pDataLoader.load_IA_tsv()\n        # IA: Basically a DataFrame with 'term_id' and 'ia_score'\n        iaMap = dict(zip(iaTsvDF['term_id'], iaTsvDF['ia_score']))\n\n        # Load Master IDs: (This defines exactly which proteins can train on).\n        # train_ids.npy DATA SHAPE: (142246,)\n        # ['P20536' 'O73864' 'O95231' ... etc]\n        train_ids = self.pDataLoader.load_train_ids_npy()\n        # Load official raw data and only keep rows where the\n        # EntryID exists in the train_ids list.\n        # train_terms.tsv dim: (537028, 3)\n        # 0  EntryID        term   aspect\n        # 1   Q5W0B1  GO:0000785        C\n        # 2   Q5W0B1  GO:0004842        F\n        # 3   Q5W0B1  GO:0051865        P\n        # Long-format table where each row maps a protein ID to a single GO term.\n        # A critical step in the pipeline involves grouping these terms by\n        # protein ID and subsequently mapping them to a fixed-length vector.\n        train_terms = self.pDataLoader.load_train_terms_tsv()\n\n        # Loaded 40122 nodes and 77229 edges.\n        # First 5 nodes of go_graph_:\n        # ['GO:0000001', 'GO:0000002', 'GO:0000006', 'GO:0000007', 'GO:0000009']\n        # ====================================================\n        # First 5 edges of go_graph_:  [('GO:0000001', 'GO:0048308'),\n        # ('GO:0000001', 'GO:0048311'), ('GO:0000002', 'GO:0007005'),\n        # ('GO:0000006', 'GO:0005385'), ('GO:0000007', 'GO:0005385')]\n        # ====================================================\n        # First 5 items of term_aspects_:\n        # [('GO:0000001', 'BPO'), ('GO:0000002', 'BPO'), ('GO:0000006', 'MFO'),\n        #  ('GO:0000007', 'MFO'), ('GO:0000009', 'MFO')]\n        go_graph, term_aspects = self.pDataLoader.load_GO_Onto()\n\n\n        ##############################\n        # 2. LABEL ENGINEERING #######\n        ##############################\n        # Define label space (go_terms_list)\n        # Pick the top 1500 most frequent terms to keep the model manageable\n        # So they are 1500 for each subontology so 4500\n        print(\"PROCESSING: Generating Multi-Hot Labels (4500 dimensions)...\")\n\n        # Filter terms to only those in our training protein set\n        train_terms = train_terms[train_terms['protein'].isin(train_ids)]\n\n        # A. Split by Ontology\n        # We need top 1500 for *each* subontology to match model heads\n        mf_terms = train_terms[train_terms['ontology'] ==\n                               'F']['term'].value_counts().index[:config.NUM_GO_TERMS].tolist()\n        bp_terms = train_terms[train_terms['ontology'] ==\n                               'P']['term'].value_counts().index[:config.NUM_GO_TERMS].tolist()\n        cc_terms = train_terms[train_terms['ontology'] ==\n                               'C']['term'].value_counts().index[:config.NUM_GO_TERMS].tolist()\n\n        print(f\"Selected Terms: MF={len(mf_terms)}, BP={len(bp_terms)}, CC={len(cc_terms)}\")\n\n        # B. Create Master Target List (Order matters: MF -> BP -> CC)\n        # This aligns with torch.cat([mf_head, bp_head, cc_head], dim=1) in the model\n        target_terms = mf_terms + bp_terms + cc_terms\n\n        # Map term to index (0 to 4499)\n        term_to_idx = {term: i for i, term in enumerate(target_terms)}\n\n        # C. Align IA Weights to this specific term list\n        # If a term is missing in IA file, default to 1.0 (or 0.0)\n        ia_weights = np.array([iaMap.get(t, 1.0) for t in target_terms], dtype=np.float32)\n        # Load IA weights and GO terms\n        # ia_weights_tensor must match the column order of train_labels\n        ia_weights_tensor = torch.tensor(ia_weights).to(config.DEVICE)\n\n        #@-LM.1-@ Initialize zero matrix (Proteins x 4500)\n        # D. Populate the Multi-Hot Label Matrix\n        n_proteins = len(train_ids)\n        n_labels = len(target_terms)\n        train_labels = np.zeros((n_proteins, n_labels), dtype=np.float32)\n\n        #@-LM.2-@#Faster Filling (Vectorized)\n        print(\"MAPPING: Populating label matrix (this may take a minute)...\")\n        # Create a map of Protein ID -> Matrix Row Index\n        pid_to_row = {pid: i for i, pid in enumerate(train_ids)}\n        term_to_col = {term: i for i, term in enumerate(target_terms)}\n        # Fill the matrix: intersection of protein terms and our target terms\n        target_terms_set = set(target_terms)\n\n        # Iterate via grouping for speed (much faster than looping rows)\n        # Group by protein and collect terms\n        # grouped = train_terms.groupby('protein')['term'].apply(set)\n        # Iterate through the dataframe directly instead of grouping (often safer/faster)\n        # Filter only relevant rows first\n        relevant_terms = train_terms[\n            (train_terms['protein'].isin(pid_to_row)) &\n            (train_terms['term'].isin(target_terms_set))\n        ]\n\n        # Use numpy advanced indexing for speed and safety\n        row_indices = [pid_to_row[p] for p in relevant_terms['protein']]\n        col_indices = [term_to_col[t] for t in relevant_terms['term']]\n\n        # Set 1s\n        train_labels[row_indices, col_indices] = 1.0\n        print(f\">>>>>>>>>>>>Final Train Labels Shape: {train_labels.shape}\") # Should be (142246, 4500)\n\n        ###################################\n        # 3. DATASET & DATALOADER - H5 ####\n        ###################################\n        # H5 path\n        path_ = config.H5_PATH + config.EMBEDS_H5_CONSOLIDATED\n\n        # Verify H5 file exists before proceeding\n        if os.path.exists(path_):\n             # Delete last run\n             print(\"DELETING .h5 files from LAST RUN\")\n             #delete_h5_dirs_files(config.H5_PATH) // TODO check this deletes 0k\n             #delete_h5_dirs_files(path_)\n\n        print(f\"MSG: H5 file not found at {path_}. Exporting EMBEDDS now...\")\n        convert_parallel_npy_to_h5(config.FILE_CONFIG, path_)\n\n        # TEST EMBEDDS to H5 conversion\n        # TODO: I need to pack Test files in .h5 too.\n        #print(f\"MSG: H5 file not found at {path_}. Exporting TEST now...\")\n        #path_ = config.H5_PATH + config.EMBEDS_H5_TEST_CONSOLI\n        #build_test_h5(config.FILE_TEST_CONFIG, path_)\n\n        dataset = MultiProteinDataset(path_,\n                                      train_ids,\n                                      list(config.EMBEDDING_MAP.keys()),\n                                      train_labels)\n        loader = DataLoader(dataset, batch_size=config.BATCH_SIZE, shuffle=True)\n\n        ###################################\n        # 4. MODEL & TRAINING #############\n        ###################################\n        model = CAFA6MultiTowerModel(config).to(config.DEVICE)\n        optimizer = torch.optim.AdamW(model.parameters(), lr=config.LEARNING_RATE)\n\n        # Use the IA weights we calculated above\n        # OPTIMIZATION STRATEGY: Hybrid Loss (Focal and F1 all in one)\n        criterion = IAWeightedHybridLoss(ia_weights_tensor)\n\n        # --- THE SECRET SAUCE: SCHEDULER ---\n        # This gently lowers the learning rate to 0 by the end of the run\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.EPOCHS)\n\n        print(\"TRAINING: Start...\")\n        for epoch in range(config.EPOCHS):\n            model.train()\n            total_loss = 0\n\n            # Use TQDM to visualize progress\n            pbar = tqdm(loader, desc=f\"Epoch {epoch}\")\n            for embeddings, y in pbar:\n                # Move data to GPU\n                embeddings = [e.to(config.DEVICE) for e in embeddings]\n                y = y.to(config.DEVICE)\n\n                optimizer.zero_grad()\n                # Forward Pass\n                logits = model(embeddings)\n                # Verify shapes match exactly (Safety check)\n                if logits.shape != y.shape:\n                    raise ValueError(f\"Shape Mismatch! Logits: {logits.shape}, Labels: {y.shape}\")\n                # Loss & Backward\n                loss = criterion(logits, y)\n                loss.backward()\n                optimizer.step()\n\n                total_loss += loss.item()\n                # Show LR in progress bar (Dashboard)\n                current_lr = scheduler.get_last_lr()[0]\n                pbar.set_postfix({'loss': loss.item(), 'lr': f\"{current_lr:.1e}\"})    \n                #pbar.set_postfix({'loss': loss.item()})\n\n            # Step the scheduler at the end of every epoch\n            scheduler.step()\n            # Logging        \n            avg_loss = total_loss / len(loader)\n            print(f\"Epoch {epoch} complete. Avg Loss: {avg_loss:.4f} | LR: {current_lr:.2e}\")\n            #print(f\"Epoch {epoch} complete. Avg Loss: {avg_loss:.4f}\")\n\n            # Save checkpoint every epoch\n            torch.save(model.state_dict(), f\"cafa6_model_epoch_{epoch}.pth\")\n\n        # final_preds = apply_true_path_rule(raw_preds, go_graph, go_terms_list)\n        #torch.save(model.state_dict(), \"cafa6_final_model.pth\")\n        print(\"SUCCESS: Training Complete.\")\n\n        return model, target_terms, go_graph\n\n        #---------------------------------------------------------#\n        # END\n\n\n","metadata":{"id":"EgUJwLLK8A1b","executionInfo":{"status":"ok","timestamp":1769983014244,"user_tz":480,"elapsed":58,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.905090Z","iopub.execute_input":"2026-02-02T19:48:57.905500Z","iopub.status.idle":"2026-02-02T19:48:57.926138Z","shell.execute_reply.started":"2026-02-02T19:48:57.905477Z","shell.execute_reply":"2026-02-02T19:48:57.925432Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Taxonomic Filtering","metadata":{}},{"cell_type":"code","source":"# ---------------------------------------------------------\n# STEP 4: AUTOMATIC CONSTRAINT PARSING (USING YOUR FULL OBO)\n# ---------------------------------------------------------\nprint(\"\\n--- STEP 4: PARSING TAXONOMY RULES ---\")\n\n# 1. DEFINE THE PARSER FUNCTION\ndef parse_obo_constraints(obo_path):\n    \"\"\"\n    Scans full go.obo for 'never_in_taxon' lines.\n    \"\"\"\n    print(f\"PARSER: Scanning {obo_path}...\")\n    rules = {}\n    current_go_id = None\n    count = 0\n        \n    with open(obo_path, 'r') as f:\n        for line in f:\n            line = line.strip()\n            if line.startswith(\"[Term]\"):\n                current_go_id = None\n                \n            # Capture GO ID\n            elif line.startswith(\"id: GO:\"):\n                current_go_id = line.split()[1] # \"GO:0005634\"\n                \n            # Capture Constraint\n            # Look for lines like: \"relationship: never_in_taxon NCBITaxon:562\"\n            elif current_go_id and \"never_in_taxon\" in line:\n                parts = line.split()\n                for p in parts:\n                    if \"NCBITaxon:\" in p:\n                        try:\n                            taxon_id = int(p.split(\":\")[1])\n                            if taxon_id not in rules: rules[taxon_id] = []\n                            rules[taxon_id].append(current_go_id)\n                            count += 1\n                        except:\n                            continue\n                                \n    print(f\"PARSER: Success! Found {count} constraints for {len(rules)} taxa.\")\n    return rules\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.927135Z","iopub.execute_input":"2026-02-02T19:48:57.927427Z","iopub.status.idle":"2026-02-02T19:48:57.944098Z","shell.execute_reply.started":"2026-02-02T19:48:57.927381Z","shell.execute_reply":"2026-02-02T19:48:57.943298Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Main():**","metadata":{"id":"tdDZcZikJ2DH"}},{"cell_type":"code","source":"# =============================================================================\n# 3. MAIN EXECUTION FLOW\n# =============================================================================\nif __name__ == \"__main__\": \n    # 1. SETUP\n    conf = CAFA6Config()\n    #conf.EPOCHS = 3  # Emergency limit for time\n\n    # 2. RUN TRAINING (And catch the returned model!)\n    print(f\"--- STEP 1: TRAINING ({conf.EPOCHS} Epochs) ---\")\n    pipeline = ProteinFunctionPipeline(conf.CAFA6_PATH)\n    trained_model, final_terms, go_graph = pipeline.run_cafa6_training()\n\n    # 3. BUILD TEST DATA\n    test_h5_path = conf.H5_PATH + conf.EMBEDS_H5_TEST_CONSOLI\n\n    print(\"\\n--- STEP 2: PREPARING TEST H5 ---\")\n    print(f\"Target: {test_h5_path}\")\n    build_test_h5(conf.FILE_TEST_CONFIG, test_h5_path)\n\n    # 4. INFERENCE\n    print(\"\\n--- STEP 3: PREDICTION ---\")\n\n    # Load the Master List of IDs (Just to know who to predict for)\n    # Why loading IDs from \"prott5\"?\n    # This is just to get the \"Guest List\" (Protein Names).\n    # The code loads the names from the ProtT5 file just to know which proteins to process.\n    # Crucial: Once it has the name (e.g., \"P12345\"), it goes into the loop and\n    #  pulls data for ALL models (ProtT5 AND ESM2 AND Ankh) using that name.\n    test_ids_source = conf.FILE_TEST_CONFIG[\"prott5\"][\"ids\"]\n    all_test_ids = np.load(test_ids_source)\n\n    # Get the keys the model actually learned (e.g. ['prott5', 'esm2', 'ankh'])\n    model_input_keys = list(conf.EMBEDDING_MAP.keys())\n\n # ----------------------------------------------------------------------------\n    ##############################################################################\n    # CONSTRAINTS::\n    ##############################################################################\n    # 3. Define Constraints (The Rules)\n    #  knowing a protein is '9606' (Human) is useless \n    # if you don't have a rule saying what Humans can/can't do.\n    # For now, keep this empty if you don't have the rules file ready.\n\n    # 2. RUN PARSER ON YOUR FILE\n    # <--- YOUR SPECIFIC PATH\n    obo_file_path = \"/kaggle/input/protein-function-prediction/Protein Function Prediction/go.obo\"\n    \n    constraints_dict = {}\n    \n    if os.path.exists(obo_file_path):\n        constraints_dict = parse_obo_constraints(obo_file_path)\n    else:\n        print(f\"WARNING: OBO file not found at {obo_file_path}\")\n\n    # 3. FAIL-SAFE (THE \"DEAD ZOO\" BACKUP)\n    # If parsing yields 0 results (e.g. file format diff), use the manual list so we don't crash.\n    if len(constraints_dict) < 10:\n        print(\"PARSER: File yielded low/zero constraints. ACTIVATING BACKUP RULES.\")\n        \n        # Bacteria (No Nucleus/Mito)\n        for tax in [2, 562, 511145, 208964]: \n             constraints_dict[tax] = [\"GO:0005634\", \"GO:0005739\", \"GO:0005783\", \"GO:0005794\"]\n\n        # Plants (No Nervous System)\n        for tax in [33090, 3702, 39947, 4577]:\n             constraints_dict[tax] = [\"GO:0007268\", \"GO:0007610\", \"GO:0008015\", \"GO:0040011\"]\n\n        # Animals (No Photosynthesis)\n        for tax in [33208, 9606, 10090, 10116, 7227]:\n             constraints_dict[tax] = [\"GO:0015979\", \"GO:0019684\"]\n             \n        # Fungi\n        for tax in [4751, 4932, 284812]:\n             constraints_dict[tax] = [\"GO:0015979\", \"GO:0007268\"]\n\n    ##############################################################################\n    # ---------------------------------------------------------------------------\n    # 4. LOAD TAXONOMY MAP (Test Protein -> Species)\n    # ---------------------------------------------------------------------------\n    print(\"SETUP: Loading Taxonomy Map for Test Proteins...\")\n    # This is the Kaggle Test List\n    taxon_file_path = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\"\n    taxonomy_map = {}\n    \n    if os.path.exists(taxon_file_path):\n        print(f\"LOADING: Taxonomy map from {taxon_file_path}...\")\n        with open(taxon_file_path, 'r') as f:\n            for line in tqdm(f, desc=\"Loading Taxa\"):\n                parts = line.strip().split('\\t')\n                if len(parts) >= 2:\n                    try:\n                        # Try to convert to int. If it's a header (\"Species\"), it jumps to except.\n                        taxonomy_map[parts[0]] = int(parts[1])\n                    except ValueError:\n                        continue # Skip the header line safely\n                    \n    # 5. CONVERT TO MODEL INDICES\n    print(\"MAPPING: Converting Rules to Model Indices...\")\n    term_to_idx = {t: i for i, t in enumerate(final_terms)}\n    final_constraints = {}\n    \n    for taxon, go_list in constraints_dict.items():\n        forbidden_indices = []\n        for term_str in go_list:\n            if term_str in term_to_idx:\n                forbidden_indices.append(term_to_idx[term_str])\n        \n        if forbidden_indices:\n            final_constraints[taxon] = forbidden_indices\n\n    print(f\"READY: Engine configured with constraints for {len(final_constraints)} taxa.\")\n\n    # ---------------------------------------------------------\n    # STEP 5: RUN THE ENGINE\n    # ---------------------------------------------------------    \n    \"\"\"\n    run_inference_h5(\n        model=trained_model,\n        h5_path=test_h5_path,\n        test_ids=all_test_ids,\n        target_terms=final_terms,\n        model_keys=model_input_keys,\n        go_graph=go_graph\n    )\n\n    \"\"\"\n    #  NEW WAY (Ensemble + True Path + Taxon Filter)\n    run_ensemble_inference(\n        h5_path=test_h5_path,\n        test_ids=all_test_ids,\n        target_terms=final_terms,\n        model_keys=model_input_keys,\n        go_graph=go_graph,\n        taxonomy_map=taxonomy_map,\n        constraints_dict=final_constraints\n    )\n     \n\n","metadata":{"id":"0DbC3KcoJ0-J","executionInfo":{"status":"error","timestamp":1769940190651,"user_tz":480,"elapsed":1792333,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"outputId":"c5458a3e-8857-42d5-983f-3e663fcc1148","trusted":true,"execution":{"iopub.status.busy":"2026-02-02T19:48:57.945174Z","iopub.execute_input":"2026-02-02T19:48:57.945915Z","iopub.status.idle":"2026-02-02T20:06:03.931140Z","shell.execute_reply.started":"2026-02-02T19:48:57.945889Z","shell.execute_reply":"2026-02-02T20:06:03.930126Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## **Some Test**","metadata":{"id":"x2M07JDqMV7b"}},{"cell_type":"code","source":"# PRE-FLIGHT CHECK\ndef check_system():\n    c = CAFA6Config()\n    print(\"CHECKING PATHS...\")\n\n    # 1. Check Train Files\n    for key, paths in c.FILE_CONFIG.items():\n        if os.path.exists(paths['embeds']): print(f\"[OK] Train {key} Embeds found\")\n        else: print(f\"[FAIL] Train {key} Embeds MISSING: {paths['embeds']}\")\n\n    # 2. Check Test Files\n    for key, paths in c.FILE_TEST_CONFIG.items():\n        if os.path.exists(paths['embeds']): print(f\"[OK] Test {key} Embeds found\")\n        else: print(f\"[FAIL] Test {key} Embeds MISSING: {paths['embeds']}\")\n\n    # 3. Check H5 Write Permissions\n    test_h5 = c.H5_PATH + c.EMBEDS_H5_TEST_CONSOLI\n    try:\n        with open(test_h5, 'w') as f:\n            f.write(\"test\")\n        os.remove(test_h5)\n        print(f\"[OK] H5 Write Permission to {test_h5}\")\n    except Exception as e:\n        print(f\"[FAIL] Cannot write H5: {e}\")\n\ncheck_system()","metadata":{"id":"EDapSWB2MZhm","executionInfo":{"status":"ok","timestamp":1769983385245,"user_tz":480,"elapsed":38,"user":{"displayName":"ANDRES HERNAN PITYLA C.","userId":"08511245577899731113"}},"outputId":"29c8fb19-d4c3-476a-98e5-07772c93e440","trusted":true,"execution":{"iopub.status.busy":"2026-02-02T20:06:03.933835Z","iopub.execute_input":"2026-02-02T20:06:03.934402Z","iopub.status.idle":"2026-02-02T20:06:04.372931Z","shell.execute_reply.started":"2026-02-02T20:06:03.934372Z","shell.execute_reply":"2026-02-02T20:06:04.371951Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Ensemble Inference**","metadata":{}}]}