{"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":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"},{"sourceId":611663,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":459479,"modelId":475349}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"pip install obonet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T10:17:03.059343Z","iopub.execute_input":"2025-10-19T10:17:03.059623Z","iopub.status.idle":"2025-10-19T10:17:07.081423Z","shell.execute_reply.started":"2025-10-19T10:17:03.059602Z","shell.execute_reply":"2025-10-19T10:17:07.080502Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nProtein Function Prediction using Deep Learning\nCAFA Challenge - Advanced Solution with Detailed Evaluation\n\"\"\"\n\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 torch.cuda.amp import autocast, GradScaler\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import f1_score, precision_score, recall_score\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom collections import defaultdict\nimport warnings\nimport gc\nimport obonet\nimport networkx as nx\nfrom tqdm.auto import tqdm\nimport os\n\nwarnings.filterwarnings('ignore')\n\n# Set device\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# ============================================================================\n# 1. DATA LOADING AND PREPROCESSING\n# ============================================================================\n\nclass DataLoader_CAFA:\n    \"\"\"Load and preprocess all competition data\"\"\"\n\n    def __init__(self):\n        self.amino_acids = 'ACDEFGHIKLMNPQRSTVWY'\n        self.aa_to_idx = {aa: idx for idx, aa in enumerate(self.amino_acids)}\n        self.aa_to_idx['X'] = len(self.amino_acids)  # Unknown\n\n    def load_fasta(self, filepath):\n        \"\"\"Load FASTA file\"\"\"\n        sequences = {}\n        current_id = None\n        current_seq = []\n\n        with open(filepath, 'r') as f:\n            for line in f:\n                line = line.strip()\n                if line.startswith('>'):\n                    if current_id:\n                        sequences[current_id] = ''.join(current_seq)\n                    # Extract UniProt ID\n                    parts = line[1:].split('|')\n                    current_id = parts[1] if len(parts) > 1 else parts[0].split()[0]\n                    current_seq = []\n                else:\n                    current_seq.append(line)\n\n            if current_id:\n                sequences[current_id] = ''.join(current_seq)\n\n        return sequences\n\n    def load_terms(self, filepath):\n        \"\"\"Load GO term annotations\"\"\"\n        df = pd.read_csv(filepath, sep='\\t', header=None,\n                         names=['EntryID', 'term', 'aspect'])\n        return df\n\n    def load_taxonomy(self, filepath):\n        \"\"\"Load taxonomy information\"\"\"\n        df = pd.read_csv(filepath, sep='\\t', header=None,\n                         names=['EntryID', 'taxon'])\n        return df\n\n    def load_ia_weights(self, filepath):\n        \"\"\"Load information accretion weights\"\"\"\n        df = pd.read_csv(filepath, sep='\\t', header=None,\n                         names=['GO_term', 'IA'])\n        return dict(zip(df['GO_term'], df['IA']))\n\n    def load_go_graph(self, filepath):\n        \"\"\"Load GO ontology graph\"\"\"\n        graph = obonet.read_obo(filepath)\n        return graph\n\n    def encode_sequence(self, seq, max_len=1024):\n        \"\"\"Encode amino acid sequence to indices\"\"\"\n        encoded = [self.aa_to_idx.get(aa, self.aa_to_idx['X']) for aa in seq[:max_len]]\n        # Pad if necessary\n        if len(encoded) < max_len:\n            encoded += [self.aa_to_idx['X']] * (max_len - len(encoded))\n        return np.array(encoded, dtype=np.int32)\n\n# ============================================================================\n# 2. GO HIERARCHY HANDLER\n# ============================================================================\n\nclass GOHierarchy:\n    \"\"\"Handle GO term hierarchy and propagation\"\"\"\n\n    def __init__(self, go_graph):\n        self.graph = go_graph\n        self.subontology_roots = {\n            'P': 'GO:0008150',  # Biological Process\n            'C': 'GO:0005575',  # Cellular Component\n            'F': 'GO:0003674'   # Molecular Function\n        }\n        self.term_to_ancestors = {}\n        self._build_ancestor_map()\n\n    def _build_ancestor_map(self):\n        \"\"\"Build mapping of each term to all its ancestors\"\"\"\n        for node in self.graph.nodes():\n            try:\n                ancestors = nx.ancestors(self.graph, node)\n                self.term_to_ancestors[node] = ancestors\n            except:\n                self.term_to_ancestors[node] = set()\n\n    def get_ancestors(self, term):\n        \"\"\"Get all ancestor terms\"\"\"\n        return self.term_to_ancestors.get(term, set())\n\n    def propagate_predictions(self, term_scores):\n        \"\"\"Propagate predictions to parent terms\"\"\"\n        propagated = term_scores.copy()\n        for term, score in term_scores.items():\n            ancestors = self.get_ancestors(term)\n            for ancestor in ancestors:\n                if ancestor in propagated:\n                    propagated[ancestor] = max(propagated[ancestor], score)\n                else:\n                    propagated[ancestor] = score\n        return propagated\n\n# ============================================================================\n# 3. DATASET CLASS\n# ============================================================================\n\nclass ProteinDataset(Dataset):\n    \"\"\"PyTorch Dataset for protein sequences\"\"\"\n\n    def __init__(self, sequences, labels=None, data_loader=None, max_len=1024):\n        self.sequences = sequences\n        self.labels = labels\n        self.protein_ids = list(sequences.keys())\n        self.data_loader = data_loader\n        self.max_len = max_len\n\n    def __len__(self):\n        return len(self.protein_ids)\n\n    def __getitem__(self, idx):\n        protein_id = self.protein_ids[idx]\n        seq = self.sequences[protein_id]\n\n        # Encode sequence\n        encoded_seq = self.data_loader.encode_sequence(seq, self.max_len)\n\n        item = {\n            'protein_id': protein_id,\n            'sequence': torch.tensor(encoded_seq, dtype=torch.long),\n            'length': min(len(seq), self.max_len)\n        }\n\n        if self.labels is not None:\n            item['labels'] = torch.tensor(self.labels[protein_id], dtype=torch.float32)\n\n        return item\n\n# ============================================================================\n# 4. ADVANCED PROTEIN ENCODER MODEL\n# ============================================================================\n\nclass ProteinEncoder(nn.Module):\n    \"\"\"Advanced protein sequence encoder with CNN + Transformer\"\"\"\n\n    def __init__(self, vocab_size=21, embed_dim=128, num_layers=4,\n                 num_heads=8, ff_dim=512, max_len=1024, dropout=0.1):\n        super().__init__()\n\n        # Embedding layer\n        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=20)\n        self.pos_encoding = self.create_positional_encoding(max_len, embed_dim)\n\n        # CNN layers for local pattern extraction\n        self.conv1 = nn.Conv1d(embed_dim, 256, kernel_size=3, padding=1)\n        self.conv2 = nn.Conv1d(256, 256, kernel_size=5, padding=2)\n        self.conv3 = nn.Conv1d(256, 256, kernel_size=7, padding=3)\n        self.bn1 = nn.BatchNorm1d(256)\n        self.bn2 = nn.BatchNorm1d(256)\n        self.bn3 = nn.BatchNorm1d(256)\n\n        # Transformer encoder for global context\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=256, nhead=num_heads, dim_feedforward=ff_dim,\n            dropout=dropout, activation='gelu', batch_first=True\n        )\n        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n\n        # Attention pooling\n        self.attention_pool = nn.Linear(256, 1)\n\n        self.dropout = nn.Dropout(dropout)\n\n    def create_positional_encoding(self, max_len, embed_dim):\n        \"\"\"Create sinusoidal positional encoding\"\"\"\n        pe = torch.zeros(max_len, embed_dim)\n        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, embed_dim, 2).float() *\n                            (-np.log(10000.0) / embed_dim))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        return pe.unsqueeze(0)\n\n    def forward(self, x, lengths):\n        batch_size, seq_len = x.shape\n\n        # Embedding\n        x = self.embedding(x)  # [B, L, E]\n        x = x + self.pos_encoding[:, :seq_len, :].to(x.device)\n\n        # CNN path\n        x_cnn = x.transpose(1, 2)  # [B, E, L]\n        x1 = F.relu(self.bn1(self.conv1(x_cnn)))\n        x2 = F.relu(self.bn2(self.conv2(x1)))\n        x3 = F.relu(self.bn3(self.conv3(x2)))\n        x_cnn = x3.transpose(1, 2)  # [B, L, 256]\n\n        # Create attention mask\n        mask = torch.arange(seq_len, device=x.device)[None, :] >= lengths[:, None]\n\n        # Transformer path\n        x_trans = self.transformer(x_cnn, src_key_padding_mask=mask)\n\n        # Attention pooling\n        attn_weights = self.attention_pool(x_trans).squeeze(-1)  # [B, L]\n\n        # --- START MODIFICATION ---\n        # Replace the hardcoded -1e9 with a value that is safe for the current dtype (float16 or float32)\n        # This prevents the overflow error when using autocast (AMP)\n        fill_value = torch.finfo(attn_weights.dtype).min\n        attn_weights = attn_weights.masked_fill(mask, fill_value)\n        # --- END MODIFICATION ---\n\n        attn_weights = F.softmax(attn_weights, dim=1).unsqueeze(1)  # [B, 1, L]\n\n        # Weighted sum\n        pooled = torch.bmm(attn_weights, x_trans).squeeze(1)  # [B, 256]\n\n        return self.dropout(pooled)\n# ============================================================================\n# 5. MULTI-LABEL CLASSIFICATION HEAD\n# ============================================================================\n\nclass MultiLabelClassifier(nn.Module):\n    \"\"\"Multi-label classifier for GO term prediction\"\"\"\n\n    def __init__(self, input_dim, num_classes, hidden_dims=[512, 256]):\n        super().__init__()\n\n        layers = []\n        prev_dim = input_dim\n\n        for hidden_dim in hidden_dims:\n            layers.extend([\n                nn.Linear(prev_dim, hidden_dim),\n                nn.BatchNorm1d(hidden_dim),\n                nn.ReLU(),\n                nn.Dropout(0.3)\n            ])\n            prev_dim = hidden_dim\n\n        layers.append(nn.Linear(prev_dim, num_classes))\n        self.classifier = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.classifier(x)\n\n# ============================================================================\n# 6. COMPLETE MODEL\n# ============================================================================\n\nclass ProteinFunctionPredictor(nn.Module):\n    \"\"\"Complete model for protein function prediction\"\"\"\n\n    def __init__(self, num_classes, vocab_size=21):\n        super().__init__()\n        self.encoder = ProteinEncoder(vocab_size=vocab_size)\n        self.classifier = MultiLabelClassifier(256, num_classes)\n\n    def forward(self, x, lengths):\n        features = self.encoder(x, lengths)\n        logits = self.classifier(features)\n        return logits\n\n# ============================================================================\n# 7. TRAINING UTILITIES\n# ============================================================================\n\nclass FocalLoss(nn.Module):\n    \"\"\"Focal Loss for handling class imbalance\"\"\"\n\n    def __init__(self, alpha=0.25, gamma=2.0):\n        super().__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n\n    def forward(self, inputs, targets):\n        bce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')\n        pt = torch.exp(-bce_loss)\n        focal_loss = self.alpha * (1 - pt) ** self.gamma * bce_loss\n        return focal_loss.mean()\n\nclass Trainer:\n    \"\"\"Training and evaluation handler\"\"\"\n\n    def __init__(self, model, train_loader, val_loader,\n                 ia_weights, device, num_epochs=10):\n        self.model = model.to(device)\n        self.train_loader = train_loader\n        self.val_loader = val_loader\n        self.ia_weights = ia_weights\n        self.device = device\n        self.num_epochs = num_epochs\n\n        self.optimizer = torch.optim.AdamW(\n            model.parameters(), lr=1e-4, weight_decay=0.01\n        )\n        self.scheduler = torch.optim.lr_scheduler.OneCycleLR(\n            self.optimizer, max_lr=1e-3,\n            steps_per_epoch=len(train_loader),\n            epochs=num_epochs\n        )\n        self.criterion = FocalLoss()\n        self.scaler = GradScaler()\n\n        self.history = {'train_loss': [], 'val_loss': []}\n\n    def train_epoch(self):\n        self.model.train()\n        total_loss = 0\n\n        pbar = tqdm(self.train_loader, desc='Training')\n        for batch in pbar:\n            sequences = batch['sequence'].to(self.device)\n            lengths = batch['length'].to(self.device)\n            labels = batch['labels'].to(self.device)\n\n            self.optimizer.zero_grad()\n\n            with autocast():\n                logits = self.model(sequences, lengths)\n                loss = self.criterion(logits, labels)\n\n            self.scaler.scale(loss).backward()\n            self.scaler.step(self.optimizer)\n            self.scaler.update()\n            self.scheduler.step()\n\n            total_loss += loss.item()\n            pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n\n        return total_loss / len(self.train_loader)\n\n    def validate(self):\n        self.model.eval()\n        total_loss = 0\n\n        with torch.no_grad():\n            for batch in tqdm(self.val_loader, desc='Validation'):\n                sequences = batch['sequence'].to(self.device)\n                lengths = batch['length'].to(self.device)\n                labels = batch['labels'].to(self.device)\n\n                with autocast():\n                    logits = self.model(sequences, lengths)\n                    loss = self.criterion(logits, labels)\n\n                total_loss += loss.item()\n\n        return total_loss / len(self.val_loader)\n\n    def train(self):\n        best_val_loss = float('inf')\n\n        for epoch in range(self.num_epochs):\n            print(f\"\\nEpoch {epoch+1}/{self.num_epochs}\")\n\n            train_loss = self.train_epoch()\n            val_loss = self.validate()\n\n            self.history['train_loss'].append(train_loss)\n            self.history['val_loss'].append(val_loss)\n\n            print(f\"Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}\")\n\n            if val_loss < best_val_loss:\n                best_val_loss = val_loss\n                torch.save(self.model.state_dict(), 'best_model.pth')\n                print(\"✓ Model saved!\")\n\n        return self.history\n\n# ============================================================================\n# 8. EVALUATION METRICS\n# ============================================================================\n\n# ============================================================================\n# 8. EVALUATION METRICS (with Progress Bar)\n# ============================================================================\n\nclass Evaluator:\n    \"\"\"Comprehensive evaluation with detailed metrics\"\"\"\n\n    def __init__(self, ia_weights, go_hierarchy):\n        self.ia_weights = ia_weights\n        self.go_hierarchy = go_hierarchy\n\n    def calculate_weighted_f1(self, y_true, y_pred, term_names):\n        \"\"\"Calculate weighted F1 score using IA weights\"\"\"\n        precisions = []\n        recalls = []\n\n        for i in range(len(y_true)):\n            true_terms = set([term_names[j] for j, val in enumerate(y_true[i]) if val == 1])\n            pred_terms = set([term_names[j] for j, val in enumerate(y_pred[i]) if val == 1])\n\n            if not pred_terms:\n                if true_terms:\n                    precisions.append(0)\n                    recalls.append(0)\n                else:\n                    precisions.append(1)\n                    recalls.append(1)\n                continue\n\n            # Calculate weighted precision\n            intersection = true_terms & pred_terms\n            weighted_prec_num = sum([self.ia_weights.get(t, 0) for t in intersection])\n            weighted_prec_den = sum([self.ia_weights.get(t, 0) for t in pred_terms])\n            precision = weighted_prec_num / weighted_prec_den if weighted_prec_den > 0 else 0\n\n            # Calculate weighted recall\n            weighted_rec_den = sum([self.ia_weights.get(t, 0) for t in true_terms])\n            recall = weighted_prec_num / weighted_rec_den if weighted_rec_den > 0 else 0\n\n            precisions.append(precision)\n            recalls.append(recall)\n\n        avg_precision = np.mean(precisions) if precisions else 0\n        avg_recall = np.mean(recalls) if recalls else 0\n\n        if avg_precision + avg_recall == 0:\n            return 0, avg_precision, avg_recall\n\n        f1 = 2 * avg_precision * avg_recall / (avg_precision + avg_recall)\n        return f1, avg_precision, avg_recall\n\n    # --- START MODIFICATION ---\n    def find_optimal_threshold(self, y_true, y_scores, term_names, thresholds=None, aspect=''):\n        \"\"\"Find optimal threshold for each aspect with a progress bar\"\"\"\n        if thresholds is None:\n            thresholds = np.arange(0.1, 0.9, 0.05)\n\n        best_f1 = 0\n        best_threshold = 0.5\n\n        # Added tqdm for the progress bar\n        pbar_desc = f\"Optimizing threshold for Aspect {aspect}\"\n        for threshold in tqdm(thresholds, desc=pbar_desc):\n            y_pred = (y_scores > threshold).astype(int)\n            f1, _, _ = self.calculate_weighted_f1(y_true, y_pred, term_names)\n\n            if f1 > best_f1:\n                best_f1 = f1\n                best_threshold = threshold\n\n        return best_threshold, best_f1\n\n    def evaluate_by_aspect(self, y_true, y_scores, term_names, aspects):\n        \"\"\"Evaluate separately for each GO aspect\"\"\"\n        results = {}\n\n        for aspect in ['F', 'P', 'C']:\n            aspect_mask = np.array([aspects.get(t) == aspect for t in term_names])\n\n            if aspect_mask.sum() == 0:\n                continue\n\n            y_true_aspect = y_true[:, aspect_mask]\n            y_scores_aspect = y_scores[:, aspect_mask]\n            aspect_terms = [t for t, m in zip(term_names, aspect_mask) if m]\n\n            # Find optimal threshold (passing the aspect for the progress bar description)\n            threshold, best_f1 = self.find_optimal_threshold(\n                y_true_aspect, y_scores_aspect, aspect_terms, aspect=aspect\n            )\n\n            # Calculate metrics at optimal threshold\n            y_pred_aspect = (y_scores_aspect > threshold).astype(int)\n            f1, precision, recall = self.calculate_weighted_f1(\n                y_true_aspect, y_pred_aspect, aspect_terms\n            )\n\n            results[aspect] = {\n                'threshold': threshold,\n                'f1': f1,\n                'precision': precision,\n                'recall': recall\n            }\n        # --- END MODIFICATION ---\n\n        # Calculate average F1\n        avg_f1 = np.mean([results[asp]['f1'] for asp in results if asp in results])\n\n        return results, avg_f1\n# ============================================================================\n# 9. MAIN EXECUTION\n# ============================================================================\n\ndef main(pretrained_model_path=None):\n    \"\"\"Main execution pipeline\"\"\"\n\n    print(\"=\"*80)\n    print(\"PROTEIN FUNCTION PREDICTION - CAFA CHALLENGE\")\n    print(\"=\"*80)\n\n    # Initialize data loader\n    print(\"\\n[1/9] Loading data...\")\n    data_loader = DataLoader_CAFA()\n\n    # Load all data\n    train_seqs = data_loader.load_fasta('/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta')\n    train_terms = data_loader.load_terms('/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv')\n    ia_weights = data_loader.load_ia_weights('/kaggle/input/cafa-6-protein-function-prediction/IA.tsv')\n    go_graph = data_loader.load_go_graph('/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo')\n    test_seqs = data_loader.load_fasta('/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta')\n\n    print(f\"  • Loaded {len(train_seqs):,} training sequences\")\n    print(f\"  • Loaded {len(test_seqs):,} test sequences\")\n    print(f\"  • Loaded {len(train_terms):,} term annotations\")\n\n    # Initialize GO hierarchy\n    print(\"\\n[2/9] Building GO hierarchy...\")\n    go_hierarchy = GOHierarchy(go_graph)\n\n    # Prepare labels\n    print(\"\\n[3/9] Preparing multi-label matrix...\")\n\n    # Get unique terms and create term-to-index mapping\n    unique_terms = sorted(train_terms['term'].unique())\n    term_to_idx = {term: idx for idx, term in enumerate(unique_terms)}\n    aspect_map = dict(zip(train_terms['term'], train_terms['aspect']))\n\n    # Filter to most common terms (for computational efficiency)\n    term_counts = train_terms['term'].value_counts()\n    selected_terms = term_counts[term_counts >= 20].index.tolist()[:2000]\n    term_to_idx = {term: idx for idx, term in enumerate(selected_terms)}\n\n    print(f\"  • Selected {len(selected_terms)} GO terms for prediction\")\n\n    # --- START MODIFICATION ---\n    # Optimization: Pre-process train_terms into a dictionary for fast lookup\n    print(\"  • Creating protein-to-term map for fast lookup...\")\n    protein_to_terms_map = defaultdict(list)\n    for entry_id, term in zip(train_terms['EntryID'], train_terms['term']):\n        protein_to_terms_map[entry_id].append(term)\n\n    # Create multi-label matrix with progress bar\n    print(\"  • Building multi-label matrix (this might take a few minutes)...\")\n    label_matrix = {}\n\n    # Added tqdm for the progress bar\n    for protein_id in tqdm(train_seqs.keys(), desc='Building label matrix'):\n        labels = np.zeros(len(selected_terms), dtype=np.float32)\n\n        # Optimized lookup using the pre-built map\n        protein_terms = protein_to_terms_map.get(protein_id, [])\n\n        for term in protein_terms:\n            if term in term_to_idx:\n                labels[term_to_idx[term]] = 1.0\n\n        label_matrix[protein_id] = labels\n    # --- END MODIFICATION ---\n    # Split data\n    print(\"\\n[4/9] Creating train/validation split...\")\n    protein_ids = list(train_seqs.keys())\n    train_ids, val_ids = train_test_split(protein_ids, test_size=0.1, random_state=42)\n\n    train_seqs_split = {pid: train_seqs[pid] for pid in train_ids}\n    val_seqs_split = {pid: train_seqs[pid] for pid in val_ids}\n    train_labels = {pid: label_matrix[pid] for pid in train_ids}\n    val_labels = {pid: label_matrix[pid] for pid in val_ids}\n\n    print(f\"  • Train: {len(train_ids):,} proteins\")\n    print(f\"  • Validation: {len(val_ids):,} proteins\")\n\n    # Create datasets\n    print(\"\\n[5/9] Creating PyTorch datasets...\")\n    train_dataset = ProteinDataset(train_seqs_split, train_labels, data_loader, max_len=1024)\n    val_dataset = ProteinDataset(val_seqs_split, val_labels, data_loader, max_len=1024)\n\n    # Fix for multiprocessing issue\n    num_workers = 0 if os.name == 'nt' else 2\n\n    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=num_workers)\n    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=num_workers)\n\n    # Initialize model\n    print(\"\\n[6/9] Initializing model...\")\n    model = ProteinFunctionPredictor(num_classes=len(selected_terms))\n    total_params = sum(p.numel() for p in model.parameters())\n    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n    print(f\"  • Total parameters: {total_params:,}\")\n    print(f\"  • Trainable parameters: {trainable_params:,}\")\n    \n    history = None\n    if pretrained_model_path:\n        print(f\"\\n[7/9] Loading pretrained model from {pretrained_model_path}...\")\n        model.load_state_dict(torch.load(pretrained_model_path))\n    else:\n        # Train model\n        print(\"\\n[7/9] Training model...\")\n        trainer = Trainer(model, train_loader, val_loader, ia_weights, device, num_epochs=5)\n        history = trainer.train()\n\n        # Plot training history\n        print(\"\\n[8/9] Plotting training history...\")\n        plt.figure(figsize=(10, 5))\n        plt.plot(history['train_loss'], label='Train Loss')\n        plt.plot(history['val_loss'], label='Val Loss')\n        plt.xlabel('Epoch')\n        plt.ylabel('Loss')\n        plt.title('Training History')\n        plt.legend()\n        plt.savefig('training_history.png', dpi=300, bbox_inches='tight')\n        print(\"  • Saved training_history.png\")\n\n    # Evaluate on validation set\n    print(\"\\n[9/9] Evaluating on validation set...\")\n    model.to(device)\n    model.eval()\n\n    # Get predictions\n    all_preds = []\n    all_labels = []\n\n    with torch.no_grad():\n        for batch in tqdm(val_loader, desc='Generating predictions'):\n            sequences = batch['sequence'].to(device)\n            lengths = batch['length'].to(device)\n            labels = batch['labels']\n\n            logits = model(sequences, lengths)\n            probs = torch.sigmoid(logits).cpu().numpy()\n\n            all_preds.append(probs)\n            all_labels.append(labels.numpy())\n\n    all_preds = np.vstack(all_preds)\n    all_labels = np.vstack(all_labels)\n\n    # Detailed evaluation\n    evaluator = Evaluator(ia_weights, go_hierarchy)\n    aspect_results, avg_f1 = evaluator.evaluate_by_aspect(\n        all_labels, all_preds, selected_terms, aspect_map\n    )\n\n    print(\"\\n\" + \"=\"*80)\n    print(\"EVALUATION RESULTS\")\n    print(\"=\"*80)\n\n    for aspect_name, aspect_code in [('Molecular Function', 'F'),\n                                      ('Biological Process', 'P'),\n                                      ('Cellular Component', 'C')]:\n        if aspect_code in aspect_results:\n            res = aspect_results[aspect_code]\n            print(f\"\\n{aspect_name} ({aspect_code}):\")\n            print(f\"  • Optimal Threshold: {res['threshold']:.3f}\")\n            print(f\"  • Weighted F1: {res['f1']:.4f}\")\n            print(f\"  • Weighted Precision: {res['precision']:.4f}\")\n            print(f\"  • Weighted Recall: {res['recall']:.4f}\")\n\n    print(f\"\\n{'='*80}\")\n    print(f\"AVERAGE F1 SCORE: {avg_f1:.4f}\")\n    print(f\"{'='*80}\\n\")\n\n    # Generate submission file\n    # 10. EFFICIENT SUBMISSION FILE GENERATION\n    # ============================================================================\n    print(\"\\n[10/9] Generating submission file...\")\n    test_dataset = ProteinDataset(test_seqs, None, data_loader, max_len=1024)\n    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=num_workers)\n\n    # File path for the submission\n    submission_filepath = 'submission.tsv'\n    total_predictions = 0\n\n    # Open the file once and write to it iteratively\n    with open(submission_filepath, 'w') as f_out:\n        with torch.no_grad():\n            for batch in tqdm(test_loader, desc='Predicting test set'):\n                sequences = batch['sequence'].to(device)\n                lengths = batch['length'].to(device)\n                protein_ids = batch['protein_id']\n\n                logits = model(sequences, lengths)\n                probs = torch.sigmoid(logits).cpu().numpy()\n\n                # Process predictions for each protein in the batch\n                for i, protein_id in enumerate(protein_ids):\n                    # Step 1: Get initial predictions above a low threshold\n                    term_scores = {}\n                    for j, term in enumerate(selected_terms):\n                        score = probs[i, j]\n                        if score > 0.01: # Use a reasonable base threshold\n                            term_scores[term] = score\n\n                    if not term_scores:\n                        continue # No predictions for this protein\n\n                    # Step 2: Apply GO hierarchy propagation\n                    propagated_scores = go_hierarchy.propagate_predictions(term_scores)\n\n                    # Step 3: Create a DataFrame for this single protein to sort and limit\n                    protein_preds_df = pd.DataFrame(propagated_scores.items(), columns=['GO_term', 'score'])\n\n                    # Step 4: Limit to the top 1500 predictions\n                    if len(protein_preds_df) > 1500:\n                        protein_preds_df = protein_preds_df.nlargest(1500, 'score')\n                    \n                    # Step 5: Write this protein's predictions to the file\n                    for _, row in protein_preds_df.iterrows():\n                        # Format the output line\n                        line = f\"{protein_id}\\t{row['GO_term']}\\t{row['score']:.3g}\\n\"\n\n                        f_out.write(line)\n                        total_predictions += 1\n\n    print(f\"\\n✓ Submission generation complete.\")\n    print(f\"  • Saved submission.tsv with {total_predictions:,} predictions.\")    \n\n# ============================================================================\n    # 11. LOAD SUBMISSION FOR ANALYSIS & VISUALIZATION\n    # ============================================================================\n    print(\"\\n[11/9] Loading submission file for analysis...\")\n    \n    # Check if the submission file exists and is not empty before trying to read it\n    if os.path.exists(submission_filepath) and os.path.getsize(submission_filepath) > 0:\n        final_submission = pd.read_csv(submission_filepath, sep='\\t', names=['protein_id', 'GO_term', 'score'])\n        \n        # Additional visualizations\n        print(\"\\n[12/9] Creating visualizations...\")\n\n        # 1. Predictions per aspect\n        aspect_counts = []\n        for term in final_submission['GO_term'].unique():\n            aspect = aspect_map.get(term, 'Unknown')\n            aspect_counts.append(aspect)\n\n        aspect_series = pd.Series(aspect_counts)\n\n        fig, axes = plt.subplots(2, 2, figsize=(15, 12))\n\n        # Plot 1: Aspect distribution\n        aspect_series.value_counts().plot(kind='bar', ax=axes[0, 0])\n        axes[0, 0].set_title('Distribution of Predictions by GO Aspect')\n        axes[0, 0].set_xlabel('Aspect')\n        axes[0, 0].set_ylabel('Number of Predictions')\n\n        # Plot 2: Score distribution\n        axes[0, 1].hist(final_submission['score'], bins=50, edgecolor='black')\n        axes[0, 1].set_title('Distribution of Prediction Scores')\n        axes[0, 1].set_xlabel('Score')\n        axes[0, 1].set_ylabel('Frequency')\n\n        # Plot 3: Predictions per protein\n        preds_per_protein = final_submission.groupby('protein_id').size()\n        axes[1, 0].hist(preds_per_protein, bins=50, edgecolor='black')\n        axes[1, 0].set_title('Number of Predictions per Protein')\n        axes[1, 0].set_xlabel('Number of GO Terms')\n        axes[1, 0].set_ylabel('Number of Proteins')\n        axes[1, 0].axvline(preds_per_protein.mean(), color='r',\n                           linestyle='--', label=f'Mean: {preds_per_protein.mean():.1f}')\n        axes[1, 0].legend()\n\n        # Plot 4: Top predicted terms\n        top_terms = final_submission['GO_term'].value_counts().head(20)\n        top_terms.plot(kind='barh', ax=axes[1, 1])\n        axes[1, 1].set_title('Top 20 Most Frequently Predicted GO Terms')\n        axes[1, 1].set_xlabel('Number of Proteins')\n\n        plt.tight_layout()\n        plt.savefig('prediction_analysis.png', dpi=300, bbox_inches='tight')\n        print(\"  • Saved prediction_analysis.png\")\n\n        # Summary statistics\n        print(\"\\n\" + \"=\"*80)\n        print(\"SUBMISSION SUMMARY\")\n        print(\"=\"*80)\n        print(f\"Total predictions: {len(final_submission):,}\")\n        print(f\"Unique proteins: {final_submission['protein_id'].nunique():,}\")\n        print(f\"Unique GO terms: {final_submission['GO_term'].nunique():,}\")\n        print(f\"Avg predictions per protein: {len(final_submission) / final_submission['protein_id'].nunique():.1f}\")\n        print(f\"Score range: [{final_submission['score'].min():.3f}, {final_submission['score'].max():.3f}]\")\n        print(f\"Mean score: {final_submission['score'].mean():.3f}\")\n        print(\"=\"*80)\n\n    else:\n        print(\"\\n⚠️ WARNING: submission.tsv is empty or was not created. Skipping analysis and summary.\")\n        # Create an empty DataFrame so the script can complete without error\n        final_submission = pd.DataFrame(columns=['protein_id', 'GO_term', 'score'])\n\n\n    # Create detailed evaluation report\n    print(\"\\n[13/9] Creating evaluation report...\")\n\n    report = f\"\"\"\nPROTEIN FUNCTION PREDICTION - EVALUATION REPORT\n{'='*80}\n\n1. MODEL ARCHITECTURE\n   - Encoder: CNN + Transformer (ProteinEncoder)\n   - Embedding dimension: 128\n   - CNN layers: 3 (with kernels 3, 5, 7)\n   - Transformer layers: 4\n   - Attention heads: 8\n   - Total parameters: {total_params:,}\n   - Trainable parameters: {trainable_params:,}\n\n2. TRAINING CONFIGURATION\n   - Training samples: {len(train_ids):,}\n   - Validation samples: {len(val_ids):,}\n   - Batch size: 16 (train), 32 (val/test)\n   - Number of epochs: 5\n   - Optimizer: AdamW\n   - Learning rate: 1e-4 to 1e-3 (OneCycle)\n   - Loss function: Focal Loss (α=0.25, γ=2.0)\n   - Max sequence length: 1024\n\n3. GO TERM SELECTION\n   - Total unique GO terms: {len(unique_terms):,}\n   - Selected GO terms: {len(selected_terms)}\n   - Selection criteria: frequency ≥ 20\n\n4. VALIDATION PERFORMANCE\n\"\"\"\n\n    for aspect_name, aspect_code in [('Molecular Function', 'F'),\n                                      ('Biological Process', 'P'),\n                                      ('Cellular Component', 'C')]:\n        if aspect_code in aspect_results:\n            res = aspect_results[aspect_code]\n            report += f\"\"\"\n   {aspect_name} ({aspect_code}):\n   - Optimal threshold: {res['threshold']:.3f}\n   - Weighted F1 score: {res['f1']:.4f}\n   - Weighted precision: {res['precision']:.4f}\n   - Weighted recall: {res['recall']:.4f}\n\"\"\"\n\n    report += f\"\"\"\n   Overall Average F1: {avg_f1:.4f}\n\n5. TEST SET PREDICTIONS\n   - Total predictions: {len(final_submission):,}\n   - Unique proteins: {final_submission['protein_id'].nunique():,}\n   - Unique GO terms predicted: {final_submission['GO_term'].nunique():,}\n   - Avg predictions per protein: {len(final_submission) / final_submission['protein_id'].nunique() if final_submission['protein_id'].nunique() > 0 else 0:.1f}\n   - Score range: [{final_submission['score'].min() if not final_submission.empty else 0:.3f}, {final_submission['score'].max() if not final_submission.empty else 0:.3f}]\n\n6. KEY FEATURES\n   ✓ Multi-label classification with hierarchical GO structure\n   ✓ Focal loss for handling class imbalance\n   ✓ Aspect-specific threshold optimization\n   ✓ GO hierarchy propagation (children → parents)\n   ✓ Mixed precision training (AMP)\n   ✓ Attention-based sequence pooling\n   ✓ Information accretion weighted evaluation\n\n7. FILES GENERATED\n   - best_model.pth: Trained model weights\n   - submission.tsv: Competition submission file\n   - training_history.png: Loss curves\n   - prediction_analysis.png: Prediction statistics\n   - evaluation_report.txt: This report\n\n{'='*80}\n\"\"\"\n\n    with open('evaluation_report.txt', 'w') as f:\n        f.write(report)\n\n    print(report)\n    print(\"\\n✓ All tasks completed successfully!\")\n    print(\"  Files ready for submission: submission.tsv\")\n\n    return model, final_submission, history, aspect_results\n# ============================================================================\n# RUN MAIN\n# ============================================================================\n\nif __name__ == \"__main__\":\n    # To train a new model, run without arguments:\n    # model, submission, history, results = main()\n\n    # To evaluate a pretrained model, provide the path:\n    model, submission, history, results = main(pretrained_model_path='/kaggle/input/intial-cafa/pytorch/default/1/best_model.pth')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-10-19T10:18:23.724921Z","iopub.execute_input":"2025-10-19T10:18:23.725794Z","iopub.status.idle":"2025-10-19T10:19:06.030496Z","shell.execute_reply.started":"2025-10-19T10:18:23.725733Z","shell.execute_reply":"2025-10-19T10:19:06.029428Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}