{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# \ud83e\uddec CAFA 6 Protein Function Prediction\n## ESM-2 Embeddings + Multi-Label Classification\n\nThis notebook uses **ESM-2** (Evolutionary Scale Modeling) protein language model to generate embeddings for protein sequences, then trains a multi-label classifier to predict GO terms.\n\n### \ud83d\udccb Strategy\n1. Load protein sequences and GO term labels\n2. Generate ESM-2 embeddings for all proteins\n3. Train separate classifiers for MF, BP, CC subontologies\n4. Generate predictions with confidence scores\n5. Create submission file\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 1\ufe0f\u20e3 Environment 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\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Install ESM if not available\ntry:\n    import esm\n    print(\"ESM already installed\")\nexcept ImportError:\n    print(\"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\nfrom sklearn.model_selection import train_test_split\n\nprint(f\"PyTorch version: {torch.__version__}\")\nprint(f\"CUDA available: {torch.cuda.is_available()}\")\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 2\ufe0f\u20e3 Data Loading\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(\"Loading training terms...\")\ntrain_terms_df = pd.read_csv(TRAIN_TERMS, sep='\\t')\nprint(f\"Total annotations: {len(train_terms_df):,}\")\nprint(f\"Unique proteins: {train_terms_df['EntryID'].nunique():,}\")\nprint(f\"Unique GO terms: {train_terms_df['term'].nunique():,}\")\n\n# Sample EntryIDs\nprint(\"\\nSample EntryIDs from train_terms:\")\nprint(train_terms_df['EntryID'].head(10).tolist())\n"}, {"cell_type": "markdown", "metadata": {}, "source": "### Parse FASTA files (Fixed ID parsing)\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def parse_fasta(fasta_path, max_seq_len=1000):\n    \"\"\"Parse FASTA file and return dict of {id: sequence}\n    \n    Handles formats like:\n    - >sp|A0A0C5B5G6|MOTSC_HUMAN ...  -> extracts A0A0C5B5G6\n    - >A0A0C5B5G6 ...                  -> extracts A0A0C5B5G6\n    \"\"\"\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                    seq = ''.join(current_seq)\n                    sequences[current_id] = seq[:max_seq_len]\n                \n                # Parse ID - handle sp|ID|NAME format\n                header = line[1:].split()[0]  # Get first part before space\n                if '|' in header:\n                    # sp|A0A0C5B5G6|MOTSC_HUMAN -> A0A0C5B5G6\n                    parts = header.split('|')\n                    if len(parts) >= 2:\n                        current_id = parts[1]\n                    else:\n                        current_id = header\n                else:\n                    current_id = header\n                current_seq = []\n            else:\n                current_seq.append(line)\n        \n        # Last sequence\n        if current_id is not None:\n            seq = ''.join(current_seq)\n            sequences[current_id] = seq[:max_seq_len]\n    \n    return sequences\n\nprint(\"Parsing training sequences...\")\ntrain_sequences = parse_fasta(TRAIN_SEQ)\nprint(f\"Loaded {len(train_sequences):,} training sequences\")\nprint(f\"Sample IDs: {list(train_sequences.keys())[:5]}\")\n\nprint(\"\\nParsing test sequences...\")\ntest_sequences = parse_fasta(TEST_SEQ)\nprint(f\"Loaded {len(test_sequences):,} test sequences\")\nprint(f\"Sample IDs: {list(test_sequences.keys())[:5]}\")\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\"\\nSequence IDs: {len(train_seq_ids):,}\")\nprint(f\"Label IDs: {len(label_ids):,}\")\nprint(f\"Overlap: {len(overlap):,}\")\n\nif len(overlap) == 0:\n    print(\"\\n\u26a0\ufe0f WARNING: No overlap! Checking ID formats...\")\n    print(f\"Sample seq IDs: {list(train_seq_ids)[:3]}\")\n    print(f\"Sample label IDs: {list(label_ids)[:3]}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "### Prepare multi-label targets\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Get proteins that have both sequences and labels\ntrain_proteins = list(set(train_sequences.keys()) & set(train_terms_df['EntryID'].unique()))\nprint(f\"Proteins with both sequences and labels: {len(train_proteins):,}\")\n\n# Create multi-label format for each aspect\ndef prepare_labels(df, proteins, aspect, top_n=300):\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    print(f\"  Aspect {aspect}: Using top {len(top_terms)} terms\")\n    \n    # Filter to top terms only\n    aspect_df = aspect_df[aspect_df['term'].isin(top_terms)]\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 = []\n    for p in proteins:\n        if p in protein_terms:\n            labels.append(protein_terms[p])\n        else:\n            labels.append([])\n    \n    y = mlb.fit_transform(labels)\n    return y, mlb\n\n# Prepare labels for each aspect\nprint(\"\\nPreparing labels...\")\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\"\\nLabel shapes: MF={y_mf.shape}, BP={y_bp.shape}, CC={y_cc.shape}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 3\ufe0f\u20e3 Generate ESM-2 Embeddings\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load ESM model\nprint(\"Loading ESM-2 model...\")\nmodel, alphabet = esm.pretrained.esm2_t6_8M_UR50D()\nmodel = model.to(device)\nmodel.eval()\nbatch_converter = alphabet.get_batch_converter()\nprint(\"Model loaded!\")\n\ndef get_esm_embeddings(sequences_dict, protein_ids, batch_size=16):\n    \"\"\"Generate ESM-2 embeddings for proteins\"\"\"\n    embeddings = []\n    \n    for i in tqdm(range(0, len(protein_ids), batch_size), desc=\"Generating embeddings\"):\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_embeddings = results[\"representations\"][6]\n            # Mean pooling (exclude BOS/EOS)\n            mean_embedding = token_embeddings[:, 1:-1, :].mean(dim=1)\n            embeddings.append(mean_embedding.cpu().numpy())\n        \n        del batch_tokens, results\n        torch.cuda.empty_cache()\n    \n    if len(embeddings) == 0:\n        raise ValueError(\"No embeddings generated!\")\n    return np.vstack(embeddings)\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Sample for faster training\nSAMPLE_SIZE = min(20000, len(train_proteins))\nprint(f\"\\nSampling {SAMPLE_SIZE:,} proteins for training...\")\nnp.random.seed(42)\n\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_mf\n    y_bp_sample = y_bp\n    y_cc_sample = y_cc\n\nprint(f\"Training on {len(train_proteins_sample):,} proteins\")\n\n# Generate training embeddings\nprint(\"\\nGenerating training embeddings...\")\nX_train = get_esm_embeddings(train_sequences, train_proteins_sample, batch_size=16)\nprint(f\"Training embeddings shape: {X_train.shape}\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 4\ufe0f\u20e3 Train Multi-Label Classifiers\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\"\\nTraining {name} classifier...\")\n    print(f\"  X shape: {X.shape}, y shape: {y.shape}\")\n    \n    # Remove columns with no positive samples\n    col_sums = y.sum(axis=0)\n    valid_cols = col_sums > 5  # At least 5 samples\n    y_valid = y[:, valid_cols]\n    print(f\"  Valid labels (with >5 samples): {valid_cols.sum()}\")\n    \n    if y_valid.shape[1] == 0:\n        print(\"  No valid labels!\")\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\"  Training complete!\")\n    \n    return clf, valid_cols\n\nclf_mf, valid_mf = train_classifier(X_train, y_mf_sample, \"MF\")\nclf_bp, valid_bp = train_classifier(X_train, y_bp_sample, \"BP\")\nclf_cc, valid_cc = train_classifier(X_train, y_cc_sample, \"CC\")\n\ndel X_train\ngc.collect()\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 5\ufe0f\u20e3 Generate Test Predictions\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Generate test embeddings in chunks\ntest_protein_ids = list(test_sequences.keys())\nprint(f\"\\nGenerating predictions for {len(test_protein_ids):,} test proteins...\")\n\nCHUNK_SIZE = 3000\nTHRESHOLD = 0.05  # Only keep predictions above this\n\nwith open('submission.tsv', 'w') as f:\n    for chunk_start in tqdm(range(0, len(test_protein_ids), CHUNK_SIZE), desc=\"Processing\"):\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        # 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        if clf_mf:\n            proba_mf = clf_mf.predict_proba(X_chunk)\n        if clf_bp:\n            proba_bp = clf_bp.predict_proba(X_chunk)\n        if clf_cc:\n            proba_cc = clf_cc.predict_proba(X_chunk)\n        \n        # Write predictions\n        for i, protein_id in enumerate(chunk_ids):\n            if clf_mf:\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 clf_bp:\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 clf_cc:\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(\"\\nPredictions complete!\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## 6\ufe0f\u20e3 Verify Submission\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\n\nfile_size = os.path.getsize('submission.tsv') / (1024 * 1024)\nprint(f\"Submission file size: {file_size:.1f} MB\")\n\nprint(\"\\nSample predictions:\")\n!head -20 submission.tsv\n\nprint(\"\\nTotal predictions:\")\n!wc -l submission.tsv\n"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}}, "nbformat": 4, "nbformat_minor": 4}