{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CAFA-6 Ultimate v3 - ESM-Cambrian (1280 dim)\n\n**Objetivo: Top 10 (>0.40)**\n\nMejoras:\n- ESM-Cambrian embeddings (1280 dim vs 320)\n- LightGBM multi-target optimizado\n- GOA completo\n- Ensemble con existentes","metadata":{}},{"cell_type":"code","source":"import os, gc, numpy as np, pandas as pd\nfrom collections import defaultdict\nfrom tqdm.auto import tqdm\nimport warnings; warnings.filterwarnings('ignore')\n\ntry:\n    import lightgbm as lgb\n    HAS_LGBM = True\nexcept: HAS_LGBM = False\n\nprint('Ready!')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CFG:\n    BASE = '/kaggle/input/cafa-6-protein-function-prediction'\n    TRAIN_TERMS = f'{BASE}/Train/train_terms.tsv'\n    GO_OBO = f'{BASE}/Train/go-basic.obo'\n    \n    # ESM-Cambrian embeddings (archivos separados train/test)\n    EMB_PATH = '/kaggle/input/cafa6-esm-cambrian-embeddings'\n    \n    GOA_PATH = '/kaggle/input/protein-go-annotations/goa_uniprot_all.csv'\n    \n    EXISTING_SUBS = [\n        ('/kaggle/input/cafa-6-proven-ensemble-submission/submission.tsv', 0.373),\n        ('/kaggle/input/cafa-quick-submission-0-358-score/submission.tsv', 0.358),\n    ]\n    \n    ASPECTS = ['C', 'F', 'P']\n    TOP_N_TERMS = 4500\n    MIN_SCORE = 0.01\n    TARGET_MAX = 0.95\n    ROOTS = {'GO:0003674', 'GO:0008150', 'GO:0005575'}\n    \n    # LightGBM\n    LGB_PARAMS = {\n        'objective': 'binary', 'metric': 'auc',\n        'boosting_type': 'gbdt', 'num_leaves': 63,\n        'learning_rate': 0.05, 'feature_fraction': 0.8,\n        'bagging_fraction': 0.8, 'bagging_freq': 5,\n        'verbose': -1, 'n_jobs': -1,\n    }\n    NUM_ROUNDS = 200\n\nprint('Config ready!')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# GO Utilities\ndef parse_obo(obo_path):\n    parents, children = defaultdict(set), 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]': cur_id = None\n            elif line.startswith('id: '): cur_id = line.split('id: ')[1].strip()\n            elif line.startswith('is_a: '):\n                pid = line.split()[1].strip()\n                if cur_id: parents[cur_id].add(pid); children[pid].add(cur_id)\n            elif line.startswith('relationship: part_of '):\n                parts = line.split()\n                if len(parts) >= 3 and cur_id:\n                    pid = parts[2].strip()\n                    parents[cur_id].add(pid); children[pid].add(cur_id)\n    print(f'[OBO] {len(parents)} terms')\n    return parents, children\n\ndef get_ancestors(term, pm, cache):\n    if term in cache: return cache[term]\n    anc = set(); stack = [term]\n    while stack:\n        cur = stack.pop()\n        for p in pm.get(cur, []):\n            if p not in anc: anc.add(p); stack.append(p)\n    cache[term] = anc\n    return anc\n\ndef get_descendants(term, cm, cache):\n    if term in cache: return cache[term]\n    desc = set(); stack = [term]\n    while stack:\n        cur = stack.pop()\n        for c in cm.get(cur, []):\n            if c not in desc: desc.add(c); stack.append(c)\n    cache[term] = desc\n    return desc\n\nprint('Utilities ready!')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load data\nprint('Loading GO ontology...')\nparents_map, children_map = parse_obo(CFG.GO_OBO)\n\nprint('\\nLoading ESM-Cambrian embeddings...')\n# Archivos separados: train y test\ntrain_emb = np.load(f'{CFG.EMB_PATH}/protein_embeddings_train.npy')\ntest_emb = np.load(f'{CFG.EMB_PATH}/protein_embeddings_test.npy')\ntrain_ids = np.load(f'{CFG.EMB_PATH}/train_ids.npy', allow_pickle=True)\ntest_ids = np.load(f'{CFG.EMB_PATH}/test_ids.npy', allow_pickle=True)\n\nprint(f'Train embeddings: {train_emb.shape}')\nprint(f'Test embeddings: {test_emb.shape}')\nprint(f'Embedding dim: {train_emb.shape[1]}')\n\n# Crear diccionarios\ntrain_emb_dict = {str(pid): emb for pid, emb in zip(train_ids, train_emb)}\ntest_emb_dict = {str(pid): emb for pid, emb in zip(test_ids, test_emb)}\nemb_dict = {**train_emb_dict, **test_emb_dict}  # Combinar\n\nprint('\\nLoading training terms...')\ntrain_df = pd.read_csv(CFG.TRAIN_TERMS, sep='\\t')\nprint(f'Annotations: {len(train_df)}, Proteins: {train_df[\"EntryID\"].nunique()}')\n\n# Test proteins son los que estan en test_emb_dict\ntest_proteins = set(test_emb_dict.keys())\nprint(f'Train proteins: {len(train_emb_dict)}')\nprint(f'Test proteins: {len(test_proteins)}')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load GOA\nprint('Loading GOA...')\nnegative_keys = set()\npositive_df = None\n\nif os.path.exists(CFG.GOA_PATH):\n    goa = pd.read_csv(CFG.GOA_PATH).drop_duplicates()\n    print(f'GOA: {len(goa)}')\n    \n    if 'qualifier' in goa.columns:\n        neg_goa = goa[goa['qualifier'].str.contains('NOT', na=False)]\n        neg_by_prot = neg_goa.groupby('protein_id')['go_term'].apply(list).to_dict()\n        desc_cache = {}\n        for protein, terms in tqdm(neg_by_prot.items(), desc='Neg propagation'):\n            all_neg = set(terms)\n            for term in terms:\n                all_neg |= get_descendants(term, children_map, desc_cache)\n            for term in all_neg:\n                negative_keys.add(f'{protein}_{term}')\n        print(f'Negative keys: {len(negative_keys)}')\n    \n    pos_goa = goa[~goa['qualifier'].str.contains('NOT', na=False)] if 'qualifier' in goa.columns else goa\n    positive_df = pd.DataFrame({\n        'protein_id': pos_goa['protein_id'].astype(str),\n        'go_term': pos_goa['go_term'].astype(str),\n        'score': 1.0\n    }).drop_duplicates()\n    positive_df['key'] = positive_df['protein_id'] + '_' + positive_df['go_term']\n    positive_df = positive_df[~positive_df['key'].isin(negative_keys)].drop(columns=['key'])\n    print(f'Positive: {len(positive_df)}')\n    del goa; gc.collect()\nelse:\n    print('[WARN] GOA not found')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train models per aspect\nprint('\\nTraining models per aspect...')\n\nfrom sklearn.linear_model import RidgeClassifier\nfrom sklearn.multiclass import OneVsRestClassifier\n\naspect_predictions = {}\n\nfor aspect in CFG.ASPECTS:\n    print(f'\\n=== Aspect: {aspect} ===')\n    \n    aspect_df = train_df[train_df['aspect'] == aspect]\n    term_counts = aspect_df['term'].value_counts()\n    top_terms = term_counts.head(CFG.TOP_N_TERMS).index.tolist()\n    term_to_idx = {t: i for i, t in enumerate(top_terms)}\n    idx_to_term = {i: t for t, i in term_to_idx.items()}\n    \n    print(f'Top {len(top_terms)} terms')\n    \n    # Build training data\n    protein_terms = aspect_df[aspect_df['term'].isin(top_terms)].groupby('EntryID')['term'].apply(list).to_dict()\n    train_pids = [p for p in protein_terms.keys() if p in emb_dict]\n    \n    X_train = np.array([emb_dict[p] for p in train_pids])\n    Y_train = np.zeros((len(train_pids), len(top_terms)), dtype=np.float32)\n    for i, pid in enumerate(train_pids):\n        for term in protein_terms[pid]:\n            if term in term_to_idx:\n                Y_train[i, term_to_idx[term]] = 1.0\n    \n    print(f'X: {X_train.shape}, Y: {Y_train.shape}')\n    \n    # Train\n    if HAS_LGBM:\n        # Use LightGBM in chunks\n        n_val = int(0.1 * len(X_train))\n        X_tr, X_vl = X_train[:-n_val], X_train[-n_val:]\n        Y_tr, Y_vl = Y_train[:-n_val], Y_train[-n_val:]\n        \n        models = []\n        CHUNK = 500\n        for start in tqdm(range(0, len(top_terms), CHUNK), desc=f'LGB {aspect}'):\n            end = min(start + CHUNK, len(top_terms))\n            chunk_models = []\n            for j in range(start, end):\n                if Y_tr[:, j].sum() < 5:\n                    chunk_models.append(None)\n                    continue\n                ds = lgb.Dataset(X_tr, label=Y_tr[:, j])\n                vs = lgb.Dataset(X_vl, label=Y_vl[:, j], reference=ds)\n                m = lgb.train(CFG.LGB_PARAMS, ds, CFG.NUM_ROUNDS, valid_sets=[vs],\n                              callbacks=[lgb.early_stopping(30, verbose=False)])\n                chunk_models.append(m)\n            models.extend(chunk_models)\n        \n        # Predict\n        test_pids = [p for p in test_proteins if p in emb_dict]\n        X_test = np.array([emb_dict[p] for p in test_pids])\n        preds = np.zeros((len(test_pids), len(top_terms)))\n        for j, m in enumerate(models):\n            if m: preds[:, j] = m.predict(X_test)\n    else:\n        # Ridge fallback\n        model = OneVsRestClassifier(RidgeClassifier(alpha=1.0), n_jobs=-1)\n        model.fit(X_train, Y_train)\n        \n        test_pids = [p for p in test_proteins if p in emb_dict]\n        X_test = np.array([emb_dict[p] for p in test_pids])\n        dec = model.decision_function(X_test)\n        preds = 1 / (1 + np.exp(-dec))\n    \n    # To DataFrame\n    rows = []\n    for i, pid in enumerate(tqdm(test_pids, desc=f'Proc {aspect}')):\n        for j in np.where(preds[i] >= CFG.MIN_SCORE)[0]:\n            rows.append({'protein_id': pid, 'go_term': idx_to_term[j], 'score': float(preds[i, j])})\n    \n    aspect_predictions[aspect] = pd.DataFrame(rows)\n    print(f'Predictions: {len(aspect_predictions[aspect])}')\n    \n    del X_train, Y_train, X_test, preds\n    gc.collect()\n\nprint('\\nAll aspects done!')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Post-processing\nprint('Post-processing...')\n\nall_preds = pd.concat(aspect_predictions.values(), ignore_index=True)\nprint(f'Combined: {len(all_preds)}')\n\n# Propagate\nanc_cache = {}\nrows = []\nfor pid, group in tqdm(all_preds.groupby('protein_id'), desc='Propagating'):\n    scores = dict(zip(group['go_term'], group['score']))\n    final = scores.copy()\n    for term, score in list(scores.items()):\n        for anc in get_ancestors(term, parents_map, anc_cache):\n            final[anc] = max(final.get(anc, 0), score)\n    for root in CFG.ROOTS: final[root] = 1.0\n    for term, score in final.items():\n        if score >= CFG.MIN_SCORE:\n            rows.append((pid, term, score))\n\nfinal_df = pd.DataFrame(rows, columns=['protein_id', 'go_term', 'score'])\nprint(f'After propagation: {len(final_df)}')\n\n# Remove negatives\nif negative_keys:\n    final_df['key'] = final_df['protein_id'] + '_' + final_df['go_term']\n    before = len(final_df)\n    final_df = final_df[~final_df['key'].isin(negative_keys)].drop(columns=['key'])\n    print(f'Removed {before - len(final_df)} negatives')\n\n# Add positives\nif positive_df is not None:\n    tps = set(final_df['protein_id'].unique())\n    pos_rel = positive_df[positive_df['protein_id'].isin(tps)].copy()\n    pos_rel.columns = ['protein_id', 'go_term', 'score']\n    final_df = pd.concat([final_df, pos_rel], ignore_index=True)\n    final_df = final_df.groupby(['protein_id', 'go_term'])['score'].max().reset_index()\n    print(f'After GOA: {len(final_df)}')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Normalize\nprint('Normalizing...')\nrows = []\nfor pid, group in tqdm(final_df.groupby('protein_id')):\n    scores = dict(zip(group['go_term'], group['score']))\n    max_nr = max((s for t, s in scores.items() if t not in CFG.ROOTS), default=0)\n    if 0 < max_nr < CFG.TARGET_MAX:\n        scale = CFG.TARGET_MAX / max_nr\n        for t in scores:\n            if t not in CFG.ROOTS: scores[t] = min(1.0, scores[t] * scale)\n    for root in CFG.ROOTS: scores[root] = 1.0\n    for term, score in scores.items():\n        if score >= CFG.MIN_SCORE: rows.append((pid, term, round(score, 6)))\n\nfinal_df = pd.DataFrame(rows, columns=['protein_id', 'go_term', 'score'])\nprint(f'After normalization: {len(final_df)}')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Ensemble\nprint('\\nEnsembling...')\nvalid_subs = [(p, s) for p, s in CFG.EXISTING_SUBS if os.path.exists(p)]\nprint(f'Found {len(valid_subs)} existing')\n\nif valid_subs:\n    dfs = []\n    for path, lb in valid_subs:\n        df = pd.read_csv(path, sep='\\t', header=None, names=['protein_id', 'go_term', 'score'])\n        df['w'] = np.exp(lb * 10)\n        dfs.append(df)\n        print(f'  LB={lb}')\n    \n    final_df['w'] = np.exp(0.40 * 10)  # Our expected\n    all_s = pd.concat([final_df] + dfs, ignore_index=True)\n    all_s['ws'] = all_s['score'] * all_s['w']\n    agg = all_s.groupby(['protein_id', 'go_term']).agg({'ws': 'sum', 'w': 'sum', 'score': 'count'}).reset_index()\n    agg.columns = ['protein_id', 'go_term', 'ws', 'w', 'n']\n    agg['score'] = agg['ws'] / agg['w']\n    agg.loc[agg['n'] >= 2, 'score'] += 0.02\n    agg['score'] = agg['score'].clip(0, 1.0)\n    final_df = agg[['protein_id', 'go_term', 'score']].copy()\n    print(f'After ensemble: {len(final_df)}')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Save\nfinal_df = final_df.sort_values(['protein_id', 'score'], ascending=[True, False])\nfinal_df.to_csv('submission.tsv', sep='\\t', index=False, header=False)\n\nprint('\\n' + '='*60)\nprint('CAFA-6 Ultimate v3 - ESM-Cambrian 1280 dim')\nprint('='*60)\nprint(f'Predictions: {len(final_df):,}')\nprint(f'Proteins: {final_df[\"protein_id\"].nunique():,}')\nprint(f'GO terms: {final_df[\"go_term\"].nunique():,}')\nprint('='*60)\n\n!head submission.tsv","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}