{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"colab":{"provenance":[],"gpuType":"T4"},"accelerator":"GPU","kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":116062,"databundleVersionId":14875579},{"sourceType":"datasetVersion","sourceId":15724143,"datasetId":10074403,"databundleVersionId":16665060},{"sourceType":"datasetVersion","sourceId":15762116,"datasetId":10101285,"databundleVersionId":16706288}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Install Dependencies & Setup Environment","metadata":{"id":"nOTkIoElG1ev"}},{"cell_type":"code","source":"# 1. Install Python dependencies\n!pip install -q biopython obonet networkx transformers torch tqdm scipy scikit-learn pandas\n\nimport os\nimport sys\nimport shutil\nimport subprocess\nfrom pathlib import Path\nimport torch\n\n# !mkdir -p /kaggle/working/mmseqs\n# !wget -q https://mmseqs.com/latest/mmseqs-linux-gpu-cuda13.tar.gz -O /kaggle/working/mmseqs/mmseqs.tar.gz\n# !tar xvfz /kaggle/working/mmseqs/mmseqs.tar.gz -C /kaggle/working/mmseqs\n# !ln -sf /kaggle/working/mmseqs/mmseqs/bin/mmseqs /usr/local/bin/mmseqs\n\n# 3. Create Project Directory Structure\nBASE_DIR = Path(\"cafa_project\")\nDATA_DIR = BASE_DIR / \"data\"\nRAW_DIR = DATA_DIR / \"raw\"\nDEBUG_DIR = DATA_DIR / \"debug\"\nEMBED_DIR = DATA_DIR / \"embeddings\"\nMODELS_DIR = BASE_DIR / \"models\"\nRESULTS_DIR = BASE_DIR / \"results\"\n\nfor d in [RAW_DIR, DEBUG_DIR, EMBED_DIR, MODELS_DIR, RESULTS_DIR]:\n d.mkdir(parents=True, exist_ok=True)\n\nsys.path.append(str(BASE_DIR))\nprint(f\"Project environment set up at {BASE_DIR.resolve()}\")","metadata":{"id":"M7PqOLvqIWub","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Kaggle-specific file upload\n\n# !tar xf /kaggle/input/datasets/hatimjazouli/cafa6-esm2-650m-embeddings/embeddings.tar.xz\n# !mv /kaggle/working/embeddings/train/ /kaggle/working/embeddings/embeddings\n# !cp -r /kaggle/working/embeddings/embeddings /kaggle/working/cafa_project/data/\n# !rm -rf /kaggle/working/embeddings\n\n# !tar xf /kaggle/input/datasets/hatimjazouli/cafa6-esmc-600m-embeddings/esmc_embeddings.tar.gz\n# !mv /kaggle/working/embeddings/train/ /kaggle/working/embeddings/embeddings\n# !cp -r /kaggle/working/embeddings/embeddings /kaggle/working/cafa_project/data/\n# !rm -rf /kaggle/working/embeddings","metadata":{"id":"Z8NK9eKDDLWr","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Configuration","metadata":{"id":"lX8LzPyeHXyn"}},{"cell_type":"code","source":"import torch\n\nclass Config:\n # Environment\n DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n NUM_GPUS = torch.cuda.device_count()\n\n # Data Paths (Adjusted for Kaggle vs Colab)\n # If raw data isn't found in standard Kaggle location, we look in local dir\n KAGGLE_INPUT = Path(\"/kaggle/input/competitions/cafa-6-protein-function-prediction\")\n\n if KAGGLE_INPUT.exists():\n TRAIN_FASTA = KAGGLE_INPUT / \"Train/train_sequences.fasta\"\n TRAIN_LABELS = KAGGLE_INPUT / \"Train/train_terms.tsv\"\n TRAIN_OBO = KAGGLE_INPUT / \"Train/go-basic.obo\"\n TEST_FASTA = KAGGLE_INPUT / \"Test/testsuperset.fasta\"\n IA_PATH = KAGGLE_INPUT / \"IA.tsv\"\n else:\n # Fallback for Colab/Local (Assuming you upload or download small subsets)\n # If these don't exist, the 'Data Prep' cell will generate synthetic data\n TRAIN_FASTA = RAW_DIR / \"train_sequences.fasta\"\n TRAIN_LABELS = RAW_DIR / \"train_terms.tsv\"\n TRAIN_OBO = RAW_DIR / \"go-basic.obo\"\n TEST_FASTA = RAW_DIR / \"testsuperset.fasta\"\n IA_PATH = RAW_DIR / \"IA.tsv\"\n\n # Model Hyperparameters (Phase 2)\n ESM_MODEL_NAME = \"Synthyra/ESMplusplus_large\"\n ESM_DIM = 1152\n HIDDEN_DIM = 512\n NUM_EPOCHS = 50 # 10 for debug, increase to 50 for Prod\n BATCH_SIZE = 128 if NUM_GPUS < 2 else 256\n LEARNING_RATE = 1e-4\n\n # Feature Engineering\n MAX_SEQ_LEN = 1022\n USE_HOMOLOGY = True\n\n # Hierarchy\n LAMBDA_HIERARCHY = 1.0 # Weight for hierarchical consistency loss\n\n # Debug Mode\n # If True, creates a small subset of data to verify pipeline\n DEBUG = False\n DEBUG_SIZE = 2000\n\nprint(f\"Running on {Config.DEVICE} with {Config.NUM_GPUS} GPUs\")\nprint(f\"Debug Mode: {Config.DEBUG}\")","metadata":{"id":"aOZ52bgdHZrp","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data Utilities (GOLabeler & Features)","metadata":{"id":"rtt2ArAgHr1q"}},{"cell_type":"code","source":"import numpy as np\nimport obonet\nimport networkx as nx\nimport torch\nfrom collections import defaultdict\nfrom Bio.SeqUtils.ProtParam import ProteinAnalysis\n\nclass GOLabeler:\n def __init__(self, obo_path):\n self.obo_path = obo_path\n print(f\"Loading Ontology from {obo_path}...\")\n self.graph = obonet.read_obo(obo_path)\n \n # Split attributes by aspect\n self.aspects = ['BPO', 'MFO', 'CCO']\n self.ns_map = {\n 'biological_process': 'BPO',\n 'molecular_function': 'MFO',\n 'cellular_component': 'CCO'\n }\n self.valid_terms = {a: [] for a in self.aspects}\n self.term_to_idx = {a: {} for a in self.aspects}\n \n # NEW: Add cache to eliminate graph traversal redundancy\n self.ancestor_cache = {}\n\n def build_vocabulary(self, annotations, min_freq=50):\n print(f\"Building vocabulary (min_freq={min_freq})...\")\n counts = defaultdict(int)\n for _, term in annotations:\n for anc in self._get_ancestors(term):\n counts[anc] += 1\n\n valid_raw = sorted([t for t, c in counts.items() if c >= min_freq])\n \n # Assign valid terms to their respective aspects\n for t in valid_raw:\n if t in self.graph:\n ns = self.graph.nodes[t].get('namespace', '')\n aspect = self.ns_map.get(ns)\n if aspect:\n self.valid_terms[aspect].append(t)\n \n for aspect in self.aspects:\n self.term_to_idx[aspect] = {t: i for i, t in enumerate(self.valid_terms[aspect])}\n print(f\"{aspect} Vocabulary size: {len(self.valid_terms[aspect])}\")\n\n def _get_ancestors(self, term):\n if term not in self.graph: return set()\n ancestors = {term}\n try: ancestors.update(nx.descendants(self.graph, term))\n except: pass\n return ancestors\n\n def get_sparse_vector(self, terms, aspect):\n active = set()\n for t in terms: active.update(self._get_ancestors(t))\n indices = [self.term_to_idx[aspect][t] for t in active if t in self.term_to_idx[aspect]]\n return list(set(indices))\n\n def get_edge_index(self, aspect):\n edges = []\n for child_term in self.valid_terms[aspect]:\n if child_term not in self.graph: continue\n child_idx = self.term_to_idx[aspect][child_term]\n for parent_term in self.graph.successors(child_term):\n if parent_term in self.term_to_idx[aspect]:\n parent_idx = self.term_to_idx[aspect][parent_term]\n edges.append([child_idx, parent_idx])\n return torch.tensor(edges, dtype=torch.long).t() if edges else torch.empty((2,0), dtype=torch.long)\n\n def get_ia_weights(self, ia_path):\n \"\"\"Loads Information Accretion values and returns a tensor of weights for valid_terms.\"\"\"\n import pandas as pd\n import torch\n \n # Load the IA values\n ia_df = pd.read_csv(ia_path, sep='\\t', names=['term', 'ia'])\n ia_dict = dict(zip(ia_df['term'], ia_df['ia']))\n \n # Map IA values to the vocabulary in the correct index order\n weights =[]\n for term in self.valid_terms:\n # Default to 1.0 if a term is somehow missing\n weights.append(ia_dict.get(term, 1.0))\n \n weights = torch.tensor(weights, dtype=torch.float32)\n \n # Normalize weights to have a mean of 1.0.\n # This preserves overall loss scale while adjusting relative class importance.\n if weights.sum() > 0:\n weights = weights / weights.mean()\n \n return weights\n\n def _get_ancestors(self, term):\n if term in self.ancestor_cache:\n return self.ancestor_cache[term]\n \n if term not in self.graph: \n return set()\n \n ancestors = {term}\n try: \n ancestors.update(nx.descendants(self.graph, term))\n except: \n pass\n \n self.ancestor_cache[term] = ancestors\n return ancestors\n\n\ndef extract_physio_features(sequence):\n \"\"\"\n Extracts basic physicochemical properties for Path B (Local).\n Returns normalized tensor of shape (5,).\n \"\"\"\n seq = str(sequence).upper().replace('U', 'C').replace('X', 'A').replace('Z', 'E').replace('B', 'D')\n if len(seq) == 0: return torch.zeros(5)\n\n try:\n analyser = ProteinAnalysis(seq)\n mol_weight = analyser.molecular_weight()\n aromaticity = analyser.aromaticity()\n instability = analyser.instability_index()\n isoelectric = analyser.isoelectric_point()\n helix_frac = analyser.secondary_structure_fraction()[0]\n\n # Simple normalization based on typical protein ranges\n feats = [\n np.log10(mol_weight + 1) / 6.0, # ~0-1 range\n aromaticity,\n min(instability / 100.0, 1.0),\n isoelectric / 14.0,\n helix_frac\n ]\n return torch.tensor(feats, dtype=torch.float32)\n except:\n return torch.zeros(5)","metadata":{"id":"--A4qWw_Htpd","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data Preparation Pipeline","metadata":{"id":"a8ufcMmhH1S1"}},{"cell_type":"code","source":"import pandas as pd\nfrom Bio import SeqIO\nfrom tqdm import tqdm\nimport os\nimport shutil\nfrom pathlib import Path\nimport random\n\ndef clean_id(header):\n \"\"\"Normalizes FASTA headers like 'sp|P12345|NAME' to 'P12345'.\"\"\"\n header = str(header).strip()\n if header.count('|') >= 2:\n return header.split('|')[1]\n return header\n\ndef check_and_move_files(config):\n files_to_check = {\n \"go-basic.obo\": config.TRAIN_OBO,\n \"train_sequences.fasta\": config.TRAIN_FASTA,\n \"train_terms.tsv\": config.TRAIN_LABELS,\n \"testsuperset.fasta\": config.TEST_FASTA\n }\n colab_root = Path(\"/content\")\n for filename, target_path in files_to_check.items():\n source = colab_root / filename\n if source.exists():\n if not target_path.exists():\n print(f\"šŸ“¦ Moving {filename} to {target_path}...\")\n shutil.move(str(source), str(target_path))\n elif source.stat().st_size > target_path.stat().st_size:\n print(f\"šŸ“¦ Overwriting {target_path.name} with uploaded version...\")\n shutil.move(str(source), str(target_path))\n else:\n print(f\"ā„¹ļø {filename} already exists in project folder.\")\n\ndef prepare_data(config):\n # 0. Move uploaded files\n check_and_move_files(config)\n\n # 1. Sync Synthetic Data (Ensure FASTA IDs match Label IDs)\n sync_synthetic_data(config)\n\n # 2. Process Annotations\n print(\"Loading Annotations...\")\n df = pd.read_csv(config.TRAIN_LABELS, sep='\\t', names=['protein_id', 'term', 'aspect'])\n # Handle Header\n if len(df) > 0 and str(df.iloc[0]['protein_id']).strip().lower() in ['proteinid', 'protein_id', 'entryid']:\n df = df.iloc[1:]\n\n # Filter for debug\n if config.DEBUG:\n unique_prots = df['protein_id'].unique()[:config.DEBUG_SIZE]\n df = df[df['protein_id'].isin(unique_prots)]\n print(f\"Debug Mode: Filtered to {len(unique_prots)} proteins.\")\n\n annotations = list(zip(df['protein_id'], df['term']))\n\n # 3. Build Labeler\n try:\n labeler = GOLabeler(config.TRAIN_OBO)\n except ValueError as e:\n print(f\"āŒ OBO Parse Error: {e}\")\n # If OBO is bad, delete it and let user know\n if config.TRAIN_OBO.exists(): config.TRAIN_OBO.unlink()\n raise ValueError(\"OBO file was corrupted. Deleted. Please re-run or re-upload.\")\n\n labeler.build_vocabulary(annotations, min_freq=10)\n\n # 4. Filter Sequences\n print(\"Filtering FASTA...\")\n valid_ids = set(df['protein_id'].unique())\n sequences = []\n\n # Check if FASTA exists\n if not config.TRAIN_FASTA.exists():\n raise FileNotFoundError(f\"FASTA file missing at {config.TRAIN_FASTA}\")\n\n with open(config.TRAIN_FASTA) as f:\n for record in SeqIO.parse(f, \"fasta\"):\n # Normalize ID\n pid = clean_id(record.id)\n if pid in valid_ids:\n record.id = pid # Ensure record uses cleaned ID\n sequences.append(record)\n\n print(f\"Kept {len(sequences)} sequences matching annotations.\")\n\n if len(sequences) == 0:\n print(\"āŒ CRITICAL ERROR: 0 Sequences matched.\")\n print(\" IDs in Labels (Sample):\", list(valid_ids)[:5])\n print(\" IDs in FASTA (Sample):\", [clean_id(r.id) for r in list(SeqIO.parse(config.TRAIN_FASTA, \"fasta\"))[:5]])\n raise ValueError(\"ID mismatch between FASTA and TSV.\")\n\n return sequences, annotations, labeler\n\ndef sync_synthetic_data(config):\n \"\"\"\n Smart generation:\n 1. If OBO missing -> Generate dummy OBO.\n 2. If Labels missing -> Generate dummy Labels.\n 3. If FASTA missing BUT Labels exist -> Generate dummy FASTA using Label IDs.\n \"\"\"\n config.TRAIN_FASTA.parent.mkdir(parents=True, exist_ok=True)\n\n # 1. OBO\n if not config.TRAIN_OBO.exists():\n print(\"āš ļø OBO missing. Generating synthetic OBO...\")\n with open(config.TRAIN_OBO, 'w') as f:\n f.write(\"format-version: 1.2\\n\\n[Term]\\nid: GO:0000001\\nname: root\\nnamespace: biological_process\\n\\n[Term]\\nid: GO:0000002\\nname: test\\nnamespace: biological_process\\nis_a: GO:0000001\\n\")\n\n # 2. LABELS\n if not config.TRAIN_LABELS.exists():\n print(\"āš ļø Labels missing. Generating synthetic labels...\")\n with open(config.TRAIN_LABELS, 'w') as f_lab:\n f_lab.write(\"ProteinID\\tTermID\\tAspect\\n\")\n for i in range(100):\n pid = f\"P{i:05d}\"\n f_lab.write(f\"{pid}\\tGO:0000002\\tP\\n\")\n\n # 3. FASTA (The critical part)\n if not config.TRAIN_FASTA.exists():\n print(\"āš ļø FASTA missing. Generating synthetic FASTA based on Labels...\")\n\n # Read IDs from the existing labels file to ensure match\n df = pd.read_csv(config.TRAIN_LABELS, sep='\\t', names=['protein_id', 'term', 'aspect'])\n if len(df) > 0 and str(df.iloc[0]['protein_id']).strip().lower() in ['proteinid', 'protein_id', 'entryid']:\n df = df.iloc[1:]\n\n unique_ids = df['protein_id'].unique()\n if config.DEBUG: unique_ids = unique_ids[:config.DEBUG_SIZE]\n\n with open(config.TRAIN_FASTA, 'w') as f_seq:\n for pid in unique_ids:\n # Generate random sequence\n seq = \"\".join(random.choices(\"ACDEFGHIKLMNPQRSTVWY\", k=50))\n f_seq.write(f\">{pid}\\n{seq}\\n\")\n print(f\" Generated {len(unique_ids)} synthetic sequences matching label IDs.\")\n\n# Run Prep\nsequences, annotations, labeler = prepare_data(Config)","metadata":{"id":"m3nzjRHWH2YD","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Feature Extraction (Homology & Embeddings)","metadata":{"id":"zCBmLGLtLaHd"}},{"cell_type":"code","source":"import subprocess\nimport shutil\nfrom pathlib import Path\nfrom collections import defaultdict\nfrom tqdm import tqdm\nfrom transformers import AutoTokenizer, AutoModel\nimport torch\nfrom Bio import SeqIO\n\ndef run_homology_search(sequences, config):\n \"\"\"\n Runs MMseqs2 to find homologs using the GPU binary.\n \"\"\"\n if not config.USE_HOMOLOGY: return None\n\n # Paths\n debug_dir = DEBUG_DIR.resolve()\n temp_fasta = debug_dir / \"current_seqs.fasta\"\n db_path = debug_dir / \"targetDB\"\n gpu_db_path = debug_dir / \"targetDB_gpu\"\n result_db = debug_dir / \"resultDB\"\n result_path = debug_dir / \"homology.m8\"\n tmp_dir = debug_dir / \"tmp\"\n\n # 1. Check Cache\n if result_path.exists() and result_path.stat().st_size > 0:\n print(f\"āœ… Found existing homology results at {result_path}. Loading...\")\n return parse_homology_results(result_path)\n\n # 2. Cleanup & Setup\n if tmp_dir.exists(): shutil.rmtree(tmp_dir)\n tmp_dir.mkdir(parents=True, exist_ok=True)\n\n print(f\"Running MMseqs2 Homology Search (GPU)...\")\n SeqIO.write(sequences, temp_fasta, \"fasta\")\n\n try:\n # 3. Create DB (CPU step)\n print(\" Creating Sequence DB...\")\n subprocess.run(\n [\"mmseqs\", \"createdb\", str(temp_fasta), str(db_path)],\n check=True, stdout=subprocess.DEVNULL\n )\n\n # 4. Make Padded DB (Required for GPU)\n # This aligns sequences for GPU memory patterns\n subprocess.run(\n [\"mmseqs\", \"makepaddedseqdb\", str(db_path), str(gpu_db_path)],\n check=True, stdout=subprocess.DEVNULL\n )\n\n # 5. Search (GPU step)\n print(\" Executing All-vs-All Search...\")\n # --gpu 1: Enable GPU\n # --gpu-mem-agg 0: Conservative memory mode to prevent T4 OOM\n # -s 5.0: High sensitivity\n cmd = [\n \"mmseqs\", \"search\",\n str(gpu_db_path), str(gpu_db_path),\n str(result_db), str(tmp_dir),\n \"-s\", \"5.0\",\n \"--gpu\", \"1\"\n ]\n\n result = subprocess.run(cmd, capture_output=True, text=True)\n\n if result.returncode != 0:\n print(\"\\nāŒ MMSEQS2 SEARCH FAILED!\")\n print(result.stderr)\n raise RuntimeError(\"MMseqs2 GPU search failed.\")\n\n # 6. Convert Results\n subprocess.run(\n [\"mmseqs\", \"convertalis\", str(gpu_db_path), str(gpu_db_path), str(result_db), str(result_path),\n \"--format-output\", \"query,target,evalue,bits\"],\n check=True, stdout=subprocess.DEVNULL\n )\n\n except subprocess.CalledProcessError as e:\n print(f\"\\nāŒ Subprocess Failed: {e}\")\n raise e\n\n print(\" Parsing results...\")\n return parse_homology_results(result_path)\n\ndef clean_id(header):\n header = str(header).strip()\n if header.count('|') >= 2:\n return header.split('|')[1]\n return header\n\ndef parse_homology_results(path):\n homologs = defaultdict(list)\n with open(path) as f:\n for line in f:\n parts = line.strip().split('\\t')\n if len(parts) < 4: continue\n q, t, e, b = parts\n q, t = clean_id(q), clean_id(t)\n if q == t: continue\n homologs[q].append((t, float(b)))\n print(f\"āœ… Loaded homologs for {len(homologs)} sequences.\")\n return homologs\n\ndef extract_embeddings(sequences, config):\n print(f\"Checking embeddings for {len(sequences)} sequences...\")\n cached_ids = set()\n for p in EMBED_DIR.glob(\"*.pt\"):\n cached_ids.add(clean_id(p.stem))\n cached_ids.add(p.stem)\n\n to_process = [s for s in sequences if s.id not in cached_ids]\n\n if not to_process:\n print(\"āœ… All embeddings cached!\")\n return\n\n print(f\"Extracting {len(to_process)} missing embeddings...\")\n tokenizer = AutoTokenizer.from_pretrained(config.ESM_MODEL_NAME)\n model = AutoModel.from_pretrained(config.ESM_MODEL_NAME).to(config.DEVICE)\n model.eval()\n\n batch_size = 8\n for i in tqdm(range(0, len(to_process), batch_size)):\n batch = to_process[i:i+batch_size]\n seqs = [str(r.seq)[:config.MAX_SEQ_LEN] for r in batch]\n inputs = tokenizer(seqs, return_tensors=\"pt\", padding=True, truncation=True, max_length=config.MAX_SEQ_LEN+2)\n inputs = {k: v.to(config.DEVICE) for k, v in inputs.items()}\n\n with torch.no_grad():\n outputs = model(**inputs)\n embeddings = outputs.last_hidden_state\n for j, record in enumerate(batch):\n emb = embeddings[j, 1:-1, :].mean(dim=0).cpu()\n torch.save(emb, EMBED_DIR / f\"{record.id}.pt\")\n\n# Execute\n# homology_dict = run_homology_search(sequences, Config)\nresult_path = \"/kaggle/input/datasets/hatimjazouli/cafa6-homology-matrix/homology.m8\"\nhomology_dict = parse_homology_results(result_path)\nextract_embeddings(sequences, Config)","metadata":{"id":"lhuAICmVLClj","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Clustering & Splitting","metadata":{"id":"Chv9v66WK3Gm"}},{"cell_type":"code","source":"import subprocess\nfrom collections import defaultdict\nimport random\n\ndef run_clustering_split(sequences, config, val_ratio=0.2):\n print(\"\\nRunning MMseqs2 Clustering for Leak-Proof Split...\")\n\n # Paths\n fasta_path = DEBUG_DIR / \"clustering_input.fasta\"\n cluster_out = DEBUG_DIR / \"clusters\"\n cluster_tsv = DEBUG_DIR / \"clusters_cluster.tsv\"\n\n # 1. Write FASTA\n # (We write it again to ensure the input to mmseqs matches our memory objects exactly)\n SeqIO.write(sequences, fasta_path, \"fasta\")\n\n # 2. Run MMseqs2 easy-linclust\n # --min-seq-id 0.3 ensures strict non-homologous split (30% identity)\n # --cov-mode 1 -c 0.8 ensures 80% coverage\n cmd = f\"mmseqs easy-linclust {fasta_path} {cluster_out} {DEBUG_DIR}/tmp --min-seq-id 0.3 --cov-mode 1 -c 0.8\"\n\n # Only run if output doesn't exist or is empty\n if not cluster_tsv.exists() or cluster_tsv.stat().st_size == 0:\n print(\"Executing MMseqs2...\")\n subprocess.run(cmd, shell=True, check=True)\n else:\n print(\"Found existing cluster file, skipping MMseqs2 run.\")\n\n # 3. Parse Clusters\n clusters = defaultdict(list)\n with open(cluster_tsv) as f:\n for line in f:\n rep, member = line.strip().split('\\t')\n # IDs in FASTA were already cleaned in Cell 4, but let's be safe\n rep, member = clean_id(rep), clean_id(member)\n clusters[rep].append(member)\n\n print(f\"Found {len(clusters)} clusters from {len(sequences)} sequences.\")\n\n # 4. Split Clusters (Stratified by cluster, not sequence)\n cluster_reps = list(clusters.keys())\n random.shuffle(cluster_reps)\n\n num_val = int(len(cluster_reps) * val_ratio)\n val_reps = set(cluster_reps[:num_val])\n\n train_ids = set()\n val_ids = set()\n\n for rep, members in clusters.items():\n if rep in val_reps:\n val_ids.update(members)\n else:\n train_ids.update(members)\n\n print(f\"Split Summary: Train={len(train_ids)}, Val={len(val_ids)} sequences.\")\n\n # Verify no overlap (Sanity Check)\n intersection = train_ids.intersection(val_ids)\n if len(intersection) > 0:\n raise ValueError(f\"CRITICAL: Data Leakage detected! {len(intersection)} sequences in both sets.\")\n\n return train_ids, val_ids\n\n# Execute Split\nrandom.seed(27)\ntrain_ids, val_ids = run_clustering_split(sequences, Config)","metadata":{"id":"_sw4toGpK4B1","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset & Dual-Path Model Architecture","metadata":{"id":"8Kx_x4ZCQGMe"}},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass CAFA2Dataset(Dataset):\n def __init__(self, sequences, labeler, annotations_dict, config, homology_dict=None, esm_embeddings=None):\n self.sequences = sequences\n self.labeler = labeler\n self.annotations_dict = annotations_dict\n self.homology_dict = homology_dict\n self.esm_embeddings = esm_embeddings or {} \n self.config = config\n self.aa_to_idx = {aa: i+1 for i, aa in enumerate(\"ACDEFGHIKLMNPQRSTVWY\")}\n\n def __len__(self): return len(self.sequences)\n\n def __getitem__(self, idx):\n record = self.sequences[idx]\n pid = record.id \n\n # 1. Path A: ESMC Multi-Layer Pooling (Read from pre-loaded RAM)\n if pid in self.esm_embeddings:\n emb_dict = self.esm_embeddings[pid]\n \n # Verify if it is in the dict format extracted by your ESMC notebook\n if isinstance(emb_dict, dict) and \"layer_34\" in emb_dict:\n esm_emb = torch.stack([\n emb_dict[\"layer_34\"],\n emb_dict[\"layer_35\"],\n emb_dict[\"layer_36\"]\n ], dim=0) # Yields Shape: (3, 1152)\n else:\n # Robust fallback for un-pooled or legacy embeddings\n single_emb = emb_dict if isinstance(emb_dict, torch.Tensor) else torch.zeros(self.config.ESM_DIM)\n if single_emb.shape[-1] != self.config.ESM_DIM:\n single_emb = torch.zeros(self.config.ESM_DIM)\n esm_emb = single_emb.unsqueeze(0).repeat(3, 1) # Replicates to shape (3, 1152)\n else:\n esm_emb = torch.zeros((3, self.config.ESM_DIM)) # Fallback shape: (3, 1152)\n\n # 2. Path B: Local Sequence Features (remain unchanged)\n seq_str = str(record.seq)[:self.config.MAX_SEQ_LEN]\n seq_idx = [self.aa_to_idx.get(aa, 0) for aa in seq_str]\n target_len = self.config.MAX_SEQ_LEN\n if len(seq_idx) < target_len:\n seq_idx = seq_idx + [0] * (target_len - len(seq_idx))\n else:\n seq_idx = seq_idx[:target_len]\n seq_tensor = torch.tensor(seq_idx, dtype=torch.long)\n\n # Physio Features (remain unchanged)\n physio = extract_physio_features(seq_str)\n\n # Setup multi-branch dictionaries (remain unchanged)\n homology_vecs = {a: torch.zeros(len(self.labeler.valid_terms[a])) for a in self.labeler.aspects}\n labels = {a: torch.zeros(len(self.labeler.valid_terms[a])) for a in self.labeler.aspects}\n\n # 3. Homology (remain unchanged)\n if self.homology_dict and pid in self.homology_dict:\n hits = self.homology_dict[pid] \n for target_id, score in hits[:10]:\n if target_id in self.annotations_dict:\n target_terms = self.annotations_dict[target_id]\n for aspect in self.labeler.aspects:\n indices = self.labeler.get_sparse_vector(target_terms, aspect)\n if indices:\n homology_vecs[aspect][indices] += np.log1p(score)\n\n for aspect in self.labeler.aspects:\n if homology_vecs[aspect].max() > 0:\n homology_vecs[aspect] = homology_vecs[aspect] / homology_vecs[aspect].max()\n\n # 4. Target Labels (remain unchanged)\n if pid in self.annotations_dict:\n true_terms = self.annotations_dict[pid]\n for aspect in self.labeler.aspects:\n true_indices = self.labeler.get_sparse_vector(true_terms, aspect)\n labels[aspect][true_indices] = 1.0\n\n return {\n \"esm\": esm_emb,\n \"seq\": seq_tensor,\n \"physio\": physio,\n \"homology_BPO\": homology_vecs[\"BPO\"],\n \"homology_MFO\": homology_vecs[\"MFO\"],\n \"homology_CCO\": homology_vecs[\"CCO\"],\n \"labels_BPO\": labels[\"BPO\"],\n \"labels_MFO\": labels[\"MFO\"],\n \"labels_CCO\": labels[\"CCO\"]\n }\n\nclass MultiHeadDualPathModel(nn.Module):\n def __init__(self, num_classes_dict, esm_dim=1152, cnn_vocab=22):\n super().__init__()\n \n # --- Shared Feature Extractor ---\n # Define 3 parallel MLPs to process layers 34, 35, and 36 individually\n self.parallel_mlps = nn.ModuleList([\n nn.Sequential(\n nn.Linear(esm_dim, 256),\n nn.BatchNorm1d(256),\n nn.ReLU(),\n nn.Dropout(0.3)\n ) for _ in range(3)\n ])\n \n # Takes the concatenated representations (3 * 256 = 768) and maps to 512\n self.path_a_mlp_combined = nn.Sequential(\n nn.Linear(3 * 256, 512),\n nn.BatchNorm1d(512),\n nn.ReLU(),\n nn.Dropout(0.4)\n )\n \n # Local Sequence Features (Path B)\n self.embedding = nn.Embedding(cnn_vocab, 64, padding_idx=0)\n self.conv1 = nn.Conv1d(64, 128, kernel_size=9, padding=4)\n self.conv2 = nn.Conv1d(128, 256, kernel_size=5, padding=2)\n self.dropout1d = nn.Dropout1d(0.2)\n \n self.path_b_mlp = nn.Sequential(\n nn.Linear(256 + 5, 256),\n nn.BatchNorm1d(256),\n nn.ReLU(),\n nn.Dropout(0.3)\n )\n self.fusion = nn.Sequential(\n nn.Linear(512 + 256, 512),\n nn.BatchNorm1d(512),\n nn.ReLU(),\n nn.Dropout(0.4)\n )\n\n # --- 3 Independent Classification Heads ---\n self.classifiers = nn.ModuleDict({\n aspect: nn.Linear(512, num_classes) \n for aspect, num_classes in num_classes_dict.items()\n })\n \n # --- 3 Independent Homology Gates ---\n self.homology_gates = nn.ParameterDict({\n aspect: nn.Parameter(torch.tensor([0.1])) \n for aspect in num_classes_dict.keys()\n })\n\n def forward(self, esm, seq, physio, hom_dict):\n # Path A: Multi-Layer Parallel Processing (esm shape: [Batch, 3, 1152])\n layer_feats = []\n for i in range(3):\n layer_emb = esm[:, i, :] # Extract layer index i: shape [Batch, 1152]\n layer_feat = self.parallel_mlps[i](layer_emb) # shape [Batch, 256]\n layer_feats.append(layer_feat)\n \n # Concatenate evolutionary representations across all 3 layers\n combined_layers = torch.cat(layer_feats, dim=1) # shape [Batch, 768]\n global_feat = self.path_a_mlp_combined(combined_layers) # shape [Batch, 512]\n\n # Path B: Local Structural and Chemical properties (remain unchanged)\n x = self.embedding(seq).permute(0, 2, 1)\n x = F.relu(self.conv1(x))\n x = self.dropout1d(x)\n x = F.relu(self.conv2(x))\n x = F.adaptive_max_pool1d(x, 1).squeeze(-1) \n local_feat = self.path_b_mlp(torch.cat([x, physio], dim=1))\n\n # Fusion (remain unchanged)\n fused = self.fusion(torch.cat([global_feat, local_feat], dim=1))\n\n # Process each head separately (remain unchanged)\n logits = {}\n for aspect in self.classifiers.keys():\n base_logits = self.classifiers[aspect](fused)\n logits[aspect] = base_logits + (self.homology_gates[aspect] * hom_dict[aspect])\n\n return logits","metadata":{"id":"6lh2JzrkQHFx","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Loss Functions (ASL + Hierarchy)","metadata":{"id":"fTdqPaEtQO3d"}},{"cell_type":"code","source":"class AsymmetricLossOptimized(nn.Module):\n \"\"\"\n Asymmetric Loss for Multi-Label Classification (reduces weight of easy negatives).\n \"\"\"\n def __init__(self, gamma_neg=4, gamma_pos=0, clip=0.05, eps=1e-8, disable_torch_grad_focal_loss=False, weight=None):\n super().__init__()\n self.gamma_neg = gamma_neg\n self.gamma_pos = gamma_pos\n self.clip = clip\n self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss\n self.eps = eps\n self.weight = weight\n\n def forward(self, x, y):\n # x: logits, y: targets\n xs_pos = x.sigmoid()\n xs_neg = 1 - xs_pos\n\n # Asymmetric Clipping\n if self.clip is not None and self.clip > 0:\n xs_neg = (xs_neg + self.clip).clamp(max=1)\n\n # Basic Cross Entropy components\n los_pos = y * torch.log(xs_pos.clamp(min=self.eps))\n los_neg = (1 - y) * torch.log(xs_neg.clamp(min=self.eps))\n\n # Asymmetric Focusing\n loss = -1 * (los_pos + los_neg) # Standard BCE\n\n # Apply focusing\n if self.gamma_neg > 0 or self.gamma_pos > 0:\n pt0 = xs_pos * y\n pt1 = xs_neg * (1 - y) # pt = p if t=1 else 1-p\n pt = pt0 + pt1\n one_sided_gamma = self.gamma_pos * y + self.gamma_neg * (1 - y)\n one_sided_w = torch.pow(1 - pt, one_sided_gamma)\n\n if self.disable_torch_grad_focal_loss:\n torch.set_grad_enabled(False)\n loss *= one_sided_w\n\n if self.weight is not None:\n class_weights = self.weight.to(x.device)\n loss *= class_weights \n\n return loss.mean()\n\nclass HierarchicalLoss(nn.Module):\n \"\"\"\n Penalizes violations where P(Child) > P(Parent).\n L = sum(max(0, P_child - P_parent))\n \"\"\"\n def __init__(self, edge_index):\n super().__init__()\n # edge_index is (2, E) where row 0 is child, row 1 is parent\n self.register_buffer('edge_index', edge_index)\n\n def forward(self, logits):\n probs = torch.sigmoid(logits)\n child_probs = probs[:, self.edge_index[0]]\n parent_probs = probs[:, self.edge_index[1]]\n\n # Violation: child > parent\n # We want child <= parent, so parent - child >= 0\n # Violation = max(0, child - parent)\n diff = child_probs - parent_probs\n violations = torch.clamp(diff, min=0)\n return violations.mean()","metadata":{"id":"V1flldrAQPs3","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training Loop","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nfrom torch.amp import autocast, GradScaler\nfrom tqdm import tqdm\nfrom pathlib import Path\nfrom collections import defaultdict\n\ndef train_phase2(config, sequences, labeler, annotations, homology_dict, train_ids, val_ids):\n # --- 0. PRE-LOAD EMBEDDINGS (THE SPEED HACK) ---\n print(\"Pre-loading ESM embeddings into RAM to eliminate I/O bottleneck...\")\n esm_embeddings = {}\n for p in tqdm(list(EMBED_DIR.glob(\"*.pt\")), desc=\"Loading embeddings\"):\n # Load directly to CPU memory to prevent GPU RAM overflow\n esm_embeddings[p.stem] = torch.load(p, map_location='cpu')\n\n # --- 1. Setup Data ---\n print(\"Initializing Datasets with Cluster Splits...\")\n ann_dict = defaultdict(list)\n for p, t in annotations: ann_dict[p].append(t)\n\n train_seqs =[s for s in sequences if s.id in train_ids]\n val_seqs =[s for s in sequences if s.id in val_ids]\n\n train_ds = CAFA2Dataset(train_seqs, labeler, ann_dict, config, homology_dict, esm_embeddings)\n val_ds = CAFA2Dataset(val_seqs, labeler, ann_dict, config, homology_dict, esm_embeddings)\n\n train_loader = DataLoader(\n train_ds, \n batch_size=config.BATCH_SIZE, \n shuffle=True, \n num_workers=2, \n pin_memory=True,\n persistent_workers=True\n )\n val_loader = DataLoader(\n val_ds, \n batch_size=config.BATCH_SIZE, \n shuffle=False,\n num_workers=2,\n pin_memory=True,\n persistent_workers=True\n )\n\n # --- 2. Setup Model & Optimization ---\n num_classes_dict = {a: len(labeler.valid_terms[a]) for a in labeler.aspects}\n model = MultiHeadDualPathModel(num_classes_dict).to(config.DEVICE)\n\n if config.NUM_GPUS > 1:\n model = nn.DataParallel(model)\n\n optimizer = torch.optim.AdamW(model.parameters(), lr=config.LEARNING_RATE, weight_decay=1e-2)\n scaler = GradScaler()\n\n scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n optimizer, mode='min', factor=0.5, patience=5, min_lr=1e-6\n )\n\n if config.IA_PATH.exists():\n print(f\"Loading IA weights from {config.IA_PATH}...\")\n ia_weights = labeler.get_ia_weights(config.IA_PATH)\n else:\n print(\"IA weights not found, defaulting to unweighted.\")\n ia_weights = None\n \n # Setup hierarchical edges for each aspect\n edges_dict = {a: labeler.get_edge_index(a).to(config.DEVICE) for a in labeler.aspects}\n criterion_hier = {a: HierarchicalLoss(edges_dict[a]) for a in labeler.aspects}\n criterion_asl = AsymmetricLossOptimized(gamma_neg=3, gamma_pos=1) # Applied per-aspect\n \n # --- 3. Resume Logic ---\n start_epoch = 0\n best_val_loss = float('inf')\n last_ckpt_path = MODELS_DIR / \"last_checkpoint.pt\"\n best_model_path = MODELS_DIR / \"best_model_phase2.pt\"\n \n # Initialize Tracking Dictionary\n history = {'train_loss':[], 'val_loss': [], 'weight_norm':[]}\n\n if last_ckpt_path.exists():\n checkpoint = torch.load(last_ckpt_path, map_location=config.DEVICE, weights_only=False)\n if isinstance(model, nn.DataParallel):\n model.module.load_state_dict(checkpoint['model_state_dict'])\n else:\n model.load_state_dict(checkpoint['model_state_dict'])\n\n optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n scaler.load_state_dict(checkpoint['scaler_state_dict'])\n start_epoch = checkpoint['epoch'] + 1\n best_val_loss = checkpoint['best_val_loss']\n if 'history' in checkpoint:\n history = checkpoint['history']\n print(f\"šŸ”„ Resuming from Epoch {start_epoch+1}\")\n\n # --- 4. Training Loop ---\n for epoch in range(start_epoch, config.NUM_EPOCHS):\n model.train()\n train_loss = 0\n for batch in tqdm(train_loader, desc=f\"Epoch {epoch+1} Train\", leave=False):\n esm = batch['esm'].to(config.DEVICE)\n seq = batch['seq'].to(config.DEVICE)\n phys = batch['physio'].to(config.DEVICE)\n \n # Load separate dicts\n hom_dict = {a: batch[f'homology_{a}'].to(config.DEVICE) for a in labeler.aspects}\n y_dict = {a: batch[f'labels_{a}'].to(config.DEVICE) for a in labeler.aspects}\n \n optimizer.zero_grad()\n \n with autocast('cuda'):\n logits_dict = model(esm, seq, phys, hom_dict)\n \n # Calculate loss per aspect and sum them\n total_loss = 0\n for aspect in labeler.aspects:\n loss_main = criterion_asl(logits_dict[aspect], y_dict[aspect])\n loss_h = criterion_hier[aspect](logits_dict[aspect])\n total_loss += loss_main + (config.LAMBDA_HIERARCHY * loss_h)\n \n scaler.scale(total_loss).backward()\n scaler.step(optimizer)\n scaler.update()\n train_loss += total_loss.item()\n \n model.eval()\n val_loss = 0\n with torch.no_grad():\n for batch in val_loader:\n esm = batch['esm'].to(config.DEVICE)\n seq = batch['seq'].to(config.DEVICE)\n phys = batch['physio'].to(config.DEVICE)\n \n # Load multi-head dictionaries for validation\n hom_dict = {a: batch[f'homology_{a}'].to(config.DEVICE) for a in labeler.aspects}\n y_dict = {a: batch[f'labels_{a}'].to(config.DEVICE) for a in labeler.aspects}\n\n with autocast('cuda'):\n logits_dict = model(esm, seq, phys, hom_dict)\n \n # Compute validation loss across all 3 heads\n batch_val_loss = 0\n for aspect in labeler.aspects:\n # We usually only track the main ASL loss for validation metrics\n loss_main = criterion_asl(logits_dict[aspect], y_dict[aspect])\n batch_val_loss += loss_main\n \n val_loss += batch_val_loss.item()\n\n avg_train = train_loss / len(train_loader)\n avg_val = val_loss / len(val_loader)\n \n scheduler.step(avg_val)\n\n total_norm = 0.0\n for p in model.parameters():\n if p.requires_grad:\n total_norm += p.data.norm(2).item() ** 2\n total_norm = total_norm ** 0.5\n \n history['train_loss'].append(avg_train)\n history['val_loss'].append(avg_val)\n history['weight_norm'].append(total_norm)\n\n print(f\"Epoch {epoch+1}: Train Loss {avg_train:.4f} | Val Loss {avg_val:.4f} | Weights Norm {total_norm:.2f}\")\n\n raw_model_state = model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict()\n\n if avg_val < best_val_loss:\n best_val_loss = avg_val\n torch.save({'state_dict': raw_model_state, 'config': config}, best_model_path)\n print(f\" āœ“ Best model saved\")\n\n if (epoch + 1) % 5 == 0 or (epoch + 1) == config.NUM_EPOCHS:\n torch.save({\n 'epoch': epoch,\n 'model_state_dict': raw_model_state,\n 'optimizer_state_dict': optimizer.state_dict(),\n 'scaler_state_dict': scaler.state_dict(),\n 'best_val_loss': best_val_loss,\n 'history': history\n }, last_ckpt_path)\n print(f\" šŸ’¾ Checkpoint saved at Epoch {epoch+1}\")\n \n best_ckpt = torch.load(best_model_path, map_location=config.DEVICE, weights_only=False)\n if isinstance(model, nn.DataParallel):\n model.module.load_state_dict(best_ckpt['state_dict'])\n else:\n model.load_state_dict(best_ckpt['state_dict'])\n\n return model, val_loader, history\n\n# Execute\ntrained_model, val_loader, history = train_phase2(Config, sequences, labeler, annotations, homology_dict, train_ids, val_ids)","metadata":{"id":"0z3klKhGQUyL","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ndef plot_training_curves(history):\n epochs = range(1, len(history['train_loss']) + 1)\n \n fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\n \n # 1. Loss Plot\n ax1.plot(epochs, history['train_loss'], 'b-', label='Training Loss', linewidth=2)\n ax1.plot(epochs, history['val_loss'], 'r-', label='Validation Loss', linewidth=2)\n ax1.set_title('Training and Validation Loss', fontsize=14)\n ax1.set_xlabel('Epochs', fontsize=12)\n ax1.set_ylabel('Loss', fontsize=12)\n ax1.grid(True, linestyle='--', alpha=0.7)\n ax1.legend(fontsize=12)\n \n # 2. Total Weights Plot\n if 'weight_norm' in history:\n ax2.plot(epochs, history['weight_norm'], 'g-', label='Total Weights (L2 Norm)', linewidth=2)\n ax2.set_title('Evolution of Total Model Weights', fontsize=14)\n ax2.set_xlabel('Epochs', fontsize=12)\n ax2.set_ylabel('L2 Norm Magnitude', fontsize=12)\n ax2.grid(True, linestyle='--', alpha=0.7)\n ax2.legend(fontsize=12)\n \n plt.tight_layout()\n plt.savefig(RESULTS_DIR / \"training_curves.png\", dpi=300, bbox_inches='tight')\n plt.show()\n\n# Run the Plot\nplot_training_curves(history)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Per-Term Threshold Optimization","metadata":{"id":"kxjXDAX6QYpn"}},{"cell_type":"code","source":"from sklearn.metrics import precision_recall_curve\nimport numpy as np\n\ndef optimize_thresholds(model, val_loader, config, labeler):\n print(\"Optimizing thresholds per aspect and per term...\")\n model.eval()\n\n # 1. Collect all predictions per aspect\n y_true = {a: [] for a in labeler.aspects}\n y_scores = {a: [] for a in labeler.aspects}\n\n with torch.no_grad():\n for batch in tqdm(val_loader, desc=\"Collecting Val Preds\"):\n esm = batch['esm'].to(config.DEVICE)\n seq = batch['seq'].to(config.DEVICE)\n phys = batch['physio'].to(config.DEVICE)\n \n hom_dict = {a: batch[f'homology_{a}'].to(config.DEVICE) for a in labeler.aspects}\n y_dict = {a: batch[f'labels_{a}'] for a in labeler.aspects} # Keep on CPU\n\n logits_dict = model(esm, seq, phys, hom_dict)\n \n for aspect in labeler.aspects:\n probs = torch.sigmoid(logits_dict[aspect]).cpu()\n y_true[aspect].append(y_dict[aspect])\n y_scores[aspect].append(probs)\n\n optimal_thresholds = {}\n \n # 2. Optimize per column for each aspect\n for aspect in labeler.aspects:\n print(f\"\\nEvaluating {aspect}...\")\n y_t = torch.cat(y_true[aspect], dim=0).numpy()\n y_s = torch.cat(y_scores[aspect], dim=0).numpy()\n \n num_classes = len(labeler.valid_terms[aspect])\n thresh_array = np.full(num_classes, 0.2) # Default start\n\n # We only optimize terms that actually appear in validation set\n active_cols = np.where(y_t.sum(axis=0) > 0)[0]\n\n for i in tqdm(active_cols, desc=f\"Tuning {aspect} thresholds\"):\n prec, rec, thresh = precision_recall_curve(y_t[:, i], y_s[:, i])\n \n # Calculate F1\n f1 = 2 * (prec * rec) / (prec + rec + 1e-9)\n best_idx = np.argmax(f1)\n\n # If best F1 is decent, update threshold. Else stick to default.\n if f1[best_idx] > 0.0:\n thresh_array[i] = thresh[best_idx]\n \n optimal_thresholds[aspect] = thresh_array\n\n # Save multiple arrays into a single compressed format\n np.savez(RESULTS_DIR / \"optimal_thresholds.npz\", **optimal_thresholds)\n print(\"\\nThresholds saved to optimal_thresholds.npz.\")\n return optimal_thresholds\n\n# Run Optimization\nthresholds_dict = optimize_thresholds(trained_model, val_loader, Config, labeler)","metadata":{"id":"l5IXvhGCQZYr","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Final Evaluation","metadata":{"id":"4vMBTBn9Qg6a"}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport torch\nfrom pathlib import Path\n\ndef load_ia_weights(ia_path, valid_terms):\n \"\"\"\n Loads raw (unnormalized) Information Accretion weights from IA.tsv \n exactly as required for the CAFA evaluation metric.\n \"\"\"\n if not Path(ia_path).exists():\n print(f\"āš ļø IA file not found at {ia_path}. Using uniform weights.\")\n return np.ones(len(valid_terms))\n \n ia_df = pd.read_csv(ia_path, sep='\\t', names=['term', 'ia'])\n ia_dict = dict(zip(ia_df['term'], ia_df['ia']))\n \n weights = np.zeros(len(valid_terms))\n for i, term in enumerate(valid_terms):\n weights[i] = ia_dict.get(term, 1.0)\n return weights\n\ndef evaluate_cafa_wfmax(model, val_loader, config, labeler):\n model.eval()\n \n scores = {a: [] for a in labeler.aspects}\n targets = {a: [] for a in labeler.aspects}\n\n with torch.no_grad():\n for batch in val_loader:\n esm = batch['esm'].to(config.DEVICE)\n seq = batch['seq'].to(config.DEVICE)\n phys = batch['physio'].to(config.DEVICE)\n \n hom_dict = {a: batch[f'homology_{a}'].to(config.DEVICE) for a in labeler.aspects}\n y_dict = {a: batch[f'labels_{a}'] for a in labeler.aspects}\n\n logits_dict = model(esm, seq, phys, hom_dict)\n \n for aspect in labeler.aspects:\n probs = torch.sigmoid(logits_dict[aspect]).cpu().numpy()\n scores[aspect].append(probs)\n targets[aspect].append(y_dict[aspect].numpy())\n\n print(f\"=== Final Validation Results (CAFA wFmax) ===\")\n \n global_wFmax = []\n \n # Evaluate Per-Aspect wFmax\n for aspect in labeler.aspects:\n y_score = np.vstack(scores[aspect])\n y_true = np.vstack(targets[aspect])\n \n # Load raw IA weights mapping to this aspect's vocabulary\n ia_weights = load_ia_weights(config.IA_PATH, labeler.valid_terms[aspect])\n \n # Threshold sweeping\n thresholds = np.arange(0.01, 1.0, 0.01)\n best_wf1 = 0.0\n best_t = 0.0\n best_wpr = 0.0\n best_wrc = 0.0\n \n for t in thresholds:\n y_pred = (y_score >= t).astype(int)\n \n # Element-wise weighting\n intersection_weights = (y_true * y_pred) * ia_weights\n pred_weights = y_pred * ia_weights\n true_weights = y_true * ia_weights\n \n # Sum weights per protein\n intersection_sum = np.sum(intersection_weights, axis=1)\n pred_sum = np.sum(pred_weights, axis=1)\n true_sum = np.sum(true_weights, axis=1)\n \n # CAFA isolates precision strictly to targets making >= 1 prediction \n # and isolates recall to targets possessing >= 1 ground truth term\n pred_mask = pred_sum > 0\n true_mask = true_sum > 0\n \n w_pr = np.mean(intersection_sum[pred_mask] / pred_sum[pred_mask]) if np.sum(pred_mask) > 0 else 0.0\n w_rc = np.mean(intersection_sum[true_mask] / true_sum[true_mask]) if np.sum(true_mask) > 0 else 0.0\n \n if w_pr + w_rc > 0:\n w_f1 = 2 * w_pr * w_rc / (w_pr + w_rc)\n else:\n w_f1 = 0.0\n \n if w_f1 > best_wf1:\n best_wf1 = w_f1\n best_t = t\n best_wpr = w_pr\n best_wrc = w_rc\n\n print(f\"\\n[{aspect} Performance]\")\n print(f\"Optimal Threshold: {best_t:.2f}\")\n print(f\"Weighted Precision (wPr): {best_wpr:.4f}\")\n print(f\"Weighted Recall (wRc): {best_wrc:.4f}\")\n print(f\"Weighted Fmax (wFmax): {best_wf1:.4f}\")\n \n global_wFmax.append(best_wf1)\n\n print(f\"\\n[GLOBAL Performance]\")\n print(f\"Average wFmax across all aspects: {np.mean(global_wFmax):.4f}\")\n\n# Execute evaluation\nevaluate_cafa_wfmax(trained_model, val_loader, Config, labeler)","metadata":{"id":"5nueH_pSQh2T","trusted":true},"outputs":[],"execution_count":null}]}