{"cells":[{"cell_type":"markdown","metadata":{},"source":"# CAFA-6: LightGBM + ESM2 Embeddings\n\nBased on CAFA-5 top solutions approach:\n- ESM2 embeddings (1280 dim)\n- LightGBM gradient boosting\n- GO hierarchy propagation"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"import os\nimport gc\nimport numpy as np\nimport pandas as pd\nimport lightgbm as lgb\nfrom collections import defaultdict\nfrom tqdm.auto import tqdm\nfrom sklearn.model_selection import train_test_split\nimport warnings\nwarnings.filterwarnings('ignore')\n\nprint('CAFA-6 LightGBM + ESM2')"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Configuration\nDATA_PATH = '/kaggle/input/cafa-6-protein-function-prediction'\nEMB_PATH = '/kaggle/input/cafa6-protein-embeddings-esm2'\nOBO_PATH = f'{DATA_PATH}/Train/go-basic.obo'\n\n# LightGBM params optimized for protein function prediction\nLGBM_PARAMS = {\n    'objective': 'binary',\n    'metric': 'binary_logloss',\n    'boosting_type': 'gbdt',\n    'num_leaves': 63,\n    'max_depth': 8,\n    'learning_rate': 0.05,\n    'feature_fraction': 0.8,\n    'bagging_fraction': 0.8,\n    'bagging_freq': 5,\n    'min_child_samples': 20,\n    'verbose': -1,\n    'n_jobs': -1,\n    'random_state': 42\n}\n\nN_BOOST_ROUNDS = 200\nTOP_N_TERMS = 2000  # Top GO terms to train\nPRED_THRESHOLD = 0.01\n\nprint(f'TOP_N_TERMS: {TOP_N_TERMS}')\nprint(f'N_BOOST_ROUNDS: {N_BOOST_ROUNDS}')"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Load ESM2 embeddings\nprint('[1/6] Loading ESM2 embeddings...')\n\n# Train embeddings\ntrain_emb = np.load(f'{EMB_PATH}/train_embeds.npy')\ntrain_ids = np.load(f'{EMB_PATH}/train_ids.npy', allow_pickle=True)\nprint(f'Train embeddings: {train_emb.shape}')\n\n# Test embeddings\ntest_emb = np.load(f'{EMB_PATH}/test_embeds.npy')\ntest_ids = np.load(f'{EMB_PATH}/test_ids.npy', allow_pickle=True)\nprint(f'Test embeddings: {test_emb.shape}')\n\n# Create ID to index mapping\ntrain_id2idx = {pid: i for i, pid in enumerate(train_ids)}\ntest_id2idx = {pid: i for i, pid in enumerate(test_ids)}"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Load training labels\nprint('[2/6] Loading training labels...')\n\ntrain_terms = pd.read_csv(f'{DATA_PATH}/Train/train_terms.tsv', sep='\\t')\nprint(f'Total annotations: {len(train_terms):,}')\n\n# Get top N most frequent GO terms\nterm_counts = train_terms['term'].value_counts()\ntop_terms = term_counts.head(TOP_N_TERMS).index.tolist()\nprint(f'Training on top {len(top_terms)} GO terms')\n\n# Filter to top terms only\ntrain_terms_filtered = train_terms[train_terms['term'].isin(top_terms)]\nprint(f'Filtered annotations: {len(train_terms_filtered):,}')\n\n# Create term to index mapping\nterm2idx = {t: i for i, t in enumerate(top_terms)}\nidx2term = {i: t for t, i in term2idx.items()}"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Build label matrix\nprint('[3/6] Building label matrix...')\n\n# Find proteins that are in both embeddings and have labels\nlabeled_proteins = set(train_terms_filtered['EntryID'].unique())\nembedded_proteins = set(train_ids)\nvalid_proteins = list(labeled_proteins & embedded_proteins)\nprint(f'Valid proteins (have embeddings + labels): {len(valid_proteins):,}')\n\n# Create ordered list of valid proteins\nvalid_protein_idx = [train_id2idx[p] for p in valid_proteins]\nX_train = train_emb[valid_protein_idx]\nprint(f'X_train shape: {X_train.shape}')\n\n# Build Y matrix\nprotein2idx = {p: i for i, p in enumerate(valid_proteins)}\nY_train = np.zeros((len(valid_proteins), len(top_terms)), dtype=np.float32)\n\nfor row in tqdm(train_terms_filtered.itertuples(), total=len(train_terms_filtered), desc='Building Y'):\n    if row.EntryID in protein2idx and row.term in term2idx:\n        Y_train[protein2idx[row.EntryID], term2idx[row.term]] = 1.0\n\nprint(f'Y_train shape: {Y_train.shape}')\nprint(f'Label density: {Y_train.mean():.4%}')\n\n# Free memory\ndel train_emb\ngc.collect()"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Train LightGBM models - One per GO term\nprint('[4/6] Training LightGBM models...')\n\nmodels = {}\nskipped = 0\n\n# Train/val split for early stopping\nX_tr, X_val, Y_tr, Y_val = train_test_split(X_train, Y_train, test_size=0.1, random_state=42)\n\nfor term_idx in tqdm(range(len(top_terms)), desc='Training'):\n    term = idx2term[term_idx]\n    y_tr = Y_tr[:, term_idx]\n    y_val = Y_val[:, term_idx]\n    \n    # Skip terms with too few positives\n    if y_tr.sum() < 10:\n        skipped += 1\n        continue\n    \n    # Create datasets\n    train_data = lgb.Dataset(X_tr, label=y_tr)\n    val_data = lgb.Dataset(X_val, label=y_val, reference=train_data)\n    \n    # Train with early stopping\n    model = lgb.train(\n        LGBM_PARAMS,\n        train_data,\n        num_boost_round=N_BOOST_ROUNDS,\n        valid_sets=[val_data],\n        callbacks=[lgb.early_stopping(stopping_rounds=20, verbose=False)]\n    )\n    \n    models[term] = model\n\nprint(f'\\nTrained {len(models)} models, skipped {skipped} rare terms')\n\n# Free training data\ndel X_train, Y_train, X_tr, X_val, Y_tr, Y_val\ngc.collect()"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Generate predictions\nprint('[5/6] Generating predictions...')\n\n# Predict for all test proteins\nX_test = test_emb\npredictions = {}\n\nfor term, model in tqdm(models.items(), desc='Predicting'):\n    preds = model.predict(X_test)\n    predictions[term] = preds\n\nprint(f'Generated predictions for {len(predictions)} GO terms')"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Parse GO ontology for propagation\nprint('[5.5/6] Parsing GO ontology...')\n\ndef parse_obo(obo_path):\n    parents = defaultdict(set)\n    with open(obo_path, 'r') as f:\n        cur_id = None\n        for line in f:\n            line = line.strip()\n            if line == '[Term]':\n                cur_id = None\n            elif line.startswith('id: '):\n                cur_id = line.split('id: ')[1].strip()\n            elif line.startswith('is_a: '):\n                pid = line.split()[1].strip()\n                if cur_id:\n                    parents[cur_id].add(pid)\n            elif line.startswith('relationship: part_of '):\n                parts = line.split()\n                if len(parts) >= 3:\n                    pid = parts[2].strip()\n                    if cur_id:\n                        parents[cur_id].add(pid)\n    return parents\n\ndef build_ancestors(term_parents):\n    ancestors = {}\n    def get_ancestors(term):\n        if term in ancestors:\n            return ancestors[term]\n        parents = term_parents.get(term, set())\n        all_anc = set(parents)\n        for p in parents:\n            all_anc |= get_ancestors(p)\n        ancestors[term] = all_anc\n        return all_anc\n    \n    for term in term_parents:\n        get_ancestors(term)\n    return ancestors\n\nterm_parents = parse_obo(OBO_PATH)\nancestors_map = build_ancestors(term_parents)\nROOTS = {'GO:0003674', 'GO:0008150', 'GO:0005575'}\nprint(f'Loaded {len(term_parents)} GO terms')"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Create submission with GO propagation\nprint('[6/6] Creating submission with GO propagation...')\n\nsubmission_rows = []\n\nfor i, pid in enumerate(tqdm(test_ids, desc='Processing proteins')):\n    # Get predictions for this protein\n    protein_scores = {}\n    for term, preds in predictions.items():\n        score = preds[i]\n        if score >= PRED_THRESHOLD:\n            protein_scores[term] = score\n    \n    # Apply GO propagation\n    final_scores = protein_scores.copy()\n    for term, score in protein_scores.items():\n        if term in ancestors_map:\n            for anc in ancestors_map[term]:\n                final_scores[anc] = max(final_scores.get(anc, 0), score)\n    \n    # Force roots to 1.0\n    if len(final_scores) > 0:\n        for root in ROOTS:\n            final_scores[root] = 1.0\n    \n    # Collect rows\n    for term, score in final_scores.items():\n        if score >= PRED_THRESHOLD:\n            submission_rows.append((pid, term, round(score, 6)))\n\nprint(f'Total predictions: {len(submission_rows):,}')"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Save submission\nprint('Saving submission...')\n\nsubmission_df = pd.DataFrame(submission_rows, columns=['protein_id', 'go_term', 'score'])\nsubmission_df = submission_df.sort_values(['protein_id', 'score'], ascending=[True, False])\nsubmission_df.to_csv('submission.tsv', sep='\\t', index=False, header=False)\n\nprint(f'\\n{\"=\"*60}')\nprint('CAFA-6 LightGBM + ESM2 Complete')\nprint(f'Total predictions: {len(submission_df):,}')\nprint(f'Unique proteins: {submission_df.protein_id.nunique():,}')\nprint(f'Unique GO terms: {submission_df.go_term.nunique():,}')\nprint(f'{\"=\"*60}')"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"!head submission.tsv"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.0"}},"nbformat":4,"nbformat_minor":4}