{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"},{"sourceId":5499219,"sourceType":"datasetVersion","datasetId":3167603},{"sourceId":5549164,"sourceType":"datasetVersion","datasetId":3197305},{"sourceId":5607816,"sourceType":"datasetVersion","datasetId":3225525},{"sourceId":5792099,"sourceType":"datasetVersion","datasetId":3327296},{"sourceId":6247561,"sourceType":"datasetVersion","datasetId":3590060}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ============================================================================\n# PACKAGE INSTALLATION\n# ============================================================================\n# Install any missing packages\n!pip install -q obonet biopython","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T08:33:44.099551Z","iopub.execute_input":"2025-12-14T08:33:44.099885Z","iopub.status.idle":"2025-12-14T08:33:49.885049Z","shell.execute_reply.started":"2025-12-14T08:33:44.099848Z","shell.execute_reply":"2025-12-14T08:33:49.884238Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\nimport numpy as np\nimport pandas as pd\nimport os\n\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# ============================================================================\n# CORE IMPORTS\n# ============================================================================\nimport gc\nimport sys\nimport subprocess\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\nfrom tqdm.auto import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# ============================================================================\n# CONFIG (ADVANCED IMPROVEMENTS)\n# ============================================================================\nCONFIG = {\n    'MAX_TRAIN_SAMPLES': 90000,        # all data\n    'TOP_K_LABELS': 800,\n    'RANDOM_SEED': 253,\n    'HIDDEN_DIMS': [512, 256, 128],\n    'DROPOUT_RATE': 0.27,             # IMPROVED: 0.3 → 0.27 (less dropout for more data)\n    'EPOCHS': 36,                     # IMPROVED: 12 - 36 epochs\n    'BATCH_SIZE': 64,                 # IMPROVED: 32 → 64 (faster training)\n    'LEARNING_RATE': 1e-3,\n    'PREDICT_BATCH_SIZE': 128,        # IMPROVED: 64 → 128 (faster prediction)\n    'MIN_CONFIDENCE': 0.05,           # IMPROVED: 0.10 → 0.05 (even more predictions)\n    'MAX_PREDS_PER_PROTEIN': 200,    # IMPROVED: 150 → 200\n    'LABEL_SMOOTHING': 0.08,          # NEW: Add label smoothing\n    'TEMPERATURE': 1.4,               # NEW: Temperature scaling for predictions\n    'MERGE_WITH_BLAST': True,\n    'BLAST_WEIGHT': 0.60,             \n    'DL_WEIGHT': 0.40,                \n    'BASE_PATH': '/kaggle/input/cafa-6-protein-function-prediction',\n    'ESM2_PATH': 'cafa-5-ems-2-embeddings-numpy',\n    'BLAST_PATH': '/kaggle/input/blast-quick-sprof-zero-pred/submission.tsv',\n}\n\nprint(\"=\"*80)\nprint(\"🚀 CAFA 6 - ADVANCED OPTIMIZATION\")\nprint(\"=\"*80)\n\ndef install(package):\n    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", package])\n\ntry:\n    import torch\n    import torch.nn as nn\n    import torch.nn.functional as F\nexcept:\n    install('torch')\n    import torch\n    import torch.nn as nn\n    import torch.nn.functional as F\n\nnp.random.seed(CONFIG['RANDOM_SEED'])\ntorch.manual_seed(CONFIG['RANDOM_SEED'])\nif torch.cuda.is_available():\n    torch.cuda.manual_seed(CONFIG['RANDOM_SEED'])\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Device: {device}\")\n\n# ============================================================================\n# LOAD DATA\n# ============================================================================\nBASE = Path(CONFIG['BASE_PATH'])\nTRAIN_DIR = BASE / 'Train'\n\nprint(\"\\n[1/5] Loading annotations...\")\ntrain_terms_df = pd.read_csv(TRAIN_DIR / 'train_terms.tsv', sep='\\t', \n                              header=None, names=['protein', 'term', 'ontology'])\nprotein_to_terms = train_terms_df.groupby('protein')['term'].apply(list).to_dict()\n\nterm_counts = Counter()\nfor terms in protein_to_terms.values():\n    term_counts.update(terms)\ntop_terms = [t for t, _ in term_counts.most_common(CONFIG['TOP_K_LABELS'])]\n\nfor protein in protein_to_terms:\n    protein_to_terms[protein] = [t for t in protein_to_terms[protein] if t in top_terms]\n\nterm_to_idx = {term: idx for idx, term in enumerate(top_terms)}\nprint(f\"   ✓ {len(protein_to_terms)} proteins, {len(top_terms)} terms\")\n\nprint(\"\\n[2/5] Loading ESM2 embeddings...\")\nesm2_base = f\"/kaggle/input/{CONFIG['ESM2_PATH']}\"\ntrain_ids = np.load(f\"{esm2_base}/train_ids.npy\", allow_pickle=True)\ntrain_embeds = np.load(f\"{esm2_base}/train_embeddings.npy\")\ntest_ids = np.load(f\"{esm2_base}/test_ids.npy\", allow_pickle=True)\ntest_embeds = np.load(f\"{esm2_base}/test_embeddings.npy\")\n\ntrain_dict = {str(pid): emb for pid, emb in zip(train_ids, train_embeds)}\ntest_dict = {str(pid): emb for pid, emb in zip(test_ids, test_embeds)}\nembed_dim = train_embeds.shape[1]\nprint(f\"   ✓ Dim: {embed_dim}\")\n\ndel train_ids, train_embeds, test_ids, test_embeds\ngc.collect()\n\nprint(\"\\n[3/5] Preparing training data...\")\nvalid_proteins = [p for p in protein_to_terms.keys() if p in train_dict][:CONFIG['MAX_TRAIN_SAMPLES']]\n\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.model_selection import train_test_split\n\nmlb = MultiLabelBinarizer(classes=range(len(top_terms)))\ny_labels = [[term_to_idx[t] for t in protein_to_terms.get(p, []) if t in term_to_idx] \n            for p in valid_proteins]\ny_encoded = mlb.fit_transform(y_labels)\n\n# Apply label smoothing\nif CONFIG['LABEL_SMOOTHING'] > 0:\n    y_encoded = y_encoded.astype(float)\n    y_encoded = y_encoded * (1 - CONFIG['LABEL_SMOOTHING']) + CONFIG['LABEL_SMOOTHING'] / len(top_terms)\n\ntrain_proteins, val_proteins, y_train, y_val = train_test_split(\n    valid_proteins, y_encoded, test_size=0.15, random_state=CONFIG['RANDOM_SEED']\n)\nprint(f\"   ✓ Train: {len(train_proteins)}, Val: {len(val_proteins)}\")\n\n# ============================================================================\n# MODEL\n# ============================================================================\nprint(\"\\n[4/5] Building model...\")\n\nclass ImprovedProteinModel(nn.Module):\n    \"\"\"Improved 3-layer architecture with BatchNorm\"\"\"\n    def __init__(self, input_dim, output_dim, hidden_dims, dropout):\n        super().__init__()\n        layers = []\n        prev_dim = input_dim\n        for dim in hidden_dims:\n            layers.extend([\n                nn.Linear(prev_dim, dim),\n                nn.BatchNorm1d(dim),\n                nn.ReLU(),\n                nn.Dropout(dropout)\n            ])\n            prev_dim = dim\n        self.encoder = nn.Sequential(*layers)\n        self.output = nn.Linear(prev_dim, output_dim)\n    \n    def forward(self, x):\n        return torch.sigmoid(self.output(self.encoder(x)))\n\nmodel = ImprovedProteinModel(embed_dim, len(top_terms), CONFIG['HIDDEN_DIMS'], CONFIG['DROPOUT_RATE']).to(device)\nprint(f\"   ✓ Parameters: {sum(p.numel() for p in model.parameters()):,}\")\n\n# ============================================================================\n# TRAINING\n# ============================================================================\nprint(\"\\n\" + \"=\"*80)\nprint(\"TRAINING (12 EPOCHS WITH LABEL SMOOTHING)\")\nprint(\"=\"*80)\n\ncriterion = nn.BCELoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=CONFIG['LEARNING_RATE'], weight_decay=1e-5)  # Added L2 regularization\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2, verbose=False)\n\nbest_val_loss = float('inf')\n\nfor epoch in range(CONFIG['EPOCHS']):\n    model.train()\n    indices = np.random.permutation(len(train_proteins))\n    epoch_loss = 0\n    n_batches = 0\n    \n    for i in range(0, len(indices), CONFIG['BATCH_SIZE']):\n        batch_idx = indices[i:i + CONFIG['BATCH_SIZE']]\n        batch_proteins = [train_proteins[j] for j in batch_idx]\n        \n        X_batch = torch.FloatTensor([train_dict[p] for p in batch_proteins]).to(device)\n        y_batch = torch.FloatTensor(y_train[batch_idx]).to(device)\n        \n        optimizer.zero_grad()\n        outputs = model(X_batch)\n        loss = criterion(outputs, y_batch)\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # Gradient clipping\n        optimizer.step()\n        \n        epoch_loss += loss.item()\n        n_batches += 1\n    \n    model.eval()\n    val_loss = 0\n    val_batches = 0\n    \n    with torch.no_grad():\n        for i in range(0, len(val_proteins), CONFIG['BATCH_SIZE']):\n            batch_proteins = val_proteins[i:i + CONFIG['BATCH_SIZE']]\n            X_batch = torch.FloatTensor([train_dict[p] for p in batch_proteins]).to(device)\n            y_batch = torch.FloatTensor(y_val[i:i + CONFIG['BATCH_SIZE']]).to(device)\n            \n            outputs = model(X_batch)\n            loss = criterion(outputs, y_batch)\n            val_loss += loss.item()\n            val_batches += 1\n    \n    train_loss_avg = epoch_loss/n_batches\n    val_loss_avg = val_loss/val_batches\n    \n    scheduler.step(val_loss_avg)\n    \n    if val_loss_avg < best_val_loss:\n        best_val_loss = val_loss_avg\n        print(f\"Epoch {epoch+1}: Train={train_loss_avg:.4f}, Val={val_loss_avg:.4f} ⭐ NEW BEST\")\n    else:\n        print(f\"Epoch {epoch+1}: Train={train_loss_avg:.4f}, Val={val_loss_avg:.4f}\")\n    \n    gc.collect()\n\nprint(f\"\\n✅ Best Validation Loss: {best_val_loss:.4f}\")\n\n# ============================================================================\n# PREDICTIONS (WITH TEMPERATURE SCALING)\n# ============================================================================\nprint(\"\\n\" + \"=\"*80)\nprint(\"PREDICTIONS (WITH TEMPERATURE SCALING)\")\nprint(\"=\"*80)\n\nmodel.eval()\ntest_protein_ids = list(test_dict.keys())\n\nn_predictions = 0\nwith open('temp_dl.tsv', 'w', newline='') as f:\n    with torch.no_grad():\n        for start in tqdm(range(0, len(test_protein_ids), CONFIG['PREDICT_BATCH_SIZE']), desc=\"Predicting\"):\n            batch_ids = test_protein_ids[start:start + CONFIG['PREDICT_BATCH_SIZE']]\n            X_batch = torch.FloatTensor([test_dict[p] for p in batch_ids]).to(device)\n            \n            # Get raw logits before sigmoid\n            logits = model.encoder(X_batch)\n            logits = model.output(logits)\n            \n            # Apply temperature scaling\n            outputs = torch.sigmoid(logits / CONFIG['TEMPERATURE']).cpu().numpy()\n            \n            for i, pid in enumerate(batch_ids):\n                probs = outputs[i]\n                top_indices = np.argsort(probs)[::-1][:CONFIG['MAX_PREDS_PER_PROTEIN']]\n                confident_indices = [idx for idx in top_indices if probs[idx] > CONFIG['MIN_CONFIDENCE']]\n                \n                for idx in confident_indices:\n                    line = f\"{pid}\\t{top_terms[idx]}\\t{min(probs[idx], 0.999):.3f}\\n\"\n                    f.write(line)\n                    n_predictions += 1\n            \n            del X_batch, outputs, logits\n            if start % 1000 == 0:\n                gc.collect()\n\nprint(f\"✓ Generated {n_predictions:,} predictions\")\n\ndel model, train_dict\ngc.collect()\n\n# ============================================================================\n# ENSEMBLE\n# ============================================================================\nprint(\"\\n\" + \"=\"*80)\nprint(\"ENSEMBLE\")\nprint(\"=\"*80)\n\nif CONFIG['MERGE_WITH_BLAST'] and os.path.exists(CONFIG['BLAST_PATH']):\n    print(\"Merging with BLAST...\")\n    \n    dl_df = pd.read_csv('temp_dl.tsv', sep='\\t', header=None, names=['Id', 'GO term', 'Confidence'])\n    dl_df['Confidence'] = dl_df['Confidence'].astype(float) * CONFIG['DL_WEIGHT']\n    \n    blast_chunks = []\n    for chunk in pd.read_csv(CONFIG['BLAST_PATH'], sep='\\t', header=None,\n                             names=['Id', 'GO term', 'Confidence'], chunksize=1000000):\n        chunk['Confidence'] = chunk['Confidence'].astype(float) * CONFIG['BLAST_WEIGHT']\n        blast_chunks.append(chunk)\n        if len(blast_chunks) >= 10:\n            blast_df = pd.concat(blast_chunks, ignore_index=True)\n            blast_chunks = [blast_df]\n            gc.collect()\n    \n    blast_df = pd.concat(blast_chunks, ignore_index=True)\n    print(f\"  DL={len(dl_df):,}, BLAST={len(blast_df):,}\")\n    \n    ensemble_df = pd.concat([dl_df, blast_df], ignore_index=True)\n    del dl_df, blast_df, blast_chunks\n    gc.collect()\n    \n    ensemble_df = ensemble_df.groupby(['Id', 'GO term'], as_index=False)['Confidence'].sum()\n    ensemble_df = ensemble_df.sort_values(['Id', 'Confidence'], ascending=[True, False])\n    ensemble_df = ensemble_df.groupby('Id').head(1500)\n    \n    import csv\n    ensemble_df.to_csv('submission.tsv', sep='\\t', header=False, index=False,\n                      quoting=csv.QUOTE_NONE, escapechar='\\\\', lineterminator='\\n')\n    print(f\"✓ Saved {len(ensemble_df):,} predictions\")\n    \n    del ensemble_df\nelse:\n    os.rename('temp_dl.tsv', 'submission.tsv')\n\ngc.collect()\n\n# ============================================================================\n# CLEANUP\n# ============================================================================\nprint(\"\\n[CLEANUP]\")\nfor fname in os.listdir('.'):\n    if 'submission' in fname.lower() and fname != 'submission.tsv':\n        try:\n            os.remove(fname)\n            print(f\"  ✗ Removed {fname}\")\n        except:\n            pass\n\nif os.path.exists('temp_dl.tsv'):\n    os.remove('temp_dl.tsv')\n    print(f\"  ✗ Removed temp_dl.tsv\")\n\n# ============================================================================\n# VERIFICATION\n# ============================================================================\nprint(\"\\n\" + \"=\"*80)\nprint(\"✅ VERIFICATION\")\nprint(\"=\"*80)\n\nif os.path.exists('submission.tsv'):\n    with open('submission.tsv', 'r') as f:\n        lines = f.readlines()\n    \n    n_lines = len(lines)\n    file_size_mb = os.path.getsize('submission.tsv') / (1024*1024)\n    \n    print(f\"📋 File: submission.tsv\")\n    print(f\"📊 Lines: {n_lines:,}\")\n    print(f\"💾 Size: {file_size_mb:.1f} MB\")\n    \n    print(\"\\n✅ First 5 lines:\")\n    for i in range(min(5, len(lines))):\n        line = lines[i].strip()\n        parts = line.split('\\t')\n        if len(parts) == 3:\n            print(f\"  ✓ {parts[0]:<15} {parts[1]:<12} {parts[2]}\")\n        else:\n            print(f\"  ✗ Line {i+1}: {len(parts)} columns (expected 3)\")\n    \n    first_line = lines[0].strip()\n    parts = first_line.split('\\t')\n    \n    print(\"\\n🔍 Format check:\")\n    all_good = True\n    \n    if len(parts) == 3:\n        print(f\"  ✓ Columns: 3\")\n    else:\n        print(f\"  ✗ Columns: {len(parts)} (expected 3)\")\n        all_good = False\n    \n    if parts[1].startswith('GO:'):\n        print(f\"  ✓ GO term format\")\n    else:\n        print(f\"  ✗ GO term: {parts[1]}\")\n        all_good = False\n    \n    try:\n        conf = float(parts[2])\n        if 0 < conf <= 1:\n            print(f\"  ✓ Confidence: {conf}\")\n        else:\n            print(f\"  ✗ Confidence out of range: {conf}\")\n            all_good = False\n    except:\n        print(f\"  ✗ Confidence not a number: {parts[2]}\")\n        all_good = False\n    \n    if '\\n' in lines[0]:\n        print(f\"  ✓ Newlines present\")\n    else:\n        print(f\"  ✗ No newlines!\")\n        all_good = False\n    \n    if all_good:\n        print(\"\\n✅✅✅ FORMAT IS CORRECT! ✅✅✅\")\n    else:\n        print(\"\\n⚠️  FORMAT HAS ISSUES!\")\n        \n    print(\"\\n📁 Files in working directory:\")\n    for fname in sorted(os.listdir('.')):\n        if not fname.startswith('.'):\n            size = os.path.getsize(fname) / (1024*1024)\n            print(f\"  {fname} ({size:.1f} MB)\")\nelse:\n    print(\"❌ ERROR: submission.tsv not found!\")\n\nprint(\"\\n\" + \"=\"*80)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T08:33:49.886996Z","iopub.execute_input":"2025-12-14T08:33:49.887234Z","iopub.status.idle":"2025-12-14T08:47:57.983088Z","shell.execute_reply.started":"2025-12-14T08:33:49.887213Z","shell.execute_reply":"2025-12-14T08:47:57.982373Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}