{"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":14875579,"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":"import numpy as np\nimport pandas as pd\nimport os\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\n# ============================================================================\nCONFIG = {\n    # Data sampling\n    'MAX_TRAIN_SAMPLES': 120000,\n    'TOP_K_FREQUENT': 1000, # từ 650 và tăng điểm\n    'TOP_K_RARE': 1200, # từ 800 và tăng điểm\n    'MIN_FREQ_COMMON': 20,\n    'MIN_FREQ_RARE': 3,\n    \n    # Focal Loss parameters\n    'FOCAL_LOSS_GAMMA': 2.5,\n    'FOCAL_LOSS_ALPHA': 0.3,\n    'USE_IA_WEIGHTS': True,\n    \n    # Multi-embedding ensemble\n    'USE_MULTI_EMBEDDING': True,\n    'EMBEDDING_SOURCES': ['esm2', 'protbert', 't5'],\n    'EMBEDDING_FUSION': 'attention',\n    \n    # High-IA term prioritization\n    'IA_SAMPLING_RATIO': 0.4,\n    'HIGH_IA_THRESHOLD': 2.5,\n    'IA_CLIP_MIN': None,        # ✅ BỎ clipping!\n    'IA_CLIP_MAX': None,\n    'IA_TRANSFORM': 'log1p',\n    \n    # Co-occurrence modeling\n    'USE_COOCCURRENCE': True,\n    'COOCCUR_TOP_K': 50,\n    \n    # Model architecture\n    'HIDDEN_DIMS': [512, 256],\n    'DROPOUT_RATE': 0.3,\n    'USE_BATCH_NORM': True,\n    \n    # Training\n    'EPOCHS': 35,\n    'BATCH_SIZE': 64,\n    'LEARNING_RATE': 1e-3,\n    'WEIGHT_DECAY': 1e-4,\n    'GRAD_CLIP': 1.0,\n    'WARMUP_EPOCHS': 5,\n    \n    # Prediction\n    'MAX_PREDS_PER_PROTEIN': 2000,\n    'MIN_CONFIDENCE': 0.005,# từ 0.005\n    'TEMPERATURE': 0.8,\n    \n    # GO propagation\n    'USE_GO_PROPAGATION': True,\n    'PROPAGATION_DECAY': 0.70,\n    \n    # Ontology calibration\n    'ONTOLOGY_CALIBRATION': {\n        'MFO': 1.10,\n        'BPO': 1.00,\n        'CCO': 1.05\n    },\n    \n    # Ensemble weights\n    'BASE_BLAST_WEIGHT': 0.55, # best\n    'BASE_DL_WEIGHT': 0.35, # best\n    'BASE_FREQ_WEIGHT': 0.10, # best\n    \n    # Paths\n    'BASE_PATH': '/kaggle/input/cafa-6-protein-function-prediction',\n    'ESM2_PATH': 'cafa-5-ems-2-embeddings-numpy',\n    'PROTBERT_PATH': 'protbert-embeddings-for-cafa5',\n    'T5_PATH': 't5embeds',\n    'BLAST_PATH': '/kaggle/input/blast-quick-sprof-zero-pred/submission.tsv',\n    'RANDOM_SEED': 42,\n}\n\nprint(\"=\"*80)\nprint(\"🚀 CAFA 6 - FIXED VERSION\")\nprint(\"=\"*80)\n\n# ============================================================================\n# SETUP\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\ntry:\n    import networkx as nx\n    import obonet\nexcept:\n    install('networkx')\n    install('obonet')\n    import networkx as nx\n    import obonet\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'])\n    \ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"🎯 Device: {device}\")\n\n# ============================================================================\n# FOCAL LOSS\n# ============================================================================\nclass FocalLoss(nn.Module):\n    def __init__(self, gamma=2.0, alpha=0.25, reduction='mean'):\n        super().__init__()\n        self.gamma = gamma\n        self.alpha = alpha\n        self.reduction = reduction\n    \n    def forward(self, inputs, targets):\n        bce_loss = F.binary_cross_entropy(inputs, targets, reduction='none')\n        p_t = inputs * targets + (1 - inputs) * (1 - targets)\n        focal_weight = (1 - p_t) ** self.gamma\n        \n        if self.alpha is not None:\n            alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)\n            focal_loss = alpha_t * focal_weight * bce_loss\n        else:\n            focal_loss = focal_weight * bce_loss\n        \n        if self.reduction == 'mean':\n            return focal_loss.mean()\n        elif self.reduction == 'sum':\n            return focal_loss.sum()\n        else:\n            return focal_loss\n\n# ============================================================================\n# LOAD GO ONTOLOGY\n# ============================================================================\nprint(\"\\n[1/12] Loading GO ontology...\")\nBASE = Path(CONFIG['BASE_PATH'])\nTRAIN_DIR = BASE / 'Train'\n\ngo_graph = None\nancestor_cache = {}\n\nif CONFIG['USE_GO_PROPAGATION']:\n    try:\n        go_graph = obonet.read_obo(TRAIN_DIR / 'go-basic.obo')\n        print(f\"   ✓ Loaded {len(go_graph):,} GO terms\")\n        \n        def get_ancestors(term_id, graph):\n            ancestors = set()\n            try:\n                for node in nx.descendants(graph, term_id):\n                    ancestors.add(node)\n            except:\n                pass\n            return ancestors\n    except Exception as e:\n        print(f\"   ⚠️ Could not load GO graph: {e}\")\n        CONFIG['USE_GO_PROPAGATION'] = False\n\n# ============================================================================\n# 🔧 FIX #2: LOAD IA WEIGHTS (CORRECT FILENAME)\n# ============================================================================\nprint(\"\\n[2/12] Loading IA weights...\")\nia_weights_dict = {}\n\ntry:\n    # Try .txt first (correct format), then .tsv\n    ia_file = BASE / 'IA.txt'\n    if not ia_file.exists():\n        ia_file = BASE / 'IA.tsv'\n    \n    ia_df = pd.read_csv(ia_file, sep='\\t', header=None, names=['term', 'ia_weight'])\n    ia_weights_dict = dict(zip(ia_df['term'], ia_df['ia_weight']))\n    \n    ia_values = list(ia_weights_dict.values())\n    print(f\"   ✓ Loaded IA weights for {len(ia_weights_dict):,} terms\")\n    print(f\"   📊 Mean IA: {np.mean(ia_values):.2f}\")\n    print(f\"   📊 Max IA: {np.max(ia_values):.2f}\")\n    \n    high_ia_terms = [t for t, ia in ia_weights_dict.items() \n                     if ia >= CONFIG['HIGH_IA_THRESHOLD']]\n    print(f\"   🔥 High-IA terms: {len(high_ia_terms):,}\")\n    \nexcept Exception as e:\n    print(f\"   ⚠️ Could not load IA weights: {e}\")\n    print(f\"   → Using uniform weights\")\n    CONFIG['USE_IA_WEIGHTS'] = False\n    high_ia_terms = []\n\n# ============================================================================\n# TWO-STAGE TERM SELECTION\n# ============================================================================\nprint(\"\\n[3/12] Loading annotations with two-stage selection...\")\ntrain_terms_df = pd.read_csv(TRAIN_DIR / 'train_terms.tsv', sep='\\t', \n                              header=None, names=['protein', 'term', 'aspect'])\n\nterm_freq = train_terms_df['term'].value_counts()\nprint(f\"   Total unique terms: {len(term_freq):,}\")\n\n# Stage 1: Frequent terms\nfrequent_terms = term_freq[term_freq >= CONFIG['MIN_FREQ_COMMON']].index.tolist()\nfrequent_terms = frequent_terms[:CONFIG['TOP_K_FREQUENT']]\nprint(f\"   📊 Stage 1 - Frequent terms: {len(frequent_terms):,}\")\n\n# Stage 2: High-IA rare terms\nrare_candidates = term_freq[\n    (term_freq >= CONFIG['MIN_FREQ_RARE']) & \n    (term_freq < CONFIG['MIN_FREQ_COMMON'])\n].index.tolist()\n\nif CONFIG['USE_IA_WEIGHTS'] and high_ia_terms:\n    rare_with_ia = [(t, ia_weights_dict.get(t, 0)) for t in rare_candidates \n                    if t in ia_weights_dict]\n    rare_with_ia.sort(key=lambda x: x[1], reverse=True)\n    rare_high_ia = [t for t, ia in rare_with_ia[:CONFIG['TOP_K_RARE']]]\n    print(f\"   📊 Stage 2 - High-IA rare terms: {len(rare_high_ia):,}\")\nelse:\n    rare_high_ia = []\n    print(f\"   📊 Stage 2 - High-IA rare terms: 0 (no IA data)\")\n\n# Combine\ntop_terms = list(set(frequent_terms + rare_high_ia))\n\nprint(f\"   ✅ Total selected terms: {len(top_terms):,}\")\n# ✅ BUILD CACHE SAU KHI CÓ top_terms\nif CONFIG['USE_GO_PROPAGATION'] and go_graph:\n    print(\"   Building ancestor cache for selected terms...\")\n    for term in tqdm(top_terms, desc=\"   Caching\", leave=False):\n        if term in go_graph:\n            ancestor_cache[term] = get_ancestors(term, go_graph)\n    print(f\"   ✓ Cached {len(ancestor_cache):,}/{len(top_terms):,} term ancestors\")\ntrain_terms_df = train_terms_df[train_terms_df['term'].isin(top_terms)]\n\n# Create mappings\nprotein_to_terms = train_terms_df.groupby('protein')['term'].apply(list).to_dict()\nterm_to_aspect = dict(zip(train_terms_df['term'], train_terms_df['aspect']))\nterm_to_idx = {term: idx for idx, term in enumerate(top_terms)}\nidx_to_term = {idx: term for term, idx in term_to_idx.items()}\n\nprint(f\"   ✓ {len(protein_to_terms):,} proteins with annotations\")\n\n# Create IA weight vector\nia_weight_vector = np.ones(len(top_terms))\n\nif CONFIG['USE_IA_WEIGHTS']:\n    for idx, term in enumerate(top_terms):\n        ia_weight_vector[idx] = ia_weights_dict.get(term, 1.0)\n    \n    # ✅ Sử dụng log1p để soft-scale (giữ được tín hiệu)\n    ia_weight_vector = np.log1p(ia_weight_vector)  # log(1+x)\n    \n    # ✅ Normalize về mean=1 (không clip!)\n    ia_weight_vector = ia_weight_vector / ia_weight_vector.mean()\n    \n    # ✅ Chỉ clip extreme outliers nếu cần (rất rộng)\n    ia_weight_vector = np.clip(ia_weight_vector, 0.1, 10.0)  # Thay vì 0.5-2.0\n    \n    print(f\"   ✓ IA weights (log-scaled): mean={ia_weight_vector.mean():.2f}, \"\n          f\"range=[{ia_weight_vector.min():.2f}, {ia_weight_vector.max():.2f}]\")\n\n# ============================================================================\n# CO-OCCURRENCE MATRIX\n# ============================================================================\nprint(\"\\n[4/12] Building co-occurrence matrix...\")\n\nif CONFIG['USE_COOCCURRENCE']:\n    cooccur_matrix = np.zeros((len(top_terms), len(top_terms)))\n    \n    for terms in tqdm(protein_to_terms.values(), desc=\"   Computing\", leave=False):\n        term_indices = [term_to_idx[t] for t in terms if t in term_to_idx]\n        \n        for i in term_indices:\n            for j in term_indices:\n                if i != j:\n                    cooccur_matrix[i, j] += 1\n    \n    row_sums = cooccur_matrix.sum(axis=1, keepdims=True)\n    row_sums[row_sums == 0] = 1\n    cooccur_matrix = cooccur_matrix / row_sums\n    \n    print(f\"   ✓ Co-occurrence matrix: {cooccur_matrix.shape}\")\n    print(f\"   📊 Sparsity: {(cooccur_matrix > 0).sum() / cooccur_matrix.size * 100:.1f}%\")\nelse:\n    cooccur_matrix = None\n\n# ============================================================================\n# 🔧 FIX #3: LOAD EMBEDDINGS WITH FALLBACK\n# ============================================================================\nprint(\"\\n[5/12] Loading embeddings...\")\n\nembedding_dicts = {}\nembedding_dims = []\n\n# ESM2 (always available)\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\nembedding_dicts['esm2'] = {\n    'train': {str(pid): emb for pid, emb in zip(train_ids, train_embeds)},\n    'test': {str(pid): emb for pid, emb in zip(test_ids, test_embeds)}\n}\nembedding_dims.append(train_embeds.shape[1])\nprint(f\"   ✓ ESM2: dim={train_embeds.shape[1]}\")\n\ndel train_ids, train_embeds, test_ids, test_embeds\ngc.collect()\n\n# Try other embeddings\nif CONFIG['USE_MULTI_EMBEDDING']:\n    for emb_name, path_key in [('protbert', 'PROTBERT_PATH'), ('t5', 'T5_PATH')]:\n        try:\n            emb_path = f\"/kaggle/input/{CONFIG[path_key]}\"\n            \n            # ✅ Kiểm tra filename khác nhau\n            if emb_name == 't5':\n                # T5 dùng \"embeds\" thay vì \"embeddings\"\n                train_ids = np.load(f\"{emb_path}/train_ids.npy\", allow_pickle=True)\n                train_embeds = np.load(f\"{emb_path}/train_embeds.npy\")\n                test_ids = np.load(f\"{emb_path}/test_ids.npy\", allow_pickle=True)\n                test_embeds = np.load(f\"{emb_path}/test_embeds.npy\")\n            else:\n                # ProtBERT dùng \"embeddings\"\n                train_ids = np.load(f\"{emb_path}/train_ids.npy\", allow_pickle=True)\n                train_embeds = np.load(f\"{emb_path}/train_embeddings.npy\")\n                test_ids = np.load(f\"{emb_path}/test_ids.npy\", allow_pickle=True)\n                test_embeds = np.load(f\"{emb_path}/test_embeddings.npy\")\n            \n            embedding_dicts[emb_name] = {\n                'train': {str(pid): emb for pid, emb in zip(train_ids, train_embeds)},\n                'test': {str(pid): emb for pid, emb in zip(test_ids, test_embeds)}\n            }\n                \n            embedding_dims.append(train_embeds.shape[1])\n            print(f\"   ✓ {emb_name.upper()}: dim={train_embeds.shape[1]}\")\n            \n            del train_ids, train_embeds, test_ids, test_embeds\n            gc.collect()\n            \n        except Exception as e:\n            print(f\"   ⚠️ Could not load {emb_name}: {e}\")  # ← Debug\n            pass\n\nprint(f\"\\n   ✅ Total sources: {len(embedding_dicts)}\")\nprint(f\"   📊 Embedding dims: {embedding_dims}\")\n\n# ============================================================================\n# PREPARE TRAINING DATA\n# ============================================================================\nprint(\"\\n[6/12] Preparing training data...\")\n\nvalid_proteins = []\nfor p in protein_to_terms.keys():\n    has_all = all(p in embedding_dicts[src]['train'] \n                  for src in embedding_dicts.keys())\n    if has_all:\n        valid_proteins.append(p)\n\nvalid_proteins = valid_proteins[:CONFIG['MAX_TRAIN_SAMPLES']]\nprint(f\"   ✓ Valid proteins: {len(valid_proteins):,}\")\n\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.model_selection import StratifiedShuffleSplit\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).astype(float)\n\nprotein_categories = []\nfor p in valid_proteins:\n    n_terms = len(protein_to_terms.get(p, []))\n    if n_terms < 5:\n        category = 'sparse'\n    elif n_terms < 15:\n        category = 'medium'\n    else:\n        category = 'rich'\n    protein_categories.append(category)\n\nsplitter = StratifiedShuffleSplit(n_splits=1, test_size=0.15, \n                                  random_state=CONFIG['RANDOM_SEED'])\ntrain_idx, val_idx = next(splitter.split(valid_proteins, protein_categories))\n\ntrain_proteins = [valid_proteins[i] for i in train_idx]\nval_proteins = [valid_proteins[i] for i in val_idx]\ny_train = y_encoded[train_idx]\ny_val = y_encoded[val_idx]\n\nprint(f\"   ✓ Train: {len(train_proteins):,} proteins\")\nprint(f\"   ✓ Val: {len(val_proteins):,} proteins\")\n\n# ============================================================================\n# 🔧 FIX #1: IMPROVED MODEL WITH PROPER MULTI-EMBEDDING HANDLING\n# ============================================================================\nprint(\"\\n[7/12] Building model...\")\n\nclass MultiEmbeddingFusion(nn.Module):\n    \"\"\"Fuse multiple embeddings using attention\"\"\"\n    def __init__(self, embedding_dims, output_dim):\n        super().__init__()\n        self.n_sources = len(embedding_dims)\n        self.projections = nn.ModuleList([\n            nn.Linear(dim, output_dim) for dim in embedding_dims\n        ])\n        if self.n_sources > 1:\n            self.attention = nn.Linear(output_dim, 1)\n    \n    def forward(self, embeddings):\n        # Handle both single and multiple embeddings\n        if not isinstance(embeddings, list):\n            embeddings = [embeddings]\n        \n        # Project all to same dimension\n        projected = [proj(emb) for proj, emb in zip(self.projections, embeddings)]\n        \n        if len(projected) == 1:\n            return projected[0]\n        \n        # Attention fusion\n        stacked = torch.stack(projected, dim=1)\n        attn_scores = self.attention(stacked).squeeze(-1)\n        attn_weights = F.softmax(attn_scores, dim=1).unsqueeze(-1)\n        fused = (stacked * attn_weights).sum(dim=1)\n        \n        return fused\n\nclass AdvancedProteinClassifier(nn.Module):\n    def __init__(self, embedding_dims, num_terms, hidden_dims, \n                 dropout, use_cooccurrence=False):\n        super().__init__()\n        \n        # Multi-embedding fusion\n        fusion_dim = hidden_dims[0]\n        self.fusion = MultiEmbeddingFusion(embedding_dims, fusion_dim)\n        \n        # Deep encoder\n        layers = []\n        in_dim = fusion_dim\n        \n        for hidden_dim in hidden_dims:\n            layers.extend([\n                nn.Linear(in_dim, hidden_dim),\n                nn.BatchNorm1d(hidden_dim),\n                nn.ReLU(),\n                nn.Dropout(dropout)\n            ])\n            in_dim = hidden_dim\n        \n        self.encoder = nn.Sequential(*layers)\n        self.output = nn.Linear(hidden_dims[-1], num_terms)\n        \n        # Initialize\n        for module in self.modules():\n            if isinstance(module, nn.Linear):\n                nn.init.xavier_uniform_(module.weight)\n                if module.bias is not None:\n                    nn.init.zeros_(module.bias)\n    \n    def forward(self, embeddings, cooccur_matrix=None, return_logits=False):\n        # Fuse embeddings (handles both list and single tensor)\n        fused = self.fusion(embeddings)\n        \n        # Encode\n        features = self.encoder(fused)\n        \n        # Output\n        logits = self.output(features)\n        \n        if return_logits:\n            return logits\n        return torch.sigmoid(logits)\n\nmodel = AdvancedProteinClassifier(\n    embedding_dims,\n    len(top_terms),\n    CONFIG['HIDDEN_DIMS'],\n    CONFIG['DROPOUT_RATE'],\n    use_cooccurrence=CONFIG['USE_COOCCURRENCE']\n).to(device)\n\ntotal_params = sum(p.numel() for p in model.parameters())\nprint(f\"   ✓ Parameters: {total_params:,}\")\nprint(f\"   ✓ Embedding sources: {len(embedding_dims)}\")\n\n# ============================================================================\n# TRAINING\n# ============================================================================\nprint(\"\\n\" + \"=\"*80)\nprint(f\"TRAINING ({CONFIG['EPOCHS']} EPOCHS)\")\nprint(\"=\"*80)\n\ncriterion = FocalLoss(\n    gamma=CONFIG['FOCAL_LOSS_GAMMA'],\n    alpha=CONFIG['FOCAL_LOSS_ALPHA'],\n    reduction='none'\n)\n\nia_weight_tensor = torch.FloatTensor(ia_weight_vector).to(device)\n\noptimizer = torch.optim.AdamW(\n    model.parameters(),\n    lr=CONFIG['LEARNING_RATE'],\n    weight_decay=CONFIG['WEIGHT_DECAY']\n)\n\ndef get_lr_lambda(epoch):\n    if epoch < CONFIG['WARMUP_EPOCHS']:\n        return (epoch + 1) / CONFIG['WARMUP_EPOCHS']\n    else:\n        progress = (epoch - CONFIG['WARMUP_EPOCHS']) / (CONFIG['EPOCHS'] - CONFIG['WARMUP_EPOCHS'])\n        return 0.5 * (1 + np.cos(np.pi * progress))\n\nscheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, get_lr_lambda)\n\nbest_val_loss = float('inf')\ncooccur_tensor = torch.FloatTensor(cooccur_matrix).to(device) if cooccur_matrix is not None else None\n\nfor epoch in range(CONFIG['EPOCHS']):\n    # Training\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        # Gather embeddings\n        if len(embedding_dicts) == 1:\n            # Single source - pass as tensor\n            src = list(embedding_dicts.keys())[0]\n            embeddings = torch.FloatTensor([\n                embedding_dicts[src]['train'][p] for p in batch_proteins\n            ]).to(device)\n        else:\n            # Multiple sources - pass as list\n            embeddings = []\n            for src in embedding_dicts.keys():\n                emb_batch = torch.FloatTensor([\n                    embedding_dicts[src]['train'][p] for p in batch_proteins\n                ]).to(device)\n                embeddings.append(emb_batch)\n        \n        y_batch = torch.FloatTensor(y_train[batch_idx]).to(device)\n        \n        optimizer.zero_grad()\n        outputs = model(embeddings, cooccur_tensor)\n        \n        loss_per_sample = criterion(outputs, y_batch)\n        weighted_loss = (loss_per_sample * ia_weight_tensor).mean()\n        \n        weighted_loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), CONFIG['GRAD_CLIP'])\n        optimizer.step()\n        \n        epoch_loss += weighted_loss.item()\n        n_batches += 1\n    \n    # Validation\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            \n            if len(embedding_dicts) == 1:\n                src = list(embedding_dicts.keys())[0]\n                embeddings = torch.FloatTensor([\n                    embedding_dicts[src]['train'][p] for p in batch_proteins\n                ]).to(device)\n            else:\n                embeddings = []\n                for src in embedding_dicts.keys():\n                    emb_batch = torch.FloatTensor([\n                        embedding_dicts[src]['train'][p] for p in batch_proteins\n                    ]).to(device)\n                    embeddings.append(emb_batch)\n            \n            y_batch = torch.FloatTensor(y_val[i:i + CONFIG['BATCH_SIZE']]).to(device)\n            \n            outputs = model(embeddings, cooccur_tensor)\n            loss_per_sample = criterion(outputs, y_batch)\n            weighted_loss = (loss_per_sample * ia_weight_tensor).mean()\n            \n            val_loss += weighted_loss.item()\n            val_batches += 1\n    \n    train_loss_avg = epoch_loss / n_batches\n    val_loss_avg = val_loss / val_batches\n    current_lr = optimizer.param_groups[0]['lr']\n    \n    scheduler.step()\n    gc.collect()\n    torch.cuda.empty_cache() \n    \n    if val_loss_avg < best_val_loss:\n        best_val_loss = val_loss_avg\n        marker = \"⭐\"\n    else:\n        marker = \"\"\n    \n    if (epoch + 1) % 5 == 0 or marker:\n        print(f\"Epoch {epoch+1:2d}: Train={train_loss_avg:.4f}, Val={val_loss_avg:.4f}, \"\n              f\"LR={current_lr:.6f} {marker}\")\n    \n    gc.collect()\n\nprint(f\"\\n✅ Best Val Loss: {best_val_loss:.4f}\")\n\n# ============================================================================\n# PREDICTIONS WITH ALL IMPROVEMENTS\n# ============================================================================\nprint(\"\\n\" + \"=\"*80)\nprint(\"GENERATING PREDICTIONS\")\nprint(\"=\"*80)\n\nmodel.eval()\n\n# Get all test protein IDs (from first embedding source)\ntest_protein_ids = list(embedding_dicts[list(embedding_dicts.keys())[0]]['test'].keys())\nprint(f\"Total test proteins: {len(test_protein_ids):,}\")\n\n# [Step 1/5] Generate DL predictions with multi-embedding\ndl_predictions = {}\n\nprint(\"\\n[Step 1/5] Generating multi-embedding DL predictions...\")\nwith torch.no_grad():\n    for start in tqdm(range(0, len(test_protein_ids), CONFIG['BATCH_SIZE']), \n                     desc=\"DL Inference\"):\n        batch_ids = test_protein_ids[start:start + CONFIG['BATCH_SIZE']]\n        \n        # Gather embeddings (match training logic)\n        if len(embedding_dicts) == 1:\n            src = list(embedding_dicts.keys())[0]\n            embeddings = torch.FloatTensor([\n                embedding_dicts[src]['test'][p] for p in batch_ids\n            ]).to(device)\n        else:\n            embeddings = []\n            for src in embedding_dicts.keys():\n                emb_batch = torch.FloatTensor([\n                    embedding_dicts[src]['test'][p] for p in batch_ids\n                ]).to(device)\n                embeddings.append(emb_batch)\n        \n        logits = model(embeddings, cooccur_tensor, return_logits=True)\n        outputs = torch.sigmoid(logits / CONFIG['TEMPERATURE']).cpu().numpy()\n        \n        for i, pid in enumerate(batch_ids):\n            dl_predictions[pid] = outputs[i]\n        \n        del embeddings, outputs, logits\n        if start % 5000 == 0:\n            gc.collect()\n\nprint(f\"✓ Generated DL predictions for {len(dl_predictions):,} proteins\")\n\n# [Step 2/5] Load BLAST predictions\nblast_dict = {}\n\nprint(\"\\n[Step 2/5] Loading BLAST predictions...\")\nif os.path.exists(CONFIG['BLAST_PATH']):\n    blast_data = defaultdict(lambda: np.zeros(len(top_terms)))\n    \n    for chunk in pd.read_csv(CONFIG['BLAST_PATH'], sep='\\t', header=None,\n                             names=['Id', 'GO term', 'Confidence'], \n                             chunksize=100000):\n        for _, row in chunk.iterrows():\n            pid = row['Id']\n            term = row['GO term']\n            conf = float(row['Confidence'])\n            \n            if term in term_to_idx:\n                idx = term_to_idx[term]\n                blast_data[pid][idx] = max(blast_data[pid][idx], conf)\n    \n    blast_dict = dict(blast_data)\n    print(f\"  ✓ Loaded BLAST for {len(blast_dict):,} proteins\")\nelse:\n    print(\"  ⚠️ BLAST file not found\")\n\n# [Step 3/5] Compute frequency baseline\nprint(\"\\n[Step 3/5] Computing frequency baseline...\")\nterm_frequencies = np.zeros(len(top_terms))\ntotal_proteins = len(protein_to_terms)\n\nfor idx, term in enumerate(top_terms):\n    count = term_freq.get(term, 0)\n    term_frequencies[idx] = min(count / total_proteins, 0.5)\n\nprint(f\"✓ Frequency baseline computed\")\n\n# [Step 4/5] Adaptive ensemble merging\nprint(\"\\n[Step 4/5] Adaptive ensemble merging...\")\n\ndef get_knowledge_level(protein_id, train_terms_dict, term_to_aspect_dict):\n    \"\"\"Detect knowledge level\"\"\"\n    if protein_id not in train_terms_dict:\n        return 'no'\n    \n    terms = train_terms_dict[protein_id]\n    aspects = set(term_to_aspect_dict.get(t, '') for t in terms)\n    aspect_map = {'F': 'MFO', 'P': 'BPO', 'C': 'CCO'}\n    mapped_aspects = set(aspect_map.get(a, a) for a in aspects)\n    \n    if len(mapped_aspects) == 0:\n        return 'no'\n    elif len(mapped_aspects) < 3:\n        return 'limited'\n    else:\n        return 'partial'\n\nmerged_predictions = {}\nknowledge_stats = Counter()\n\nfor pid, dl_probs in tqdm(dl_predictions.items(), desc=\"Merging\"):\n    knowledge = get_knowledge_level(pid, protein_to_terms, term_to_aspect)\n    knowledge_stats[knowledge] += 1\n    \n    blast_probs = blast_dict.get(pid, np.zeros(len(top_terms)))\n    \n    # Adaptive weighting\n    if knowledge == 'no':\n        dl_w = 0.35\n        freq_w = 0.15\n        blast_w = 0.50\n    elif knowledge == 'limited':\n        dl_w = 0.40\n        freq_w = 0.10\n        blast_w = 0.50\n    else:\n        dl_w = 0.50\n        freq_w = 0.05\n        blast_w = 0.45\n    \n    # Weighted combination\n    merged = (dl_probs * dl_w + \n              term_frequencies * freq_w + \n              blast_probs * blast_w)\n    \n    merged_predictions[pid] = merged\n\nprint(f\"✓ Merged predictions for {len(merged_predictions):,} proteins\")\nprint(\"\\n  Knowledge distribution:\")\nfor level, count in sorted(knowledge_stats.items()):\n    pct = count / len(merged_predictions) * 100\n    print(f\"    {level:8s}: {count:6,} ({pct:5.1f}%)\")\n\n# [Step 5/5] GO propagation + Ontology calibration\nprint(\"\\n[Step 5/5] GO propagation + Ontology calibration...\")\n\nfinal_predictions = {}\n\nfor pid, probs in tqdm(merged_predictions.items(), desc=\"Processing\"):\n    new_probs = probs.copy()\n    \n    # GO propagation\n    # Thêm threshold cao hơn cho propagation\n    PROPAGATION_MIN_CONF = 0.1 # Chỉ propagate high-confidence predictions # từ 0.3\n    if CONFIG['USE_GO_PROPAGATION'] and go_graph and ancestor_cache:\n        # ✅ Optimized\n        high_conf_indices = np.where(probs > PROPAGATION_MIN_CONF)[0]\n        for term_idx in high_conf_indices:\n            term = idx_to_term[term_idx]\n            score = probs[term_idx]\n            \n            if term in ancestor_cache:\n                ancestor_indices = [term_to_idx[a] for a in ancestor_cache[term] \n                                   if a in term_to_idx]\n                if ancestor_indices:\n                    propagated = score * CONFIG['PROPAGATION_DECAY']\n                    new_probs[ancestor_indices] = np.maximum(\n                        new_probs[ancestor_indices], \n                        propagated\n                    )\n    \n    # Ontology-specific calibration\n    for term_idx, score in enumerate(new_probs):\n        if score > CONFIG['MIN_CONFIDENCE']:\n            term = idx_to_term[term_idx]\n            aspect = term_to_aspect.get(term, 'F')\n            \n            aspect_map = {'F': 'MFO', 'P': 'BPO', 'C': 'CCO'}\n            ontology = aspect_map.get(aspect, 'MFO')\n            \n            calibration_factor = CONFIG['ONTOLOGY_CALIBRATION'][ontology]\n            new_probs[term_idx] = score * calibration_factor\n    \n    final_predictions[pid] = new_probs\n\nprint(f\"✓ Final predictions ready for {len(final_predictions):,} proteins\")\n\n\n\n\n# ============================================================================\n# ADAPTIVE CUTOFF HELPERS - Thêm vào trước phần WRITE SUBMISSION\n# ============================================================================\n\ndef get_adaptive_thresholds(protein_id, knowledge_level, blast_coverage, probs):\n    \"\"\"Tính adaptive thresholds dựa trên protein characteristics\"\"\"\n    \n    # Base thresholds\n    if knowledge_level == 'no':\n        if blast_coverage > 50:\n            min_conf = 0.008\n            max_preds = 3000\n        else:\n            min_conf = 0.015\n            max_preds = 1500\n    elif knowledge_level == 'limited':\n        min_conf = 0.010\n        max_preds = 2500\n    else:  # partial\n        min_conf = 0.008\n        max_preds = 3000\n    \n    # Điều chỉnh theo high-confidence predictions\n    high_conf_count = (probs > 0.1).sum()\n    \n    if high_conf_count > 100:\n        max_preds = min(max_preds + 500, 4000)\n    elif high_conf_count < 20:\n        max_preds = min(max_preds, 2000)\n        min_conf = max(min_conf, 0.015)\n    \n    return min_conf, max_preds\n\ndef should_include_prediction(score, position, prev_score, min_conf, max_preds):\n    \"\"\"Quyết định có include prediction không\"\"\"\n    \n    # Hard limits\n    if score < min_conf or position >= max_preds:\n        return False\n    \n    # Progressive thresholds\n    if position > 2000 and score < 0.02:\n        return False\n    \n    if position > 3000 and score < 0.05:\n        return False\n    \n    # Score drop detection (cliff detection)\n    if prev_score > 0 and position > 500:\n        if score < prev_score * 0.3:  # Drop > 70%\n            return False\n    \n    return True\n\n# ============================================================================\n# WRITE SUBMISSION WITH ADAPTIVE CUTOFF\n# ============================================================================\nprint(\"\\n\" + \"=\"*80)\nprint(\"WRITING SUBMISSION WITH ADAPTIVE CUTOFF\")\nprint(\"=\"*80)\n\nn_predictions = 0\nprotein_pred_counts = []\ncutoff_stats = {'min_conf': 0, 'max_preds': 0, 'quality': 0}\n\nwith open('submission.tsv', 'w') as f:\n    for pid, probs in tqdm(final_predictions.items(), desc=\"Writing\"):\n        # Get protein characteristics\n        knowledge = get_knowledge_level(pid, protein_to_terms, term_to_aspect)\n        blast_probs = blast_dict.get(pid, np.zeros(len(top_terms)))\n        blast_coverage = (blast_probs > 0.01).sum()\n        \n        # Get adaptive thresholds\n        min_conf, max_preds = get_adaptive_thresholds(\n            pid, knowledge, blast_coverage, probs\n        )\n        \n        # Sort predictions\n        top_indices = np.argsort(probs)[::-1]\n        \n        protein_preds = 0\n        prev_score = 1.0\n        \n        for position, idx in enumerate(top_indices):\n            score = probs[idx]\n            \n            # Check if should include\n            if should_include_prediction(score, position, prev_score, \n                                        min_conf, max_preds):\n                term = idx_to_term[idx]\n                f.write(f\"{pid}\\t{term}\\t{min(score, 0.999):.3f}\\n\")\n                n_predictions += 1\n                protein_preds += 1\n                prev_score = score\n            else:\n                # Track why stopped\n                if score < min_conf:\n                    cutoff_stats['min_conf'] += 1\n                elif position >= max_preds:\n                    cutoff_stats['max_preds'] += 1\n                else:\n                    cutoff_stats['quality'] += 1\n                break\n        \n        if protein_preds > 0:\n            protein_pred_counts.append(protein_preds)\n\nprint(f\"\\n✅ Adaptive cutoff statistics:\")\nprint(f\"  Stopped by min_confidence: {cutoff_stats['min_conf']:,}\")\nprint(f\"  Stopped by max_preds: {cutoff_stats['max_preds']:,}\")\nprint(f\"  Stopped by quality filter: {cutoff_stats['quality']:,}\")\n\n\nprint(\"\\n✅ ALL DONE!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T08:52:57.867146Z","iopub.execute_input":"2025-12-14T08:52:57.867392Z","iopub.status.idle":"2025-12-14T09:28:39.485987Z","shell.execute_reply.started":"2025-12-14T08:52:57.867372Z","shell.execute_reply":"2025-12-14T09:28:39.484985Z"}},"outputs":[],"execution_count":null}]}