"""Train the final model on all of train.parquet and write weights/ (run on your computer, not by the grader). v1 (no ESM): python train.py --train ../data/train.parquet v2 (ESM-2): python train.py --train ../data/train.parquet --esm ../esm2_35M python train.py --train ../data/train.parquet --esm ../esm2_150M --quantize v3 (blend): python train.py --train ../data/train.parquet --esm ../esm2_35M --nn v4: python train.py --train ../data/train.parquet --esm ../esm2_35M --esm-layer 6 --nn v5: python train.py --train ../data/train.parquet --esm ../esm2_35M --esm-layer 6 --et """ import argparse import json import time from itertools import product from pathlib import Path import lightgbm as lgb import numpy as np import pandas as pd from scipy import sparse from scipy.sparse.csgraph import connected_components from sklearn.model_selection import GroupKFold from features import TARGETS, base_features, prior_features, prior_table LGB_PARAMS = dict(objective='regression', learning_rate=0.03, num_leaves=63, min_data_in_leaf=200, feature_fraction=0.05, bagging_fraction=0.8, bagging_freq=1, lambda_l2=10.0, max_bin=63, verbose=-1, seed=0) def enzyme_clusters(sequences, threshold=0.3): """Group enzymes whose 3-mer cosine similarity > threshold (connected components).""" aa = 'ACDEFGHIKLMNPQRSTVWY' k3 = {''.join(p): i for i, p in enumerate(product(aa, repeat=3))} useq = pd.unique(sequences) r, c, d = [], [], [] for i, s in enumerate(useq): s = s.upper() u, n = np.unique([k3[s[j:j + 3]] for j in range(len(s) - 2) if s[j:j + 3] in k3], return_counts=True) r += [i] * len(u); c += list(u); d += list(n) M = sparse.csr_matrix((d, (r, c)), shape=(len(useq), len(k3)), dtype=np.float32) nrm = np.sqrt(M.multiply(M).sum(1)).A1; nrm[nrm == 0] = 1 M = sparse.diags(1 / nrm) @ M rows, cols = [], [] for i in range(0, M.shape[0], 1000): S = (M[i:i + 1000] @ M.T).tocoo(); keep = S.data > threshold rows += list(S.row[keep] + i); cols += list(S.col[keep]) G = sparse.csr_matrix((np.ones(len(rows)), (rows, cols)), shape=(len(useq),) * 2) lab = dict(zip(useq, connected_components(G, directed=False)[1])) return np.array([lab[s] for s in sequences]) def esm_train_features(sequences, esm_dir, out, quantize, layer=None): """Embed every unique training enzyme with ESM-2 (GPU if available) and copy the model into weights/esm.""" import shutil import torch from features import esm_embed, load_esm device = 'cuda' if torch.cuda.is_available() else 'cpu' useq = list(pd.unique(sequences)) tag = f'_L{layer}' if layer else '' cache = Path(out) / '..' / f'.esm_cache_{Path(esm_dir).name}{tag}.npy' # outside weights/, not uploaded tok, model = load_esm(esm_dir, device=device, layer=layer) if cache.exists(): E = np.load(cache) else: t = time.time() E = esm_embed(useq, tok, model, device=device, batch_tokens=8000) np.save(cache, E); print(f'ESM embeddings {E.shape} on {device} in {(time.time() - t) / 60:.1f} min') # Save a clean fp32 copy (safetensors + tokenizer only) into weights/esm for the grader. dst = Path(out) / 'esm' if dst.exists(): shutil.rmtree(dst) model.to('cpu').save_pretrained(dst); tok.save_pretrained(dst) del model pos = {s: i for i, s in enumerate(useq)} return E[[pos[s] for s in sequences]], {'dir': 'esm', 'quantize': quantize, 'dim': int(E.shape[1]), 'layer': layer} ET_PARAMS = dict(n_estimators=200, max_features='sqrt', min_samples_leaf=3, random_state=0) ET_WEIGHT = 0.5 # final = 0.5 * LightGBM + 0.5 * ExtraTrees (Experiment 5, blend a) def train_et(Xall, Y, out, threads): """One multi-output ExtraTrees on the same features as LightGBM (NaN -> -999, as in Experiment 5).""" import joblib from sklearn.ensemble import ExtraTreesRegressor t = time.time() m = ExtraTreesRegressor(**ET_PARAMS, n_jobs=threads or -1).fit(np.nan_to_num(Xall, nan=-999).astype(np.float32), Y) joblib.dump(m, out / 'et.joblib') print(f'ExtraTrees {ET_PARAMS["n_estimators"]} trees in {time.time() - t:.0f}s, ' f'{(out / "et.joblib").stat().st_size / 1e6:.0f} MB') return {'file': 'et.joblib', 'weight': ET_WEIGHT, **ET_PARAMS} NN_SEEDS, NN_LGB_WEIGHT, NN_ACT, NN_DROPOUT = (0, 1, 2), 0.7, 'silu', 0.5 def train_nn(E_tr, S_tr, y_tr, E_es, S_es, y_es, seed, epochs, device, early_stop=True, run_epochs=None): """Two-tower net: AdamW, warmup + cosine LR, Huber loss, AMP, grad clipping; early stopping on held-out Spearman.""" import copy import torch from scipy.stats import spearmanr from nn_model import TwoTower torch.manual_seed(seed); np.random.seed(seed) T = lambda a: torch.tensor(a, dtype=torch.float32, device=device) mu, sd = y_tr.mean(0), y_tr.std(0) e, s, y = T(E_tr), T(S_tr), T((y_tr - mu) / sd) net = TwoTower(e.shape[1], s.shape[1], p=NN_DROPOUT, act=NN_ACT).to(device) opt = torch.optim.AdamW(net.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=1e-2) bs = 512; steps = epochs * (len(y) // bs); warm = max(1, steps // 20) sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda t: min(1, (t + 1) / warm) * 0.5 * (1 + np.cos(np.pi * min(1, t / steps)))) amp = device == 'cuda' scaler = torch.amp.GradScaler(enabled=amp) best, best_ep, best_state, bad = -1, epochs, None, 0 for ep in range(run_epochs or epochs): net.train(); perm = torch.randperm(len(y), device=device) for i in range(0, len(perm) - bs + 1, bs): b = perm[i:i + bs] with torch.autocast(device, enabled=amp): p = net(e[b], s[b]) loss = torch.nn.functional.huber_loss(p.float(), y[b]) opt.zero_grad(); scaler.scale(loss).backward() scaler.unscale_(opt); torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0) scaler.step(opt); scaler.update(); sched.step() if not early_stop: continue net.eval() with torch.no_grad(): pe = net(T(E_es), T(S_es)).float().cpu().numpy() sc = np.mean([spearmanr(y_es[:, k], pe[:, k]).statistic for k in range(2)]) if sc > best: best, best_ep, bad, best_state = sc, ep + 1, 0, copy.deepcopy(net.state_dict()) else: bad += 1 if bad >= 8: break if early_stop: net.load_state_dict(best_state) return net.eval().cpu(), mu, sd, best_ep, best def train_nn_ensemble(base, E, P, Y, es_tr, es_va, out, epochs): """Fit quantile normalisation on all rows; per seed find the epoch count on held-out clusters, refit on all rows.""" import torch from sklearn.preprocessing import QuantileTransformer from nn_model import split_blocks device = 'cuda' if torch.cuda.is_available() else 'cpu' ENZ, SUB = split_blocks(base, E, P) qts, Z = {}, {} for name, M in [('enz', ENZ), ('sub', SUB)]: qt = QuantileTransformer(n_quantiles=200, output_distribution='normal', subsample=20000, random_state=0).fit(M) qts[name] = qt Z[name] = np.nan_to_num(qt.transform(M), nan=0.0).astype(np.float32) d = out / 'nn'; d.mkdir(exist_ok=True) np.savez(d / 'quantiles.npz', enz_q=qts['enz'].quantiles_, enz_r=qts['enz'].references_, sub_q=qts['sub'].quantiles_, sub_r=qts['sub'].references_) info = [] for seed in NN_SEEDS: t = time.time() _, _, _, ep, sc = train_nn(Z['enz'][es_tr], Z['sub'][es_tr], Y[es_tr], Z['enz'][es_va], Z['sub'][es_va], Y[es_va], seed, epochs, device) net, mu, sd, _, _ = train_nn(Z['enz'], Z['sub'], Y, None, None, None, seed, epochs, device, early_stop=False, run_epochs=ep) torch.save(net.state_dict(), d / f'net_{seed}.pt') info.append({'seed': seed, 'epochs': ep, 'heldout_scc': round(float(sc), 4), 'mu': mu.tolist(), 'sd': sd.tolist()}) print(f'nn seed {seed}: {ep} epochs (held-out SCC {sc:.4f}) {time.time() - t:.0f}s on {device}') return {'seeds': info, 'de': int(ENZ.shape[1]), 'ds': int(SUB.shape[1]), 'lgb_weight': NN_LGB_WEIGHT, 'act': NN_ACT, 'p': NN_DROPOUT} def main(train_path, out_dir, threads, esm_dir=None, quantize=False, nn=False, nn_epochs=60, esm_layer=None, et=False): assert esm_dir or not nn, '--nn needs --esm' t0 = time.time() df = pd.read_parquet(train_path) Y = df[TARGETS].to_numpy(float) X = base_features(df.sequence.to_numpy(), df.smiles.to_numpy()) esm_meta = None if esm_dir: Path(out_dir).mkdir(parents=True, exist_ok=True) E, esm_meta = esm_train_features(df.sequence.to_numpy(), esm_dir, out_dir, quantize, esm_layer) base = X X = np.hstack([X, E]) groups = enzyme_clusters(df.sequence.to_numpy()) print(f'features {X.shape}, {len(np.unique(groups))} enzyme clusters, {time.time() - t0:.0f}s') # Training rows get out-of-fold priors (a row never sees its own label), exactly as in cross-validation. P = np.zeros((len(df), 11)) for a, b in GroupKFold(5).split(df, groups=groups): tab = prior_table(df.smiles.values[a], Y[a]) P[b] = prior_features(tab, df.smiles.values[b], Y[a].mean(0)) Xall = np.hstack([X, P]) params = {**LGB_PARAMS, 'num_threads': threads} es_tr, es_va = next(GroupKFold(8).split(df, groups=groups)) # held-out enzyme clusters for early stopping out = Path(out_dir); out.mkdir(parents=True, exist_ok=True) iters = {} for i, t in enumerate(TARGETS): m = lgb.train(params, lgb.Dataset(Xall[es_tr], Y[es_tr, i]), 4000, valid_sets=[lgb.Dataset(Xall[es_va], Y[es_va, i])], callbacks=[lgb.early_stopping(150, verbose=False)]) iters[t] = int(m.best_iteration * 1.1) + 1 # a bit more data -> a few more trees final = lgb.train(params, lgb.Dataset(Xall, Y[:, i]), iters[t]) final.save_model(str(out / f'lgb_{t}.txt')) print(t, 'trees', iters[t]) et_meta = train_et(Xall, Y, out, threads) if et else None nn_meta = None if nn: nn_meta = train_nn_ensemble(base, E, P, Y.astype(np.float32), es_tr, es_va, out, nn_epochs) # Inference-time priors use ALL training rows. tab = prior_table(df.smiles.values, Y) tab.reset_index(names='smiles').to_parquet(out / 'prior.parquet', index=False) (out / 'meta.json').write_text(json.dumps({'global_mean': Y.mean(0).tolist(), 'n_features': Xall.shape[1], 'trees': iters, 'esm': esm_meta, 'et': et_meta, 'nn': nn_meta}, indent=2)) print(f'saved to {out} in {(time.time() - t0) / 60:.1f} min') if __name__ == '__main__': ap = argparse.ArgumentParser() ap.add_argument('--train', default='../data/train.parquet') ap.add_argument('--out', default='weights') ap.add_argument('--threads', type=int, default=0, help='0 = all cores') ap.add_argument('--esm', default=None, help='folder with a downloaded ESM-2 model (enables v2)') ap.add_argument('--quantize', action='store_true', help='run ESM in int8 at inference (needed for 150M)') ap.add_argument('--nn', action='store_true', help='also train the two-tower network and blend it (v3)') ap.add_argument('--nn-epochs', type=int, default=60, help='max epochs for the network') ap.add_argument('--esm-layer', type=int, default=None, help='use this intermediate ESM layer (v4: 6)') ap.add_argument('--et', action='store_true', help='also train ExtraTrees and average it with LightGBM (v5)') a = ap.parse_args() main(a.train, a.out, a.threads, a.esm, a.quantize, a.nn, a.nn_epochs, a.esm_layer, a.et)