{"nbformat": 4, "nbformat_minor": 4, "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}}, "cells": [{"cell_type": "markdown", "metadata": {}, "source": "
\n

CAFA 6: Advanced Hybrid Solution

\n

ESM2-15B + DIAMOND Homology + Calibrated Ensemble

\n
\nESM2-15B (5120d)\n5-Fold CV\nCalibrated F-max\n
\n
"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "!pip install -q biopython goatools\n!apt-get update -qq && apt-get install -y -qq diamond-aligner"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import numpy as np\nimport pandas as pd\nimport os, gc, warnings\nfrom collections import Counter\nfrom tqdm.auto import tqdm\n\nfrom Bio import SeqIO\nfrom goatools.obo_parser import GODag\nfrom sklearn.model_selection import KFold\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\n\nwarnings.filterwarnings('ignore')\nprint(f\"PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "class CFG:\n COMP_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\n EMB_DIR = '/kaggle/input/cafa-6-embeddings/esm2_15B'\n OUTPUT_DIR = '/kaggle/working'\n TRAIN_TERMS = f'{COMP_DIR}/Train/train_terms.tsv'\n TRAIN_SEQ = f'{COMP_DIR}/Train/train_sequences.fasta'\n TEST_SEQ = f'{COMP_DIR}/Test/testsuperset.fasta'\n GO_OBO = f'{COMP_DIR}/Train/go-basic.obo'\n IA_FILE = f'{COMP_DIR}/IA.tsv'\n \n NUM_LABELS = 2500\n BATCH_SIZE = 32\n EPOCHS = 8\n N_FOLDS = 5\n DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n\nprint(f\"Device: {CFG.DEVICE}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load data\ntrain_terms = pd.read_csv(CFG.TRAIN_TERMS, sep='\\t')\nprint(f\"Annotations: {len(train_terms):,} | Proteins: {train_terms['EntryID'].nunique():,}\")\n\ngodag = GODag(CFG.GO_OBO, optional_attrs={'relationship'})\nprint(f\"GO terms: {len(godag):,}\")\n\nia_df = pd.read_csv(CFG.IA_FILE, sep='\\t', header=None, names=['term', 'ia'])\nia_dict = dict(zip(ia_df['term'], ia_df['ia']))\nprint(f\"IA scores: {len(ia_dict):,}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def parse_id(s):\n s = str(s)\n return s.split('|')[1] if '|' in s else s\n\ntest_ids = [parse_id(r.id) for r in SeqIO.parse(CFG.TEST_SEQ, 'fasta')]\nprint(f\"Test proteins: {len(test_ids):,}\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Part 1: DIAMOND Homology"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "print(\"Building DIAMOND database...\")\n!diamond makedb --in {CFG.TRAIN_SEQ} -d /kaggle/working/train_db --quiet\n\nprint(\"Running DIAMOND BLASTp...\")\n!diamond blastp -d /kaggle/working/train_db -q {CFG.TEST_SEQ} \\\n -o /kaggle/working/diamond_hits.tsv \\\n --very-sensitive --top 10 \\\n --outfmt 6 qseqid sseqid pident evalue bitscore \\\n --threads 4 --quiet\n\nprint(\"Done!\")\n!wc -l /kaggle/working/diamond_hits.tsv"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "diamond_hits = pd.read_csv('/kaggle/working/diamond_hits.tsv', sep='\\t',\n names=['query_id', 'subject_id', 'pident', 'evalue', 'bitscore'])\nprint(f\"Total hits: {len(diamond_hits):,}\")\n\ndiamond_hits['query_id'] = diamond_hits['query_id'].apply(parse_id)\ndiamond_hits['subject_id'] = diamond_hits['subject_id'].apply(parse_id)\ndiamond_hits = diamond_hits[diamond_hits['pident'] >= 25]\nprint(f\"Hits >= 25% identity: {len(diamond_hits):,}\")\nprint(f\"Proteins with homologs: {diamond_hits['query_id'].nunique():,}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Compute sophisticated homology score\ndef compute_score(row):\n id_score = row['pident'] / 100.0\n evalue_score = np.clip(1 - np.log10(row['evalue'] + 1e-180) / 180, 0, 1)\n bit_score = np.clip(row['bitscore'] / 500.0, 0, 1)\n return 0.5 * id_score + 0.3 * evalue_score + 0.2 * bit_score\n\ndiamond_hits['score'] = diamond_hits.apply(compute_score, axis=1)\nprint(f\"Score range: [{diamond_hits['score'].min():.3f}, {diamond_hits['score'].max():.3f}]\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Transfer labels\nprint(\"Transferring labels...\")\nhomology_merged = diamond_hits.merge(\n train_terms[['EntryID', 'term']], \n left_on='subject_id', right_on='EntryID', how='inner'\n)\n\nhomology_df = homology_merged.groupby(['query_id', 'term'])['score'].max().reset_index()\nhomology_df.columns = ['id', 'term', 'score']\nhomology_df = homology_df[homology_df['score'] >= 0.25]\n\nprint(f\"Homology predictions: {len(homology_df):,}\")\nprint(f\"Proteins covered: {homology_df['id'].nunique():,}\")\nprint(f\"GO terms: {homology_df['term'].nunique():,}\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Part 2: ESM2-15B Neural Network"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Prepare labels\nmf_terms = train_terms[train_terms['aspect'] == 'F']['term'].value_counts().head(CFG.NUM_LABELS // 3).index.tolist()\nbp_terms = train_terms[train_terms['aspect'] == 'P']['term'].value_counts().head(CFG.NUM_LABELS // 3).index.tolist()\ncc_terms = train_terms[train_terms['aspect'] == 'C']['term'].value_counts().head(CFG.NUM_LABELS // 3).index.tolist()\n\nall_terms = mf_terms + bp_terms + cc_terms\nterm2idx = {t: i for i, t in enumerate(all_terms)}\nidx2term = {i: t for t, i in term2idx.items()}\n\nprint(f\"MF: {len(mf_terms)}, BP: {len(bp_terms)}, CC: {len(cc_terms)}, Total: {len(all_terms)}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load ESM2-15B embeddings\nprint(\"Loading ESM2-15B embeddings...\")\ntrain_emb = np.load(f'{CFG.EMB_DIR}/train_sequences_emb.npy')\ntrain_ids_raw = np.load(f'{CFG.EMB_DIR}/train_sequences_ids.npy')\ntest_emb = np.load(f'{CFG.EMB_DIR}/testsuperset_emb.npy')\ntest_ids_raw = np.load(f'{CFG.EMB_DIR}/testsuperset_ids.npy')\n\ntrain_ids = [parse_id(x) for x in train_ids_raw]\ntest_ids_emb = [parse_id(x) for x in test_ids_raw]\nEMB_DIM = train_emb.shape[1]\n\nprint(f\"Train: {train_emb.shape}, Test: {test_emb.shape}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Prepare training data\nprot2labels = train_terms[train_terms['term'].isin(all_terms)].groupby('EntryID')['term'].apply(list).to_dict()\ntrain_id2idx = {pid: i for i, pid in enumerate(train_ids)}\nvalid_ids = [pid for pid in train_ids if pid in prot2labels]\n\nX_train = np.array([train_emb[train_id2idx[pid]] for pid in valid_ids])\ny_train = np.zeros((len(valid_ids), len(all_terms)), dtype=np.float32)\n\nfor i, pid in enumerate(valid_ids):\n for term in prot2labels[pid]:\n if term in term2idx:\n y_train[i, term2idx[term]] = 1.0\n\nprint(f\"X: {X_train.shape}, y: {y_train.shape}\")\ndel train_emb; gc.collect()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Model\nclass Model(nn.Module):\n def __init__(self, in_dim, out_dim):\n super().__init__()\n self.net = nn.Sequential(\n nn.Linear(in_dim, 2048), nn.LayerNorm(2048), nn.GELU(), nn.Dropout(0.2),\n nn.Linear(2048, 1024), nn.LayerNorm(1024), nn.GELU(), nn.Dropout(0.3),\n nn.Linear(1024, 512), nn.LayerNorm(512), nn.GELU(), nn.Dropout(0.3),\n nn.Linear(512, out_dim)\n )\n def forward(self, x): return self.net(x)\n\nclass FocalLoss(nn.Module):\n def __init__(self, alpha=0.25, gamma=2.0):\n super().__init__()\n self.alpha, self.gamma = alpha, gamma\n def forward(self, logits, targets):\n bce = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')\n pt = torch.exp(-bce)\n return (self.alpha * (1 - pt) ** self.gamma * bce).mean()\n\nclass ProteinDataset(Dataset):\n def __init__(self, emb, labels=None):\n self.emb = torch.tensor(emb, dtype=torch.float32)\n self.labels = torch.tensor(labels, dtype=torch.float32) if labels is not None else None\n def __len__(self): return len(self.emb)\n def __getitem__(self, i): return (self.emb[i], self.labels[i]) if self.labels is not None else self.emb[i]\n\nprint(f\"Model: {EMB_DIM} -> {len(all_terms)}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Training\nprint(f\"Training {CFG.N_FOLDS}-fold CV...\")\nkf = KFold(n_splits=CFG.N_FOLDS, shuffle=True, random_state=42)\nmodels = []\noof_preds = np.zeros_like(y_train)\n\nfor fold, (tr_idx, val_idx) in enumerate(kf.split(X_train)):\n print(f\"\\nFold {fold+1}/{CFG.N_FOLDS}\")\n \n train_ds = ProteinDataset(X_train[tr_idx], y_train[tr_idx])\n val_ds = ProteinDataset(X_train[val_idx], y_train[val_idx])\n train_loader = DataLoader(train_ds, batch_size=CFG.BATCH_SIZE, shuffle=True)\n val_loader = DataLoader(val_ds, batch_size=CFG.BATCH_SIZE)\n \n model = Model(EMB_DIM, len(all_terms)).to(CFG.DEVICE)\n optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=1e-4)\n scheduler = CosineAnnealingLR(optimizer, T_max=CFG.EPOCHS)\n criterion = FocalLoss()\n \n best_loss, best_state = float('inf'), None\n \n for epoch in range(CFG.EPOCHS):\n model.train()\n for X, y in train_loader:\n X, y = X.to(CFG.DEVICE), y.to(CFG.DEVICE)\n optimizer.zero_grad()\n loss = criterion(model(X), y)\n loss.backward()\n torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n optimizer.step()\n \n model.eval()\n val_loss = sum(criterion(model(X.to(CFG.DEVICE)), y.to(CFG.DEVICE)).item() \n for X, y in val_loader) / len(val_loader)\n scheduler.step()\n \n if val_loss < best_loss:\n best_loss, best_state = val_loss, model.state_dict().copy()\n print(f\" Epoch {epoch+1}: val_loss={val_loss:.4f}\")\n \n model.load_state_dict(best_state)\n models.append(model)\n \n model.eval()\n with torch.no_grad():\n oof_preds[val_idx] = np.vstack([torch.sigmoid(model(X.to(CFG.DEVICE))).cpu().numpy() \n for X, _ in val_loader])\n\nprint(\"\\nTraining complete!\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Calibrate thresholds using OOF\nprint(\"Calibrating thresholds...\")\noptimal_thresholds = []\nfor j in range(len(all_terms)):\n y_true_j, y_pred_j = y_train[:, j], oof_preds[:, j]\n if y_true_j.sum() < 10:\n optimal_thresholds.append(0.5)\n continue\n \n best_f1, best_t = 0, 0.5\n for t in np.arange(0.1, 0.9, 0.05):\n tp = ((y_pred_j >= t) & (y_true_j == 1)).sum()\n fp = ((y_pred_j >= t) & (y_true_j == 0)).sum()\n fn = ((y_pred_j < t) & (y_true_j == 1)).sum()\n f1 = 2 * tp / (2 * tp + fp + fn + 1e-10)\n if f1 > best_f1: best_f1, best_t = f1, t\n optimal_thresholds.append(best_t)\n\noptimal_thresholds = np.array(optimal_thresholds)\nprint(f\"Threshold: mean={optimal_thresholds.mean():.2f}, range=[{optimal_thresholds.min():.2f}, {optimal_thresholds.max():.2f}]\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Test inference\nprint(\"Running inference...\")\ntest_ds = ProteinDataset(test_emb)\ntest_loader = DataLoader(test_ds, batch_size=CFG.BATCH_SIZE)\n\ntest_preds = np.zeros((len(test_emb), len(all_terms)))\nfor model in models:\n model.eval()\n with torch.no_grad():\n preds = np.vstack([torch.sigmoid(model(X.to(CFG.DEVICE))).cpu().numpy() for X in test_loader])\n test_preds += preds / CFG.N_FOLDS\n\nprint(f\"NN predictions: {test_preds.shape}\")\ndel test_emb; gc.collect()"}, {"cell_type": "markdown", "metadata": {}, "source": "## Part 3: Ensemble & Submission"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "test_id2idx = {pid: i for i, pid in enumerate(test_ids_emb)}\nproteins_with_homology = set(homology_df['id'].unique())\nproteins_without_homology = set(test_ids_emb) - proteins_with_homology\nprint(f\"With homology: {len(proteins_with_homology):,}, Orphans: {len(proteins_without_homology):,}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Build ensemble\nprint(\"Building ensemble...\")\nensemble_rows = []\n\n# Homology predictions\nfor _, row in tqdm(homology_df.iterrows(), total=len(homology_df), desc=\"Homology\"):\n ensemble_rows.append({'id': row['id'], 'term': row['term'], 'score': row['score']})\n\n# NN for orphans\nfor pid in tqdm(proteins_without_homology, desc=\"Orphan NN\"):\n if pid not in test_id2idx: continue\n pred = test_preds[test_id2idx[pid]]\n for j in range(len(all_terms)):\n if pred[j] >= optimal_thresholds[j] * 0.7:\n ensemble_rows.append({'id': pid, 'term': idx2term[j], 'score': min(pred[j] * 1.2, 0.95)})\n\n# NN boost for low-confidence homology\nhomology_max = homology_df.groupby('id')['score'].max().to_dict()\nfor pid in tqdm(proteins_with_homology, desc=\"NN Boost\"):\n if pid not in test_id2idx or homology_max.get(pid, 1.0) > 0.7: continue\n pred = test_preds[test_id2idx[pid]]\n for j in range(len(all_terms)):\n if pred[j] >= 0.5:\n ensemble_rows.append({'id': pid, 'term': idx2term[j], 'score': pred[j] * 0.8})\n\nensemble_df = pd.DataFrame(ensemble_rows).groupby(['id', 'term'])['score'].max().reset_index()\nprint(f\"Ensemble: {len(ensemble_df):,}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# GO propagation\nprint(\"GO propagation...\")\nterm_ancestors = {}\nfor term in tqdm(ensemble_df['term'].unique(), desc=\"Ancestors\"):\n term_ancestors[term] = list(godag[term].get_all_parents()) if term in godag else []\n\nprop_rows = []\nfor _, row in tqdm(ensemble_df.iterrows(), total=len(ensemble_df), desc=\"Propagate\"):\n prop_rows.append({'id': row['id'], 'term': row['term'], 'score': row['score']})\n for anc in term_ancestors.get(row['term'], []):\n prop_rows.append({'id': row['id'], 'term': anc, 'score': row['score']})\n\nfinal_df = pd.DataFrame(prop_rows).groupby(['id', 'term'])['score'].max().reset_index()\nfinal_df = final_df[final_df['score'] >= 0.01]\nprint(f\"After propagation: {len(final_df):,}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Save\nfinal_df = final_df.sort_values(['id', 'score'], ascending=[True, False])\nfinal_df['score'] = final_df['score'].round(3)\nfinal_df[['id', 'term', 'score']].to_csv(f'{CFG.OUTPUT_DIR}/submission.tsv', sep='\\t', index=False, header=False)\n\nprint(f\"\\n{'='*60}\")\nprint(f\"SUBMISSION SAVED\")\nprint(f\"{'='*60}\")\nprint(f\"Predictions: {len(final_df):,}\")\nprint(f\"Proteins: {final_df['id'].nunique():,}\")\nprint(f\"GO terms: {final_df['term'].nunique():,}\")\nprint(f\"Coverage: {final_df['id'].nunique() / len(test_ids) * 100:.1f}%\")"}, {"cell_type": "markdown", "metadata": {}, "source": "
\n

Thank You!

\n

ESM2-15B + DIAMOND + Calibrated Ensemble

\n
"}]}