{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"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"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14207560,"sourceType":"datasetVersion","datasetId":9061916}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:18:56.791191Z","iopub.execute_input":"2026-01-27T17:18:56.791802Z","iopub.status.idle":"2026-01-27T17:18:57.445327Z","shell.execute_reply.started":"2026-01-27T17:18:56.791768Z","shell.execute_reply":"2026-01-27T17:18:57.444609Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\nimport gc\nfrom collections import defaultdict\n\nimport pandas as pd\nfrom tqdm.auto import tqdm\nimport numpy as np\n# Core Imports\nimport pandas as pd\nimport numpy as np\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\nimport warnings\nwarnings.filterwarnings('ignore')\nimport joblib\nimport networkx as nx\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.linear_model import LogisticRegression\n# import Bio\n# Visualization imports\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom matplotlib.patches import Rectangle\nimport matplotlib.patches as mpatches\nplt.style.use('seaborn-v0_8-whitegrid')\nsns.set_palette(\"husl\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:18:57.446564Z","iopub.execute_input":"2026-01-27T17:18:57.446916Z","iopub.status.idle":"2026-01-27T17:18:59.291817Z","shell.execute_reply.started":"2026-01-27T17:18:57.446892Z","shell.execute_reply":"2026-01-27T17:18:59.290972Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\nimport gc\nfrom collections import defaultdict\n\nimport pandas as pd\nfrom tqdm.auto import tqdm\nimport numpy as np\n\n# ⚙️ CONFIGURATION - ADJUST THIS FOR SPEED VS ACCURACY\nSAMPLE_PERCENT = 100  # Use 100% of data\nQUICK_MODE = True   # Enable full feature computation\n\n# Package Installation\nimport subprocess\nimport sys\n\ndef install(package):\n    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", package])\n\nprint(\"Installing required packages...\")\ntry:\n    import obonet\nexcept:\n    install('obonet')\n    import obonet\n\ntry:\n    from Bio import SeqIO\nexcept:\n    install('biopython')\n    from Bio import SeqIO\n\n# Core Imports\nimport pandas as pd\nimport numpy as np\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\nimport warnings\nwarnings.filterwarnings('ignore')\n\nimport networkx as nx\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.linear_model import LogisticRegression\nimport Bio\n# Visualization imports\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom matplotlib.patches import Rectangle\nimport matplotlib.patches as mpatches\nplt.style.use('seaborn-v0_8-whitegrid')\nsns.set_palette(\"husl\")\n\nprint(\"=\"*80)\nprint(\"CAFA 6 PROTEIN FUNCTION PREDICTION - ENHANCED STARTER\")\nprint(f\"📊 SAMPLE MODE: {SAMPLE_PERCENT}% of data\")\nprint(f\"⚡ QUICK MODE: {'ON' if QUICK_MODE else 'OFF'}\")\nprint(\"=\"*80)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:18:59.292781Z","iopub.execute_input":"2026-01-27T17:18:59.293127Z","iopub.status.idle":"2026-01-27T17:19:07.776884Z","shell.execute_reply.started":"2026-01-27T17:18:59.293106Z","shell.execute_reply":"2026-01-27T17:19:07.776170Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# CONFIG (ADVANCED IMPROVEMENTS)\n# ============================================================================\nCONFIG = {\n    'MAX_TRAIN_SAMPLES': 80000,        # IMPROVED: 5000 → 8000 (60% more data!)\n    'TOP_K_LABELS': 1500,\n    'VALIDATION_SIZE': 0.15,\n    'RANDOM_SEED': 42,\n    'HIDDEN_DIMS': [1024, 512, 256],\n    'DROPOUT_RATE': 0.25,             # IMPROVED: 0.3 → 0.25 (less dropout for more data)\n    'EPOCHS': 100,                     # IMPROVED: 10 → 12 epochs\n    'BATCH_SIZE': 64,                 # IMPROVED: 32 → 64 (faster training)\n    'LEARNING_RATE': 1e-3,\n    'PREDICT_BATCH_SIZE': 128,        # IMPROVED: 64 → 128 (faster prediction)\n    'MIN_CONFIDENCE': 0.08,           # IMPROVED: 0.10 → 0.08 (even more predictions)\n    'MAX_PREDS_PER_PROTEIN': 200,    # IMPROVED: 150 → 200\n    'LABEL_SMOOTHING': 0.1,          # NEW: Add label smoothing\n    'TEMPERATURE': 1.5,               # NEW: Temperature scaling for predictions\n    'MERGE_WITH_BLAST': True,\n    'BLAST_WEIGHT': 0.60,             # IMPROVED: 0.65 → 0.60\n    'DL_WEIGHT': 0.40,                # IMPROVED: 0.35 → 0.40 (more trust in DL)\n    'BASE_PATH': '/kaggle/input/cafa-6-protein-function-prediction',\n    'TRAIN_DIR': '/kaggle/input/cafa-6-protein-function-prediction/Train',\n    'TEST_DIR': '/kaggle/input/cafa-6-protein-function-prediction/Test',\n    'ESM1_train': '/kaggle/input/cafa-6-embeddings/esm1b_650M/train_sequences_emb.npy',\n    'ESM1_train_id': '/kaggle/input/cafa-6-embeddings/esm1b_650M/train_sequences_ids.npy',\n    'ESM1_test': '/kaggle/input/cafa-6-embeddings/esm1b_650M/testsuperset_emb.npy',\n    'ESM1_test_id': '/kaggle/input/cafa-6-embeddings/esm1b_650M/testsuperset_ids.npy',\n    'D2V_train': '/kaggle/input/cafa6-embeddings/d2v_embedding.csv',\n    'TFIDF_train': '/kaggle/input/cafa6-embeddings/tfidf_embedding.csv',\n    'PUBMED_train': '/kaggle/input/cafa6-embeddings/pubmed_embedding.csv',\n    'D2V_test': '/kaggle/input/cafa6-embeddings/d2vtest_embedding.csv',\n    'TFIDF_test': '/kaggle/input/cafa6-embeddings/test_tfidf_embedding.csv',\n    'PUBMED_test': '/kaggle/input/cafa6-embeddings/test_pubmed_embedding.csv',\n    'IA': '/kaggle/input/cafa-6-protein-function-prediction/IA.tsv',\n    'ESM1_PATH': '/kaggle/input/cafa-6-embeddings/esm1b_650M',\n    'ESM2_650': '/kaggle/input/cafa-6-embeddings/esm2_650M',\n    'Taxonomy_HL': '/kaggle/input/cafa-6-embeddings/taxonomy_highlevel',\n    \n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:07.778630Z","iopub.execute_input":"2026-01-27T17:19:07.778880Z","iopub.status.idle":"2026-01-27T17:19:07.784359Z","shell.execute_reply.started":"2026-01-27T17:19:07.778859Z","shell.execute_reply":"2026-01-27T17:19:07.783601Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_fasta = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\ntest_fasta = '/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta'\nobo_path = '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:07.785262Z","iopub.execute_input":"2026-01-27T17:19:07.785609Z","iopub.status.idle":"2026-01-27T17:19:07.801015Z","shell.execute_reply.started":"2026-01-27T17:19:07.785582Z","shell.execute_reply":"2026-01-27T17:19:07.800397Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_terms = pd.read_csv(Path(CONFIG['TRAIN_DIR'])/ 'train_terms.tsv', sep='\\t', \n                          names=['protein', 'term', 'ontology'])\n\ntrain_terms.drop(train_terms.index[0], inplace=True)\ntrain_taxonomy = pd.read_csv(Path(CONFIG['TRAIN_DIR'])/ 'train_taxonomy.tsv', sep='\\t',\n                             names=['protein', 'taxon'])\n\nia_df = pd.read_csv(Path(CONFIG['BASE_PATH'])/'IA.tsv', sep='\\t', names=['term', 'ia_score'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:07.801788Z","iopub.execute_input":"2026-01-27T17:19:07.802037Z","iopub.status.idle":"2026-01-27T17:19:08.190164Z","shell.execute_reply.started":"2026-01-27T17:19:07.802019Z","shell.execute_reply":"2026-01-27T17:19:08.189577Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_taxonomy","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:08.190990Z","iopub.execute_input":"2026-01-27T17:19:08.191211Z","iopub.status.idle":"2026-01-27T17:19:08.212804Z","shell.execute_reply.started":"2026-01-27T17:19:08.191192Z","shell.execute_reply":"2026-01-27T17:19:08.212143Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_tax_hl = pd.read_csv(Path(CONFIG['Taxonomy_HL'])/ 'train_taxonomy_highlevel.tsv', sep='\\t',\n                             )\ntrain_tax_hl.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:08.213692Z","iopub.execute_input":"2026-01-27T17:19:08.214070Z","iopub.status.idle":"2026-01-27T17:19:08.339164Z","shell.execute_reply.started":"2026-01-27T17:19:08.214043Z","shell.execute_reply":"2026-01-27T17:19:08.338576Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_tax_hl = pd.read_csv(Path(CONFIG['Taxonomy_HL'])/ 'test_taxonomy_highlevel.tsv', sep='\\t',\n                             )\ntest_tax_hl.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:08.339999Z","iopub.execute_input":"2026-01-27T17:19:08.340251Z","iopub.status.idle":"2026-01-27T17:19:08.624850Z","shell.execute_reply.started":"2026-01-27T17:19:08.340232Z","shell.execute_reply":"2026-01-27T17:19:08.624169Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_dist = test_tax_hl[['protein', 'taxon']]\ntest_dist","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:08.627661Z","iopub.execute_input":"2026-01-27T17:19:08.628416Z","iopub.status.idle":"2026-01-27T17:19:08.640123Z","shell.execute_reply.started":"2026-01-27T17:19:08.628378Z","shell.execute_reply":"2026-01-27T17:19:08.639533Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\n\ndef create_cafa6_validation_set(train_df, test_dist, val_size=1500):\n    \"\"\"\n    Samples validation data to match Test Set species proportions.\n    test_dist: The Series you shared (test_df['taxon_id'].value_counts(normalize=True))\n    \"\"\"\n    val_indices = []\n    \n    # 1. Iterate through test set species proportions\n    for taxon_id, proportion in test_dist.items():\n        # Calculate how many proteins we need for this species in our 1000-sample val set\n        n_to_sample = int(round(proportion * val_size))\n        if n_to_sample == 0: continue\n        \n        # Find available proteins in training for this taxon\n        available_pool = train_df[train_df['taxon'] == taxon_id].index.tolist()\n        \n        if len(available_pool) >= n_to_sample:\n            sampled = np.random.choice(available_pool, n_to_sample, replace=False)\n            val_indices.extend(sampled)\n        else:\n            # If training is missing enough of this species, take all we have\n            val_indices.extend(available_pool)\n            \n    # 2. Final check: if rounding left us short of 1000, fill randomly from remaining\n    remaining_pool = train_df.index.difference(val_indices)\n    if len(val_indices) < val_size:\n        extra_needed = val_size - len(val_indices)\n        extra_samples = np.random.choice(remaining_pool, extra_needed, replace=False)\n        val_indices.extend(extra_samples)\n        \n    # 3. Create the final split\n    val_df = train_df.loc[val_indices]\n    train_split_df = train_df.drop(val_indices)\n    \n    print(f\"✅ Species-Balanced Validation Set Created: {len(val_df)} proteins\")\n    return train_split_df, val_df\n    \n# Run the split\ntaxon_test = test_dist['taxon'].value_counts(normalize=True)\ntrain_final, val_final = create_cafa6_validation_set(train_taxonomy, taxon_test)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:08.641048Z","iopub.execute_input":"2026-01-27T17:19:08.641339Z","iopub.status.idle":"2026-01-27T17:19:08.769539Z","shell.execute_reply.started":"2026-01-27T17:19:08.641317Z","shell.execute_reply":"2026-01-27T17:19:08.768830Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_final.taxon.value_counts(normalize=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:08.770509Z","iopub.execute_input":"2026-01-27T17:19:08.770923Z","iopub.status.idle":"2026-01-27T17:19:08.782232Z","shell.execute_reply.started":"2026-01-27T17:19:08.770893Z","shell.execute_reply":"2026-01-27T17:19:08.781613Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val_final.taxon.value_counts(normalize=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:08.783433Z","iopub.execute_input":"2026-01-27T17:19:08.783821Z","iopub.status.idle":"2026-01-27T17:19:08.797258Z","shell.execute_reply.started":"2026-01-27T17:19:08.783801Z","shell.execute_reply":"2026-01-27T17:19:08.796649Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_sequences = list(SeqIO.parse(train_fasta, \"fasta\"))\nnum_sequences_train = sum(1 for seq in train_sequences)\n\nprint(\"Number of train sequences:\", num_sequences_train)\n\ntest_sequences = list(SeqIO.parse(test_fasta, \"fasta\"))\nnum_sequences_test = sum(1 for seq in test_sequences)\nprint(\"Number of test sequences:\", num_sequences_test)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:08.798124Z","iopub.execute_input":"2026-01-27T17:19:08.798332Z","iopub.status.idle":"2026-01-27T17:19:12.553142Z","shell.execute_reply.started":"2026-01-27T17:19:08.798314Z","shell.execute_reply":"2026-01-27T17:19:12.552480Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\n\ndef calculate_ctd_features(sequence):\n    \"\"\"Internal helper to calculate CTD for a single sequence\"\"\"\n    group_map = {\n        'R':1, 'K':1, 'E':1, 'D':1, 'Q':1, 'N':1,\n        'G':2, 'A':2, 'S':2, 'T':2, 'P':2, 'H':2, 'Y':2,\n        'C':3, 'V':3, 'L':3, 'I':3, 'M':3, 'F':3, 'W':3\n    }\n    encoded = [group_map.get(aa, 0) for aa in sequence]\n    n = len(encoded)\n    res = {}\n\n    # 1. Composition\n    for g in [1, 2, 3]:\n        res[f'C_Group{g}'] = encoded.count(g) / n if n > 0 else 0\n        \n    # 2. Transition\n    t_counts = {'12': 0, '13': 0, '23': 0}\n    for i in range(n - 1):\n        pair = \"\".join(map(str, sorted([encoded[i], encoded[i+1]])))\n        if pair in t_counts: t_counts[pair] += 1\n    for k, v in t_counts.items():\n        res[f'T_{k}'] = v / (n - 1) if n > 1 else 0\n\n    # 3. Distribution\n    for g in [1, 2, 3]:\n        idxs = [i + 1 for i, val in enumerate(encoded) if val == g]\n        for p, lbl in zip([0, 24, 49, 74, 99], ['001', '025', '050', '075', '100']):\n            if idxs:\n                val = idxs[int(np.percentile(range(len(idxs)), p))] / n\n            else:\n                val = 0.0\n            res[f'D_Group{g}_{lbl}'] = val\n    return res","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:12.553947Z","iopub.execute_input":"2026-01-27T17:19:12.554149Z","iopub.status.idle":"2026-01-27T17:19:12.562591Z","shell.execute_reply.started":"2026-01-27T17:19:12.554132Z","shell.execute_reply":"2026-01-27T17:19:12.561922Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def create_sequence_dataframe_from_fasta(fasta_sequences, data_type: str) -> pd.DataFrame:\n    hydrophobic, polar, charged = 'AILMFWYV', 'STNQ', 'DEKR'\n    rows = []\n\n    print(f\"Extracting features for {data_type} data...\")\n    \n    for seq_record in tqdm(fasta_sequences):\n        seq_str = str(seq_record.seq)\n        seq_len = len(seq_str)\n        if seq_len == 0: continue\n\n        # 1. Create the row dictionary\n        # We start with the IDs so they are naturally at the front\n        entry = {}\n        \n        if data_type == 'train':\n            entry[\"protein\"] = seq_record.id.split('|')[1]\n            entry[\"EntryID\"] = seq_record.id\n        else: # test\n            entry[\"protein\"] = seq_record.id\n            entry[\"taxon\"] = seq_record.description.split(\" \")[1]\n\n        # 2. Add the physical features\n        entry.update({\n            \"sequence\": seq_str,\n            \"Length\": seq_len,\n            \"hydrophobic_ratio\": sum(seq_str.count(aa) for aa in hydrophobic) / seq_len,\n            \"polar_ratio\": sum(seq_str.count(aa) for aa in polar) / seq_len,\n            \"charged_ratio\": sum(seq_str.count(aa) for aa in charged) / seq_len,\n            \"molecular_weight\": 110 * seq_len\n        })\n\n        # 3. Add the CTD features\n        ctd_data = calculate_ctd_features(seq_str)\n        entry.update(ctd_data)\n\n        rows.append(entry)\n\n    # Convert to DataFrame\n    df = pd.DataFrame(rows)\n\n    # --- THE REORDERING LOGIC ---\n    # Define exactly which columns you want at the front\n    if data_type == 'train':\n        first_cols = ['protein', 'EntryID']\n    else:\n        first_cols = ['protein', 'taxon']\n\n    # Get all other columns that aren't in our \"first_cols\" list\n    remaining_cols = [c for c in df.columns if c not in first_cols]\n    \n    # Reconstruct the DataFrame with the new column order\n    df = df[first_cols + remaining_cols]\n\n    return df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:12.563387Z","iopub.execute_input":"2026-01-27T17:19:12.563561Z","iopub.status.idle":"2026-01-27T17:19:12.573039Z","shell.execute_reply.started":"2026-01-27T17:19:12.563545Z","shell.execute_reply":"2026-01-27T17:19:12.572374Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df = create_sequence_dataframe_from_fasta(train_sequences, data_type='train')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:19:12.573958Z","iopub.execute_input":"2026-01-27T17:19:12.574181Z","iopub.status.idle":"2026-01-27T17:21:52.820806Z","shell.execute_reply.started":"2026-01-27T17:19:12.574163Z","shell.execute_reply":"2026-01-27T17:21:52.820129Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df['Length_log'] = np.log1p(train_df['Length'])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:52.824227Z","iopub.execute_input":"2026-01-27T17:21:52.824452Z","iopub.status.idle":"2026-01-27T17:21:52.829718Z","shell.execute_reply.started":"2026-01-27T17:21:52.824434Z","shell.execute_reply":"2026-01-27T17:21:52.829049Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:52.830681Z","iopub.execute_input":"2026-01-27T17:21:52.830981Z","iopub.status.idle":"2026-01-27T17:21:52.855589Z","shell.execute_reply.started":"2026-01-27T17:21:52.830953Z","shell.execute_reply":"2026-01-27T17:21:52.854894Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.preprocessing import StandardScaler\n\n# Select only the numeric manual columns\nmanual_cols = ['hydrophobic_ratio', 'polar_ratio', 'charged_ratio', 'molecular_weight',\n               'C_Group1', 'C_Group2', 'C_Group3', 'T_12', 'T_13', 'T_23',\n               'D_Group1_001', 'D_Group1_025', 'D_Group1_050', 'D_Group1_075',\n               'D_Group1_100', 'D_Group2_001', 'D_Group2_025', 'D_Group2_050',\n               'D_Group2_075', 'D_Group2_100', 'D_Group3_001', 'D_Group3_025',\n               'D_Group3_050', 'D_Group3_075', 'D_Group3_100', 'Length_log']\n\nscaler = StandardScaler()\ntrain_df[manual_cols] = scaler.fit_transform(train_df[manual_cols])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:52.856655Z","iopub.execute_input":"2026-01-27T17:21:52.857235Z","iopub.status.idle":"2026-01-27T17:21:52.911257Z","shell.execute_reply.started":"2026-01-27T17:21:52.857207Z","shell.execute_reply":"2026-01-27T17:21:52.910494Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\n Loading ESM1b embeddings...\")\nesm1_base = f\"{CONFIG['ESM2_650']}\"\ntrain_ids = np.load(f\"{esm1_base}/train_sequences_ids.npy\", allow_pickle=True)\ntrain_embeds = np.load(f\"{esm1_base}/train_sequences_emb.npy\")\n\ntrain_dict = {str(pid): emb for pid, emb in zip(train_ids, train_embeds)}\nembed_dim = train_embeds.shape[1]\nprint(f\"   ✓ Dim: {embed_dim}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:52.912180Z","iopub.execute_input":"2026-01-27T17:21:52.912451Z","iopub.status.idle":"2026-01-27T17:21:57.913839Z","shell.execute_reply.started":"2026-01-27T17:21:52.912419Z","shell.execute_reply":"2026-01-27T17:21:57.913084Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"esm_emb_names = [f\"esm_{i}\" for i in range(train_embeds.shape[1])]\nesm1_train = pd.DataFrame(train_embeds, columns=esm_emb_names)\nesm1_train.insert(0, 'train_ids', train_ids)\nesm1_train.insert(1, 'protein', esm1_train['train_ids'].str.split(\"|\").str[1])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:57.914861Z","iopub.execute_input":"2026-01-27T17:21:57.915145Z","iopub.status.idle":"2026-01-27T17:21:57.984836Z","shell.execute_reply.started":"2026-01-27T17:21:57.915117Z","shell.execute_reply":"2026-01-27T17:21:57.984268Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"esm1_train","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:57.985672Z","iopub.execute_input":"2026-01-27T17:21:57.985927Z","iopub.status.idle":"2026-01-27T17:21:58.044808Z","shell.execute_reply.started":"2026-01-27T17:21:57.985908Z","shell.execute_reply":"2026-01-27T17:21:58.044161Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df = pd.merge(train_df, esm1_train, on='protein', how='left')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:58.045584Z","iopub.execute_input":"2026-01-27T17:21:58.045830Z","iopub.status.idle":"2026-01-27T17:21:59.406663Z","shell.execute_reply.started":"2026-01-27T17:21:58.045809Z","shell.execute_reply":"2026-01-27T17:21:59.405910Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_full_embedding = pd.merge(train_df, train_tax_hl, on='protein', how='left')\nprint(f\"The shape of all embeddings: {train_full_embedding.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:59.407663Z","iopub.execute_input":"2026-01-27T17:21:59.407952Z","iopub.status.idle":"2026-01-27T17:21:59.549197Z","shell.execute_reply.started":"2026-01-27T17:21:59.407932Z","shell.execute_reply":"2026-01-27T17:21:59.548566Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_full_embedding.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:59.549955Z","iopub.execute_input":"2026-01-27T17:21:59.550139Z","iopub.status.idle":"2026-01-27T17:21:59.565983Z","shell.execute_reply.started":"2026-01-27T17:21:59.550122Z","shell.execute_reply":"2026-01-27T17:21:59.565321Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_ids1 = train_full_embedding['protein'].astype(str).tolist()\nesm1_fts = train_full_embedding[esm_emb_names].values\nman_features = train_full_embedding[manual_cols].values","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:59.566915Z","iopub.execute_input":"2026-01-27T17:21:59.567205Z","iopub.status.idle":"2026-01-27T17:21:59.943549Z","shell.execute_reply.started":"2026-01-27T17:21:59.567176Z","shell.execute_reply":"2026-01-27T17:21:59.942953Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Convert to float32 BEFORE scaling\nesm1_fts = esm1_fts.astype('float32')\nman_features = man_features.astype('float32')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:21:59.947122Z","iopub.execute_input":"2026-01-27T17:21:59.947545Z","iopub.status.idle":"2026-01-27T17:22:00.312007Z","shell.execute_reply.started":"2026-01-27T17:21:59.947525Z","shell.execute_reply":"2026-01-27T17:22:00.311365Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df = create_sequence_dataframe_from_fasta(test_sequences, data_type='test')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:22:00.312841Z","iopub.execute_input":"2026-01-27T17:22:00.313127Z","iopub.status.idle":"2026-01-27T17:28:44.495539Z","shell.execute_reply.started":"2026-01-27T17:22:00.313100Z","shell.execute_reply":"2026-01-27T17:28:44.494941Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df['Length_log'] = np.log1p(test_df['Length'])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:28:44.500999Z","iopub.execute_input":"2026-01-27T17:28:44.501216Z","iopub.status.idle":"2026-01-27T17:28:44.508446Z","shell.execute_reply.started":"2026-01-27T17:28:44.501198Z","shell.execute_reply":"2026-01-27T17:28:44.507918Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.preprocessing import StandardScaler\n\n# Select only the numeric manual columns\nmanual_cols = ['hydrophobic_ratio', 'polar_ratio', 'charged_ratio', 'molecular_weight',\n               'C_Group1', 'C_Group2', 'C_Group3', 'T_12', 'T_13', 'T_23',\n               'D_Group1_001', 'D_Group1_025', 'D_Group1_050', 'D_Group1_075',\n               'D_Group1_100', 'D_Group2_001', 'D_Group2_025', 'D_Group2_050',\n               'D_Group2_075', 'D_Group2_100', 'D_Group3_001', 'D_Group3_025',\n               'D_Group3_050', 'D_Group3_075', 'D_Group3_100', 'Length_log']\n\nscaler = StandardScaler()\ntest_df[manual_cols] = scaler.fit_transform(test_df[manual_cols])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:28:44.509170Z","iopub.execute_input":"2026-01-27T17:28:44.509378Z","iopub.status.idle":"2026-01-27T17:28:44.665257Z","shell.execute_reply.started":"2026-01-27T17:28:44.509351Z","shell.execute_reply":"2026-01-27T17:28:44.664684Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_ids = np.load(f\"{esm1_base}/testsuperset_ids.npy\", allow_pickle=True)\ntest_embeds = np.load(f\"{esm1_base}/testsuperset_emb.npy\")\ntest_dict = {str(pid): emb for pid, emb in zip(test_ids, test_embeds)}\n# train_dict1 = {str(pid): embbedding for pid, embbedding in zip(train_ids1, train_embeds1)}\n# test_dict1 = {str(pid): embbedding for pid, embbedding in zip(test_ids1, test_embeds1)}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:28:44.666096Z","iopub.execute_input":"2026-01-27T17:28:44.666379Z","iopub.status.idle":"2026-01-27T17:28:51.818587Z","shell.execute_reply.started":"2026-01-27T17:28:44.666358Z","shell.execute_reply":"2026-01-27T17:28:51.818032Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"testesm_emb_names = [f\"esm_{i}\" for i in range(test_embeds.shape[1])]\nesm1_test = pd.DataFrame(test_embeds, columns=testesm_emb_names)\nesm1_test.insert(0, 'test_ids', test_ids)\nesm1_test.insert(1, 'protein', esm1_test['test_ids'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:28:51.819378Z","iopub.execute_input":"2026-01-27T17:28:51.819629Z","iopub.status.idle":"2026-01-27T17:28:51.842927Z","shell.execute_reply.started":"2026-01-27T17:28:51.819596Z","shell.execute_reply":"2026-01-27T17:28:51.842231Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df = pd.merge(test_df, esm1_test,on='protein', how='left')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:28:51.843824Z","iopub.execute_input":"2026-01-27T17:28:51.844069Z","iopub.status.idle":"2026-01-27T17:28:56.755302Z","shell.execute_reply.started":"2026-01-27T17:28:51.844051Z","shell.execute_reply":"2026-01-27T17:28:56.754743Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_full_embedding = pd.merge(test_df, test_tax_hl, on='protein', how='left')\nprint(f\"The shape of all embeddings: {test_full_embedding.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:28:56.756235Z","iopub.execute_input":"2026-01-27T17:28:56.756847Z","iopub.status.idle":"2026-01-27T17:28:57.146945Z","shell.execute_reply.started":"2026-01-27T17:28:56.756817Z","shell.execute_reply":"2026-01-27T17:28:57.146165Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_ids1 = test_full_embedding['protein'].astype(str).tolist()\nesm1_test_fts = test_full_embedding[esm_emb_names].values\nman_testfeatures = test_full_embedding[manual_cols].values","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:28:57.147965Z","iopub.execute_input":"2026-01-27T17:28:57.148201Z","iopub.status.idle":"2026-01-27T17:28:58.097262Z","shell.execute_reply.started":"2026-01-27T17:28:57.148181Z","shell.execute_reply":"2026-01-27T17:28:58.096692Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Convert to float32 BEFORE scaling\nesm1_test_fts = esm1_test_fts.astype('float32')\nman_testfeatures = man_testfeatures.astype('float32')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:28:58.098251Z","iopub.execute_input":"2026-01-27T17:28:58.098555Z","iopub.status.idle":"2026-01-27T17:28:59.086510Z","shell.execute_reply.started":"2026-01-27T17:28:58.098528Z","shell.execute_reply":"2026-01-27T17:28:59.085766Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport torch\n\n# Load Training Taxonomy Indices\ntrain_order = np.load(Path(CONFIG['Taxonomy_HL']) /'train_order_idx.npy')\ntrain_family = np.load(Path(CONFIG['Taxonomy_HL'])/'train_family_idx.npy')\ntrain_genus = np.load(Path(CONFIG['Taxonomy_HL'])/'train_genus_idx.npy')\n\n# Stack them into a single matrix: (N_proteins, 3)\ntrain_tax_indices = np.stack([train_order, train_family, train_genus], axis=1)\n\n# Load Test Taxonomy Indices\ntest_order = np.load(Path(CONFIG['Taxonomy_HL'])/'test_order_idx.npy')\ntest_family = np.load(Path(CONFIG['Taxonomy_HL'])/'test_family_idx.npy')\ntest_genus = np.load(Path(CONFIG['Taxonomy_HL'])/'test_genus_idx.npy')\n\ntest_tax_indices = np.stack([test_order, test_family, test_genus], axis=1)\n\nprint(f\"Taxonomy feature shape: {train_tax_indices.shape}\") # Should be (142246, 3) or similar","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:28:59.087551Z","iopub.execute_input":"2026-01-27T17:28:59.087943Z","iopub.status.idle":"2026-01-27T17:29:02.796690Z","shell.execute_reply.started":"2026-01-27T17:28:59.087913Z","shell.execute_reply":"2026-01-27T17:29:02.796050Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:02.797589Z","iopub.execute_input":"2026-01-27T17:29:02.798078Z","iopub.status.idle":"2026-01-27T17:29:02.854627Z","shell.execute_reply.started":"2026-01-27T17:29:02.798054Z","shell.execute_reply":"2026-01-27T17:29:02.853924Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# class ProteinDataset(Dataset):\n    # def __init__(self, embeddings, manual_fts, taxonomy_indices, sparse_labels):\n    #     self.embeddings = torch.tensor(embeddings, dtype=torch.float32)\n    #     self.manual_ftrs = torch.tensor(manual_fts, dtype=torch.float32)\n    #     # taxonomy_indices is a Nx3 matrix: [order_idx, family_idx, genus_idx]\n    #     self.tax_indices = torch.tensor(taxonomy_indices, dtype=torch.long)\n    #     self.labels = sparse_labels\n        \n    # def __getitem__(self, idx):\n    #     return (self.embeddings[idx], \n    #             self.manual_ftrs[idx], \n    #             self.tax_indices[idx], \n    #             torch.tensor(self.labels[idx].toarray().flatten(), dtype=torch.float32))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:02.855617Z","iopub.execute_input":"2026-01-27T17:29:02.855943Z","iopub.status.idle":"2026-01-27T17:29:02.868158Z","shell.execute_reply.started":"2026-01-27T17:29:02.855911Z","shell.execute_reply":"2026-01-27T17:29:02.867643Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ProteinDataset(Dataset):\n    def __init__(self, esm_embeddings, manual_fts, tax_indices, labels):\n        self.esm = torch.tensor(esm_embeddings, dtype=torch.float32)\n        self.manual = torch.tensor(manual_fts, dtype=torch.float32)\n        # taxonomy_indices is a Nx3 matrix: [order_idx, family_idx, genus_idx]\n        self.tax = torch.tensor(tax_indices, dtype=torch.long) # Indices must be Long\n        self.labels = labels # Keep sparse if necessary\n        \n    def __len__(self):\n        return len(self.esm)\n    \n    def __getitem__(self, idx):\n        label_row = torch.tensor(self.labels[idx].toarray().flatten(), dtype=torch.float32)\n        \n        # Return 4 items now\n        return self.esm[idx], self.manual[idx], self.tax[idx], label_row","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:02.868883Z","iopub.execute_input":"2026-01-27T17:29:02.869116Z","iopub.status.idle":"2026-01-27T17:29:02.880553Z","shell.execute_reply.started":"2026-01-27T17:29:02.869094Z","shell.execute_reply":"2026-01-27T17:29:02.879988Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\nclass MLP(nn.Module):\n    def __init__(self, input_dim=3304, output_dim=1500):\n        super(MLP, self).__init__()\n        self.input_norm = nn.BatchNorm1d(input_dim)\n        self.feature_extractor = nn.Sequential(\n            # Layer 1: 3304 -> 2048\n            nn.Linear(input_dim, 2048),\n            nn.BatchNorm1d(2048),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            \n            # Layer 2: 2048 -> 1024\n            nn.Linear(2048, 1024),\n            nn.BatchNorm1d(1024),\n            nn.ReLU(),\n            nn.Dropout(0.2)\n        )\n        \n        # Final Output Layer: 1024 -> 1500\n        self.classifier = nn.Linear(1024, output_dim)\n        \n    def forward(self, x):\n        x = self.input_norm(x)\n        x = self.feature_extractor(x)\n        # We don't use Sigmoid here because we use BCEWithLogitsLoss\n        logits = self.classifier(x)\n        return","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:02.881377Z","iopub.execute_input":"2026-01-27T17:29:02.881618Z","iopub.status.idle":"2026-01-27T17:29:02.898142Z","shell.execute_reply.started":"2026-01-27T17:29:02.881599Z","shell.execute_reply":"2026-01-27T17:29:02.897590Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CAFA6HybridModel(nn.Module):\n    def __init__(self, num_labels, esm_dim=1280, manual_dim=26, \n                 max_order=200, max_family=1000, max_genus=2000):\n        super().__init__()\n        \n        # Taxonomy Embeddings (small vectors are enough)\n        self.order_emb = nn.Embedding(max_order, 8)\n        self.family_emb = nn.Embedding(max_family, 16)\n        self.genus_emb = nn.Embedding(max_genus, 16)\n        \n        # self.manual_norm = nn.BatchNorm1d(manual_dim)\n        \n        # Total Input = 1280 + 26 + (8 + 16 + 16) = 1346\n        combined_dim = esm_dim + manual_dim + 8 + 16 + 16\n        self.input_norm = nn.BatchNorm1d(combined_dim)\n        self.feature_extractor = nn.Sequential(\n            nn.Linear(combined_dim, 2048),\n            nn.BatchNorm1d(2048),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            \n            # Layer 2: 2048 -> 1024\n            nn.Linear(2048, 1024),\n            nn.BatchNorm1d(1024),\n            nn.ReLU(),\n            nn.Dropout(0.2)\n        )\n\n        # Final Output Layer: 1024 -> 1500\n        self.classifier = nn.Linear(1024, num_labels)\n        \n    def forward(self, esm_feat, manual_feat, tax_indices):\n        # 1. Process Taxonomy\n        o = self.order_emb(tax_indices[:, 0])\n        f = self.family_emb(tax_indices[:, 1])\n        g = self.genus_emb(tax_indices[:, 2])\n        \n        # 2. Normalize Manual\n        # m = self.manual_norm(manual_feat)\n        \n        # 3. Concatenate Everything\n        combined = torch.cat([esm_feat, manual_feat, o, f, g], dim=1)\n        x = self.input_norm(combined)\n        x = self.feature_extractor(x)\n        # We don't use Sigmoid here because we use BCEWithLogitsLoss\n        logits = self.classifier(x)\n        \n        return logits","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:02.898860Z","iopub.execute_input":"2026-01-27T17:29:02.899103Z","iopub.status.idle":"2026-01-27T17:29:02.911447Z","shell.execute_reply.started":"2026-01-27T17:29:02.899072Z","shell.execute_reply":"2026-01-27T17:29:02.910760Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"val_final_list = list(val_final.protein)\ntrain_final_list = list(train_final.protein)\nprint(len(set(train_final_list)), len(set(val_final_list)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:02.912250Z","iopub.execute_input":"2026-01-27T17:29:02.912493Z","iopub.status.idle":"2026-01-27T17:29:02.943518Z","shell.execute_reply.started":"2026-01-27T17:29:02.912473Z","shell.execute_reply":"2026-01-27T17:29:02.942767Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from collections import Counter\n\ndef pr_ont(train_terms,train_final, val_final):\n    val_final_list = list(val_final.protein)\n    train_final_list = list(train_final.protein)\n    val_list_ont, val_list_C, val_list_F = [], [], []\n    \n    for ontology in ['P', 'C', 'F']:   \n        ont_terms_df = train_terms[train_terms['ontology'] == ontology]\n        pr_on = list(ont_terms_df.protein)\n        result = list(set(pr_on) & set(val_final_list))\n        \n        # val_list_ont.append(result) \n        val_list_C.append(len(result))\n    \n    return val_list_C\npr_ont(train_terms,train_final, val_final)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:02.944277Z","iopub.execute_input":"2026-01-27T17:29:02.944454Z","iopub.status.idle":"2026-01-27T17:29:03.150904Z","shell.execute_reply.started":"2026-01-27T17:29:02.944438Z","shell.execute_reply":"2026-01-27T17:29:03.150228Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def parse_obo_parents(go_obo_path):\n    print(f\"[1/5] Parsing OBO Ontology...\")\n    term_parents = defaultdict(set)\n    roots = set(['GO:0003674', 'GO:0008150', 'GO:0005575'])\n    \n    with open(go_obo_path, \"r\") as f:\n        cur_id = None\n        for line in f:\n            line = line.strip()\n            if line == \"[Term]\":\n                cur_id = None\n            elif line.startswith(\"id: \"):\n                cur_id = line.split(\"id: \")[1].strip()\n            elif line.startswith(\"is_a: \"):\n                pid = line.split()[1].strip()\n                if cur_id: term_parents[cur_id].add(pid)\n            elif line.startswith(\"relationship: part_of \"):\n                parts = line.split()\n                if len(parts) >= 3:\n                    pid = parts[2].strip()\n                    if cur_id: term_parents[cur_id].add(pid)\n    return term_parents, roots","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:03.151750Z","iopub.execute_input":"2026-01-27T17:29:03.151996Z","iopub.status.idle":"2026-01-27T17:29:03.157901Z","shell.execute_reply.started":"2026-01-27T17:29:03.151975Z","shell.execute_reply":"2026-01-27T17:29:03.157300Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def parse_obo_parents(go_obo_path):\n    print(f\"[1/5] Parsing OBO Ontology...\")\n    term_parents = defaultdict(set)\n    roots = set(['GO:0003674', 'GO:0008150', 'GO:0005575'])\n    \n    with open(go_obo_path, \"r\") as f:\n        cur_id = None\n        for line in f:\n            line = line.strip()\n            if line == \"[Term]\":\n                cur_id = None\n            elif line.startswith(\"id: \"):\n                cur_id = line.split(\"id: \")[1].strip()\n            elif line.startswith(\"is_a: \"):\n                pid = line.split()[1].strip()\n                if cur_id: term_parents[cur_id].add(pid)\n            elif line.startswith(\"relationship: part_of \"):\n                parts = line.split()\n                if len(parts) >= 3:\n                    pid = parts[2].strip()\n                    if cur_id: term_parents[cur_id].add(pid)\n    return term_parents, roots\n\ndef get_ancestors_map(term_parents):\n    print(\"[1/5] Building Ancestor Map...\")\n    ancestors = {}\n    def get_all_ancestors(term):\n        if term in ancestors: return ancestors[term]\n        parents = term_parents.get(term, set())\n        all_anc = set(parents)\n        for p in parents:\n            all_anc |= get_all_ancestors(p)\n        ancestors[term] = all_anc\n        return all_anc\n    \n    for term in tqdm(list(term_parents.keys())):\n        get_all_ancestors(term)\n    return ancestors\n# Create a cache to store results of get_ancestors\nancestor_cache = {}\n\ndef get_ancestors_cached(go_id, parents_map):\n    if go_id in ancestor_cache:\n        return ancestor_cache[go_id]\n    \n    ans = {go_id}\n    stack = [go_id]\n    while stack:\n        cur = stack.pop()\n        for p in parents_map.get(cur, []):\n            if p not in ans:\n                ans.add(p)\n                stack.append(p)\n    \n    ancestor_cache[go_id] = ans\n    return ans","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:03.158674Z","iopub.execute_input":"2026-01-27T17:29:03.158947Z","iopub.status.idle":"2026-01-27T17:29:03.171723Z","shell.execute_reply.started":"2026-01-27T17:29:03.158918Z","shell.execute_reply":"2026-01-27T17:29:03.171206Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"parents_map, children_map = parse_obo_parents(obo_path)\n# parents_map","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:03.172606Z","iopub.execute_input":"2026-01-27T17:29:03.172895Z","iopub.status.idle":"2026-01-27T17:29:03.636172Z","shell.execute_reply.started":"2026-01-27T17:29:03.172870Z","shell.execute_reply":"2026-01-27T17:29:03.635589Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def rebuild_propagated_matrix(protein_mapping, ontology_terms, parents_map):\n    \"\"\"\n    protein_mapping: {protein_id: [list of leaf GO terms]}\n    ontology_terms: set of GO terms allowed for this ontology (from your MLB)\n    \"\"\"\n    propagated_rows = []\n    \n    # We iterate through the proteins in the EXACT order of your train_ids\n    for pid in tqdm(train_ids1):\n        # Get original terms for this protein\n        original_terms = protein_mapping.get(pid, [])\n        \n        full_set = set()\n        for term in original_terms:\n            # 1. Get all ancestors\n            ancestors = get_ancestors_cached(term, parents_map)\n            # 2. Only keep ancestors that are in our allowed ontology list\n            valid_ancestors = ancestors.intersection(ontology_terms)\n            full_set.update(valid_ancestors)\n            \n        propagated_rows.append(list(full_set))\n        \n    return propagated_rows\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:03.637041Z","iopub.execute_input":"2026-01-27T17:29:03.637370Z","iopub.status.idle":"2026-01-27T17:29:03.642060Z","shell.execute_reply.started":"2026-01-27T17:29:03.637349Z","shell.execute_reply":"2026-01-27T17:29:03.641457Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from collections import Counter\nfrom sklearn.preprocessing import MultiLabelBinarizer\nimport pickle\nmapping = defaultdict(list)\nfor _, r in tqdm(train_terms.iterrows(), total=len(train_terms)): \n    mapping[r.protein].append(r.term)\n# Define your limits clearly\n\ny_trains = dict()\nmlb_dict = dict()\nmodels = dict()\nproteins_dict = dict()\n\nLIMITS = {'P': 1000, 'C': 1500, 'F': 1500}\n\nfor ontology in ['P', 'C', 'F']:\n    \n    # 1. Filter and Group\n    ont_terms_df = train_terms[train_terms['ontology'] == ontology]\n    protein_terms_dict = ont_terms_df.groupby('protein')['term'].apply(list).to_dict()\n    protein_list = list(ont_terms_df.protein.unique())\n    proteins_dict[ontology] = protein_list\n    # 2. Identify Top K Terms for THIS ontology\n    # We count frequencies within the current ontology only\n    term_counts = Counter()\n    for terms in protein_terms_dict.values():\n        term_counts.update(terms)\n    \n    # Select the Top N terms (e.g., 1500 for F, 800 for P/C)\n    top_n = LIMITS[ontology]\n    top_terms = [t for t, _ in term_counts.most_common(top_n)]\n    \n    # 3. Prepare the labels for the MLB\n    # Crucial: Ensure every EntryID in your X_train has a corresponding row in Y\n    labels = [protein_terms_dict.get(entry_id, []) for entry_id in train_ids1]\n    \n    # 4. Fit the Binarizer restricted to ONLY top_terms\n    # Any term NOT in top_terms will be ignored during fit_transform\n    mlb = MultiLabelBinarizer(classes=top_terms, sparse_output=True)\n    # Example usage for Ontology F\n#     allowed_terms = set(mlb.classes)\n    allowed_terms = set(mlb.classes)\n    \n    propagated_list = rebuild_propagated_matrix(mapping, allowed_terms, parents_map)\n\n#     Convert to Sparse Matrix immediately\n    y_trains[ontology] = mlb.fit_transform(propagated_list)\n    \n#     propagated_list = rebuild_propagated_matrix(mapping, allowed_terms, parents_map)\n\n#     Convert to Sparse Matrix immediately\n    # y_trains[ontology] = mlb.fit_transform(labels)\n    \n#     y_trains[ontology] = mlb.fit_transform(labels)\n    \n    labels_train = y_trains[ontology]\n    mlb_dict[ontology] = mlb\n    \n    print(f\"ontology {ontology}: Output shape {labels_train.shape} (Limited to Top {top_n})\")\n    # 2. SAVE IT IMMEDIATELY THIS TIME\n    with open('mlb_dict.pkl', 'wb') as f:\n        pickle.dump(mlb_dict, f)\n    # 5. Initialize Model\n    # input_dim is 3304, output_dim is now exactly 1500 or 800\n    model = CAFA6HybridModel(num_labels=labels_train.shape[1], esm_dim=esm1_fts.shape[1], manual_dim=man_features.shape[1], \n                 max_order=200, max_family=1000, max_genus=2000)\n    models[ontology] = model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:03.642853Z","iopub.execute_input":"2026-01-27T17:29:03.643119Z","iopub.status.idle":"2026-01-27T17:29:31.302100Z","shell.execute_reply.started":"2026-01-27T17:29:03.643089Z","shell.execute_reply":"2026-01-27T17:29:31.301351Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"models","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:31.302958Z","iopub.execute_input":"2026-01-27T17:29:31.303224Z","iopub.status.idle":"2026-01-27T17:29:31.308335Z","shell.execute_reply.started":"2026-01-27T17:29:31.303203Z","shell.execute_reply":"2026-01-27T17:29:31.307618Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pp = proteins_dict['P']\npc = proteins_dict['C']\npf = proteins_dict['F']\nprint(len(pp), len(pc), len(pf))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:31.309287Z","iopub.execute_input":"2026-01-27T17:29:31.309528Z","iopub.status.idle":"2026-01-27T17:29:31.323920Z","shell.execute_reply.started":"2026-01-27T17:29:31.309509Z","shell.execute_reply":"2026-01-27T17:29:31.323246Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for ontology in ['P', 'C', 'F']:\n    print(f\"\\n--- Training ontology: {ontology} ---\")\n    \n    # --- RESET for each ontology ---\n    best_val_loss = float('inf')\n    counter = 0\n    patience = 5 \n    train_history = []\n    val_history = []\n    scores = []\n    \n    # 1. Generate indices for the split\n    train_indices = proteins_dict[ontology]\n    indices = np.arange(len(train_indices)) \n    val_prot = list(set(train_indices) & set(val_final_list))\n    index_map = {protein: i for i, protein in enumerate(train_indices)}\n    val_idx = [index_map[p] for p in val_prot if p in index_map]\n    print(val_idx[:5])\n    train_prot = [item for item in train_indices if item not in val_idx]\n    train_idx = [index_map[p] for p in train_prot if p in index_map]\n    print(train_idx[:5])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:31.324789Z","iopub.execute_input":"2026-01-27T17:29:31.325147Z","iopub.status.idle":"2026-01-27T17:29:33.701793Z","shell.execute_reply.started":"2026-01-27T17:29:31.325117Z","shell.execute_reply":"2026-01-27T17:29:33.701182Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for ontology in ['P', 'C', 'F']:\n    print(f\"\\n--- Training ontology: {ontology} ---\")\n    \n    # --- RESET for each ontology ---\n    best_val_loss = float('inf')\n    counter = 0\n    patience = 5 \n    train_history = []\n    val_history = []\n    scores = []\n    \n    # 1. Generate indices for the split\n    train_indices = proteins_dict[ontology]\n    indices = np.arange(len(train_indices)) \n    val_idx = list(set(train_indices) & set(val_final_list))\n    print(val_idx[:5])\n    train_idx = [item for item in train_indices if item not in val_idx]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:33.702743Z","iopub.execute_input":"2026-01-27T17:29:33.703061Z","iopub.status.idle":"2026-01-27T17:29:35.980041Z","shell.execute_reply.started":"2026-01-27T17:29:33.703028Z","shell.execute_reply":"2026-01-27T17:29:35.979452Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n\ndef get_f1_max_with_ia(y_true, y_probs, ia_map, mlb_obj):\n    \"\"\"\n    Correct Protein-Centric IA-weighted F1-Max calculation for CAFA5.\n    \"\"\"\n    # 1. Align weights with the columns of y_probs\n    cls = mlb_obj.classes_\n    weights = np.array([ia_map.get(c, 0.0) for c in cls])\n    \n    # 2. Pre-calculate true-label weights per protein (for Recall denominator)\n    # Shape: (num_proteins,)\n    true_weighted_sum = (y_true * weights).sum(axis=1)\n    \n    best_f1 = 0\n    best_threshold = 0\n\n    # 3. Threshold Sweep\n    for threshold in np.linspace(0.01, 0.9, 50):\n        y_pred_bin = (y_probs >= threshold).astype(int)\n        \n        # Calculate intersection of Pred and True weighted by IA\n        # Shape: (num_proteins, num_labels) -> sum(axis=1) -> (num_proteins,)\n        intersection_weighted_sum = ((y_true * y_pred_bin) * weights).sum(axis=1)\n        \n        # Predicted weights per protein (for Precision denominator)\n        pred_weighted_sum = (y_pred_bin * weights).sum(axis=1)\n        \n        # Precision per protein: (Intersection / Predicted)\n        # We use a mask to handle proteins with zero predictions\n        prec_n = np.zeros_like(intersection_weighted_sum)\n        mask_p = pred_weighted_sum > 0\n        prec_n[mask_p] = intersection_weighted_sum[mask_p] / pred_weighted_sum[mask_p]\n        \n        # Recall per protein: (Intersection / True)\n        rec_n = np.zeros_like(intersection_weighted_sum)\n        mask_r = true_weighted_sum > 0\n        rec_n[mask_r] = intersection_weighted_sum[mask_r] / true_weighted_sum[mask_r]\n        \n        # 4. Average across all proteins (Sample-centric)\n        avg_precision = np.mean(prec_n)\n        avg_recall = np.mean(rec_n)\n        \n        # 5. Calculate F1 for this threshold\n        if avg_precision + avg_recall > 0:\n            f1 = (2 * avg_precision * avg_recall) / (avg_precision + avg_recall)\n        else:\n            f1 = 0\n            \n        if f1 > best_f1:\n            best_f1 = f1\n            best_threshold = threshold\n            \n    return best_f1, best_threshold","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:35.980947Z","iopub.execute_input":"2026-01-27T17:29:35.981240Z","iopub.status.idle":"2026-01-27T17:29:35.988604Z","shell.execute_reply.started":"2026-01-27T17:29:35.981218Z","shell.execute_reply":"2026-01-27T17:29:35.987938Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ia_map = ia_df.set_index('term')['ia_score'].to_dict() # Assuming you loaded IA.tsv\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:35.989547Z","iopub.execute_input":"2026-01-27T17:29:35.989843Z","iopub.status.idle":"2026-01-27T17:29:36.019594Z","shell.execute_reply.started":"2026-01-27T17:29:35.989824Z","shell.execute_reply":"2026-01-27T17:29:36.018958Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom torch.cuda.amp import autocast, GradScaler\nfrom sklearn.model_selection import train_test_split\nfrom torch.optim.lr_scheduler import LambdaLR, ReduceLROnPlateau\n\n# 1. Initialize result storage\nall_val_predictions = {}\nall_train_loss = {}\nall_val_loss = {}\nall_val_true = {} # New dictionary for ground truth\n\nscaler = GradScaler() # For Mixed Precision\nLIMITS = {'P': 1000, 'C': 1000, 'F': 1000}\n\nfor ontology in ['P', 'C', 'F']:\n    print(f\"\\n--- Training ontology: {ontology} ---\")\n    # --- Parameters ---\n    WARMUP_EPOCHS = 3\n    # --- RESET for each ontology ---\n    best_val_loss = float('inf')\n    counter = 0\n    patience = 5 \n    train_history = []\n    val_history = []\n    scores = []\n    \n    # 1. Generate indices for the split\n    train_indices = proteins_dict[ontology]\n    indices = np.arange(len(train_indices)) \n    val_prot = list(set(train_indices) & set(val_final_list))\n    train_prot = [item for item in train_indices if item not in val_prot]\n\n    index_map = {protein: i for i, protein in enumerate(train_indices)}\n    val_idx = [index_map[p] for p in val_prot if p in index_map]\n    train_idx = [index_map[p] for p in train_prot if p in index_map]\n\n    \n    train_ds = ProteinDataset(esm1_fts[train_idx], man_features[train_idx],train_tax_indices[train_idx] , y_trains[ontology][train_idx])\n    val_ds = ProteinDataset(esm1_fts[val_idx], man_features[val_idx],train_tax_indices[val_idx], y_trains[ontology][val_idx])\n    \n    train_loader = DataLoader(train_ds, batch_size=CONFIG['BATCH_SIZE'], shuffle=True)\n    val_loader = DataLoader(val_ds, batch_size=CONFIG['BATCH_SIZE'], shuffle=False)\n    \n    # f1_score = MultilabelF1Score(num_labels=LIMITS[ontology]).to(device)\n    \n    model = models[ontology].to(device)\n    optimizer = torch.optim.Adam(model.parameters(), lr=1e-5, weight_decay=1e-5)\n    # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n#     optimizer, \n#     mode='min', \n#     factor=0.5, \n#     patience=5, \n#     # verbose=True\n# )\n    # 2. Define Warmup Scheduler\n    # Increases LR linearly from 10% to 100% of TARGET_LR over WARMUP_EPOCHS\n    # warmup_lambda = lambda epoch: 0.1 + 0.9 * (epoch / WARMUP_EPOCHS) if epoch < WARMUP_EPOCHS else 1.0\n    # warmup_scheduler = LambdaLR(optimizer, lr_lambda=warmup_lambda)\n    # 3. Define Plateau Scheduler\n    # plateau_scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)\n#     criterion = torch.nn.BCELoss()\n    criterion = torch.nn.BCEWithLogitsLoss()\n    \n#     pos_weight = compute_pos_weight(y_trains[ontology]).to(device)\n#     pos_weight = torch.clamp(pos_weight, max=50.0)\n#     criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\n    \n#     criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight_tensor)\n    for epoch in range(CONFIG['EPOCHS']):\n        # --- TRAINING PHASE (with Mixed Precision) ---\n        model.train()\n        total_train_loss = 0\n        \n        for batch_x,mn_ft_x,tax_x,batch_y in train_loader:\n            batch_x, mn_ft_x, tax_x, batch_y = batch_x.to(device), mn_ft_x.to(device),tax_x.to(device) ,batch_y.to(device)\n            optimizer.zero_grad()\n            logits = model(batch_x, mn_ft_x, tax_x)\n            loss = criterion(logits, batch_y)\n            \n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            total_train_loss += loss.item()\n            \n        avg_train_loss = total_train_loss / len(train_loader)\n        train_history.append(avg_train_loss)\n\n        # --- VALIDATION PHASE ---\n        model.eval()\n        total_val_loss = 0\n        with torch.no_grad():\n            for batch_x,mn_ft_x,tax_x, batch_y in val_loader:\n                batch_x,mn_ft_x, tax_x, batch_y = batch_x.to(device), mn_ft_x.to(device),tax_x.to(device) ,batch_y.to(device)\n                logits = model(batch_x, mn_ft_x, tax_x)\n                \n                # USE LOGITS HERE, NOT SIGMOID PROBS\n                loss = criterion(logits, batch_y) \n                total_val_loss += loss.item()\n#                 score=f1_score(logits, batch_y)\n#                 scores.append(score.item())\n#         avg_score = np.mean(scores)\n        avg_val_loss = total_val_loss / len(val_loader)\n        val_history.append(avg_val_loss)\n        # scheduler.step(avg_val_loss)\n    \n        # current_lr = optimizer.param_groups[0]['lr']\n#         print(f\"Epoch {epoch+1}: TRAIN F1-Score : , {avg_score}\")\n        print(f\"Epoch {epoch+1}: Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}\")\n\n       \n        # --- EARLY STOPPING LOGIC (Inside Epoch Loop) ---\n        if avg_val_loss < best_val_loss:\n            best_val_loss = avg_val_loss\n            torch.save(model.state_dict(), f'best_model_ESM_Model_{ontology}.pth')\n            print(f\"  ➜ Best model saved (Loss: {best_val_loss:.4f})\")\n            counter = 0\n        else:\n            counter += 1\n            if counter >= patience:\n                print(f\"  ➜ Early stopping triggered for {ontology} at epoch {epoch+1}\")\n                break\n\n         # --- Step the Schedulers ---\n        # if epoch < WARMUP_EPOCHS:\n        #     warmup_scheduler.step()\n        #     print(f\"🔥 Warmup Phase: LR set to {optimizer.param_groups[0]['lr']:.6f}\")\n        # else:\n        #     # After warmup, let the plateau scheduler take over based on val_loss\n        #     plateau_scheduler.step(avg_val_loss)\n        #     print(f\"❄️ Plateau Phase: LR is {optimizer.param_groups[0]['lr']:.6f}\")\n            \n                \n    y_probs_list = []\n    y_true_list = []\n    # --- FINAL PROBABILITY RECORDING (Using Best Model) ---\n    print(f\"Recording final probabilities for {ontology}...\")\n    model.load_state_dict(torch.load(f'best_model_ESM_Model_{ontology}.pth'))\n    model.eval()\n    final_probs = []\n    with torch.no_grad():\n        for batch_x,mn_ft_x, tax_x, batch_y in val_loader:\n            logits = model(batch_x.to(device), mn_ft_x.to(device), tax_x.to(device))\n            probs = torch.sigmoid(logits).cpu().numpy()\n            y_true_list.append(batch_y.numpy())\n            final_probs.append(probs)\n            \n    all_val_true[ontology] = np.vstack(y_true_list)\n    all_val_predictions[ontology] = np.vstack(final_probs)\n    all_train_loss[ontology] = train_history\n    all_val_loss[ontology] = val_history","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:29:36.020623Z","iopub.execute_input":"2026-01-27T17:29:36.020958Z","iopub.status.idle":"2026-01-27T17:57:06.597593Z","shell.execute_reply.started":"2026-01-27T17:29:36.020930Z","shell.execute_reply":"2026-01-27T17:57:06.596919Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ndef plot_learning_curves(all_train_loss, all_val_loss):\n    fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n    for i, ont in enumerate(['P', 'C', 'F']):\n        axes[i].plot(all_train_loss[ont], label='Train Loss')\n        axes[i].plot(all_val_loss[ont], label='Val Loss')\n        axes[i].set_title(f'Ontology {ont} - Learning Curve')\n        axes[i].set_xlabel('Epochs')\n        axes[i].set_ylabel('BCE Loss')\n        axes[i].legend()\n    plt.tight_layout()\n    plt.show()\n\nplot_learning_curves(all_train_loss, all_val_loss)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:57:06.598665Z","iopub.execute_input":"2026-01-27T17:57:06.599227Z","iopub.status.idle":"2026-01-27T17:57:07.059740Z","shell.execute_reply.started":"2026-01-27T17:57:06.599203Z","shell.execute_reply":"2026-01-27T17:57:07.059003Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"final_scores = {}\n# ia_map = ia_df.set_index('term')['ia'].to_dict() # Assuming you loaded IA.tsv\n\nprint(\"\\n=== Final Evaluation: IA-weighted F1-Max ===\")\n\nfor ontology in ['P', 'C', 'F']:\n    y_true = all_val_true[ontology]\n    y_probs = all_val_predictions[ontology]\n    mlb_obj = mlb_dict[ontology]\n    \n    # Calculate the score using your function\n    f1_max_ia, final_tau = get_f1_max_with_ia(y_true, y_probs, ia_map, mlb_obj)\n    \n    final_scores[ontology] = f1_max_ia\n    print(f\"Ontology {ontology}: F1-Max (IA) = {f1_max_ia:.4f}\")\n\n# Optional: Calculate an overall weighted average (CAFA5 Style)\n# BPO (P) is usually weighted more heavily in official rankings\n# overall_score = (final_scores['P'] * 0.3) + (final_scores['C'] * 0.3) + (final_scores['F'] * 0.4)\noverall_score = (final_scores['P'] * 0.33) + (final_scores['C'] * 0.33) + (final_scores['F'] * 0.33)\nprint(f\"\\nEstimated Weighted CAFA6 Score: {overall_score:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:57:07.060755Z","iopub.execute_input":"2026-01-27T17:57:07.061063Z","iopub.status.idle":"2026-01-27T17:57:08.414413Z","shell.execute_reply.started":"2026-01-27T17:57:07.061034Z","shell.execute_reply":"2026-01-27T17:57:08.413781Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:57:08.415237Z","iopub.execute_input":"2026-01-27T17:57:08.415508Z","iopub.status.idle":"2026-01-27T17:57:08.435408Z","shell.execute_reply.started":"2026-01-27T17:57:08.415481Z","shell.execute_reply":"2026-01-27T17:57:08.434528Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ProteinTestDataset(Dataset):\n    def __init__(self, embeddings, manual_fts, tax_indices):\n        self.embeddings = torch.tensor(embeddings, dtype=torch.float32)\n        self.manual_ftrs = torch.tensor(manual_fts, dtype=torch.float32)\n        self.tax_indices = torch.tensor(tax_indices, dtype=torch.long)\n        \n    def __len__(self):\n        return len(self.embeddings)\n    \n    def __getitem__(self, idx):\n        return self.embeddings[idx], self.manual_ftrs[idx], self.tax_indices[idx]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:57:08.436431Z","iopub.execute_input":"2026-01-27T17:57:08.436650Z","iopub.status.idle":"2026-01-27T17:57:08.441465Z","shell.execute_reply.started":"2026-01-27T17:57:08.436630Z","shell.execute_reply":"2026-01-27T17:57:08.440651Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import csv\n\n# 1. Setup\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nSUBMISSION_FILE = 'submission_ESM.tsv'\n\ntest_idx = np.arange(len(test_ids1))\ntest_ds = ProteinTestDataset(esm1_test_fts[test_idx],man_testfeatures[test_idx], test_tax_indices[test_idx] ) \ntest_loader = DataLoader(test_ds, batch_size=1000, shuffle=False, num_workers=0)\n\n# We open the file once and write results as we find them\nwith open(SUBMISSION_FILE, 'w') as f:\n    # Use a standard tab-separated writer\n    writer = csv.writer(f, delimiter='\\t')\n    \n    for ont, model in models.items():\n        print(f\"\\n--- Processing Ontology: {ont} ---\")\n       \n        checkpoint_path = f'/kaggle/working/best_model_ESM_Model_{ont}.pth'\n        model.load_state_dict(torch.load(checkpoint_path, map_location=device))\n        model.to(device).eval()\n        \n        mlb = mlb_dict[ont]\n        go_terms = mlb.classes_ if hasattr(mlb, 'classes_') else mlb.classes\n        \n        current_idx = 0\n        with torch.no_grad():\n            for batch_x,mn_ft_x, tax_x in tqdm(test_loader, desc=f\"Inference {ont}\"):\n                batch_x = batch_x.to(device)\n                mn_ft_x = mn_ft_x.to(device)\n                tax_x = tax_x.to(device)\n                logits = model(batch_x, mn_ft_x, tax_x)\n                probs = torch.sigmoid(logits).cpu().numpy()\n                \n                for b in range(probs.shape[0]):\n                    protein_id = test_ids1[current_idx]\n                    \n                    # 1. Get indices and scores that pass threshold\n                    # We use a slightly higher threshold (e.g., 0.01) if 0.001 is too heavy\n                    mask = probs[b] >= 0.005 \n                    indices = np.where(mask)[0]\n                    scores = probs[b][mask]\n                    \n                    # 2. Sort this specific protein's results by confidence\n                    # This ensures we don't need a massive global sort later\n                    sorted_map = sorted(zip(indices, scores), key=lambda x: x[1], reverse=True)\n                    \n                    # 3. Write top 1500 (or however many pass threshold) to disk immediately\n                    for idx, score in sorted_map[:1500]:\n                        writer.writerow([protein_id, go_terms[idx], f\"{score:.3f}\"])\n                    \n                    current_idx += 1\n        \n        # Free up GPU memory after each ontology\n        del model\n        torch.cuda.empty_cache()\n\nprint(f\"✅ {SUBMISSION_FILE} created successfully using stream-to-disk!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T17:57:08.443199Z","iopub.execute_input":"2026-01-27T17:57:08.443439Z","iopub.status.idle":"2026-01-27T17:59:57.992955Z","shell.execute_reply.started":"2026-01-27T17:57:08.443419Z","shell.execute_reply":"2026-01-27T17:59:57.992283Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}