{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport gc\nimport numpy as np\nimport pandas as pd\nfrom collections import defaultdict\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import T5EncoderModel, T5Tokenizer\nfrom sklearn.model_selection import train_test_split\nfrom tqdm.auto import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# ========================= CONFIG =========================\nclass Config:\n    # Paths\n    DATA_PATH = '/kaggle/input/cafa-6-protein-function-prediction'\n    \n    # Model settings - FULL POWER MODE\n    MAX_LENGTH = 512          # Full length\n    BATCH_SIZE = 4            # Small batch for memory\n    GRADIENT_ACCUMULATION = 4  # Effective batch = 16\n    HIDDEN_DIM = 512\n    NUM_EPOCHS = 1            # Single epoch\n    LEARNING_RATE = 1e-4\n    \n    # Data - USE ALL\n    MAX_TRAIN_SAMPLES = None  # ALL 82K\n    MIN_TERM_FREQUENCY = 10   # Filter very rare terms\n    \n    # Device\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nconfig = Config()\n\nprint(f\"🚀 FULL POWER MODE ACTIVATED\")\nprint(f\"Device: {config.DEVICE}\")\nprint(f\"Using ALL training data (~82K proteins)\")\nprint(f\"Using ALL GO terms (filtered by frequency)\")\nprint(f\"Encoder will be fine-tuned\")\nprint(f\"Max sequence length: {config.MAX_LENGTH}\")\n\n# ========================= DATA LOADING =========================\n\ndef load_fasta(filepath):\n    \"\"\"Load FASTA file and return dict of sequences\"\"\"\n    sequences = {}\n    with open(filepath, 'r') as f:\n        current_id = None\n        current_seq = []\n        for line in f:\n            line = line.strip()\n            if line.startswith('>'):\n                if current_id is not None:\n                    sequences[current_id] = ''.join(current_seq)\n                current_id = line.split('|')[1] if '|' in line else line[1:].split()[0]\n                current_seq = []\n            else:\n                current_seq.append(line)\n        if current_id is not None:\n            sequences[current_id] = ''.join(current_seq)\n    return sequences\n\ndef load_go_terms(filepath):\n    \"\"\"Load GO term annotations\"\"\"\n    df = pd.read_csv(filepath, sep='\\t', header=None, names=['protein_id', 'go_term', 'ontology'])\n    return df\n\ndef load_ia_weights(filepath):\n    \"\"\"Load information accretion weights\"\"\"\n    df = pd.read_csv(filepath, sep='\\t', header=None, names=['go_term', 'ia_weight'])\n    return dict(zip(df['go_term'], df['ia_weight']))\n\nprint(\"\\n📁 Loading data...\")\ntrain_sequences = load_fasta(f'{config.DATA_PATH}/Train/train_sequences.fasta')\ntest_sequences = load_fasta(f'{config.DATA_PATH}/Test/testsuperset.fasta')\ntrain_terms_df = load_go_terms(f'{config.DATA_PATH}/Train/train_terms.tsv')\nia_weights = load_ia_weights(f'{config.DATA_PATH}/IA.tsv')\n\nprint(f\"✓ Loaded {len(train_sequences)} training sequences\")\nprint(f\"✓ Loaded {len(test_sequences)} test sequences\")\nprint(f\"✓ Loaded {len(train_terms_df)} annotations\")\n\n# ========================= PREPROCESSING =========================\n\n# Filter to only proteins we have sequences for\ntrain_terms_df = train_terms_df[train_terms_df['protein_id'].isin(train_sequences.keys())]\n\n# Get ALL GO terms with minimum frequency\ngo_term_counts = train_terms_df['go_term'].value_counts()\ntop_go_terms = go_term_counts[go_term_counts >= config.MIN_TERM_FREQUENCY].index.tolist()\ngo_term_to_idx = {term: idx for idx, term in enumerate(top_go_terms)}\nidx_to_go_term = {idx: term for term, idx in go_term_to_idx.items()}\n\nprint(f\"\\n🎯 Predicting {len(top_go_terms)} GO terms (frequency >= {config.MIN_TERM_FREQUENCY})\")\n\n# Create multi-label matrix for ALL proteins\nprotein_ids = list(train_sequences.keys())\nprint(f\"📊 Building label matrix for {len(protein_ids)} proteins...\")\n\nlabels_matrix = np.zeros((len(protein_ids), len(top_go_terms)), dtype=np.float32)\nprotein_id_to_idx = {pid: idx for idx, pid in enumerate(protein_ids)}\n\nfor _, row in tqdm(train_terms_df.iterrows(), total=len(train_terms_df), desc='Processing labels'):\n    pid = row['protein_id']\n    go_term = row['go_term']\n    if pid in protein_id_to_idx and go_term in go_term_to_idx:\n        protein_idx = protein_id_to_idx[pid]\n        term_idx = go_term_to_idx[go_term]\n        labels_matrix[protein_idx, term_idx] = 1.0\n\nprint(f\"✓ Label matrix: {labels_matrix.shape}\")\nprint(f\"✓ Label density: {labels_matrix.mean():.6f}\")\nprint(f\"✓ Avg labels per protein: {labels_matrix.sum(axis=1).mean():.1f}\")\n\n# ========================= PROTEIN DATASET =========================\n\nclass ProteinDataset(Dataset):\n    def __init__(self, protein_ids, sequences, labels=None, tokenizer=None, max_length=512):\n        self.protein_ids = protein_ids\n        self.sequences = sequences\n        self.labels = labels\n        self.tokenizer = tokenizer\n        self.max_length = max_length\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        sequence = self.sequences[protein_id]\n        \n        # Add spaces between amino acids for ProtT5\n        sequence = ' '.join(list(sequence[:self.max_length]))\n        \n        # Tokenize\n        encoded = self.tokenizer(\n            sequence,\n            max_length=self.max_length,\n            padding='max_length',\n            truncation=True,\n            return_tensors='pt'\n        )\n        \n        item = {\n            'protein_id': protein_id,\n            'input_ids': encoded['input_ids'].squeeze(0),\n            'attention_mask': encoded['attention_mask'].squeeze(0),\n        }\n        \n        if self.labels is not None:\n            item['labels'] = torch.tensor(self.labels[idx], dtype=torch.float32)\n        \n        return item\n\n# ========================= FULL MODEL =========================\n\nclass ProteinFunctionPredictor(nn.Module):\n    def __init__(self, num_labels, hidden_dim=512, dropout=0.3):\n        super().__init__()\n        self.num_labels = num_labels\n        \n        print(\"⚡ Loading ProtT5 model (will be fine-tuned)...\")\n        self.encoder = T5EncoderModel.from_pretrained(\"Rostlab/prot_t5_xl_half_uniref50-enc\")\n        \n        # Fine-tune last 6 layers only (for speed)\n        for param in list(self.encoder.parameters())[:-18]:\n            param.requires_grad = False\n        \n        # Classification head\n        self.classifier = nn.Sequential(\n            nn.Linear(1024, hidden_dim),\n            nn.LayerNorm(hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, hidden_dim // 2),\n            nn.LayerNorm(hidden_dim // 2),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim // 2, num_labels)\n        )\n        \n    def forward(self, input_ids, attention_mask):\n        # Get ProtT5 embeddings (will have gradients for last layers)\n        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)\n        embeddings = outputs.last_hidden_state\n        \n        # Mean pooling\n        mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float()\n        sum_embeddings = torch.sum(embeddings * mask_expanded, 1)\n        sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)\n        pooled = sum_embeddings / sum_mask\n        \n        # Classification\n        logits = self.classifier(pooled)\n        return logits\n\n# ========================= TRAINING =========================\n\ndef train_model(model, train_loader, val_loader, num_epochs, device):\n    optimizer = torch.optim.AdamW(model.parameters(), lr=config.LEARNING_RATE)\n    criterion = nn.BCEWithLogitsLoss()\n    \n    best_val_loss = float('inf')\n    \n    for epoch in range(num_epochs):\n        model.train()\n        train_loss = 0\n        optimizer.zero_grad()\n        \n        pbar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs}')\n        for i, batch in enumerate(pbar):\n            input_ids = batch['input_ids'].to(device)\n            attention_mask = batch['attention_mask'].to(device)\n            labels = batch['labels'].to(device)\n            \n            # Forward pass\n            logits = model(input_ids, attention_mask)\n            loss = criterion(logits, labels)\n            \n            # Scale loss for gradient accumulation\n            loss = loss / config.GRADIENT_ACCUMULATION\n            loss.backward()\n            \n            # Update weights every N batches\n            if (i + 1) % config.GRADIENT_ACCUMULATION == 0:\n                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n                optimizer.step()\n                optimizer.zero_grad()\n            \n            train_loss += loss.item() * config.GRADIENT_ACCUMULATION\n            \n            if i % 10 == 0:\n                pbar.set_postfix({\n                    'loss': f'{loss.item() * config.GRADIENT_ACCUMULATION:.4f}',\n                    'avg': f'{train_loss / (i + 1):.4f}'\n                })\n            \n            # Clear cache periodically\n            if i % 100 == 0:\n                torch.cuda.empty_cache()\n        \n        avg_train_loss = train_loss / len(train_loader)\n        \n        # Validation\n        print(\"Running validation...\")\n        model.eval()\n        val_loss = 0\n        with torch.no_grad():\n            for batch in tqdm(val_loader, desc='Validation'):\n                input_ids = batch['input_ids'].to(device)\n                attention_mask = batch['attention_mask'].to(device)\n                labels = batch['labels'].to(device)\n                \n                logits = model(input_ids, attention_mask)\n                loss = criterion(logits, labels)\n                val_loss += loss.item()\n        \n        avg_val_loss = val_loss / len(val_loader)\n        print(f'\\n✓ Epoch {epoch+1}: Train Loss={avg_train_loss:.4f}, Val Loss={avg_val_loss:.4f}')\n        \n        if avg_val_loss < best_val_loss:\n            best_val_loss = avg_val_loss\n            torch.save(model.state_dict(), 'best_model.pth')\n            print(f'  💾 Best model saved!')\n        \n        torch.cuda.empty_cache()\n        gc.collect()\n    \n    return model\n\n# ========================= MAIN EXECUTION =========================\n\nprint(\"\\n🔧 Initializing tokenizer...\")\ntokenizer = T5Tokenizer.from_pretrained(\"Rostlab/prot_t5_xl_half_uniref50-enc\", do_lower_case=False)\n\n# Split train/val (90/10)\ntrain_ids, val_ids, train_labels, val_labels = train_test_split(\n    protein_ids, labels_matrix, test_size=0.1, random_state=42\n)\n\nprint(f\"\\n📊 Dataset split:\")\nprint(f\"  Train: {len(train_ids):,} proteins\")\nprint(f\"  Val: {len(val_ids):,} proteins\")\n\n# Create datasets\nprint(\"\\n🔨 Creating datasets...\")\ntrain_dataset = ProteinDataset(train_ids, train_sequences, train_labels, tokenizer, config.MAX_LENGTH)\nval_dataset = ProteinDataset(val_ids, train_sequences, val_labels, tokenizer, config.MAX_LENGTH)\n\ntrain_loader = DataLoader(\n    train_dataset, \n    batch_size=config.BATCH_SIZE, \n    shuffle=True, \n    num_workers=2, \n    pin_memory=True\n)\nval_loader = DataLoader(\n    val_dataset, \n    batch_size=config.BATCH_SIZE * 2, \n    shuffle=False, \n    num_workers=2, \n    pin_memory=True\n)\n\nprint(f\"  Train batches: {len(train_loader):,}\")\nprint(f\"  Val batches: {len(val_loader):,}\")\n\n# Initialize model\nprint(\"\\n🤖 Initializing model...\")\nmodel = ProteinFunctionPredictor(num_labels=len(top_go_terms), hidden_dim=config.HIDDEN_DIM)\nmodel = model.to(config.DEVICE)\n\ntotal_params = sum(p.numel() for p in model.parameters())\ntrainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\nprint(f\"  Total parameters: {total_params:,}\")\nprint(f\"  Trainable parameters: {trainable_params:,} ({100*trainable_params/total_params:.1f}%)\")\n\n# Train\nprint(\"\\n🏋️ Starting training...\")\nimport time\nstart_time = time.time()\n\nmodel = train_model(model, train_loader, val_loader, config.NUM_EPOCHS, config.DEVICE)\n\nelapsed = time.time() - start_time\nprint(f\"\\n⏱️ Training completed in {elapsed/60:.1f} minutes\")\n\n# Load best model\nprint(\"\\n📥 Loading best model...\")\nmodel.load_state_dict(torch.load('best_model.pth'))\n\n# ========================= PREDICTION =========================\n\nprint(\"\\n🔮 Generating predictions on test set...\")\n\ntest_protein_ids = list(test_sequences.keys())\nprint(f\"  Test proteins: {len(test_protein_ids):,}\")\n\ntest_dataset = ProteinDataset(test_protein_ids, test_sequences, labels=None, \n                              tokenizer=tokenizer, max_length=config.MAX_LENGTH)\ntest_loader = DataLoader(\n    test_dataset, \n    batch_size=config.BATCH_SIZE * 2, \n    shuffle=False, \n    num_workers=2, \n    pin_memory=True\n)\n\nmodel.eval()\nall_predictions = []\nall_protein_ids = []\n\nwith torch.no_grad():\n    for batch in tqdm(test_loader, desc='Predicting'):\n        input_ids = batch['input_ids'].to(config.DEVICE)\n        attention_mask = batch['attention_mask'].to(config.DEVICE)\n        \n        logits = model(input_ids, attention_mask)\n        probs = torch.sigmoid(logits)\n        \n        all_predictions.append(probs.cpu().numpy())\n        all_protein_ids.extend(batch['protein_id'])\n        \n        if len(all_predictions) % 50 == 0:\n            torch.cuda.empty_cache()\n\nall_predictions = np.vstack(all_predictions)\nprint(f\"✓ Predictions shape: {all_predictions.shape}\")\n\n# ========================= CREATE SUBMISSION =========================\n\nprint(\"\\n📝 Creating submission file...\")\n\nsubmission_rows = []\nthreshold = 0.05  # Confidence threshold\n\nfor i, protein_id in enumerate(tqdm(all_protein_ids, desc='Building submission')):\n    pred_probs = all_predictions[i]\n    \n    # Get predictions above threshold\n    top_indices = np.where(pred_probs > threshold)[0]\n    top_probs = pred_probs[top_indices]\n    \n    # Sort by probability\n    sorted_idx = np.argsort(top_probs)[::-1]\n    top_indices = top_indices[sorted_idx]\n    top_probs = top_probs[sorted_idx]\n    \n    # Limit to top 150 per protein\n    for idx, prob in zip(top_indices[:150], top_probs[:150]):\n        go_term = idx_to_go_term[idx]\n        submission_rows.append(f\"{protein_id}\\t{go_term}\\t{prob:.3f}\")\n\n# Write submission\nsubmission_path = 'submission.tsv'\nwith open(submission_path, 'w') as f:\n    f.write('\\n'.join(submission_rows))\n\nprint(f\"\\n✅ SUBMISSION COMPLETE!\")\nprint(f\"  Total predictions: {len(submission_rows):,}\")\nprint(f\"  Average per protein: {len(submission_rows)/len(all_protein_ids):.1f}\")\nprint(f\"  File: {submission_path}\")\n\nprint(f\"\\n📊 Summary:\")\nprint(f\"  Training time: {elapsed/60:.1f} min\")\nprint(f\"  Proteins trained: {len(train_ids):,}\")\nprint(f\"  GO terms predicted: {len(top_go_terms):,}\")\nprint(f\"  Test proteins: {len(all_protein_ids):,}\")\n\nprint(\"\\n🎯 Sample predictions:\")\nfor row in submission_rows[:10]:\n    print(f\"  {row}\")\n\nprint(\"\\n🚀 Ready to submit!\")\n\n# Cleanup\ndel model\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-23T09:21:24.435228Z","iopub.execute_input":"2026-01-23T09:21:24.435544Z"}},"outputs":[],"execution_count":null}]}