File size: 6,268 Bytes
b145940 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """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}
# ---------------- local validation run (gives you a trustworthy score) ----
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})")
# ---------------- final training on ALL training rows ---------------------
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()
|