{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# CAFA-6: ESM2 Embeddings -> GO Term Prediction\n\n**Strategy:**\n1. Train MLP on ESM2 protein embeddings to predict GO terms\n2. Combine with proven ensemble for proteins with weak ESM2 predictions\n3. GO ontology propagation for complete submission"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\nimport gc\nimport numpy as np\nimport pandas as pd\nfrom collections import defaultdict\nfrom sklearn.model_selection import train_test_split\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm.auto import tqdm\n\nprint('CAFA-6 ESM2 GO Predictor v1')\nprint(f'PyTorch: {torch.__version__}')\nprint(f'CUDA: {torch.cuda.is_available()}')\nif torch.cuda.is_available():\n    print(f'GPU: {torch.cuda.get_device_name(0)}')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Paths\nEMBEDDINGS_PATH = '/kaggle/input/cafa6-protein-embeddings-esm2/protein_embeddings.npy'\nPROTEIN_IDS_PATH = '/kaggle/input/cafa6-protein-embeddings-esm2/protein_ids.csv'\nTRAIN_TERMS_PATH = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv'\nGO_OBO_PATH = '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo'\n\n# Ensemble kernels\nQUICK_KERNEL = '/kaggle/input/cafa-quick-submission-0-358-score/submission.tsv'\nPROVEN_KERNEL = '/kaggle/input/cafa-6-proven-ensemble-submission/submission.tsv'\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f'Device: {device}')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 1. LOAD EMBEDDINGS\n# ==========================================\nprint('[1/8] Loading ESM2 embeddings...')\nembeddings = np.load(EMBEDDINGS_PATH)\nprotein_ids_df = pd.read_csv(PROTEIN_IDS_PATH)\nprotein_ids = protein_ids_df['protein_id'].tolist()\nprotein_to_idx = {pid: idx for idx, pid in enumerate(protein_ids)}\n\nprint(f'Embeddings shape: {embeddings.shape}')\nprint(f'Proteins: {len(protein_ids)}')\nEMBEDDING_DIM = embeddings.shape[1]"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 2. LOAD TRAINING GO TERMS\n# ==========================================\nprint('[2/8] Loading training GO terms...')\ntrain_terms = pd.read_csv(TRAIN_TERMS_PATH, sep='\\t', header=None,\n                          names=['protein', 'go_term', 'aspect'])\nprint(f'Total annotations: {len(train_terms):,}')\n\n# Filter to proteins we have embeddings for\ntrain_terms = train_terms[train_terms['protein'].isin(protein_to_idx)]\nprint(f'Annotations with embeddings: {len(train_terms):,}')\n\n# Create protein -> GO terms mapping\nprotein_to_go = defaultdict(set)\nfor _, row in tqdm(train_terms.iterrows(), total=len(train_terms), desc='Building GO map'):\n    protein_to_go[row['protein']].add(row['go_term'])\n\n# Filter to most common GO terms for efficiency\ngo_term_counts = train_terms['go_term'].value_counts()\nMIN_ANNOTATIONS = 50\ngo_terms_filtered = go_term_counts[go_term_counts >= MIN_ANNOTATIONS].index.tolist()\nprint(f'GO terms with >= {MIN_ANNOTATIONS} annotations: {len(go_terms_filtered)}')\n\ngo_term_to_idx = {term: idx for idx, term in enumerate(go_terms_filtered)}\nNUM_CLASSES = len(go_terms_filtered)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 3. PREPARE TRAINING DATA\n# ==========================================\nprint('[3/8] Preparing training data...')\n\ntrain_proteins = []\ntrain_labels = []\n\nfor protein, go_set in tqdm(protein_to_go.items(), desc='Building labels'):\n    if protein not in protein_to_idx:\n        continue\n    label_indices = [go_term_to_idx[term] for term in go_set if term in go_term_to_idx]\n    if label_indices:\n        train_proteins.append(protein)\n        train_labels.append(label_indices)\n\nprint(f'Training proteins: {len(train_proteins)}')\n\n# Create multi-label targets\ny_train = np.zeros((len(train_proteins), NUM_CLASSES), dtype=np.float32)\nfor i, labels in enumerate(train_labels):\n    y_train[i, labels] = 1.0\n\nX_indices = np.array([protein_to_idx[p] for p in train_proteins])\n\n# Split\nX_train_idx, X_val_idx, y_train_split, y_val_split = train_test_split(\n    X_indices, y_train, test_size=0.1, random_state=42\n)\nprint(f'Train: {len(X_train_idx)}, Val: {len(X_val_idx)}')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 4. CREATE MODEL\n# ==========================================\nprint('[4/8] Creating model...')\n\nclass ProteinDataset(Dataset):\n    def __init__(self, indices, labels, embeddings):\n        self.indices = indices\n        self.labels = labels\n        self.embeddings = embeddings\n    \n    def __len__(self):\n        return len(self.indices)\n    \n    def __getitem__(self, idx):\n        emb_idx = self.indices[idx]\n        x = torch.tensor(self.embeddings[emb_idx], dtype=torch.float32)\n        y = torch.tensor(self.labels[idx], dtype=torch.float32)\n        return x, y\n\nclass GOPredictor(nn.Module):\n    def __init__(self, input_dim, hidden_dim, num_classes, dropout=0.3):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, hidden_dim // 2),\n            nn.BatchNorm1d(hidden_dim // 2),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim // 2, num_classes)\n        )\n    \n    def forward(self, x):\n        return self.net(x)\n\nBATCH_SIZE = 512\nHIDDEN_DIM = 2048\n\ntrain_dataset = ProteinDataset(X_train_idx, y_train_split, embeddings)\nval_dataset = ProteinDataset(X_val_idx, y_val_split, embeddings)\n\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)\n\nmodel = GOPredictor(EMBEDDING_DIM, HIDDEN_DIM, NUM_CLASSES).to(device)\nprint(f'Parameters: {sum(p.numel() for p in model.parameters()):,}')\n\ncriterion = nn.BCEWithLogitsLoss()\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=2, factor=0.5)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 5. TRAINING\n# ==========================================\nprint('[5/8] Training...')\n\nNUM_EPOCHS = 15  # Reduced for Kaggle time limits\nbest_val_loss = float('inf')\n\nfor epoch in range(NUM_EPOCHS):\n    model.train()\n    train_loss = 0\n    for x, y in tqdm(train_loader, desc=f'Epoch {epoch+1}/{NUM_EPOCHS}', leave=False):\n        x, y = x.to(device), y.to(device)\n        optimizer.zero_grad()\n        outputs = model(x)\n        loss = criterion(outputs, y)\n        loss.backward()\n        optimizer.step()\n        train_loss += loss.item()\n    \n    train_loss /= len(train_loader)\n    \n    model.eval()\n    val_loss = 0\n    with torch.no_grad():\n        for x, y in val_loader:\n            x, y = x.to(device), y.to(device)\n            outputs = model(x)\n            loss = criterion(outputs, y)\n            val_loss += loss.item()\n    \n    val_loss /= len(val_loader)\n    scheduler.step(val_loss)\n    \n    print(f'Epoch {epoch+1}: Train={train_loss:.4f}, Val={val_loss:.4f}, LR={optimizer.param_groups[0][\"lr\"]:.6f}')\n    \n    if val_loss < best_val_loss:\n        best_val_loss = val_loss\n        torch.save(model.state_dict(), 'best_model.pt')\n        print('  -> Best model saved!')\n\nprint(f'Best val loss: {best_val_loss:.4f}')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 6. GENERATE ESM2 PREDICTIONS\n# ==========================================\nprint('[6/8] Generating ESM2 predictions...')\n\nmodel.load_state_dict(torch.load('best_model.pt'))\nmodel.eval()\n\nesm2_predictions = defaultdict(dict)\nbatch_size = 1024\n\nwith torch.no_grad():\n    for start_idx in tqdm(range(0, len(protein_ids), batch_size), desc='Predicting'):\n        end_idx = min(start_idx + batch_size, len(protein_ids))\n        batch_emb = torch.tensor(embeddings[start_idx:end_idx], dtype=torch.float32).to(device)\n        \n        logits = model(batch_emb)\n        probs = torch.sigmoid(logits).cpu().numpy()\n        \n        for i, idx in enumerate(range(start_idx, end_idx)):\n            protein = protein_ids[idx]\n            for j, prob in enumerate(probs[i]):\n                if prob > 0.1:  # Threshold\n                    go_term = go_terms_filtered[j]\n                    esm2_predictions[protein][go_term] = float(prob)\n\nprint(f'ESM2 predictions for {len(esm2_predictions):,} proteins')\n\n# Clean up GPU memory\ndel model, train_dataset, val_dataset, train_loader, val_loader\ngc.collect()\ntorch.cuda.empty_cache()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 7. LOAD AND COMBINE WITH ENSEMBLE\n# ==========================================\nprint('[7/8] Loading ensemble and combining...')\n\ndef load_predictions(path):\n    df = pd.read_csv(path, sep='\\t', header=None, names=['protein', 'go_term', 'score'])\n    preds = defaultdict(dict)\n    for _, row in tqdm(df.iterrows(), total=len(df), desc=f'Loading {os.path.basename(path)}'):\n        preds[row['protein']][row['go_term']] = max(preds[row['protein']].get(row['go_term'], 0), row['score'])\n    return preds\n\nquick_preds = load_predictions(QUICK_KERNEL)\nproven_preds = load_predictions(PROVEN_KERNEL)\n\n# Merge quick + proven ensemble\nensemble_preds = defaultdict(dict)\nall_proteins = set(quick_preds.keys()) | set(proven_preds.keys())\n\nfor protein in tqdm(all_proteins, desc='Merging ensemble'):\n    all_terms = set(quick_preds.get(protein, {}).keys()) | set(proven_preds.get(protein, {}).keys())\n    for term in all_terms:\n        s1 = quick_preds.get(protein, {}).get(term, 0)\n        s2 = proven_preds.get(protein, {}).get(term, 0)\n        ensemble_preds[protein][term] = (s1 + s2) / 2\n\nprint(f'Ensemble: {len(ensemble_preds):,} proteins')\nprint(f'ESM2: {len(esm2_predictions):,} proteins')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Combine ESM2 + Ensemble\n# Strategy: ESM2 for proteins with embeddings, ensemble as fallback\n# Weight ESM2 predictions higher when confident\n\nfinal_predictions = defaultdict(dict)\nESM2_WEIGHT = 0.6\nENSEMBLE_WEIGHT = 0.4\n\nall_final_proteins = set(esm2_predictions.keys()) | set(ensemble_preds.keys())\n\nfor protein in tqdm(all_final_proteins, desc='Combining'):\n    esm2_terms = esm2_predictions.get(protein, {})\n    ens_terms = ensemble_preds.get(protein, {})\n    all_terms = set(esm2_terms.keys()) | set(ens_terms.keys())\n    \n    for term in all_terms:\n        e_score = esm2_terms.get(term, 0)\n        s_score = ens_terms.get(term, 0)\n        \n        if e_score > 0 and s_score > 0:\n            # Both have predictions - weighted average\n            final_predictions[protein][term] = ESM2_WEIGHT * e_score + ENSEMBLE_WEIGHT * s_score\n        elif e_score > 0:\n            final_predictions[protein][term] = e_score\n        else:\n            final_predictions[protein][term] = s_score\n\nprint(f'Combined: {len(final_predictions):,} proteins')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 8. GO PROPAGATION AND SUBMISSION\n# ==========================================\nprint('[8/8] GO propagation and submission...')\n\ndef parse_obo(obo_path):\n    parents = defaultdict(set)\n    with open(obo_path, 'r') as f:\n        cur_id = None\n        for line in f:\n            line = line.strip()\n            if line == '[Term]':\n                cur_id = None\n            elif line.startswith('id: '):\n                cur_id = line.split('id: ')[1].strip()\n            elif line.startswith('is_a: '):\n                pid = line.split()[1].strip()\n                if cur_id:\n                    parents[cur_id].add(pid)\n            elif line.startswith('relationship: part_of '):\n                parts = line.split()\n                if len(parts) >= 3:\n                    pid = parts[2].strip()\n                    if cur_id:\n                        parents[cur_id].add(pid)\n    return parents\n\ndef get_ancestors(term, parents_map, cache):\n    if term in cache:\n        return cache[term]\n    ancestors = set()\n    stack = [term]\n    while stack:\n        cur = stack.pop()\n        for p in parents_map.get(cur, []):\n            if p not in ancestors:\n                ancestors.add(p)\n                stack.append(p)\n    cache[term] = ancestors\n    return ancestors\n\nparents_map = parse_obo(GO_OBO_PATH)\nprint(f'Parsed {len(parents_map)} GO terms')\n\nROOTS = {'GO:0003674', 'GO:0008150', 'GO:0005575'}\nancestor_cache = {}\nfinal_rows = []\n\nfor protein, terms_dict in tqdm(final_predictions.items(), desc='Propagating'):\n    final_scores = dict(terms_dict)\n    \n    for term, score in list(terms_dict.items()):\n        for ancestor in get_ancestors(term, parents_map, ancestor_cache):\n            final_scores[ancestor] = max(final_scores.get(ancestor, 0), score)\n    \n    for root in ROOTS:\n        final_scores[root] = 1.0\n    \n    for term, score in final_scores.items():\n        if score >= 0.01:\n            final_rows.append((protein, term, round(score, 6)))\n\nprint(f'Total predictions: {len(final_rows):,}')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Save submission\nfinal_df = pd.DataFrame(final_rows, columns=['protein_id', 'go_term', 'score'])\nfinal_df = final_df.sort_values(['protein_id', 'score'], ascending=[True, False])\nfinal_df.to_csv('submission.tsv', sep='\\t', index=False, header=False)\n\nprint(f'\\n{\"=\"*60}')\nprint('DONE - CAFA6 ESM2 GO Predictor v1')\nprint(f'Predictions: {len(final_df):,}')\nprint(f'Proteins: {final_df[\"protein_id\"].nunique():,}')\nprint(f'{\"=\"*60}')\nprint(final_df.head(10))"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.0"}}, "nbformat": 4, "nbformat_minor": 4}