| """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' |
| 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') |
| |
| 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 |
|
|
|
|
| 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') |
|
|
| |
| 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)) |
| 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 |
| 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) |
|
|
| |
| 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) |
|
|
