{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-05T12:13:18.118681Z","iopub.execute_input":"2025-12-05T12:13:18.119007Z","iopub.status.idle":"2025-12-05T12:13:18.581358Z","shell.execute_reply.started":"2025-12-05T12:13:18.118977Z","shell.execute_reply":"2025-12-05T12:13:18.579911Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import KFold\nfrom sklearn.metrics import f1_score\nfrom tqdm import tqdm\nimport networkx as nx\nimport gc\nimport re\n\n# Try to import transformers for Real Data Embedding\ntry:\n    from transformers import AutoTokenizer, AutoModel, EsmTokenizer\n    TRANSFORMERS_AVAILABLE = True\nexcept ImportError:\n    TRANSFORMERS_AVAILABLE = False\n    print(\"Warning: 'transformers' library not found. Install it for real embedding generation.\")\n\n# ==========================================\n# CONFIGURATION & HYPERPARAMETERS\n# ==========================================\nCONFIG = {\n    'MOCK_DATA': False, # Set to False to use real logic\n    'PATHS': {\n        'train_sequences': '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta',\n        'train_taxonomy': '/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv',\n        'train_terms': '/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv',\n        'test_sequences': '/kaggle/input/cafa-6-protein-function-prediction/Test/test_sequences.fasta',\n        'ia_file': '/kaggle/input/cafa-6-protein-function-prediction/IA.tsv',\n        'go_obo': '/kaggle/input/cafa-6-protein-function-prediction/go-basic.obo', \n        'text_data': '/kaggle/input/cafa-6-protein-function-prediction/text_data.tsv'\n    },\n    'MODEL': {\n        'model_name': 'facebook/esm2_t6_8M_UR50D', # Small model for speed. Use 'facebook/esm2_t33_650M_UR50D' for high score\n        'seq_emb_dim': 320, # Matches esm2_t6_8M_UR50D output (650M has 1280 dim)\n        'text_emb_dim': 768, \n        'hidden_dim': 512,\n        'dropout': 0.3,\n        'learning_rate': 2e-4, \n        'batch_size': 32,\n        'epochs': 5,           \n        'num_classes': 500,    # Target top 500 most frequent terms\n        'n_folds': 3           \n    },\n    'LOSS': {\n        'gamma': 2.0,          \n        'alpha': 0.25          \n    }\n}\n\n# ==========================================\n# PHASE 1: DATA ENGINEERING & ONTOLOGY\n# ==========================================\n\nclass GODAG:\n    \"\"\"Handles the Gene Ontology Graph structure.\"\"\"\n    def __init__(self):\n        self.graph = nx.DiGraph()\n        self.term_to_idx = {}\n        self.idx_to_term = {}\n\n    def build_mock_graph(self, terms):\n        print(\"Building GO DAG...\")\n        for term in terms:\n            self.graph.add_node(term)\n        sorted_terms = sorted(list(terms))\n        for i in range(len(sorted_terms) - 1):\n            if np.random.rand() > 0.7:\n                self.graph.add_edge(sorted_terms[i], sorted_terms[i+1])\n        self.update_indices(terms)\n\n    def load_real_graph(self, obo_path, relevant_terms):\n        print(f\"Loading GO Graph from {obo_path}...\")\n        \n        # Simple OBO Parser\n        term_pattern = re.compile(r'\\[Term\\]')\n        id_pattern = re.compile(r'^id: (GO:\\d+)')\n        is_a_pattern = re.compile(r'^is_a: (GO:\\d+)')\n        \n        current_term = None\n        \n        # Initialize graph with relevant terms\n        for term in relevant_terms:\n            self.graph.add_node(term)\n\n        if not os.path.exists(obo_path):\n            print(\"OBO file not found. Skipping graph edges.\")\n            self.update_indices(relevant_terms)\n            return\n\n        with open(obo_path, 'r') as f:\n            for line in f:\n                line = line.strip()\n                if term_pattern.match(line):\n                    current_term = None\n                \n                id_match = id_pattern.match(line)\n                if id_match:\n                    current_term = id_match.group(1)\n                    if current_term not in self.graph:\n                        self.graph.add_node(current_term)\n                \n                is_a_match = is_a_pattern.match(line)\n                if is_a_match and current_term:\n                    parent = is_a_match.group(1)\n                    # Edge direction: Child -> Parent (standard for propagation)\n                    # OR Parent -> Child (standard for graph traversal)\n                    # For consistency propagation: if Child is true, Parent is true.\n                    # We add Child -> Parent edges.\n                    self.graph.add_edge(current_term, parent)\n        \n        print(f\"Graph loaded with {self.graph.number_of_nodes()} nodes.\")\n        self.update_indices(relevant_terms)\n\n    def update_indices(self, terms):\n        # Only index the terms we are actually predicting\n        self.idx_to_term = {i: t for i, t in enumerate(sorted(list(terms)))}\n        self.term_to_idx = {t: i for i, t in enumerate(sorted(list(terms)))}\n\n    def propagate_probs(self, preds_dict):\n        \"\"\"Enforce Hierarchy: P(parent) >= P(child)\"\"\"\n        # In a DAG, efficient propagation follows topological sort.\n        # Here we do iterative updates for simplicity and robustness.\n        # We iterate \"up\" the graph.\n        changed = True\n        pass_count = 0\n        while changed and pass_count < 3:\n            changed = False\n            pass_count += 1\n            # Iterate over all edges (child -> parent)\n            for child, parent in self.graph.edges():\n                if child in preds_dict and parent in preds_dict:\n                    if preds_dict[child] > preds_dict[parent]:\n                        preds_dict[parent] = preds_dict[child]\n                        changed = True\n        return preds_dict\n\n# ==========================================\n# PHASE 3 & 4: DATASET & ADVANCED MODEL\n# ==========================================\n\nclass ProteinTextDataset(Dataset):\n    def __init__(self, ids, seq_embeddings, text_embeddings, labels=None, num_classes=100):\n        self.ids = ids\n        self.seq_embeddings = seq_embeddings\n        self.text_embeddings = text_embeddings\n        self.labels = labels\n        self.num_classes = num_classes\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, idx):\n        seq_emb = torch.tensor(self.seq_embeddings[idx], dtype=torch.float32)\n        text_emb = torch.tensor(self.text_embeddings[idx], dtype=torch.float32)\n        \n        item = {'seq_emb': seq_emb, 'text_emb': text_emb, 'id': self.ids[idx]}\n\n        if self.labels is not None:\n            label_vec = torch.zeros(self.num_classes, dtype=torch.float32)\n            # Create Multi-hot vector\n            for lbl_idx in self.labels[idx]:\n                if lbl_idx < self.num_classes:\n                    label_vec[lbl_idx] = 1.0\n            item['labels'] = label_vec\n            \n        return item\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, in_features, out_features, dropout=0.3):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Linear(in_features, out_features),\n            nn.BatchNorm1d(out_features),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(out_features, out_features),\n            nn.BatchNorm1d(out_features)\n        )\n        self.relu = nn.ReLU()\n        self.shortcut = nn.Sequential()\n        if in_features != out_features:\n            self.shortcut = nn.Linear(in_features, out_features)\n\n    def forward(self, x):\n        return self.relu(self.block(x) + self.shortcut(x))\n\nclass AdvancedMultimodalModel(nn.Module):\n    def __init__(self, seq_dim, text_dim, hidden_dim, num_classes):\n        super().__init__()\n        \n        # 1. Sequence Branch with Residuals\n        self.seq_encoder = nn.Sequential(\n            nn.Linear(seq_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(),\n            ResidualBlock(hidden_dim, hidden_dim)\n        )\n        \n        # 2. Text Branch with Residuals\n        self.text_encoder = nn.Sequential(\n            nn.Linear(text_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(),\n            ResidualBlock(hidden_dim, hidden_dim)\n        )\n        \n        # 3. Gated Fusion Mechanism\n        self.fusion_gate = nn.Sequential(\n            nn.Linear(hidden_dim * 2, 1),\n            nn.Sigmoid()\n        )\n        \n        self.classifier = nn.Sequential(\n            nn.Linear(hidden_dim * 2, hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(0.4),\n            nn.Linear(hidden_dim, num_classes)\n        )\n\n    def forward(self, seq_emb, text_emb):\n        s_feat = self.seq_encoder(seq_emb)\n        t_feat = self.text_encoder(text_emb)\n        \n        combined_raw = torch.cat([s_feat, t_feat], dim=1)\n        gate = self.fusion_gate(combined_raw)\n        \n        # Weighted combination + Skip connection\n        weighted_feat = torch.cat([s_feat * gate, t_feat * (1 - gate)], dim=1)\n        \n        logits = self.classifier(weighted_feat)\n        return torch.sigmoid(logits)\n\nclass FocalLoss(nn.Module):\n    def __init__(self, alpha=0.25, gamma=2.0):\n        super().__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.bce = nn.BCELoss(reduction='none')\n\n    def forward(self, inputs, targets):\n        bce_loss = self.bce(inputs, targets)\n        pt = torch.exp(-bce_loss)\n        focal_loss = self.alpha * (1-pt)**self.gamma * bce_loss\n        return focal_loss.mean()\n\n# ==========================================\n# DATA PREPARATION (REAL DATA LOGIC)\n# ==========================================\n\ndef parse_fasta(path, limit=None):\n    \"\"\"Parses FASTA file into a dictionary {id: sequence}.\"\"\"\n    sequences = {}\n    if not os.path.exists(path):\n        return sequences\n    \n    print(f\"Parsing FASTA: {path}\")\n    with open(path, 'r') as f:\n        current_id = None\n        current_seq = []\n        count = 0\n        for line in f:\n            if line.startswith('>'):\n                if current_id:\n                    sequences[current_id] = ''.join(current_seq)\n                    count += 1\n                    if limit and count >= limit: break\n                current_id = line.strip().split()[0][1:]\n                current_seq = []\n            else:\n                current_seq.append(line.strip())\n        if current_id and (not limit or count < limit):\n            sequences[current_id] = ''.join(current_seq)\n    return sequences\n\ndef parse_terms(path, relevant_ids):\n    \"\"\"Parses Train Terms TSV into dictionary {id: [go_terms]}.\"\"\"\n    labels = {}\n    print(f\"Parsing Labels: {path}\")\n    if not os.path.exists(path): return labels\n    \n    df = pd.read_csv(path, sep='\\t', header=None, names=['Id', 'Term', 'Aspect'])\n    # Filter for IDs we actually have sequences for\n    df = df[df['Id'].isin(relevant_ids)]\n    \n    # Group by ID\n    for protein_id, group in tqdm(df.groupby('Id')):\n        labels[protein_id] = group['Term'].tolist()\n        \n    return labels\n\ndef generate_embeddings(sequences, model_name, batch_size=32):\n    \"\"\"\n    Generates embeddings using a pre-trained ESM model.\n    \"\"\"\n    if not TRANSFORMERS_AVAILABLE:\n        print(\"Transformers not installed. Returning random embeddings.\")\n        return np.random.randn(len(sequences), CONFIG['MODEL']['seq_emb_dim']), list(sequences.keys())\n\n    print(f\"Loading Model: {model_name}\")\n    \n    tokenizer = None\n    model = None\n\n    # --- ERROR FIX: Robust Tokenizer Loading ---\n    # 1. Try EsmTokenizer (Preferred)\n    try:\n        tokenizer = EsmTokenizer.from_pretrained(model_name)\n    except Exception as e_esm:\n        print(f\"Warning: EsmTokenizer failed ({e_esm}). Attempting AutoTokenizer...\")\n        # 2. Try AutoTokenizer (Fallback)\n        try:\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n        except Exception as e_auto:\n            print(f\"Error: AutoTokenizer also failed ({e_auto}).\")\n            print(\"Fallback: Returning random embeddings to prevent pipeline crash.\")\n            return np.random.randn(len(sequences), CONFIG['MODEL']['seq_emb_dim']), list(sequences.keys())\n    # -------------------------------------------\n\n    try:\n        model = AutoModel.from_pretrained(model_name)\n    except Exception as e:\n        print(f\"Error: Failed to load AutoModel ({e}). Returning random embeddings.\")\n        return np.random.randn(len(sequences), CONFIG['MODEL']['seq_emb_dim']), list(sequences.keys())\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model = model.to(device)\n    model.eval()\n    \n    seq_list = list(sequences.values())\n    ids_list = list(sequences.keys())\n    embeddings = []\n    \n    print(\"Generating Embeddings (this may take time)...\")\n    with torch.no_grad():\n        for i in tqdm(range(0, len(seq_list), batch_size)):\n            batch_seqs = seq_list[i:i+batch_size]\n            \n            # Tokenize with safety\n            try:\n                inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True, max_length=1024)\n                inputs = {k: v.to(device) for k, v in inputs.items()}\n                \n                # Forward pass\n                outputs = model(**inputs)\n                \n                # Take the mean of the last hidden state\n                batch_embeddings = outputs.last_hidden_state.mean(dim=1).cpu().numpy()\n                embeddings.append(batch_embeddings)\n            except Exception as e_batch:\n                print(f\"Error processing batch {i}: {e_batch}\")\n                embeddings.append(np.zeros((len(batch_seqs), CONFIG['MODEL']['seq_emb_dim'])))\n\n    if len(embeddings) > 0:\n        return np.concatenate(embeddings, axis=0), ids_list\n    else:\n        return np.array([]), []\n\ndef generate_mock_data(n_samples=1000, n_classes=100, is_test=False):\n    prefix = \"TEST\" if is_test else \"TRAIN\"\n    print(f\"--- Generating Mock {prefix} Data ({n_samples} samples) ---\")\n    ids = [f\"{prefix}_{str(i).zfill(5)}\" for i in range(n_samples)]\n    seq_embs = np.random.randn(n_samples, CONFIG['MODEL']['seq_emb_dim']).astype(np.float32)\n    text_embs = np.random.randn(n_samples, CONFIG['MODEL']['text_emb_dim']).astype(np.float32)\n    labels = []\n    if not is_test:\n        for _ in range(n_samples):\n            n_lbls = np.random.randint(1, 10)\n            lbls = np.random.choice(range(n_classes), size=n_lbls, replace=False)\n            labels.append(list(lbls))\n    go_term_ids = [f\"GO:{str(i).zfill(7)}\" for i in range(n_classes)]\n    return ids, seq_embs, text_embs, labels, go_term_ids\n\ndef load_data():\n    if CONFIG['MOCK_DATA']:\n        # Fallback to Mock\n        tr_ids, tr_seq, tr_text, tr_lbls, go_terms = generate_mock_data(n_samples=800, n_classes=CONFIG['MODEL']['num_classes'], is_test=False)\n        te_ids, te_seq, te_text, _, _ = generate_mock_data(n_samples=200, n_classes=CONFIG['MODEL']['num_classes'], is_test=True)\n        return (tr_ids, tr_seq, tr_text, tr_lbls), (te_ids, te_seq, te_text), go_terms\n    \n    else:\n        # --- REAL DATA LOADING ---\n        print(\"--- Loading Real Data ---\")\n        \n        # 1. Parse Sequences\n        # Limit set for demonstration speed. Set limit=None for full run.\n        train_seqs_dict = parse_fasta(CONFIG['PATHS']['train_sequences'], limit=5000) \n        test_seqs_dict = parse_fasta(CONFIG['PATHS']['test_sequences'], limit=1000)\n        \n        if not train_seqs_dict or not test_seqs_dict:\n            print(\"Error: No sequences found (Train or Test). Reverting to Mock Data.\")\n            CONFIG['MOCK_DATA'] = True\n            return load_data()\n\n        # 2. Generate/Load Embeddings\n        # Check if saved embeddings exist to save time\n        if os.path.exists('train_embs.npy') and os.path.exists('train_ids.npy'):\n            print(\"Loading cached training embeddings...\")\n            tr_seq_embs = np.load('train_embs.npy')\n            tr_ids = np.load('train_ids.npy').tolist()\n        else:\n            print(\"Generating training embeddings...\")\n            tr_seq_embs, tr_ids = generate_embeddings(train_seqs_dict, CONFIG['MODEL']['model_name'])\n            np.save('train_embs.npy', tr_seq_embs)\n            np.save('train_ids.npy', np.array(tr_ids))\n            \n        if os.path.exists('test_embs.npy') and os.path.exists('test_ids.npy'):\n            print(\"Loading cached test embeddings...\")\n            te_seq_embs = np.load('test_embs.npy')\n            te_ids = np.load('test_ids.npy').tolist()\n        else:\n            print(\"Generating test embeddings...\")\n            te_seq_embs, te_ids = generate_embeddings(test_seqs_dict, CONFIG['MODEL']['model_name'])\n            np.save('test_embs.npy', te_seq_embs)\n            np.save('test_ids.npy', np.array(te_ids))\n\n        # 3. Text Embeddings (Placeholder for actual text processing)\n        # In a real run, you'd load the 'text_data' file and run PubMedBERT similar to above.\n        # Here we use random noise for the text channel to allow the pipeline to complete.\n        print(\"Simulating Text Embeddings (Implement PubMedBERT here)...\")\n        tr_text_embs = np.random.randn(len(tr_ids), CONFIG['MODEL']['text_emb_dim']).astype(np.float32)\n        te_text_embs = np.random.randn(len(te_ids), CONFIG['MODEL']['text_emb_dim']).astype(np.float32)\n\n        # 4. Process Labels\n        raw_labels = parse_terms(CONFIG['PATHS']['train_terms'], set(tr_ids))\n        \n        # Frequency Analysis: Select Top N Terms\n        all_terms = [t for sublist in raw_labels.values() for t in sublist]\n        term_counts = pd.Series(all_terms).value_counts()\n        top_terms = term_counts.head(CONFIG['MODEL']['num_classes']).index.tolist()\n        \n        # Map IDs to Top Terms only\n        final_labels = []\n        term_to_idx = {t: i for i, t in enumerate(top_terms)}\n        \n        for pid in tr_ids:\n            if pid in raw_labels:\n                # Convert terms to indices, filtering out non-top terms\n                idxs = [term_to_idx[t] for t in raw_labels[pid] if t in term_to_idx]\n                final_labels.append(idxs)\n            else:\n                final_labels.append([])\n                \n        return (tr_ids, tr_seq_embs, tr_text_embs, final_labels), (te_ids, te_seq_embs, te_text_embs), top_terms\n\n# ==========================================\n# PIPELINE EXECUTION (ENSEMBLE)\n# ==========================================\n\ndef run_pipeline():\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Running on: {device}\")\n\n    # 1. Load Data\n    (tr_ids, tr_seq, tr_text, tr_lbls), (te_ids, te_seq, te_text), go_terms = load_data()\n    \n    # 2. Ontology\n    dag = GODAG()\n    if CONFIG['MOCK_DATA']:\n        dag.build_mock_graph(go_terms) \n    else:\n        dag.load_real_graph(CONFIG['PATHS']['go_obo'], go_terms)\n\n    # 3. K-Fold Ensemble Setup\n    kf = KFold(n_splits=CONFIG['MODEL']['n_folds'], shuffle=True, random_state=42)\n    \n    # Placeholder for test predictions (summing probabilities)\n    test_preds_sum = np.zeros((len(te_ids), CONFIG['MODEL']['num_classes']))\n    \n    # Convert Test Data to Dataset once\n    test_dataset = ProteinTextDataset(te_ids, te_seq, te_text, labels=None, num_classes=CONFIG['MODEL']['num_classes'])\n    test_loader = DataLoader(test_dataset, batch_size=CONFIG['MODEL']['batch_size'], shuffle=False)\n    \n    # --- ADDED SCORE TRACKING ---\n    fold_scores = []\n    # ----------------------------\n\n    print(f\"\\n--- Starting {CONFIG['MODEL']['n_folds']}-Fold Cross Validation ---\")\n\n    for fold, (train_idx, val_idx) in enumerate(kf.split(tr_ids)):\n        print(f\"\\n[FOLD {fold+1}/{CONFIG['MODEL']['n_folds']}]\")\n        \n        # Split Data\n        train_ds = ProteinTextDataset(\n            [tr_ids[i] for i in train_idx], tr_seq[train_idx], tr_text[train_idx], [tr_lbls[i] for i in train_idx], \n            num_classes=CONFIG['MODEL']['num_classes']\n        )\n        val_ds = ProteinTextDataset(\n            [tr_ids[i] for i in val_idx], tr_seq[val_idx], tr_text[val_idx], [tr_lbls[i] for i in val_idx], \n            num_classes=CONFIG['MODEL']['num_classes']\n        )\n        \n        train_loader = DataLoader(train_ds, batch_size=CONFIG['MODEL']['batch_size'], shuffle=True)\n        val_loader = DataLoader(val_ds, batch_size=CONFIG['MODEL']['batch_size'], shuffle=False)\n\n        # Initialize Model (Fresh for each fold)\n        model = AdvancedMultimodalModel(\n            seq_dim=CONFIG['MODEL']['seq_emb_dim'],\n            text_dim=CONFIG['MODEL']['text_emb_dim'],\n            hidden_dim=CONFIG['MODEL']['hidden_dim'],\n            num_classes=CONFIG['MODEL']['num_classes']\n        ).to(device)\n\n        optimizer = torch.optim.AdamW(model.parameters(), lr=CONFIG['MODEL']['learning_rate'], weight_decay=1e-4)\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=CONFIG['MODEL']['epochs'])\n        criterion = FocalLoss(gamma=CONFIG['LOSS']['gamma'], alpha=CONFIG['LOSS']['alpha'])\n\n        # Training\n        for epoch in range(CONFIG['MODEL']['epochs']):\n            model.train()\n            train_loss = 0\n            for batch in tqdm(train_loader, desc=f\"Ep {epoch+1}\", leave=False):\n                s_emb, t_emb, targets = batch['seq_emb'].to(device), batch['text_emb'].to(device), batch['labels'].to(device)\n                \n                optimizer.zero_grad()\n                outputs = model(s_emb, t_emb)\n                loss = criterion(outputs, targets)\n                loss.backward()\n                optimizer.step()\n                train_loss += loss.item()\n            \n            scheduler.step()\n            if epoch == CONFIG['MODEL']['epochs'] - 1:\n                print(f\"  Final Train Loss: {train_loss/len(train_loader):.4f}\")\n\n        # --- VALIDATION & SCORING ---\n        model.eval()\n        val_targets = []\n        val_preds = []\n        with torch.no_grad():\n            for batch in val_loader:\n                s_emb, t_emb, targets = batch['seq_emb'].to(device), batch['text_emb'].to(device), batch['labels'].to(device)\n                outputs = model(s_emb, t_emb)\n                val_preds.append(outputs.cpu().numpy())\n                val_targets.append(targets.cpu().numpy())\n        \n        if len(val_preds) > 0:\n            val_preds = np.concatenate(val_preds)\n            val_targets = np.concatenate(val_targets)\n            # Calculate Micro-F1 (common proxy for protein function)\n            fold_f1 = f1_score(val_targets, (val_preds > 0.3).astype(int), average='micro')\n            print(f\"  Fold {fold+1} F1 Score (Micro, th=0.3): {fold_f1:.4f}\")\n            fold_scores.append(fold_f1)\n        # ----------------------------\n\n        # Inference on Test Set for this fold\n        fold_test_preds = []\n        with torch.no_grad():\n            for batch in test_loader:\n                s_emb, t_emb = batch['seq_emb'].to(device), batch['text_emb'].to(device)\n                outputs = model(s_emb, t_emb)\n                fold_test_preds.append(outputs.cpu().numpy())\n        \n        # Add to ensemble sum\n        if len(fold_test_preds) > 0:\n            test_preds_sum += np.concatenate(fold_test_preds, axis=0)\n        else:\n            print(f\"Warning: No predictions generated for fold {fold+1} (Test set empty?).\")\n\n    # 4. Average Predictions (Ensemble)\n    print(\"\\n--- Generating Final Ensemble Predictions ---\")\n    if len(fold_scores) > 0:\n        print(f\"Average CV F1 Score: {np.mean(fold_scores):.4f}\")\n        \n    avg_preds = test_preds_sum / CONFIG['MODEL']['n_folds']\n    \n    # 5. Post-Processing & Submission\n    submission_rows = []\n    print(\"Applying Hierarchical Consistency & Filtering...\")\n    \n    for i, protein_id in enumerate(te_ids):\n        pred_dict = {\n            dag.idx_to_term[idx]: float(avg_preds[i][idx])\n            for idx in range(CONFIG['MODEL']['num_classes'])\n        }\n        \n        # Propagate consistency\n        consistent_preds = dag.propagate_probs(pred_dict)\n        \n        for term, prob in consistent_preds.items():\n            if prob > 0.01: # Threshold\n                submission_rows.append(f\"{protein_id}\\t{term}\\t{prob:.3f}\")\n\n    # Write File\n    with open('submission.tsv', 'w') as f:\n        for row in submission_rows:\n            f.write(row + \"\\n\")\n            \n    print(f\"Success! Generated {len(submission_rows)} predictions from {CONFIG['MODEL']['n_folds']}-fold ensemble.\")\n    print(\"Sample:\", submission_rows[0] if submission_rows else \"No predictions\")\n\nif __name__ == \"__main__\":\n    run_pipeline()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-05T12:13:18.583431Z","iopub.execute_input":"2025-12-05T12:13:18.583989Z","iopub.status.idle":"2025-12-05T12:13:31.170219Z","shell.execute_reply.started":"2025-12-05T12:13:18.583958Z","shell.execute_reply":"2025-12-05T12:13:31.169218Z"}},"outputs":[],"execution_count":null}]}