{"cells":[{"cell_type":"markdown","metadata":{},"source":"# 🧬 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### 📋 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\n### 🎯 Key Features\n- Pre-trained ESM-2 (facebook/esm2_t6_8M_UR50D - lightweight version)\n- Multi-label classification with sklearn\n- GO hierarchy-aware prediction\n"},{"cell_type":"markdown","metadata":{},"source":"## 1️⃣ 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 RidgeClassifierCV\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️⃣ Data Loading\n\nWe load data from the Kaggle competition input directory.\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'\nIA_FILE = f'{INPUT_DIR}/IA.tsv'\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():,}\")\nprint(f\"Aspects: {train_terms_df['aspect'].unique()}\")\n"},{"cell_type":"markdown","metadata":{},"source":"### Parse FASTA files\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    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                    # Truncate long sequences\n                    sequences[current_id] = seq[:max_seq_len]\n                current_id = line[1:].split()[0]\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\")\n\nprint(\"\\nParsing test sequences...\")\ntest_sequences = parse_fasta(TEST_SEQ)\nprint(f\"Loaded {len(test_sequences):,} test sequences\")\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=500):\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    valid_proteins = []\n    for p in proteins:\n        if p in protein_terms:\n            labels.append(protein_terms[p])\n            valid_proteins.append(p)\n        else:\n            labels.append([])\n            valid_proteins.append(p)\n    \n    y = mlb.fit_transform(labels)\n    return y, mlb, valid_proteins\n\n# Prepare labels for each aspect\nprint(\"\\nPreparing labels...\")\ny_mf, mlb_mf, _ = prepare_labels(train_terms_df, train_proteins, 'F', top_n=300)\ny_bp, mlb_bp, _ = prepare_labels(train_terms_df, train_proteins, 'P', top_n=300)\ny_cc, mlb_cc, _ = prepare_labels(train_terms_df, train_proteins, 'C', top_n=300)\n\nprint(f\"\\nLabel shapes: MF={y_mf.shape}, BP={y_bp.shape}, CC={y_cc.shape}\")\n"},{"cell_type":"markdown","metadata":{},"source":"## 3️⃣ Generate ESM-2 Embeddings\n\nWe use the lightweight ESM-2 model (8M parameters) for efficiency.\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=32):\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]\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            # Mean pooling over sequence length\n            token_embeddings = results[\"representations\"][6]\n            # Exclude BOS and EOS tokens\n            mean_embedding = token_embeddings[:, 1:-1, :].mean(dim=1)\n            embeddings.append(mean_embedding.cpu().numpy())\n        \n        # Clear GPU memory\n        del batch_tokens, results\n        torch.cuda.empty_cache()\n    \n    return np.vstack(embeddings)\n\n# Sample for faster training (full dataset may take too long)\nSAMPLE_SIZE = min(30000, len(train_proteins))\nif len(train_proteins) > SAMPLE_SIZE:\n    print(f\"\\nSampling {SAMPLE_SIZE:,} proteins for training...\")\n    np.random.seed(42)\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"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# 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\n# Free memory\ndel model\ngc.collect()\ntorch.cuda.empty_cache()\n"},{"cell_type":"markdown","metadata":{},"source":"## 4️⃣ Train Multi-Label Classifiers\n"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"from sklearn.linear_model import LogisticRegression\n\ndef 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 > 0\n    y_valid = y[:, valid_cols]\n    print(f\"  Valid labels (with positive samples): {valid_cols.sum()}\")\n    \n    # Simple logistic regression with OvR\n    clf = OneVsRestClassifier(\n        LogisticRegression(max_iter=500, 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\n# Train classifiers\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# Free training data\ndel X_train, y_mf_sample, y_bp_sample, y_cc_sample\ngc.collect()\n"},{"cell_type":"markdown","metadata":{},"source":"## 5️⃣ Generate Test Predictions\n"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Reload ESM model for test embeddings\nprint(\"Reloading ESM-2 model for test predictions...\")\nmodel, alphabet = esm.pretrained.esm2_t6_8M_UR50D()\nmodel = model.to(device)\nmodel.eval()\nbatch_converter = alphabet.get_batch_converter()\n\n# Generate test embeddings in batches\ntest_protein_ids = list(test_sequences.keys())\nprint(f\"\\nGenerating embeddings for {len(test_protein_ids):,} test proteins...\")\n\nBATCH_SIZE = 16\nCHUNK_SIZE = 5000  # Process in chunks to save memory\n\n# Open submission file\nwith open('submission.tsv', 'w') as f:\n    for chunk_start in tqdm(range(0, len(test_protein_ids), CHUNK_SIZE), desc=\"Processing chunks\"):\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 for chunk\n        X_chunk = get_esm_embeddings(test_sequences, chunk_ids, batch_size=BATCH_SIZE)\n        \n        # Predict probabilities\n        proba_mf = clf_mf.predict_proba(X_chunk)\n        proba_bp = clf_bp.predict_proba(X_chunk)\n        proba_cc = clf_cc.predict_proba(X_chunk)\n        \n        # Get term names\n        mf_terms = np.array(mlb_mf.classes_)[valid_mf]\n        bp_terms = np.array(mlb_bp.classes_)[valid_bp]\n        cc_terms = np.array(mlb_cc.classes_)[valid_cc]\n        \n        # Write predictions\n        for i, protein_id in enumerate(chunk_ids):\n            # MF predictions\n            for j, (term, prob) in enumerate(zip(mf_terms, proba_mf[i])):\n                if prob > 0.01:  # Threshold to reduce file size\n                    f.write(f\"{protein_id}\\t{term}\\t{prob:.3f}\\n\")\n            \n            # BP predictions\n            for j, (term, prob) in enumerate(zip(bp_terms, proba_bp[i])):\n                if prob > 0.01:\n                    f.write(f\"{protein_id}\\t{term}\\t{prob:.3f}\\n\")\n            \n            # CC predictions\n            for j, (term, prob) in enumerate(zip(cc_terms, proba_cc[i])):\n                if prob > 0.01:\n                    f.write(f\"{protein_id}\\t{term}\\t{prob:.3f}\\n\")\n        \n        # Free memory\n        del X_chunk, proba_mf, proba_bp, proba_cc\n        gc.collect()\n        torch.cuda.empty_cache()\n\nprint(\"\\nPredictions complete!\")\n"},{"cell_type":"markdown","metadata":{},"source":"## 6️⃣ Verify Submission\n"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"import os\n\n# Check file\nfile_size = os.path.getsize('submission.tsv') / (1024 * 1024)\nprint(f\"Submission file size: {file_size:.1f} MB\")\n\n# Sample\nprint(\"\\nSample predictions:\")\n!head -20 submission.tsv\n\n# Count predictions\nprint(\"\\nCounting predictions...\")\n!wc -l submission.tsv\n"},{"cell_type":"markdown","metadata":{},"source":"## 📊 Summary\n\nThis notebook implements:\n1. **ESM-2 Protein Language Model** for sequence embeddings\n2. **Multi-label classification** with Logistic Regression\n3. **Separate models** for MF, BP, CC subontologies\n4. **Efficient batch processing** for large test set\n\n### 🚀 Next Steps for Improvement\n- Use larger ESM model (esm2_t33_650M)\n- Ensemble multiple models\n- Add GO hierarchy propagation\n- Use threshold optimization on validation set\n- Include taxonomy features\n"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.8.5"}},"nbformat":4,"nbformat_minor":4}