{"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},{"sourceId":291180104,"sourceType":"kernelVersion"}],"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":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:12.543364Z","iopub.execute_input":"2026-02-11T19:01:12.543627Z","iopub.status.idle":"2026-02-11T19:01:13.666283Z","shell.execute_reply.started":"2026-02-11T19:01:12.543596Z","shell.execute_reply":"2026-02-11T19:01:13.665349Z"}},"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-02-11T19:01:13.668040Z","iopub.execute_input":"2026-02-11T19:01:13.668646Z","iopub.status.idle":"2026-02-11T19:01:15.965533Z","shell.execute_reply.started":"2026-02-11T19:01:13.668601Z","shell.execute_reply":"2026-02-11T19:01:15.964727Z"}},"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-02-11T19:01:15.970213Z","iopub.execute_input":"2026-02-11T19:01:15.970922Z","iopub.status.idle":"2026-02-11T19:01:24.605497Z","shell.execute_reply.started":"2026-02-11T19:01:15.970883Z","shell.execute_reply":"2026-02-11T19:01:24.604781Z"}},"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 'LOSS' : 'INTERLABEL',\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': 50, # IMPROVED: 10 ā†’ 12 epochs\n 'BATCH_SIZE': 64, # IMPROVED: 32 ā†’ 64 (faster training)\n 'LEARNING_RATE': 5e-5 ,\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}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:24.606565Z","iopub.execute_input":"2026-02-11T19:01:24.606994Z","iopub.status.idle":"2026-02-11T19:01:24.612975Z","shell.execute_reply.started":"2026-02-11T19:01:24.606971Z","shell.execute_reply":"2026-02-11T19:01:24.612314Z"}},"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-02-11T19:01:24.613789Z","iopub.execute_input":"2026-02-11T19:01:24.614103Z","iopub.status.idle":"2026-02-11T19:01:24.627168Z","shell.execute_reply.started":"2026-02-11T19:01:24.614071Z","shell.execute_reply":"2026-02-11T19:01:24.626483Z"}},"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'])\ntest_tax = pd.read_csv('/kaggle/input/cafa-6-embeddings/taxonomy/test_species_proteins.tsv', sep='\\t',\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-02-11T19:01:24.628182Z","iopub.execute_input":"2026-02-11T19:01:24.628832Z","iopub.status.idle":"2026-02-11T19:01:25.116694Z","shell.execute_reply.started":"2026-02-11T19:01:24.628810Z","shell.execute_reply":"2026-02-11T19:01:25.115922Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ia_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.117764Z","iopub.execute_input":"2026-02-11T19:01:25.118025Z","iopub.status.idle":"2026-02-11T19:01:25.139934Z","shell.execute_reply.started":"2026-02-11T19:01:25.117997Z","shell.execute_reply":"2026-02-11T19:01:25.139380Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"species_counts = train_taxonomy['taxon'].value_counts()\ntop_20_species = species_counts.head(20).index.to_list()\nprint(f\"Top 3 species in training: {top_20_species[:20]}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.140713Z","iopub.execute_input":"2026-02-11T19:01:25.140978Z","iopub.status.idle":"2026-02-11T19:01:25.152882Z","shell.execute_reply.started":"2026-02-11T19:01:25.140951Z","shell.execute_reply":"2026-02-11T19:01:25.152347Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"species_to_idx = {spec: i for i, spec in enumerate(top_20_species)}\n\ndef get_taxonomic_index(taxon_id):\n return species_to_idx.get(taxon_id, 20) # Default to 20 if not in top 20\n \ntrain_taxonomy['taxon_id'] = train_taxonomy['taxon'].apply(get_taxonomic_index)\ntest_tax['taxon_id'] = test_tax['taxon'].apply(get_taxonomic_index)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.155534Z","iopub.execute_input":"2026-02-11T19:01:25.155747Z","iopub.status.idle":"2026-02-11T19:01:25.271662Z","shell.execute_reply.started":"2026-02-11T19:01:25.155729Z","shell.execute_reply":"2026-02-11T19:01:25.271126Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_taxonomy.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.272408Z","iopub.execute_input":"2026-02-11T19:01:25.272676Z","iopub.status.idle":"2026-02-11T19:01:25.279865Z","shell.execute_reply.started":"2026-02-11T19:01:25.272657Z","shell.execute_reply":"2026-02-11T19:01:25.279176Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_tax.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.280724Z","iopub.execute_input":"2026-02-11T19:01:25.281379Z","iopub.status.idle":"2026-02-11T19:01:25.292236Z","shell.execute_reply.started":"2026-02-11T19:01:25.281358Z","shell.execute_reply":"2026-02-11T19:01:25.291683Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create an empty matrix of zeros: (Total Proteins, 21)\ntax_vectors = np.zeros((len(train_taxonomy), 21))\n\n# Fill the matrix: for each row, set the column at 'tax_idx' to 1\nfor i, idx in enumerate(train_taxonomy['taxon_id']):\n tax_vectors[i, idx] = 1\n\n# Convert to a DataFrame to merge with your main features\ntax_cols = [f\"tax_{i}\" for i in range(21)]\ntrain_tax_features = pd.DataFrame(tax_vectors, columns=tax_cols)\ntrain_tax_features.insert(0, \"protein\", train_taxonomy['protein'])\ntrain_tax_features.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.292970Z","iopub.execute_input":"2026-02-11T19:01:25.293160Z","iopub.status.idle":"2026-02-11T19:01:25.352826Z","shell.execute_reply.started":"2026-02-11T19:01:25.293144Z","shell.execute_reply":"2026-02-11T19:01:25.352241Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create an empty matrix of zeros: (Total Proteins, 21)\ntest_tax_vectors = np.zeros((len(test_tax), 21))\n\n# Fill the matrix: for each row, set the column at 'tax_idx' to 1\nfor i, idx in enumerate(test_tax['taxon_id']):\n test_tax_vectors[i, idx] = 1\n\n# Convert to a DataFrame to merge with your main features\ntesttax_cols = [f\"tax_{i}\" for i in range(21)]\ntest_tax_features = pd.DataFrame(test_tax_vectors, columns=testtax_cols)\ntest_tax_features.insert(0, \"protein\", test_tax['protein'])\ntest_tax_features.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.353717Z","iopub.execute_input":"2026-02-11T19:01:25.353931Z","iopub.status.idle":"2026-02-11T19:01:25.476342Z","shell.execute_reply.started":"2026-02-11T19:01:25.353913Z","shell.execute_reply":"2026-02-11T19:01:25.475640Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ia_df.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.477237Z","iopub.execute_input":"2026-02-11T19:01:25.477459Z","iopub.status.idle":"2026-02-11T19:01:25.484119Z","shell.execute_reply.started":"2026-02-11T19:01:25.477436Z","shell.execute_reply":"2026-02-11T19:01:25.483338Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_terms.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.485117Z","iopub.execute_input":"2026-02-11T19:01:25.485395Z","iopub.status.idle":"2026-02-11T19:01:25.496966Z","shell.execute_reply.started":"2026-02-11T19:01:25.485366Z","shell.execute_reply":"2026-02-11T19:01:25.496361Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def select_top_label(train_terms, ia_df, aspect:str, topk: int):\n aspect_df = train_terms[train_terms['ontology'] == aspect]\n aspect_df_freq = aspect_df['term'].value_counts().reset_index()\n selection_df = pd.merge(aspect_df_freq, ia_df, on='term', how='left')\n selection_df['selection_metric'] = selection_df['count'] * selection_df['ia_score']\n top_terms = selection_df.sort_values(by='selection_metric', ascending=False).head(topk)\n topk_terms_list = top_terms.term.to_list()\n return top_terms, topk_terms_list","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.498008Z","iopub.execute_input":"2026-02-11T19:01:25.498262Z","iopub.status.idle":"2026-02-11T19:01:25.507332Z","shell.execute_reply.started":"2026-02-11T19:01:25.498240Z","shell.execute_reply":"2026-02-11T19:01:25.506600Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BPO_top_labels, top_BPO = select_top_label(train_terms, ia_df, aspect='P', topk=800)\nBPO_top_labels","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.508219Z","iopub.execute_input":"2026-02-11T19:01:25.508479Z","iopub.status.idle":"2026-02-11T19:01:25.612801Z","shell.execute_reply.started":"2026-02-11T19:01:25.508454Z","shell.execute_reply":"2026-02-11T19:01:25.612068Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BPO_top_labels['selection_metric'].min()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.613670Z","iopub.execute_input":"2026-02-11T19:01:25.614161Z","iopub.status.idle":"2026-02-11T19:01:25.618644Z","shell.execute_reply.started":"2026-02-11T19:01:25.614139Z","shell.execute_reply":"2026-02-11T19:01:25.618096Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BCO_top_labels, top_BCO = select_top_label(train_terms, ia_df, aspect='C', topk=800)\nBCO_top_labels['norm_metric'] = BCO_top_labels['selection_metric'] / BCO_top_labels['selection_metric'].max()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.619353Z","iopub.execute_input":"2026-02-11T19:01:25.619705Z","iopub.status.idle":"2026-02-11T19:01:25.696683Z","shell.execute_reply.started":"2026-02-11T19:01:25.619686Z","shell.execute_reply":"2026-02-11T19:01:25.695947Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"top_BCO[:3]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.697707Z","iopub.execute_input":"2026-02-11T19:01:25.697970Z","iopub.status.idle":"2026-02-11T19:01:25.702614Z","shell.execute_reply.started":"2026-02-11T19:01:25.697944Z","shell.execute_reply":"2026-02-11T19:01:25.702040Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BFO_top_labels, top_BFO = select_top_label(train_terms, ia_df, aspect='F', topk=800)\nBFO_top_labels['norm_metric'] = BFO_top_labels['selection_metric'] / BFO_top_labels['selection_metric'].max()\nBFO_top_labels","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:25.703524Z","iopub.execute_input":"2026-02-11T19:01:25.703852Z","iopub.status.idle":"2026-02-11T19:01:25.785560Z","shell.execute_reply.started":"2026-02-11T19:01:25.703825Z","shell.execute_reply":"2026-02-11T19:01:25.785012Z"}},"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-02-11T19:01:25.786418Z","iopub.execute_input":"2026-02-11T19:01:25.786719Z","iopub.status.idle":"2026-02-11T19:01:29.299220Z","shell.execute_reply.started":"2026-02-11T19:01:25.786688Z","shell.execute_reply":"2026-02-11T19:01:29.298609Z"}},"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-02-11T19:01:29.300196Z","iopub.execute_input":"2026-02-11T19:01:29.300653Z","iopub.status.idle":"2026-02-11T19:01:29.309236Z","shell.execute_reply.started":"2026-02-11T19:01:29.300630Z","shell.execute_reply":"2026-02-11T19:01:29.308546Z"}},"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-02-11T19:01:29.310042Z","iopub.execute_input":"2026-02-11T19:01:29.310520Z","iopub.status.idle":"2026-02-11T19:01:29.324882Z","shell.execute_reply.started":"2026-02-11T19:01:29.310498Z","shell.execute_reply":"2026-02-11T19:01:29.324205Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df = create_sequence_dataframe_from_fasta(train_sequences, data_type='train')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:01:29.325713Z","iopub.execute_input":"2026-02-11T19:01:29.325973Z","iopub.status.idle":"2026-02-11T19:04:14.526589Z","shell.execute_reply.started":"2026-02-11T19:01:29.325953Z","shell.execute_reply":"2026-02-11T19:04:14.525944Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df['Length_log'] = np.log1p(train_df['Length'])\ntrain_df['MW_log'] = np.log1p(train_df['molecular_weight'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:14.529606Z","iopub.execute_input":"2026-02-11T19:04:14.530000Z","iopub.status.idle":"2026-02-11T19:04:14.536715Z","shell.execute_reply.started":"2026-02-11T19:04:14.529980Z","shell.execute_reply":"2026-02-11T19:04:14.536173Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df1 = train_df.loc[train_df['Length'] < 3000]\ntrain_df1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:14.541222Z","iopub.execute_input":"2026-02-11T19:04:14.541519Z","iopub.status.idle":"2026-02-11T19:04:14.637773Z","shell.execute_reply.started":"2026-02-11T19:04:14.541500Z","shell.execute_reply":"2026-02-11T19:04:14.637149Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\n Loading ESM1b embeddings...\")\nesm1_base = f\"{CONFIG['ESM1_PATH']}\"\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-02-11T19:04:14.638686Z","iopub.execute_input":"2026-02-11T19:04:14.638925Z","iopub.status.idle":"2026-02-11T19:04:16.997324Z","shell.execute_reply.started":"2026-02-11T19:04:14.638890Z","shell.execute_reply":"2026-02-11T19:04:16.996583Z"}},"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-02-11T19:04:16.998343Z","iopub.execute_input":"2026-02-11T19:04:16.998767Z","iopub.status.idle":"2026-02-11T19:04:17.071442Z","shell.execute_reply.started":"2026-02-11T19:04:16.998736Z","shell.execute_reply":"2026-02-11T19:04:17.070864Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df1 = pd.merge(train_df1, esm1_train, on='protein', how='left')\nprint(f\"{train_df1.shape}\")\ntrain_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:17.072238Z","iopub.execute_input":"2026-02-11T19:04:17.072490Z","iopub.status.idle":"2026-02-11T19:04:18.841230Z","shell.execute_reply.started":"2026-02-11T19:04:17.072464Z","shell.execute_reply":"2026-02-11T19:04:18.840473Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df1 = pd.merge(train_df1, train_taxonomy, on='protein', how='left')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:18.842203Z","iopub.execute_input":"2026-02-11T19:04:18.842481Z","iopub.status.idle":"2026-02-11T19:04:18.987010Z","shell.execute_reply.started":"2026-02-11T19:04:18.842458Z","shell.execute_reply":"2026-02-11T19:04:18.986402Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df1 = pd.merge(train_df1, train_tax_features, on='protein', how='left')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:18.987882Z","iopub.execute_input":"2026-02-11T19:04:18.988150Z","iopub.status.idle":"2026-02-11T19:04:19.136198Z","shell.execute_reply.started":"2026-02-11T19:04:18.988125Z","shell.execute_reply":"2026-02-11T19:04:19.135433Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tfidf_embdn = pd.read_csv(CONFIG['TFIDF_train'])\nprint(f\"The shape of Train text TF-IDF embedding: {tfidf_embdn.shape}\")\ntfidf_embdn.rename(columns={'protein_id': 'protein'}, inplace=True)\ntfidf_embdn.head(2)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:19.137187Z","iopub.execute_input":"2026-02-11T19:04:19.137499Z","iopub.status.idle":"2026-02-11T19:04:26.205740Z","shell.execute_reply.started":"2026-02-11T19:04:19.137470Z","shell.execute_reply":"2026-02-11T19:04:26.205058Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_TF_features = tfidf_embdn.columns[1:]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:26.206573Z","iopub.execute_input":"2026-02-11T19:04:26.206840Z","iopub.status.idle":"2026-02-11T19:04:26.210691Z","shell.execute_reply.started":"2026-02-11T19:04:26.206808Z","shell.execute_reply":"2026-02-11T19:04:26.210071Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"traind2v_embdn = pd.read_csv(CONFIG['D2V_train'])\nprint(f\"The shape of Train text D2v embedding: {traind2v_embdn.shape}\")\ntraind2v_embdn.rename(columns={'protein_id': 'protein'}, inplace=True)\ntraind2v_embdn.head(2)\ntrain_d2v_features = traind2v_embdn.columns[1:]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:26.211559Z","iopub.execute_input":"2026-02-11T19:04:26.211829Z","iopub.status.idle":"2026-02-11T19:04:29.687766Z","shell.execute_reply.started":"2026-02-11T19:04:26.211803Z","shell.execute_reply":"2026-02-11T19:04:29.687082Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(train_d2v_features)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:29.688745Z","iopub.execute_input":"2026-02-11T19:04:29.689135Z","iopub.status.idle":"2026-02-11T19:04:29.694152Z","shell.execute_reply.started":"2026-02-11T19:04:29.689111Z","shell.execute_reply":"2026-02-11T19:04:29.693469Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_full_embedding = pd.merge(train_df1, traind2v_embdn, on='protein', how='left')\nprint(f\"The shape of train_full: {train_full_embedding .shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:29.694946Z","iopub.execute_input":"2026-02-11T19:04:29.695347Z","iopub.status.idle":"2026-02-11T19:04:30.055922Z","shell.execute_reply.started":"2026-02-11T19:04:29.695290Z","shell.execute_reply":"2026-02-11T19:04:30.055165Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_full_embedding.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:30.056790Z","iopub.execute_input":"2026-02-11T19:04:30.057039Z","iopub.status.idle":"2026-02-11T19:04:30.075677Z","shell.execute_reply.started":"2026-02-11T19:04:30.057019Z","shell.execute_reply":"2026-02-11T19:04:30.075071Z"}},"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', 'MW_log',\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_full_embedding[manual_cols] = scaler.fit_transform(train_full_embedding[manual_cols])\ntrain_full_embedding[train_d2v_features] = scaler.fit_transform(train_full_embedding[train_d2v_features])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:30.076562Z","iopub.execute_input":"2026-02-11T19:04:30.076839Z","iopub.status.idle":"2026-02-11T19:04:30.468873Z","shell.execute_reply.started":"2026-02-11T19:04:30.076811Z","shell.execute_reply":"2026-02-11T19:04:30.468073Z"}},"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\ntax_features = train_full_embedding[tax_cols].values\nd2v_features = train_full_embedding[train_d2v_features].values","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:30.469811Z","iopub.execute_input":"2026-02-11T19:04:30.470071Z","iopub.status.idle":"2026-02-11T19:04:30.841276Z","shell.execute_reply.started":"2026-02-11T19:04:30.470050Z","shell.execute_reply":"2026-02-11T19:04:30.840700Z"}},"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')\ntax_features = tax_features.astype('float32')\nd2v_features = d2v_features.astype('float32')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:30.842110Z","iopub.execute_input":"2026-02-11T19:04:30.842316Z","iopub.status.idle":"2026-02-11T19:04:31.217721Z","shell.execute_reply.started":"2026-02-11T19:04:30.842299Z","shell.execute_reply":"2026-02-11T19:04:31.216950Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data import random_split\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom torchmetrics.classification import MultilabelF1Score\nfrom torchmetrics.classification import MultilabelAccuracy\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:31.218630Z","iopub.execute_input":"2026-02-11T19:04:31.218884Z","iopub.status.idle":"2026-02-11T19:04:42.033820Z","shell.execute_reply.started":"2026-02-11T19:04:31.218861Z","shell.execute_reply":"2026-02-11T19:04:42.033120Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ProteinDataset(Dataset):\n def __init__(self, esm_embeddings, manual_fts, tax_indices, d2v_ft,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.float32) # Indices must be Long\n self.d2V = torch.tensor(d2v_ft, dtype=torch.float32)\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], dtype=torch.float32)\n \n # Return 4 items now\n return self.esm[idx], self.manual[idx], self.tax[idx], self.d2V[idx],label_row","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:42.034743Z","iopub.execute_input":"2026-02-11T19:04:42.035289Z","iopub.status.idle":"2026-02-11T19:04:42.041174Z","shell.execute_reply.started":"2026-02-11T19:04:42.035266Z","shell.execute_reply":"2026-02-11T19:04:42.040250Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class InterLabelGO_Model(nn.Module):\n def __init__(self, num_labels, esm_dim=1280, manual_dim=26, \n tax_dim=21, d2v_dim=256):\n super(InterLabelGO_Model, self).__init__()\n\n # FIXED: Sum of the OUTPUTS of each branch\n # self.final_feature_dim = 1024 + 32 + 32 +128\n self.final_feature_dim = 1024 + 64 + 64 +1024\n \n # Branch 1: ESM (Dense high-level sequence info)\n self.esm_branch = nn.Sequential(\n nn.BatchNorm1d(esm_dim),\n nn.Linear(esm_dim, 2048),\n nn.BatchNorm1d(2048),\n nn.ReLU(),\n nn.Dropout(0.4),\n nn.Linear(2048, 1024),\n nn.BatchNorm1d(1024),\n nn.ReLU(),\n nn.Dropout(0.2)\n )\n \n # Branch 2: Homology (Manual PhysChem features)\n self.homology_branch = nn.Sequential(\n nn.BatchNorm1d(manual_dim),\n nn.Linear(manual_dim, 64), # Slightly wider to capture variance\n nn.ReLU(),\n nn.Linear(64, 64),\n nn.Dropout(0.3)\n )\n \n # Branch 3: Taxonomy (One-hot or label encoded tax info)\n self.tax_branch = nn.Sequential(\n nn.Linear(tax_dim, 64),\n nn.ReLU(),\n nn.Dropout(0.3)\n )\n \n # Branch 4: D2V (Text/Literature embeddings)\n self.d2v_branch = nn.Sequential(\n nn.BatchNorm1d(d2v_dim),\n nn.Linear(d2v_dim, 1024),\n nn.ReLU(),\n nn.Dropout(0.3)\n )\n \n # Final Classifier\n self.classifier = nn.Sequential(\n nn.Linear(self.final_feature_dim, 1024), # Uses correct concatenated size\n nn.ReLU(),\n nn.BatchNorm1d(1024),\n nn.Dropout(0.5),\n nn.Linear(1024, num_labels) \n )\n\n def forward(self, esm, hom, tax, d2v):\n x1 = self.esm_branch(esm)\n x2 = self.homology_branch(hom)\n x3 = self.tax_branch(tax)\n x4 = self.d2v_branch(d2v)\n \n combined = torch.cat((x1, x2, x3, x4), dim=1)\n logits = self.classifier(combined)\n return logits","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:42.042154Z","iopub.execute_input":"2026-02-11T19:04:42.042401Z","iopub.status.idle":"2026-02-11T19:04:42.057742Z","shell.execute_reply.started":"2026-02-11T19:04:42.042382Z","shell.execute_reply":"2026-02-11T19:04:42.057110Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install iterative-stratification","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:42.058547Z","iopub.execute_input":"2026-02-11T19:04:42.058793Z","iopub.status.idle":"2026-02-11T19:04:45.394450Z","shell.execute_reply.started":"2026-02-11T19:04:42.058761Z","shell.execute_reply":"2026-02-11T19:04:45.393731Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom sklearn.metrics import roc_auc_score\n\nclass ZLPRLoss(nn.Module):\n def __init__(self):\n super(ZLPRLoss, self).__init__()\n\n def forward(self, logits, targets):\n # Masking: positive labels get a huge negative value in 'neg' branch\n # and negative labels get a huge negative value in 'pos' branch\n logits_neg = logits - targets * 1e12 \n logits_pos = -logits - (1 - targets) * 1e12 \n \n zeros = torch.zeros_like(logits[:, :1])\n logits_neg = torch.cat([logits_neg, zeros], dim=-1)\n logits_pos = torch.cat([logits_pos, zeros], dim=-1)\n \n neg_loss = torch.logsumexp(logits_neg, dim=-1)\n pos_loss = torch.logsumexp(logits_pos, dim=-1)\n return (neg_loss + pos_loss).mean()\n\nclass IAWeightedSoftF1(nn.Module):\n def __init__(self, ia_weights):\n super(IAWeightedSoftF1, self).__init__()\n self.register_buffer('ia_weights', ia_weights)\n\n def forward(self, probs, targets):\n tp = (probs * targets) * self.ia_weights\n fp = (probs * (1 - targets)) * self.ia_weights\n fn = ((1 - probs) * targets) * self.ia_weights\n \n tp_sum = tp.sum(dim=0)\n fp_sum = fp.sum(dim=0)\n fn_sum = fn.sum(dim=0)\n \n soft_f1 = (2 * tp_sum) / (2 * tp_sum + fp_sum + fn_sum + 1e-8)\n return 1 - soft_f1.mean()\n\nclass IntegratedLoss(nn.Module):\n def __init__(self, ia_weights):\n super().__init__()\n self.rank_loss = ZLPRLoss()\n self.f1_loss = IAWeightedSoftF1(ia_weights)\n \n def forward(self, logits, targets):\n probs = torch.sigmoid(logits)\n return self.rank_loss(logits, targets) * self.f1_loss(probs, targets)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:45.395773Z","iopub.execute_input":"2026-02-11T19:04:45.396132Z","iopub.status.idle":"2026-02-11T19:04:45.405603Z","shell.execute_reply.started":"2026-02-11T19:04:45.396101Z","shell.execute_reply":"2026-02-11T19:04:45.404887Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 1. SEEDING FUNCTION ---\ndef seed_everything(seed):\n import random, os\n random.seed(seed)\n os.environ['PYTHONHASHSEED'] = str(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.backends.cudnn.deterministic = True","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:45.406826Z","iopub.execute_input":"2026-02-11T19:04:45.407149Z","iopub.status.idle":"2026-02-11T19:04:45.420733Z","shell.execute_reply.started":"2026-02-11T19:04:45.407128Z","shell.execute_reply":"2026-02-11T19:04:45.420085Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MetricAccumulator:\n def __init__(self):\n self.reset()\n\n def update(self, preds, targets):\n # Move to CPU and convert to numpy to save GPU memory\n self.all_preds.append(preds.detach().cpu().numpy())\n self.all_targets.append(targets.detach().cpu().numpy())\n\n def get_full(self):\n return np.vstack(self.all_preds), np.vstack(self.all_targets)\n\n def reset(self):\n self.all_preds = []\n self.all_targets = []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:45.421631Z","iopub.execute_input":"2026-02-11T19:04:45.422281Z","iopub.status.idle":"2026-02-11T19:04:45.431305Z","shell.execute_reply.started":"2026-02-11T19:04:45.422253Z","shell.execute_reply":"2026-02-11T19:04:45.430702Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from collections import Counter\nimport numpy as np\n\nimport iterstrat\nimport pickle\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom iterstrat.ml_stratifiers import MultilabelStratifiedKFold\nimport torch\nimport 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\ny_trains = dict()\nmlb_dict = dict()\nmodels = dict()\nproteins_dict = dict()\n# 1. Initialize result storage\nall_val_predictions = {}\nall_train_loss = {}\nall_val_loss = {}\nall_val_true = {} # New dictionary for ground truth\n# Configuration\nfold_train_losses = defaultdict(list)\nfold_val_losses = defaultdict(list)\nfold_val_preds = defaultdict(list)\nfold_val_true = defaultdict(list)\nfold_models = defaultdict(list)\n\n# Storage for final ensemble\nfold_models = {ont: [] for ont in ['P', 'C', 'F']}\n\nTOP_LABELS = {'P': top_BPO, 'C': top_BCO, 'F': top_BFO}\nscaler = GradScaler() # For Mixed Precision\nSEEDS = [42, 43, 120, 90, 45] # 5 different seeds for 5 folds\nfor ontology in ['P', 'C', 'F']:\n print(f\"\\nšŸš€ Processing Ontology: {ontology}\")\n current_top_labels = TOP_LABELS[ontology]\n # label_list = sorted(current_top_labels)\n\n ont_terms_df = train_terms[\n (train_terms['ontology'] == ontology) &\n (train_terms['term'].isin(current_top_labels))\n ]\n\n protein_terms_dict = ont_terms_df.groupby('protein')['term'].apply(list).to_dict()\n \n train_ids_ont = [pid for pid in train_ids1 if pid in protein_terms_dict]\n print(train_ids_ont[:2])\n proteins_dict[ontology] = train_ids_ont\n\n labels = [protein_terms_dict[pid] for pid in train_ids_ont]\n print(f\"Labels : {labels[:2]}\")\n \n \n mlb = MultiLabelBinarizer(classes=current_top_labels,sparse_output=True)\n y_full = mlb.fit_transform(labels)\n y_trains[ontology] = y_full\n mlb_dict[ontology] = mlb\n\n # 1. Create a quick lookup for original row indices\n id_to_pos = {pid: i for i, pid in enumerate(train_ids1)}\n \n # 2. Map your labeled protein IDs to their ACTUAL row numbers in the feature matrices\n actual_row_indices = [id_to_pos[pid] for pid in train_ids_ont]\n print(f\"actual_row_indices: {actual_row_indices[:2]}\")\n # 3. Subset features using the CORRECT rows\n X_esm_lab = esm1_fts[actual_row_indices]\n X_man_lab = man_features[actual_row_indices]\n X_tax_lab = tax_features[actual_row_indices]\n X_d2v_lab = d2v_features[actual_row_indices]\n # 4. y_full is already aligned with train_ids_ont because you built 'labels' \n # using that specific list. Just convert it to a dense array for training.\n y_lab = y_full.toarray()\n print(f\"āœ… Proteins: {len(train_ids_ont)}, Labels: {y_trains[ontology].shape[1]}\")\n\n \n n_splits = 5\n mskf = MultilabelStratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)\n \n ordered_classes = mlb.classes_\n # 1. Map IA scores to your label order\n # Assume ia_map is { 'GO:000...': 5.4 }\n ia_map = ia_df.set_index('term')['ia_score'].to_dict()\n ia_weights = np.array([ia_map.get(goterm, 0.0) for goterm in mlb_dict[ontology].classes_])\n \n # 2. Normalize: This is vital. \n # If you don't normalize, your gradients will be 10x larger than normal, \n # causing the model to diverge.\n ia_weights = torch.tensor(ia_weights, dtype=torch.float32)\n ia_weights = ia_weights / ia_weights.mean() \n \n # 3. Define the Loss\n # Move weights to the same device as your model\n # criterion = nn.BCEWithLogitsLoss(weight=ia_weights.to(device))\n criterion = IntegratedLoss(ia_weights.to(device))\n \n # 3. Start 5-Fold Cross Validation\n # We split based on the labeled subset\n for fold, (t_idx, v_idx) in enumerate(mskf.split(X_esm_lab, y_lab)):\n print(t_idx, v_idx)\n # šŸ”‘ APPLY UNIQUE SEED PER FOLD\n current_seed = SEEDS[fold]\n seed_everything(current_seed)\n print(f\"\\nšŸš€ Fold {fold+1} | Seed: {current_seed} | Ontology: {ontology}\")\n print(f\" ā†’ Fold {fold+1}/{n_splits} | Train Size: {len(t_idx)} | Val Size: {len(v_idx)}\")\n \n # Create Dataset/DataLoader for this fold\n train_ds = ProteinDataset(\n X_esm_lab[t_idx], X_man_lab[t_idx], \n X_tax_lab[t_idx], X_d2v_lab[t_idx], y_lab[t_idx]\n )\n val_ds = ProteinDataset(\n X_esm_lab[v_idx], X_man_lab[v_idx], \n X_tax_lab[v_idx], X_d2v_lab[v_idx], y_lab[v_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 # 4. Initialize a FRESH model for every fold\n model = InterLabelGO_Model(\n num_labels=y_lab.shape[1], \n esm_dim=esm1_fts.shape[1], \n manual_dim=man_features.shape[1], \n tax_dim=tax_features.shape[1],\n d2v_dim=d2v_features.shape[1]\n ).to(device)\n \n # [ Insert your Training Loop here: Optimizer, Criterion, Epochs ]\n optimizer = torch.optim.Adam(model.parameters(), lr=CONFIG['LEARNING_RATE'], weight_decay=1e-4)\n \n # f1_score = MultilabelF1Score(num_labels=current_top_labels).to(config.device)\n\n counter = 0\n patience = 5 \n train_history = []\n train_score_history = []\n val_history = []\n scores = []\n best_val_loss = float('inf')\n \n for epoch in range(CONFIG['EPOCHS']):\n # --- TRAINING PHASE (with Mixed Precision) ---\n model.train()\n total_train_loss = 0\n train_scores= []\n for batch_x,mn_ft_x,tax_x,d2v_x, batch_y in train_loader:\n batch_x, mn_ft_x, tax_x, d2v_x,batch_y = batch_x.to(device), mn_ft_x.to(device),tax_x.to(device), d2v_x.to(device) ,batch_y.to(device)\n optimizer.zero_grad()\n \n logits = model(batch_x, mn_ft_x, tax_x, d2v_x)\n \n loss = criterion(logits, batch_y)\n probs = torch.sigmoid(logits)\n # score=f1_score(probs, batch_y)\n # train_scores.append(score.item())\n \n scaler.scale(loss).backward()\n scaler.step(optimizer)\n scaler.update()\n \n total_train_loss += loss.item()\n # avg_train_score = train_scores/ len(train_loader)\n avg_train_loss = total_train_loss / len(train_loader)\n train_history.append(avg_train_loss)\n # train_score_history.append(avg_train_score)\n # --- VALIDATION PHASE ---\n model.eval()\n total_val_loss = 0\n all_logits, all_targets , val_auc_history = [], [], []\n val_scores, val_score_history = [], []\n with torch.no_grad():\n for batch_x,mn_ft_x,tax_x, d2v_x, batch_y in val_loader:\n batch_x,mn_ft_x, tax_x, d2v_x, batch_y = batch_x.to(device), mn_ft_x.to(device),tax_x.to(device),d2v_x.to(device) ,batch_y.to(device)\n logits = model(batch_x, mn_ft_x, tax_x, d2v_x)\n probs = torch.sigmoid(logits)\n # USE LOGITS HERE, NOT SIGMOID PROBS\n loss = criterion(logits, batch_y) \n total_val_loss += loss.item()\n \n # score=f1_score(probs, batch_y)\n # val_scores.append(score.item())\n all_logits.append(torch.sigmoid(logits).cpu().numpy())\n all_targets.append(batch_y.cpu().numpy())\n\n # šŸ“Š Calculate AUC\n y_pred = np.vstack(all_logits)\n y_true = np.vstack(all_targets)\n try:\n # Macro AUC: Average of AUC per GO term\n val_auc = roc_auc_score(y_true, y_pred, average='macro')\n except ValueError: # Handle cases with only 1 class in validation\n val_auc = 0.0\n \n # avg_val_score = val_scores / len(val_loader)\n avg_val_loss = total_val_loss / len(val_loader)\n # val_score_history.append(avg_val_score)\n val_history.append(avg_val_loss)\n val_auc_history.append(val_auc)\n\n print(f\"Epoch {epoch+1}: Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f} | AUC: {val_auc:.4f}| \")\n # print(f\"Epoch {epoch+1}: Train score: {avg_train_score:.4f} | Val score: {avg_val_score:.4f} \")\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_{ontology}_fold{fold}.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 fold_train_losses[ontology].append(train_history)\n fold_val_losses[ontology].append(val_history)\n fold_val_preds[ontology].append(y_pred)\n fold_val_true[ontology].append(y_true)\n fold_models[ontology].append(f\"best_model_{ontology}_fold{fold}.pth\")\n \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_{ontology}_fold{fold}.pth'))\n model.eval()\n final_probs = []\n with torch.no_grad():\n for batch_x,mn_ft_x, tax_x, d2v_x, batch_y in val_loader:\n logits = model(batch_x.to(device), mn_ft_x.to(device), tax_x.to(device), d2v_x.to(device))\n probs = torch.sigmoid(logits).cpu().numpy()\n y_true_list.append(batch_y.numpy())\n final_probs.append(probs)\n # weighted_prec, weighted_rec, weighted_f1 = weighted_precision_recall_f1(y_true_list, final_probs, ia_map, mlb_obj)\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\n \n # 5. Save the Fold Model\n model_path = f\"best_model_{ontology}_fold{fold}.pth\"\n torch.save(model.state_dict(), model_path)\n fold_models[ontology].append(model_path)\n\n# Save MLB for inference later\nwith open('mlb_dict.pkl', 'wb') as f:\n pickle.dump(mlb_dict, f)\n \nwith open('all_val_true.pkl', 'wb') as f:\n pickle.dump(all_val_true, f)\n\nwith open('all_val_predictions.pkl', 'wb') as f:\n pickle.dump(all_val_predictions, f)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:04:45.432237Z","iopub.execute_input":"2026-02-11T19:04:45.432519Z","iopub.status.idle":"2026-02-11T19:54:19.804031Z","shell.execute_reply.started":"2026-02-11T19:04:45.432490Z","shell.execute_reply":"2026-02-11T19:54:19.803427Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_learning_curves_cv(train_losses, val_losses):\n fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n for i, ont in enumerate(['P', 'C', 'F']):\n min_len = min(len(l) for l in train_losses[ont])\n tr = np.mean([l[:min_len] for l in train_losses[ont]], axis=0)\n va = np.mean([l[:min_len] for l in val_losses[ont]], axis=0)\n axes[i].plot(tr, label='Train')\n axes[i].plot(va, label='Val')\n axes[i].set_title(f'{ont}')\n axes[i].legend()\n plt.show()\n\nplot_learning_curves_cv(fold_train_losses, fold_val_losses)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:54:19.804986Z","iopub.execute_input":"2026-02-11T19:54:19.805234Z","iopub.status.idle":"2026-02-11T19:54:20.168739Z","shell.execute_reply.started":"2026-02-11T19:54:19.805211Z","shell.execute_reply":"2026-02-11T19:54:20.168166Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def ensemble_predictions(pred_list):\n return np.mean(np.vstack(pred_list), axis=0)\n\nfor ont in ['P', 'C', 'F']:\n y_pred = ensemble_predictions(fold_val_preds[ont])\n y_true = np.vstack(fold_val_true[ont])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:54:20.170059Z","iopub.execute_input":"2026-02-11T19:54:20.170514Z","iopub.status.idle":"2026-02-11T19:54:20.523577Z","shell.execute_reply.started":"2026-02-11T19:54:20.170484Z","shell.execute_reply":"2026-02-11T19:54:20.522754Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_oof_score(val_preds_list, val_true_list):\n # 1. Stack the predictions and truths from all folds\n # This turns 5 arrays of (N/5, labels) into one array of (N, labels)\n oof_preds = np.vstack(val_preds_list)\n oof_true = np.vstack(val_true_list)\n \n # 2. Calculate the score on the full OOF set\n macro_auc = roc_auc_score(oof_true, oof_preds, average='macro')\n return macro_auc, oof_preds, oof_true\n\nfor ont in ['P', 'C', 'F']:\n if ont in fold_val_preds and len(fold_val_preds[ont]) > 0:\n auc, oof_p, oof_t = get_oof_score(fold_val_preds[ont], fold_val_true[ont])\n print(f\"šŸ“Š Final OOF AUC for {ont}: {auc:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:54:20.524545Z","iopub.execute_input":"2026-02-11T19:54:20.524844Z","iopub.status.idle":"2026-02-11T19:54:39.091577Z","shell.execute_reply.started":"2026-02-11T19:54:20.524822Z","shell.execute_reply":"2026-02-11T19:54:39.090912Z"}},"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 for threshold in np.linspace(0.001, 0.3, 100):\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-02-11T19:54:39.092527Z","iopub.execute_input":"2026-02-11T19:54:39.092815Z","iopub.status.idle":"2026-02-11T19:54:39.099694Z","shell.execute_reply.started":"2026-02-11T19:54:39.092792Z","shell.execute_reply":"2026-02-11T19:54:39.099073Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Evaluate & select global threshold\n# ------------------------------------------------------------\ndef weighted_precision_recall_f1(y_true, y_pred_bin, ia_map, mlb_obj):\n tp = ((y_true==1)&(y_pred_bin==1)).sum(axis=0).astype(float)\n fp = ((y_true==0)&(y_pred_bin==1)).sum(axis=0).astype(float)\n fn = ((y_true==1)&(y_pred_bin==0)).sum(axis=0).astype(float)\n eps=1e-12\n prec = tp/(tp+fp+eps); rec = tp/(tp+fn+eps)\n f1 = 2*prec*rec/(prec+rec+eps)\n cls = mlb_obj.classes_\n weights = np.array([ia_weights.get(c,1.0) for c in cls], dtype=float) if 'ia_weights' in globals() else np.ones(len(cls))\n weighted_f1 = (f1*weights).sum()/(weights.sum()+eps)\n weighted_prec = (prec*weights).sum()/(weights.sum()+eps)\n weighted_rec = (rec*weights).sum()/(weights.sum()+eps)\n return weighted_prec, weighted_rec, weighted_f1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:54:39.100631Z","iopub.execute_input":"2026-02-11T19:54:39.100884Z","iopub.status.idle":"2026-02-11T19:54:39.115783Z","shell.execute_reply.started":"2026-02-11T19:54:39.100852Z","shell.execute_reply":"2026-02-11T19:54:39.115245Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"final_scores = {}\nia_map = ia_df.set_index('term')['ia_score'].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} | final threshold :{final_tau:.4f}\")\n # weighted_prec, weighted_rec, weighted_f1 = weighted_precision_recall_f1(y_true, y_probs, ia_map, mlb_obj)\n # print(f\"Ontology {ontology}: w_prec = {weighted_prec:.4f}, w_recall = {weighted_rec:.4f}, w_F1: {weighted_f1:.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}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:54:39.116827Z","iopub.execute_input":"2026-02-11T19:54:39.117223Z","iopub.status.idle":"2026-02-11T19:55:05.829892Z","shell.execute_reply.started":"2026-02-11T19:54:39.117200Z","shell.execute_reply":"2026-02-11T19:55:05.829124Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df = create_sequence_dataframe_from_fasta(test_sequences, data_type='test')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T19:55:05.830884Z","iopub.execute_input":"2026-02-11T19:55:05.831120Z","iopub.status.idle":"2026-02-11T20:01:55.625849Z","shell.execute_reply.started":"2026-02-11T19:55:05.831099Z","shell.execute_reply":"2026-02-11T20:01:55.623631Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df['Length_log'] = np.log1p(test_df['Length'])\ntest_df['MW_log'] = np.log1p(test_df['molecular_weight'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:01:55.628871Z","iopub.execute_input":"2026-02-11T20:01:55.629111Z","iopub.status.idle":"2026-02-11T20:01:55.636329Z","shell.execute_reply.started":"2026-02-11T20:01:55.629089Z","shell.execute_reply":"2026-02-11T20:01:55.635770Z"}},"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', 'MW_log',\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-02-11T20:01:55.637282Z","iopub.execute_input":"2026-02-11T20:01:55.637546Z","iopub.status.idle":"2026-02-11T20:01:55.789152Z","shell.execute_reply.started":"2026-02-11T20:01:55.637525Z","shell.execute_reply":"2026-02-11T20:01:55.788583Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\n Loading ESM1b embeddings...\")\nesm1_base = f\"{CONFIG['ESM1_PATH']}\"\ntest_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)}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:01:55.790026Z","iopub.execute_input":"2026-02-11T20:01:55.790285Z","iopub.status.idle":"2026-02-11T20:02:02.191681Z","shell.execute_reply.started":"2026-02-11T20:01:55.790253Z","shell.execute_reply":"2026-02-11T20:02:02.191126Z"}},"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-02-11T20:02:02.192532Z","iopub.execute_input":"2026-02-11T20:02:02.192749Z","iopub.status.idle":"2026-02-11T20:02:02.215642Z","shell.execute_reply.started":"2026-02-11T20:02:02.192729Z","shell.execute_reply":"2026-02-11T20:02:02.215129Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"testd2v_embdn = pd.read_csv(CONFIG['D2V_test'])\nprint(f\"The shape of Test text D2v embedding: {testd2v_embdn.shape}\")\ntestd2v_embdn.rename(columns={'protein_id': 'protein'}, inplace=True)\ntestd2v_embdn.head(2)\ntest_d2v_features = testd2v_embdn.columns[1:]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:02:02.216433Z","iopub.execute_input":"2026-02-11T20:02:02.216984Z","iopub.status.idle":"2026-02-11T20:02:13.528516Z","shell.execute_reply.started":"2026-02-11T20:02:02.216961Z","shell.execute_reply":"2026-02-11T20:02:13.527895Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df = pd.merge(test_df, esm1_test,on='protein', how='left')\ntest_df = pd.merge(test_df, test_tax_features,on='protein', how='left')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:02:13.529363Z","iopub.execute_input":"2026-02-11T20:02:13.529687Z","iopub.status.idle":"2026-02-11T20:02:18.833620Z","shell.execute_reply.started":"2026-02-11T20:02:13.529664Z","shell.execute_reply":"2026-02-11T20:02:18.833015Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_full_embedding = pd.merge(test_df, testd2v_embdn, on='protein', how='left')\nprint(f\"The shape of all embeddings: {test_full_embedding.shape}\")\n# The shape of all embeddings: (224309, 1589)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:03:53.324219Z","iopub.execute_input":"2026-02-11T20:03:53.326173Z","iopub.status.idle":"2026-02-11T20:03:54.578222Z","shell.execute_reply.started":"2026-02-11T20:03:53.326079Z","shell.execute_reply":"2026-02-11T20:03:54.577616Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_ids1 = test_full_embedding['protein'].astype(str).tolist()\nesm1_test_fts = test_full_embedding[testesm_emb_names].values\nman_testfeatures = test_full_embedding[manual_cols].values\ntest_tax_features = test_full_embedding[tax_cols].values\ntest_d2v_features = test_full_embedding[test_d2v_features].values","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:03:54.579355Z","iopub.execute_input":"2026-02-11T20:03:54.579599Z","iopub.status.idle":"2026-02-11T20:03:55.731767Z","shell.execute_reply.started":"2026-02-11T20:03:54.579577Z","shell.execute_reply":"2026-02-11T20:03:55.731140Z"}},"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')\ntest_tax_features = test_tax_features.astype('float32')\ntest_d2v_features = test_d2v_features.astype('float32')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:03:55.732948Z","iopub.execute_input":"2026-02-11T20:03:55.733187Z","iopub.status.idle":"2026-02-11T20:03:56.899485Z","shell.execute_reply.started":"2026-02-11T20:03:55.733164Z","shell.execute_reply":"2026-02-11T20:03:56.898614Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ProteinTestDataset(Dataset):\n def __init__(self, esm_embeddings, manual_fts, tax_indices, d2v_ft):\n self.esm = torch.tensor(esm_embeddings, dtype=torch.float32)\n self.manual = torch.tensor(manual_fts, dtype=torch.float32)\n self.tax = torch.tensor(tax_indices, dtype=torch.float32)\n self.d2V = torch.tensor(d2v_ft, dtype=torch.float32)\n \n def __len__(self):\n return len(self.esm)\n \n def __getitem__(self, idx):\n return self.esm[idx], self.manual[idx], self.tax[idx], self.d2V[idx]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:06:31.871056Z","iopub.execute_input":"2026-02-11T20:06:31.871763Z","iopub.status.idle":"2026-02-11T20:06:31.876967Z","shell.execute_reply.started":"2026-02-11T20:06:31.871733Z","shell.execute_reply":"2026-02-11T20:06:31.876332Z"}},"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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:06:33.162330Z","iopub.execute_input":"2026-02-11T20:06:33.162936Z","iopub.status.idle":"2026-02-11T20:06:33.170379Z","shell.execute_reply.started":"2026-02-11T20:06:33.162904Z","shell.execute_reply":"2026-02-11T20:06:33.169684Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"term_parents, roots = parse_obo_parents(obo_path)\nancestors_map = get_ancestors_map(term_parents)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:06:34.788119Z","iopub.execute_input":"2026-02-11T20:06:34.788440Z","iopub.status.idle":"2026-02-11T20:06:35.438456Z","shell.execute_reply.started":"2026-02-11T20:06:34.788381Z","shell.execute_reply":"2026-02-11T20:06:35.437874Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def propagate_probs_safe(probs, classes, ancestors_map):\n \"\"\"\n probs: (num_labels,) raw sigmoid probabilities\n classes: mlb.classes_ (ontology-specific)\n ancestors_map: pre-filtered to SAME ontology\n \"\"\"\n new_probs = probs.copy()\n idx_map = {go:i for i, go in enumerate(classes)}\n\n for i, go in enumerate(classes):\n p = probs[i]\n for parent in ancestors_map.get(go, []):\n if parent in idx_map:\n j = idx_map[parent]\n new_probs[j] = max(new_probs[j], p)\n\n return new_probs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:06:35.770167Z","iopub.execute_input":"2026-02-11T20:06:35.770530Z","iopub.status.idle":"2026-02-11T20:06:35.775740Z","shell.execute_reply.started":"2026-02-11T20:06:35.770484Z","shell.execute_reply":"2026-02-11T20:06:35.774971Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n\ndef propagate_single(probs, classes, ancestors_map):\n out = probs.copy()\n idx = {g:i for i,g in enumerate(classes)}\n\n for i, go in enumerate(classes):\n p = probs[i]\n for parent in ancestors_map.get(go, []):\n if parent in idx:\n j = idx[parent]\n if p > out[j]:\n out[j] = p\n return out","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:06:36.920072Z","iopub.execute_input":"2026-02-11T20:06:36.920368Z","iopub.status.idle":"2026-02-11T20:06:36.931276Z","shell.execute_reply.started":"2026-02-11T20:06:36.920340Z","shell.execute_reply":"2026-02-11T20:06:36.930644Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch, csv\nimport numpy as np\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nSUBMISSION_FILE = \"submission.tsv\"\nTHRESHOLD = 0.1\nTOP_K = 1500\n\ntest_idx = np.arange(len(test_ids1))\ntest_ds = ProteinTestDataset(\n esm1_test_fts[test_idx],\n man_testfeatures[test_idx],\n test_tax_features[test_idx],\n test_d2v_features[test_idx]\n)\ntest_loader = DataLoader(test_ds, batch_size=512, shuffle=False)\n\nwith open(SUBMISSION_FILE, \"w\", newline=\"\") as f:\n writer = csv.writer(f, delimiter=\"\\t\")\n\n for ont in [\"P\", \"C\", \"F\"]:\n print(f\"\\nšŸš€ Ontology {ont}\")\n\n mlb = mlb_dict[ont]\n go_terms = mlb.classes_\n\n # šŸšØ FILTER ancestors to this ontology only\n ont_set = set(go_terms)\n \n # Load models\n models = []\n for path in fold_models[ont]:\n m = InterLabelGO_Model(\n num_labels=len(go_terms),\n esm_dim=esm1_test_fts.shape[1],\n manual_dim=man_testfeatures.shape[1],\n tax_dim=test_tax_features.shape[1],\n d2v_dim=test_d2v_features.shape[1]\n ).to(device)\n m.load_state_dict(torch.load(path, map_location=device))\n m.eval()\n models.append(m)\n\n ptr = 0\n with torch.no_grad():\n for esm, man, tax, d2v in tqdm(test_loader):\n esm, man, tax, d2v = esm.to(device), man.to(device), tax.to(device), d2v.to(device)\n\n probs = 0\n for m in models:\n probs += torch.sigmoid(m(esm, man, tax, d2v))\n probs = (probs / len(models)).cpu().numpy()\n\n for i in range(probs.shape[0]):\n protein = test_ids1[ptr]\n ptr += 1\n\n # āœ… PROPAGATE PER PROTEIN\n p = propagate_probs_safe(\n probs[i],\n go_terms,\n ancestors_map\n )\n \n\n mask = p >= THRESHOLD\n idxs = np.where(mask)[0]\n scores = p[mask]\n\n if len(scores) == 0:\n continue\n\n top = sorted(zip(idxs, scores), key=lambda x: x[1], reverse=True)[:TOP_K]\n for j,s in top:\n writer.writerow([protein, go_terms[j], f\"{s:.3f}\"])\n\n del models\n torch.cuda.empty_cache()\n\nprint(\"āœ… submission.tsv ready\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-11T20:06:38.237410Z","iopub.execute_input":"2026-02-11T20:06:38.238156Z"}},"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}]}