| """Train the kcat/Km model. |
| |
| Stage 1 (no protein language model): |
| python train.py --train ../data/train.parquet --out weights |
| Stage 2 (adds ESM2 protein embeddings; use a GPU in Colab): |
| python train.py --train ../data/train.parquet --out weights --esm |
| Stage 3 (average several models): |
| python train.py --train ../data/train.parquet --out weights --esm --seeds 3 |
| """ |
| import argparse |
| import json |
| import os |
| import time |
|
|
| import lightgbm as lgb |
| import numpy as np |
| import pandas as pd |
| from scipy.stats import spearmanr |
| from sklearn.model_selection import GroupKFold, GroupShuffleSplit |
|
|
| import features as F |
|
|
| ESM_NAME = "facebook/esm2_t12_35M_UR50D" |
|
|
|
|
| def log(*a): |
| print(time.strftime("%H:%M:%S"), *a, flush=True) |
|
|
|
|
| def te_oof(df, groups, n_splits=5): |
| """Substrate-average features for training rows, computed WITHOUT looking at |
| the row's own enzyme group (so the model learns to use them honestly).""" |
| out = np.full((len(df), 3), np.nan, dtype=np.float32) |
| out[:, 0] = 0 |
| for tr, va in GroupKFold(n_splits).split(df, groups=groups): |
| st = F.te_stats(df.iloc[tr]) |
| out[va] = F.te_apply(df.iloc[va]["smiles"].values, st) |
| return out |
|
|
|
|
| def make_model(seed, n_estimators, lr): |
| return lgb.LGBMRegressor( |
| n_estimators=n_estimators, learning_rate=lr, num_leaves=63, |
| min_child_samples=20, subsample=0.8, subsample_freq=1, |
| colsample_bytree=0.25, reg_lambda=1.0, random_state=seed, |
| n_jobs=-1, verbose=-1) |
|
|
|
|
| def report(name, y, p): |
| scc = spearmanr(y, p).correlation |
| rmse = float(np.sqrt(np.mean((y - p) ** 2))) |
| log(f" {name}: SCC={scc:.4f} RMSE={rmse:.4f}") |
| return scc |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--train", required=True) |
| ap.add_argument("--out", default="weights") |
| ap.add_argument("--esm", action="store_true", help="add ESM2 protein embeddings") |
| ap.add_argument("--seeds", type=int, default=1) |
| ap.add_argument("--lr", type=float, default=0.03) |
| ap.add_argument("--max_trees", type=int, default=3000) |
| ap.add_argument("--no_holdout", action="store_true", help="skip the local validation run") |
| args = ap.parse_args() |
| os.makedirs(args.out, exist_ok=True) |
|
|
| df = pd.read_parquet(args.train).reset_index(drop=True) |
| log("rows:", len(df), "| unique smiles:", df.smiles.nunique(), "| unique sequences:", df.sequence.nunique()) |
|
|
| usmi = df["smiles"].unique().tolist() |
| useq = df["sequence"].unique().tolist() |
| log("computing SMILES features ...") |
| smi_mat = F.smiles_features(usmi) |
| log("computing sequence features ...") |
| seq_mat = F.seq_features(useq) |
|
|
| esm_mat, esm_mean = None, None |
| if args.esm: |
| esm_dir = os.path.join(args.out, "esm") |
| if not os.path.exists(esm_dir): |
| from transformers import AutoTokenizer, EsmModel |
| log("downloading", ESM_NAME) |
| AutoTokenizer.from_pretrained(ESM_NAME).save_pretrained(esm_dir) |
| EsmModel.from_pretrained(ESM_NAME).save_pretrained(esm_dir) |
| cache = "esm_cache.npz" |
| if os.path.exists(cache) and list(np.load(cache, allow_pickle=True)["seqs"]) == useq: |
| esm_mat = np.load(cache, allow_pickle=True)["emb"] |
| log("loaded ESM embeddings from cache") |
| else: |
| import torch |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| log(f"computing ESM embeddings on {dev} (this is the slow part) ...") |
| esm_mat = F.esm_embed(useq, esm_dir, device=dev, |
| batch_tokens=16384 if dev == "cuda" else 4096, |
| n_threads=os.cpu_count() or 2) |
| np.savez(cache, seqs=np.array(useq, dtype=object), emb=esm_mat) |
| esm_mean = esm_mat.mean(0) |
| np.save(os.path.join(args.out, "esm_mean.npy"), esm_mean) |
|
|
| log("grouping similar sequences (for an honest validation split) ...") |
| leader = F.cluster_sequences(useq) |
| seq2grp = dict(zip(useq, leader)) |
| groups = df["sequence"].map(seq2grp).values |
| log("sequence groups:", len(set(groups))) |
|
|
| y_kc = df["log10_kcat"].values |
| y_km = df["log10_km"].values |
| best_iter = {"kcat": 800, "km": 800} |
|
|
| |
| if not args.no_holdout: |
| tr, va = next(GroupShuffleSplit(1, test_size=0.2, random_state=0).split(df, groups=groups)) |
| dtr, dva = df.iloc[tr].reset_index(drop=True), df.iloc[va].reset_index(drop=True) |
| te_tr = te_oof(dtr, groups[tr]) |
| te_va = F.te_apply(dva["smiles"].values, F.te_stats(dtr)) |
| Xtr = F.assemble(dtr, usmi, smi_mat, useq, seq_mat, esm_mat, te_tr) |
| Xva = F.assemble(dva, usmi, smi_mat, useq, seq_mat, esm_mat, te_va) |
| log("holdout features:", Xtr.shape, Xva.shape) |
| sccs = [] |
| for name, col in [("kcat", "log10_kcat"), ("km", "log10_km")]: |
| m = make_model(0, args.max_trees, args.lr) |
| m.fit(Xtr, dtr[col].values, eval_set=[(Xva, dva[col].values)], |
| callbacks=[lgb.early_stopping(100, verbose=False)]) |
| best_iter[name] = max(100, int(m.best_iteration_ or args.max_trees)) |
| sccs.append(report(name, dva[col].values, m.predict(Xva))) |
| log(f"LOCAL HOLDOUT mean SCC = {np.mean(sccs):.4f} (best iters: {best_iter})") |
|
|
| |
| log("training final models on all rows ...") |
| te_all = te_oof(df, groups) |
| stats = F.te_stats(df) |
| stats.to_parquet(os.path.join(args.out, "te_stats.parquet")) |
| X = F.assemble(df, usmi, smi_mat, useq, seq_mat, esm_mat, te_all) |
| for name, y in [("kcat", y_kc), ("km", y_km)]: |
| for s in range(args.seeds): |
| m = make_model(s, int(best_iter[name] * 1.15), args.lr) |
| m.fit(X, y) |
| m.booster_.save_model(os.path.join(args.out, f"{name}_{s}.txt")) |
| log(f"saved {name}_{s}.txt") |
|
|
| meta = {"use_esm": bool(args.esm), "n_seeds": args.seeds, |
| "mean_kcat": float(y_kc.mean()), "mean_km": float(y_km.mean())} |
| with open(os.path.join(args.out, "meta.json"), "w") as f: |
| json.dump(meta, f) |
| log("DONE. Weights are in", args.out) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
