File size: 4,781 Bytes
4b488f2 | 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 | """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)) # the evaluator gives us 2 vCPUs
ESM_BUDGET_S = 480 # max 8 minutes for ESM, then fall back
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()
|