{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# \ud83e\uddec CAFA 6 Protein Function Prediction\n## ESM-2 Protein Language Model + Multi-Label Classification\n\n<div style=\"background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 20px; border-radius: 10px; color: white;\">\n<h3>\ud83d\udccc Quick Summary</h3>\n<ul>\n<li><b>Approach</b>: ESM-2 embeddings \u2192 Logistic Regression multi-label classifiers</li>\n<li><b>Model</b>: facebook/esm2_t6_8M_UR50D (8M parameters, efficient)</li>\n<li><b>Targets</b>: 3 subontologies (MF, BP, CC) with top 200 GO terms each</li>\n<li><b>Key Insight</b>: Protein Language Models capture evolutionary patterns that correlate with function</li>\n</ul>\n</div>\n\n---\n\n## \ud83d\udccb Table of Contents\n1. [\ud83c\udfaf Problem Overview](#problem-overview)\n2. [\ud83d\udd27 Environment Setup](#environment-setup)\n3. [\ud83d\udcca Data Loading & EDA](#data-loading)\n4. [\ud83e\udde0 ESM-2 Embeddings](#esm-embeddings)\n5. [\ud83e\udd16 Model Training](#model-training)\n6. [\ud83d\udcc8 Predictions](#predictions)\n7. [\ud83d\udca1 Insights & Conclusions](#insights)\n\n---\n\n<a id=\"problem-overview\"></a>\n## \ud83c\udfaf Problem Overview\n\n### What is CAFA?\n**CAFA (Critical Assessment of Functional Annotation)** is a community-driven challenge to predict the biological function of proteins. This is CAFA 6, continuing a tradition since 2010.\n\n### Why is this important?\n- **~200 million proteins** in UniProt, but only a fraction have experimentally verified functions\n- Understanding protein function is crucial for **drug discovery** and **disease understanding**\n- Computational prediction can guide expensive wet-lab experiments\n\n### Gene Ontology (GO)\nProteins are annotated with GO terms across 3 subontologies:\n| Subontology | Description | Example |\n|-------------|-------------|---------|\n| **MF** (Molecular Function) | What the protein does at molecular level | `GO:0003677` (DNA binding) |\n| **BP** (Biological Process) | Larger biological objectives | `GO:0006915` (apoptotic process) |\n| **CC** (Cellular Component) | Where in the cell | `GO:0005634` (nucleus) |\n\n### Evaluation\n- **Metric**: Mean of max F1-scores across MF, BP, CC\n- **Weighting**: Information Accretion (IA) - rarer terms weighted more\n- **Prospective**: Final evaluation uses newly published experimental data\n"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n<a id=\"environment-setup\"></a>\n## \ud83d\udd27 Environment Setup\n\nWe use **ESM-2** (Evolutionary Scale Modeling), a protein language model from Meta AI trained on millions of protein sequences.\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\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Install ESM if not available\ntry:\n    import esm\n    print(\"\u2705 ESM already installed\")\nexcept ImportError:\n    print(\"\ud83d\udce6 Installing ESM...\")\n    !pip install fair-esm -q\n    import esm\n\nimport torch\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.multiclass import OneVsRestClassifier\n\nprint(f\"\ud83d\udd37 PyTorch version: {torch.__version__}\")\nprint(f\"\ud83d\udd37 CUDA available: {torch.cuda.is_available()}\")\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"\ud83d\udd37 Using device: {device}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n<a id=\"data-loading\"></a>\n## \ud83d\udcca Data Loading & Exploratory Data Analysis\n\nLet's understand our data before modeling.\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Kaggle paths\nINPUT_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'\n\n# Load training terms\nprint(\"\ud83d\udcc2 Loading training terms...\")\ntrain_terms_df = pd.read_csv(TRAIN_TERMS, sep='\\t')\n\nprint(f\"\\n\ud83d\udcca Dataset Statistics:\")\nprint(f\"  \u2022 Total annotations: {len(train_terms_df):,}\")\nprint(f\"  \u2022 Unique proteins: {train_terms_df['EntryID'].nunique():,}\")\nprint(f\"  \u2022 Unique GO terms: {train_terms_df['term'].nunique():,}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "### \ud83d\udcc8 Distribution of Annotations\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import matplotlib.pyplot as plt\n\nfig, axes = plt.subplots(1, 3, figsize=(15, 4))\n\n# 1. Annotations per aspect\naspect_counts = train_terms_df['aspect'].value_counts()\ncolors = {'F': '#2ecc71', 'P': '#3498db', 'C': '#e74c3c'}\naspect_counts.plot(kind='bar', ax=axes[0], color=[colors[x] for x in aspect_counts.index])\naxes[0].set_title('Annotations by Subontology', fontsize=12, fontweight='bold')\naxes[0].set_xlabel('Aspect')\naxes[0].set_ylabel('Count')\naxes[0].tick_params(axis='x', rotation=0)\n\n# 2. Annotations per protein\nannotations_per_protein = train_terms_df.groupby('EntryID').size()\naxes[1].hist(annotations_per_protein, bins=50, color='#9b59b6', edgecolor='white')\naxes[1].set_title('Annotations per Protein', fontsize=12, fontweight='bold')\naxes[1].set_xlabel('Number of GO terms')\naxes[1].set_ylabel('Number of proteins')\naxes[1].axvline(annotations_per_protein.median(), color='red', linestyle='--', label=f'Median: {annotations_per_protein.median():.0f}')\naxes[1].legend()\n\n# 3. Top 10 most common GO terms\ntop_terms = train_terms_df['term'].value_counts().head(10)\ntop_terms.plot(kind='barh', ax=axes[2], color='#1abc9c')\naxes[2].set_title('Top 10 Most Common GO Terms', fontsize=12, fontweight='bold')\naxes[2].set_xlabel('Count')\naxes[2].invert_yaxis()\n\nplt.tight_layout()\nplt.savefig('eda_plots.png', dpi=100, bbox_inches='tight')\nplt.show()\n\nprint(f\"\\n\ud83d\udca1 Insight: The median protein has {annotations_per_protein.median():.0f} GO term annotations.\")\nprint(f\"\ud83d\udca1 Insight: Some proteins have >50 annotations - multi-functional proteins!\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "### \ud83e\uddec Parse Protein Sequences\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def parse_fasta(fasta_path, max_seq_len=1000):\n    \"\"\"Parse FASTA file, extracting UniProt accession from sp|ID|NAME format\"\"\"\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                \n                # Parse ID - handle sp|A0A0C5B5G6|MOTSC_HUMAN format\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        \n        if current_id is not None:\n            sequences[current_id] = ''.join(current_seq)[:max_seq_len]\n    \n    return sequences\n\nprint(\"\ud83e\uddec Parsing training sequences...\")\ntrain_sequences = parse_fasta(TRAIN_SEQ)\nprint(f\"  \u2705 Loaded {len(train_sequences):,} training sequences\")\n\nprint(\"\\n\ud83e\uddec Parsing test sequences...\")\ntest_sequences = parse_fasta(TEST_SEQ)\nprint(f\"  \u2705 Loaded {len(test_sequences):,} test sequences\")\n\n# Sequence length distribution\nseq_lengths = [len(s) for s in train_sequences.values()]\nprint(f\"\\n\ud83d\udccf Sequence length stats:\")\nprint(f\"  \u2022 Min: {min(seq_lengths)}, Max: {max(seq_lengths)}, Median: {np.median(seq_lengths):.0f}\")\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Verify overlap between sequences and labels\ntrain_seq_ids = set(train_sequences.keys())\nlabel_ids = set(train_terms_df['EntryID'].unique())\noverlap = train_seq_ids & label_ids\n\nprint(f\"\\n\ud83d\udd17 ID Matching:\")\nprint(f\"  \u2022 Sequence IDs: {len(train_seq_ids):,}\")\nprint(f\"  \u2022 Label IDs: {len(label_ids):,}\")\nprint(f\"  \u2022 Overlap: {len(overlap):,} ({100*len(overlap)/len(label_ids):.1f}%)\")\n\ntrain_proteins = list(overlap)\n"}, {"cell_type": "markdown", "metadata": {}, "source": "### \ud83c\udff7\ufe0f Prepare Multi-Label Targets\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def prepare_labels(df, proteins, aspect, top_n=200):\n    \"\"\"Prepare multi-label targets for a specific aspect\"\"\"\n    aspect_df = df[df['aspect'] == aspect]\n    \n    # Get top N most frequent terms\n    term_counts = aspect_df['term'].value_counts()\n    top_terms = term_counts.head(top_n).index.tolist()\n    \n    # Group by protein\n    protein_terms = aspect_df.groupby('EntryID')['term'].apply(list).to_dict()\n    \n    # Create label matrix\n    mlb = MultiLabelBinarizer(classes=top_terms)\n    labels = [protein_terms.get(p, []) for p in proteins]\n    y = mlb.fit_transform(labels)\n    \n    return y, mlb\n\nprint(\"\ud83c\udff7\ufe0f Preparing labels for each subontology...\")\ny_mf, mlb_mf = prepare_labels(train_terms_df, train_proteins, 'F', top_n=200)\ny_bp, mlb_bp = prepare_labels(train_terms_df, train_proteins, 'P', top_n=200)\ny_cc, mlb_cc = prepare_labels(train_terms_df, train_proteins, 'C', top_n=200)\n\nprint(f\"  \u2022 MF: {y_mf.shape} (avg {y_mf.sum(axis=1).mean():.1f} labels/protein)\")\nprint(f\"  \u2022 BP: {y_bp.shape} (avg {y_bp.sum(axis=1).mean():.1f} labels/protein)\")\nprint(f\"  \u2022 CC: {y_cc.shape} (avg {y_cc.sum(axis=1).mean():.1f} labels/protein)\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n<a id=\"esm-embeddings\"></a>\n## \ud83e\udde0 ESM-2 Protein Language Model\n\n### Why ESM-2?\nESM-2 is a **transformer-based language model** trained on 250 million protein sequences. It learns:\n- Evolutionary patterns from sequence conservation\n- Structural information encoded in sequences\n- Functional motifs and domains\n\nWe use the **8M parameter version** for efficiency, but larger models (650M, 3B) could improve accuracy.\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load ESM model\nprint(\"\ud83e\udde0 Loading ESM-2 model (esm2_t6_8M_UR50D)...\")\nmodel, alphabet = esm.pretrained.esm2_t6_8M_UR50D()\nmodel = model.to(device)\nmodel.eval()\nbatch_converter = alphabet.get_batch_converter()\nprint(\"  \u2705 Model loaded!\")\nprint(f\"  \u2022 Parameters: ~8M\")\nprint(f\"  \u2022 Embedding dim: 320\")\n\ndef get_esm_embeddings(sequences_dict, protein_ids, batch_size=16):\n    \"\"\"Generate ESM-2 embeddings via mean pooling\"\"\"\n    embeddings = []\n    \n    for i in tqdm(range(0, len(protein_ids), batch_size), desc=\"\ud83d\udd04 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        _, _, 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=[6])\n            token_emb = results[\"representations\"][6]\n            # Mean pooling (excluding BOS/EOS tokens)\n            mean_emb = token_emb[:, 1:-1, :].mean(dim=1)\n            embeddings.append(mean_emb.cpu().numpy())\n        \n        del batch_tokens, results\n        torch.cuda.empty_cache()\n    \n    return np.vstack(embeddings) if embeddings else np.array([])\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Sample for training\nSAMPLE_SIZE = min(20000, len(train_proteins))\nnp.random.seed(42)\n\nprint(f\"\\n\ud83d\udcca Sampling {SAMPLE_SIZE:,} proteins for training...\")\nif len(train_proteins) > SAMPLE_SIZE:\n    sample_idx = np.random.choice(len(train_proteins), SAMPLE_SIZE, replace=False)\n    train_proteins_sample = [train_proteins[i] for i in sample_idx]\n    y_mf_sample = y_mf[sample_idx]\n    y_bp_sample = y_bp[sample_idx]\n    y_cc_sample = y_cc[sample_idx]\nelse:\n    train_proteins_sample = train_proteins\n    y_mf_sample, y_bp_sample, y_cc_sample = y_mf, y_bp, y_cc\n\nprint(f\"  \u2705 Training on {len(train_proteins_sample):,} proteins\")\n\n# Generate embeddings\nprint(\"\\n\ud83e\udde0 Generating training embeddings...\")\nX_train = get_esm_embeddings(train_sequences, train_proteins_sample, batch_size=16)\nprint(f\"  \u2705 Embedding shape: {X_train.shape}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n<a id=\"model-training\"></a>\n## \ud83e\udd16 Model Training\n\nWe train **One-vs-Rest Logistic Regression** classifiers for each subontology. This is a strong baseline that:\n- Handles multi-label prediction naturally\n- Is fast to train\n- Provides probability estimates for ranking\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def train_classifier(X, y, name):\n    \"\"\"Train a multi-label classifier\"\"\"\n    print(f\"\\n\ud83c\udfaf Training {name} classifier...\")\n    \n    # Remove terms with <10 positive samples\n    col_sums = y.sum(axis=0)\n    valid_cols = col_sums >= 10\n    y_valid = y[:, valid_cols]\n    \n    print(f\"  \u2022 Input: {X.shape[0]:,} samples, {y.shape[1]} labels\")\n    print(f\"  \u2022 Valid labels (\u226510 samples): {valid_cols.sum()}\")\n    \n    if y_valid.shape[1] == 0:\n        return None, valid_cols\n    \n    clf = OneVsRestClassifier(\n        LogisticRegression(max_iter=300, C=1.0, solver='lbfgs', n_jobs=-1),\n        n_jobs=-1\n    )\n    clf.fit(X, y_valid)\n    print(f\"  \u2705 Training complete!\")\n    \n    return clf, valid_cols\n\nclf_mf, valid_mf = train_classifier(X_train, y_mf_sample, \"MF (Molecular Function)\")\nclf_bp, valid_bp = train_classifier(X_train, y_bp_sample, \"BP (Biological Process)\")\nclf_cc, valid_cc = train_classifier(X_train, y_cc_sample, \"CC (Cellular Component)\")\n\n# Clean up\ndel X_train\ngc.collect()\nprint(\"\\n\ud83e\uddf9 Memory cleaned!\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n<a id=\"predictions\"></a>\n## \ud83d\udcc8 Generate Test Predictions\n\nNow we generate predictions for all test proteins in chunks to manage memory.\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "test_protein_ids = list(test_sequences.keys())\nprint(f\"\\n\ud83c\udfaf Generating predictions for {len(test_protein_ids):,} test proteins...\")\n\nCHUNK_SIZE = 3000\nTHRESHOLD = 0.05\n\nwith open('submission.tsv', 'w') as f:\n    for chunk_start in tqdm(range(0, len(test_protein_ids), CHUNK_SIZE), desc=\"\ud83d\udcdd 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        # Generate embeddings\n        X_chunk = get_esm_embeddings(test_sequences, chunk_ids, batch_size=16)\n        \n        if len(X_chunk) == 0:\n            continue\n        \n        # Get term names\n        mf_terms = np.array(mlb_mf.classes_)[valid_mf] if clf_mf else []\n        bp_terms = np.array(mlb_bp.classes_)[valid_bp] if clf_bp else []\n        cc_terms = np.array(mlb_cc.classes_)[valid_cc] if clf_cc else []\n        \n        # Predict\n        proba_mf = clf_mf.predict_proba(X_chunk) if clf_mf else None\n        proba_bp = clf_bp.predict_proba(X_chunk) if clf_bp else None\n        proba_cc = clf_cc.predict_proba(X_chunk) if clf_cc else None\n        \n        # Write predictions above threshold\n        for i, protein_id in enumerate(chunk_ids):\n            if proba_mf is not None:\n                for term, prob in zip(mf_terms, proba_mf[i]):\n                    if prob > THRESHOLD:\n                        f.write(f\"{protein_id}\\t{term}\\t{prob:.3f}\\n\")\n            if proba_bp is not None:\n                for term, prob in zip(bp_terms, proba_bp[i]):\n                    if prob > THRESHOLD:\n                        f.write(f\"{protein_id}\\t{term}\\t{prob:.3f}\\n\")\n            if proba_cc is not None:\n                for term, prob in zip(cc_terms, proba_cc[i]):\n                    if prob > THRESHOLD:\n                        f.write(f\"{protein_id}\\t{term}\\t{prob:.3f}\\n\")\n        \n        del X_chunk\n        gc.collect()\n        torch.cuda.empty_cache()\n\nprint(\"\\n\u2705 Predictions complete!\")\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\n\nfile_size = os.path.getsize('submission.tsv') / (1024 * 1024)\nprint(f\"\\n\ud83d\udcc1 Submission file: {file_size:.1f} MB\")\n\nprint(\"\\n\ud83d\udccb Sample predictions:\")\n!head -15 submission.tsv\n\nprint(\"\\n\ud83d\udcca Total predictions:\")\n!wc -l submission.tsv\n"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n<a id=\"insights\"></a>\n## \ud83d\udca1 Insights & Conclusions\n\n### What We Learned\n\n1. **ESM-2 captures functional information**: Protein language models encode evolutionary patterns that correlate with function, even without explicit structural or domain information.\n\n2. **Multi-label challenge**: Proteins have multiple functions - our model handles this with One-vs-Rest classification, predicting independent probabilities for each GO term.\n\n3. **Scalability matters**: With 200K+ test proteins, efficient batch processing is essential.\n\n### Limitations & Future Improvements\n\n| Current Limitation | Potential Improvement |\n|-------------------|----------------------|\n| Small ESM model (8M params) | Use ESM2-650M or ESM2-3B |\n| Logistic Regression | Neural network classifier |\n| Top 200 terms only | Predict full GO hierarchy |\n| No hierarchy propagation | Add GO DAG propagation |\n| No taxonomy features | Include species information |\n\n### Next Steps for Higher Score\n\n1. **Ensemble**: Combine multiple ESM models (6-layer, 12-layer, 33-layer)\n2. **GO Hierarchy**: Propagate predictions up the ontology DAG\n3. **Pseudo-labeling**: Use confident predictions on test set\n4. **K-fold cross-validation**: Better estimate of generalization\n5. **Threshold optimization**: Maximize F1 on validation set\n\n---\n\n## \ud83d\udcda References\n\n- [ESM-2 Paper](https://www.biorxiv.org/content/10.1101/2022.07.20.500902v1): Meta AI's protein language model\n- [Gene Ontology](http://geneontology.org/): The GO Consortium\n- [CAFA Challenge](https://www.biofunctionprediction.org/cafa/): Official CAFA website\n\n---\n\n<div style=\"text-align: center; padding: 20px; background: #f8f9fa; border-radius: 10px;\">\n<h3>\ud83d\ude4f Thank you for reading!</h3>\n<p>If this notebook was helpful, please consider giving it an <b>upvote</b>! \ud83d\udc4d</p>\n<p>Questions or suggestions? Leave a comment below!</p>\n</div>\n"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}}, "nbformat": 4, "nbformat_minor": 4}