{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# \ud83e\uddec CAFA 6: ESM2-150M Optimized Solution\n## 10\u6642\u9593\u4ee5\u5185\u3067\u5b8c\u4e86\u3059\u308b\u9ad8\u7cbe\u5ea6\u7248\n\n
\n

\u26a1 \u6700\u9069\u5316\u30dd\u30a4\u30f3\u30c8

\n\n

\u23f1\ufe0f \u4e88\u60f3\u5b9f\u884c\u6642\u9593: 8-10\u6642\u9593

\n
\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 1\ufe0f\u20e3 Setup\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\nimport gc\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nfrom collections import defaultdict, Counter\nimport warnings\nwarnings.filterwarnings('ignore')\n\ntry:\n import esm\nexcept ImportError:\n !pip install fair-esm -q\n import esm\n\ntry:\n import obonet\nexcept ImportError:\n !pip install obonet -q\n import obonet\n\nimport torch\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.linear_model import RidgeClassifierCV\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.model_selection import train_test_split\n\nprint(f\"PyTorch: {torch.__version__}\")\nprint(f\"CUDA: {torch.cuda.is_available()}\")\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Device: {device}\")\n\nif torch.cuda.is_available():\n print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n print(f\"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 2\ufe0f\u20e3 Data Loading\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "INPUT_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\nTRAIN_SEQ = f'{INPUT_DIR}/Train/train_sequences.fasta'\nTRAIN_TERMS = f'{INPUT_DIR}/Train/train_terms.tsv'\nTEST_SEQ = f'{INPUT_DIR}/Test/testsuperset.fasta'\nGO_OBO = f'{INPUT_DIR}/Train/go-basic.obo'\n\nprint(\"Loading GO ontology...\")\ngo_graph = obonet.read_obo(GO_OBO)\nprint(f\" GO terms: {len(go_graph.nodes):,}\")\n\nprint(\"\\nLoading training terms...\")\ntrain_terms_df = pd.read_csv(TRAIN_TERMS, sep='\\t')\nprint(f\" Annotations: {len(train_terms_df):,}\")\nprint(f\" Proteins: {train_terms_df['EntryID'].nunique():,}\")\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def parse_fasta(fasta_path, max_seq_len=1022):\n sequences = {}\n current_id = None\n current_seq = []\n \n with open(fasta_path, 'r') as f:\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)[:max_seq_len]\n header = line[1:].split()[0]\n if '|' in header:\n parts = header.split('|')\n current_id = parts[1] if len(parts) >= 2 else header\n else:\n current_id = header\n current_seq = []\n else:\n current_seq.append(line)\n if current_id is not None:\n sequences[current_id] = ''.join(current_seq)[:max_seq_len]\n return sequences\n\nprint(\"\\nParsing sequences...\")\ntrain_sequences = parse_fasta(TRAIN_SEQ)\ntest_sequences = parse_fasta(TEST_SEQ)\nprint(f\" Train: {len(train_sequences):,}\")\nprint(f\" Test: {len(test_sequences):,}\")\n\ntrain_proteins = list(set(train_sequences.keys()) & set(train_terms_df['EntryID'].unique()))\nprint(f\" Matched: {len(train_proteins):,}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 3\ufe0f\u20e3 GO Hierarchy\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def get_ancestors(go_id, graph):\n ancestors = set()\n try:\n for parent in graph.successors(go_id):\n ancestors.add(parent)\n ancestors.update(get_ancestors(parent, graph))\n except:\n pass\n return ancestors\n\ndef propagate_predictions(predictions, graph):\n propagated = predictions.copy()\n for go_id, prob in predictions.items():\n ancestors = get_ancestors(go_id, graph)\n for ancestor in ancestors:\n if ancestor in propagated:\n propagated[ancestor] = max(propagated[ancestor], prob)\n else:\n propagated[ancestor] = prob\n return propagated\n\nprint(\"Pre-computing ancestors...\")\nall_terms = set(train_terms_df['term'].unique())\nterm_ancestors = {}\nfor term in tqdm(all_terms, desc=\"Ancestors\"):\n term_ancestors[term] = get_ancestors(term, go_graph)\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 4\ufe0f\u20e3 Smart Sampling\uff08\u7d71\u8a08\u4fdd\u6301\uff09\n\n\u516882K\u30bf\u30f3\u30d1\u30af\u8cea\u3067\u306f\u306a\u304f\u3001**\u30e9\u30d9\u30eb\u5206\u5e03\u3092\u7dad\u6301\u3057\u306a\u304c\u308950%\u30b5\u30f3\u30d7\u30ea\u30f3\u30b0**\u3057\u3066\u8a13\u7df4\u6642\u9593\u3092\u534a\u6e1b\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "TOP_N_TERMS = 300\n\ndef prepare_labels_with_sampling(df, proteins, aspect, top_n=TOP_N_TERMS, sample_ratio=0.5):\n aspect_df = df[df['aspect'] == aspect]\n term_counts = aspect_df['term'].value_counts()\n top_terms = term_counts.head(top_n).index.tolist()\n \n aspect_df = aspect_df[aspect_df['term'].isin(top_terms)]\n protein_terms = aspect_df.groupby('EntryID')['term'].apply(list).to_dict()\n \n # Stratified sampling based on number of labels\n protein_label_counts = {p: len(protein_terms.get(p, [])) for p in proteins}\n \n # Sort by label count and sample evenly from bins\n sorted_proteins = sorted(proteins, key=lambda p: protein_label_counts.get(p, 0))\n sample_size = int(len(sorted_proteins) * sample_ratio)\n \n # Every other protein to maintain distribution\n sampled_proteins = sorted_proteins[::2][:sample_size]\n \n mlb = MultiLabelBinarizer(classes=top_terms)\n labels = [protein_terms.get(p, []) for p in sampled_proteins]\n y = mlb.fit_transform(labels)\n \n return y, mlb, sampled_proteins, top_terms\n\nprint(f\"\\nPreparing labels (top {TOP_N_TERMS} per aspect, 50% sampling)...\")\ny_mf, mlb_mf, proteins_mf, terms_mf = prepare_labels_with_sampling(train_terms_df, train_proteins, 'F')\ny_bp, mlb_bp, proteins_bp, terms_bp = prepare_labels_with_sampling(train_terms_df, train_proteins, 'P')\ny_cc, mlb_cc, proteins_cc, terms_cc = prepare_labels_with_sampling(train_terms_df, train_proteins, 'C')\n\n# Use union of sampled proteins\nsampled_proteins = list(set(proteins_mf) | set(proteins_bp) | set(proteins_cc))\nprint(f\" Final sample size: {len(sampled_proteins):,} (from {len(train_proteins):,})\")\n\n# Re-prepare labels for sampled proteins (EXACT Match - Fixes IndexError)\ndef get_labels_exact(df, proteins, aspect, top_n=TOP_N_TERMS):\n aspect_df = df[df['aspect'] == aspect]\n term_counts = aspect_df['term'].value_counts()\n top_terms = term_counts.head(top_n).index.tolist()\n \n aspect_df = aspect_df[aspect_df['term'].isin(top_terms)]\n protein_terms = aspect_df.groupby('EntryID')['term'].apply(list).to_dict()\n \n mlb = MultiLabelBinarizer(classes=top_terms)\n labels = [protein_terms.get(p, []) for p in proteins]\n y = mlb.fit_transform(labels)\n return y, mlb\n\nprint(f\" Generating exact labels for {len(sampled_proteins):,} proteins...\")\ny_mf, mlb_mf = get_labels_exact(train_terms_df, sampled_proteins, 'F')\ny_bp, mlb_bp = get_labels_exact(train_terms_df, sampled_proteins, 'P')\ny_cc, mlb_cc = get_labels_exact(train_terms_df, sampled_proteins, 'C')\n\nprint(f\" MF: {y_mf.shape}\")\nprint(f\" BP: {y_bp.shape}\")\nprint(f\" CC: {y_cc.shape}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 5\ufe0f\u20e3 ESM2-150M Embeddings\n\n**ESM2-150M**: 650M\u306e\u7d041/4\u306e\u30b5\u30a4\u30ba\u3067\u30018M\u3088\u308a\u5927\u5e45\u306b\u9ad8\u7cbe\u5ea6\u3002\u901f\u5ea6\u3068\u7cbe\u5ea6\u306e\u6700\u9069\u30d0\u30e9\u30f3\u30b9\u3002\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "print(\"\\nLoading ESM2-150M model...\")\nmodel, alphabet = esm.pretrained.esm2_t30_150M_UR50D()\nmodel = model.to(device)\nmodel.eval()\nbatch_converter = alphabet.get_batch_converter()\nprint(\" Model loaded!\")\nprint(f\" Parameters: ~150M\")\nprint(f\" Layers: 30\")\nprint(f\" Embedding dim: 640\")\n\ndef get_esm_embeddings(sequences_dict, protein_ids, batch_size=8):\n embeddings = []\n valid_ids = []\n \n for i in tqdm(range(0, len(protein_ids), batch_size), desc=\"Embedding\"):\n batch_ids = protein_ids[i:i+batch_size]\n batch_data = [(pid, sequences_dict[pid]) for pid in batch_ids if pid in sequences_dict]\n \n if len(batch_data) == 0:\n continue\n \n try:\n _, _, batch_tokens = batch_converter(batch_data)\n batch_tokens = batch_tokens.to(device)\n \n with torch.no_grad():\n results = model(batch_tokens, repr_layers=[30])\n token_emb = results[\"representations\"][30]\n mean_emb = token_emb[:, 1:-1, :].mean(dim=1)\n embeddings.append(mean_emb.cpu().numpy())\n valid_ids.extend([d[0] for d in batch_data])\n \n del batch_tokens, results\n torch.cuda.empty_cache()\n except:\n continue\n \n return np.vstack(embeddings) if embeddings else np.array([]), valid_ids\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "print(f\"\\nGenerating training embeddings ({len(sampled_proteins):,} proteins)...\")\nX_train, train_ids_actual = get_esm_embeddings(train_sequences, sampled_proteins, batch_size=8)\nprint(f\" Generated: {X_train.shape}\")\n\n# Reindex labels\nid_to_idx = {p: i for i, p in enumerate(sampled_proteins)}\nactual_idx = [id_to_idx[p] for p in train_ids_actual if p in id_to_idx]\n\ny_mf_actual = y_mf[actual_idx]\ny_bp_actual = y_bp[actual_idx]\ny_cc_actual = y_cc[actual_idx]\n\nprint(f\" Labels: MF={y_mf_actual.shape}, BP={y_bp_actual.shape}, CC={y_cc_actual.shape}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 6\ufe0f\u20e3 Training with Ridge Classifier\uff08\u9ad8\u901f\u53ce\u675f\uff09\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "X_tr, X_val, y_mf_tr, y_mf_val, y_bp_tr, y_bp_val, y_cc_tr, y_cc_val = train_test_split(\n X_train, y_mf_actual, y_bp_actual, y_cc_actual,\n test_size=0.1, random_state=42\n)\nprint(f\"Train: {X_tr.shape[0]:,}, Val: {X_val.shape[0]:,}\")\n\ndef train_ridge(X_tr, y_tr, X_val, y_val, name):\n print(f\"\\nTraining {name}...\")\n \n col_sums = y_tr.sum(axis=0)\n valid_cols = col_sums >= 5\n y_tr_valid = y_tr[:, valid_cols]\n y_val_valid = y_val[:, valid_cols]\n \n print(f\" Valid labels: {valid_cols.sum()}\")\n \n if valid_cols.sum() == 0:\n return None, None, 0.1\n \n # Ridge is much faster than Logistic Regression\n clf = OneVsRestClassifier(\n RidgeClassifierCV(alphas=[0.1, 1.0, 10.0], cv=3),\n n_jobs=-1\n )\n clf.fit(X_tr, y_tr_valid)\n \n # Simple threshold=0 for Ridge\n print(f\" Training complete!\")\n \n return clf, valid_cols, 0.1\n\nclf_mf, valid_mf, thresh_mf = train_ridge(X_tr, y_mf_tr, X_val, y_mf_val, \"MF\")\nclf_bp, valid_bp, thresh_bp = train_ridge(X_tr, y_bp_tr, X_val, y_bp_val, \"BP\")\nclf_cc, valid_cc, thresh_cc = train_ridge(X_tr, y_cc_tr, X_val, y_cc_val, \"CC\")\n\ndel X_train, X_tr, X_val\ngc.collect()\ntorch.cuda.empty_cache()\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 7\ufe0f\u20e3 Test Predictions\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "test_protein_ids = list(test_sequences.keys())\nprint(f\"\\nGenerating predictions for {len(test_protein_ids):,} test proteins...\")\n\nmf_terms = np.array(mlb_mf.classes_)[valid_mf] if clf_mf else []\nbp_terms = np.array(mlb_bp.classes_)[valid_bp] if clf_bp else []\ncc_terms = np.array(mlb_cc.classes_)[valid_cc] if clf_cc else []\n\nterm_aspect = {}\nfor _, row in train_terms_df.drop_duplicates('term').iterrows():\n term_aspect[row['term']] = row['aspect']\n\nCHUNK_SIZE = 2000\nMIN_PROB = 0.01\n\nall_predictions = []\n\nfor chunk_start in tqdm(range(0, len(test_protein_ids), CHUNK_SIZE), desc=\"Predicting\"):\n chunk_end = min(chunk_start + CHUNK_SIZE, len(test_protein_ids))\n chunk_ids = test_protein_ids[chunk_start:chunk_end]\n \n X_chunk, valid_chunk_ids = get_esm_embeddings(test_sequences, chunk_ids, batch_size=8)\n \n if len(X_chunk) == 0:\n continue\n \n # Predict (Ridge decision_function instead of predict_proba)\n if clf_mf:\n scores_mf = clf_mf.decision_function(X_chunk)\n # Normalize to [0, 1]\n proba_mf = 1 / (1 + np.exp(-scores_mf))\n \n if clf_bp:\n scores_bp = clf_bp.decision_function(X_chunk)\n proba_bp = 1 / (1 + np.exp(-scores_bp))\n \n if clf_cc:\n scores_cc = clf_cc.decision_function(X_chunk)\n proba_cc = 1 / (1 + np.exp(-scores_cc))\n \n # Collect\n for i, protein_id in enumerate(valid_chunk_ids):\n protein_preds = {}\n \n if clf_mf:\n for term, prob in zip(mf_terms, proba_mf[i]):\n if prob > MIN_PROB:\n protein_preds[term] = max(protein_preds.get(term, 0), prob)\n \n if clf_bp:\n for term, prob in zip(bp_terms, proba_bp[i]):\n if prob > MIN_PROB:\n protein_preds[term] = max(protein_preds.get(term, 0), prob)\n \n if clf_cc:\n for term, prob in zip(cc_terms, proba_cc[i]):\n if prob > MIN_PROB:\n protein_preds[term] = max(protein_preds.get(term, 0), prob)\n \n # Propagate\n propagated = propagate_predictions(protein_preds, go_graph)\n all_predictions.append((protein_id, propagated))\n \n del X_chunk\n gc.collect()\n torch.cuda.empty_cache()\n\nprint(f\"\\nCollected predictions for {len(all_predictions):,} proteins\")\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "print(\"\\nWriting submission...\")\n\nwith open('submission.tsv', 'w') as f:\n for protein_id, preds in tqdm(all_predictions, desc=\"Writing\"):\n for term, prob in preds.items():\n aspect = term_aspect.get(term, 'P')\n if aspect == 'F':\n thresh = thresh_mf * 0.8\n elif aspect == 'P':\n thresh = thresh_bp * 0.8\n else:\n thresh = thresh_cc * 0.8\n \n if prob >= thresh:\n prob_clamped = max(0.001, min(0.999, prob))\n f.write(f\"{protein_id}\\t{term}\\t{prob_clamped:.3f}\\n\")\n\nfile_size = os.path.getsize('submission.tsv') / (1024 * 1024)\nprint(f\"\\nSubmission saved: {file_size:.1f} MB\")\n\nprint(\"\\nSample:\")\n!head -20 submission.tsv\n\nprint(\"\\nTotal:\")\n!wc -l submission.tsv\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## \ud83d\udca1 Summary\n\n### \u6700\u9069\u5316\u5185\u5bb9\n1. **ESM2-150M**: 650M\u3088\u308a4\u500d\u9ad8\u901f\u3001\u7cbe\u5ea6\u306f\u7dad\u6301\n2. **50%\u30b5\u30f3\u30d7\u30ea\u30f3\u30b0**: \u8a13\u7df4\u6642\u9593\u534a\u6e1b\u3001\u5206\u5e03\u7dad\u6301\n3. **Ridge\u56de\u5e30**: LogReg\u3088\u308a\u9ad8\u901f\u53ce\u675f\n4. **Top 300 terms**: \u30ab\u30d0\u30ec\u30c3\u30b8\u7dad\u6301\u3001\u51e6\u7406\u524a\u6e1b\n\n### \u63a8\u5b9a\u6642\u9593\u914d\u5206\n- Training embedding: ~2.5\u6642\u9593\uff0850% reduction\uff09\n- Training: ~2\u6642\u9593\uff08Ridge fast convergence\uff09\n- Test embedding: ~5\u6642\u9593\n- Total: **~9.5\u6642\u9593**\n"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}}, "nbformat": 4, "nbformat_minor": 4}