| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import joblib |
| import pandas as pd |
|
|
| from src.features.rdkit_features import featurize_smiles, FeatConfig |
|
|
|
|
| def load_artifacts(artifact_dir: Path): |
| meta = json.loads((artifact_dir / "metadata.json").read_text(encoding="utf-8")) |
| pre = joblib.load(artifact_dir / "preprocess.joblib") |
| models = {t: joblib.load(artifact_dir / "models" / f"{t}.joblib") for t in meta["targets"]} |
| return meta, pre, models |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="Predict ADMET from SMILES using saved artifacts.") |
| ap.add_argument("--artifact_dir",type=str,default="artifacts/admet_smiles_merge_didb_seed202") |
| ap.add_argument("--smiles", type=str, required=True) |
| ap.add_argument("--smiles_col", type=str, default="smiles") |
| ap.add_argument("--threshold", type=float, default=0.5) |
| args = ap.parse_args() |
|
|
| meta, pre, models = load_artifacts(Path(args.artifact_dir)) |
|
|
| cfg = FeatConfig(fp_radius=meta["features"]["fp"]["radius"], fp_nbits=meta["features"]["fp"]["nbits"]) |
| df = pd.DataFrame({args.smiles_col: [args.smiles]}) |
|
|
| X, valid = featurize_smiles(df, smiles_col=args.smiles_col, config=cfg, add_physchem=True, drop_invalid=False) |
| if not valid.all(): |
| raise ValueError("Invalid SMILES.") |
|
|
| feature_names = X.columns.tolist() |
| Xt = pd.DataFrame(pre.transform(X), columns=feature_names, index=X.index) |
|
|
|
|
| out = {"smiles": args.smiles, "pred": {}, "meta": {"model_version": meta["model_version"]}} |
| for t, m in models.items(): |
| prob = float(m.predict_proba(Xt)[:, 1][0]) |
| pred = int(prob >= args.threshold) |
|
|
| inv = {v: k for k, v in meta["label_maps"][t]["mapping"].items()} |
| out["pred"][t] = {"prob": prob, "pred": pred, "label": inv[pred]} |
|
|
| print(json.dumps(out, indent=2, ensure_ascii=False)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
