| """Inference entry point: python /app/predict.py --input X.parquet --output Y.csv""" |
| import argparse |
| import json |
| import os |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
|
|
| import features as F |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| W = os.path.join(HERE, "weights") |
| N_THREADS = int(os.environ.get("NTHREADS", 2)) |
| ESM_BUDGET_S = 480 |
|
|
|
|
| def main(): |
| t0 = time.time() |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--input", required=True) |
| ap.add_argument("--output", required=True) |
| args = ap.parse_args() |
|
|
| df = pd.read_parquet(args.input).reset_index(drop=True) |
| df["smiles"] = df["smiles"].fillna("").astype(str) |
| df["sequence"] = df["sequence"].fillna("").astype(str) |
| meta = json.load(open(os.path.join(W, "meta.json"))) |
|
|
| usmi = df["smiles"].unique().tolist() |
| useq = df["sequence"].unique().tolist() |
| smi_mat = F.smiles_features(usmi) |
| seq_mat = F.seq_features(useq) |
| print(f"[predict] features done {time.time() - t0:.0f}s", flush=True) |
|
|
| esm_mat = None |
| if meta["use_esm"]: |
| fb = np.load(os.path.join(W, "esm_mean.npy")) |
| esm_mat = F.esm_embed(useq, os.path.join(W, "esm"), fallback=fb, |
| budget_s=ESM_BUDGET_S, n_threads=N_THREADS, layers=meta.get("layers")) |
| print(f"[predict] esm done {time.time() - t0:.0f}s", flush=True) |
|
|
| stats = pd.read_parquet(os.path.join(W, "te_stats.parquet")) |
| ref_bits = np.load(os.path.join(W, "ref_bits.npy")).astype(np.float32) |
| smi_index = {s: i for i, s in enumerate(usmi)} |
| bits_u = (smi_mat[:, :F.N_MORGAN] > 0).astype(np.float32) |
| te = F.te_full(df["smiles"].values, smi_index, bits_u, stats, ref_bits) |
| X0 = F.assemble(df, usmi, smi_mat, useq, seq_mat, esm_mat, te) |
| kn, er = None, None |
| q_embn = None |
| if meta.get("knn") or meta["preset"].get("res"): |
| ref_embn = np.load(os.path.join(W, "ref_embn.npy")).astype(np.float32) |
| q_embn = F.normalize_emb(esm_mat, np.load(os.path.join(W, "esm_mean.npy"))) |
| if meta["preset"].get("res"): |
| seq_pos0 = {s: i for i, s in enumerate(useq)} |
| er = F.enzyme_effect_feats(df["sequence"].map(seq_pos0).values, q_embn, ref_embn, |
| np.load(os.path.join(W, "ref_seq_res.npy"))) |
| if meta.get("knn"): |
| ref_rows = pd.read_parquet(os.path.join(W, "ref_rows.parquet")) |
| seq_pos = {s: i for i, s in enumerate(useq)} |
| kn = F.knn_feats(df["smiles"].values, df["sequence"].map(seq_pos).values, q_embn, |
| ref_rows["smiles"].values, ref_rows["seq_id"].values, ref_embn, |
| ref_rows[["kc", "km"]].values.astype(np.float32), |
| ks=tuple(meta.get("ks", [5]))) |
| pca = None |
| if meta["preset"].get("esm_pca"): |
| z = np.load(os.path.join(W, "esm_pca.npz")) |
| pca = {"mean": z["mean"], "comps": z["comps"]} |
| pf = dict(meta["preset"]) |
| if "ks" in pf: |
| pf["ks"] = tuple(pf["ks"]) |
| smi_dim, seq_dim, esm_dim = meta["dims"] |
| X = F.apply_variant(X0, kn, pf, smi_dim, seq_dim, esm_dim, ks_all=meta.get("ks", [5]), pca=pca, |
| n_blocks=len(meta["layers"]) if meta.get("layers") else 1, er=er) |
|
|
| import lightgbm as lgb |
| mlp_out = None |
| if meta.get("mlp"): |
| z = np.load(os.path.join(W, "mlp_prep.npz")) |
| st = {"mean": z["mean"], "std": z["std"]} |
| Xm = F.apply_variant(X0, None, {"nodi": True}, smi_dim, seq_dim, esm_dim) |
| Zm = F.mlp_prep(Xm, st) |
| outs = [] |
| for s in range(meta["mlp"]["n"]): |
| q = np.load(os.path.join(W, f"mlp_{s}.npz")) |
| layers = [(q[f"W{i}"], q[f"b{i}"]) for i in range(int(q["n"]))] |
| outs.append(F.mlp_forward(Zm, layers)) |
| mlp_out = np.mean(outs, axis=0) * z["ys"] + z["ym"] |
| out = pd.DataFrame({"id": df["id"]}) |
| for j, (name, col, mean_key) in enumerate([("kcat", "log10_kcat", "mean_kcat"), ("km", "log10_km", "mean_km")]): |
| preds = [] |
| for s in range(meta["n_seeds"]): |
| b = lgb.Booster(model_file=os.path.join(W, f"{name}_{s}.txt")) |
| preds.append(b.predict(X, num_threads=N_THREADS)) |
| p = np.mean(preds, axis=0) |
| if mlp_out is not None: |
| w = meta["mlp"]["w"] |
| p = (1 - w) * p + w * mlp_out[:, j] |
| p = np.nan_to_num(p, nan=meta[mean_key], posinf=meta[mean_key], neginf=meta[mean_key]) |
| out[col] = np.clip(p, -50, 50) |
|
|
| os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) |
| out.to_csv(args.output, index=False) |
| print(f"[predict] wrote {len(out)} rows in {time.time() - t0:.0f}s", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
