{"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\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 = 150\nTOP_N_TERMS = 1500\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/7] Loading ESM2 embeddings...')\n\nprotein_ids_df = pd.read_csv(f'{EMB_PATH}/protein_ids.csv')\nall_embeddings = np.load(f'{EMB_PATH}/protein_embeddings.npy')\nall_protein_ids = protein_ids_df['protein_id'].tolist()\n\nprint(f'Total embeddings: {all_embeddings.shape}')\nprint(f'Total protein IDs: {len(all_protein_ids)}')\n\n# Create mapping\nid2idx = {pid: i for i, pid in enumerate(all_protein_ids)}\nall_protein_set = set(all_protein_ids)"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Load training labels\nprint('[2/7] 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 proteins that have annotations (these are training proteins)\ntrain_protein_set = set(train_terms['EntryID'].unique())\nprint(f'Training proteins with annotations: {len(train_protein_set):,}')\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 training data\nprint('[3/7] Building training data...')\n\n# Find train proteins that have embeddings\nvalid_train = [p for p in train_protein_set if p in id2idx]\nvalid_train = list(valid_train)  # Convert to list\nprint(f'Train proteins with embeddings: {len(valid_train)}')\n\n# Get embeddings for training\ntrain_idx = [id2idx[p] for p in valid_train]\nX_train = all_embeddings[train_idx]\nprint(f'X_train shape: {X_train.shape}')\n\n# Build Y matrix\nprotein2idx = {p: i for i, p in enumerate(valid_train)}\nY_train = np.zeros((len(valid_train), 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%}')"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Train LightGBM models\nprint('[4/7] Training LightGBM models...')\n\nmodels = {}\nskipped = 0\n\n# Train/val split\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    train_data = lgb.Dataset(X_tr, label=y_tr)\n    val_data = lgb.Dataset(X_val, label=y_val, reference=train_data)\n    \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\ndel X_train, Y_train, X_tr, X_val, Y_tr, Y_val\ngc.collect()"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Identify test proteins (proteins with embeddings but no training annotations)\nprint('[5/7] Identifying test proteins...')\n\n# Test proteins are those in embeddings but NOT in training set\ntest_proteins = [p for p in all_protein_ids if p not in train_protein_set]\nprint(f'Test proteins: {len(test_proteins)}')\n\n# Get test embeddings\ntest_idx = [id2idx[p] for p in test_proteins]\nX_test = all_embeddings[test_idx]\nprint(f'X_test shape: {X_test.shape}')\n\n# Generate predictions\nprint('Generating predictions...')\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\nprint('[6/7] 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('[7/7] Creating submission...')\n\nsubmission_rows = []\n\nfor i, pid in enumerate(tqdm(test_proteins, desc='Processing')):\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    # 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\n    if len(final_scores) > 0:\n        for root in ROOTS:\n            final_scores[root] = 1.0\n    \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}