File size: 4,886 Bytes
ab9bea9 | 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 110 111 112 113 114 115 116 | """Dataset-free sampling from a downloaded maltefranke/morph bundle."""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import sys
def load_bundle(root, variant, epoch, device="cpu", steps=100):
"""Load the exact bundled runtime, statistics, current weights and EMA."""
root = Path(root).resolve()
bundle = root / "chembl3d" / variant
runtime = bundle / "runtime"
if "morph" in sys.modules:
raise RuntimeError("Load each runtime in a fresh Python process.")
sys.path.insert(0, str(runtime / "src"))
os.environ.setdefault("MORPH_COMPILE", "0")
os.environ.setdefault("MORPH_ATTENTION_BACKEND", "sdpa")
os.environ.setdefault("MORPH_FLASH_ATTN", "0")
os.environ.setdefault("MORPH_MUON_EXECUTION", "native")
import torch
from omegaconf import OmegaConf
import morph.config # registers the training config's resolvers
from morph.inference.generate import _build_module, load_checkpoint
if device == "auto":
device = "cuda" if torch.cuda.is_available() else "cpu"
cfg = OmegaConf.load(bundle / "config.yaml")
cfg.data.preprocessing.root = str(bundle)
for field, filename in {
"atom_tokens_path": "atom_tokens.txt",
"edge_tokens_path": "edge_tokens.txt",
"charge_tokens_path": "charge_tokens.txt",
"distributions_path": "distributions.pt",
}.items():
path = bundle / filename
if not path.is_file():
raise FileNotFoundError(path)
cfg.data.preprocessing[field] = str(path)
cfg.integrator.integrator.num_integration_steps = steps
module, vocab, distributions = _build_module(cfg)
checkpoint = bundle / f"epoch_{epoch:03d}.ckpt"
load_checkpoint(module, checkpoint)
module.eval().to(device)
return module, vocab, distributions, cfg
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parent)
parser.add_argument("--variant", choices=["explicit_h", "implicit_h"], required=True)
parser.add_argument("--epoch", type=int, required=True)
parser.add_argument("--num-mols", type=int, default=16)
parser.add_argument("--batch-size", type=int, default=16)
parser.add_argument("--steps", type=int, default=100)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
parser.add_argument("--output", type=Path, default=Path("molecules.sdf"))
args = parser.parse_args()
if min(args.num_mols, args.batch_size, args.steps) < 1:
parser.error("num-mols, batch-size, and steps must be positive")
import torch
torch.manual_seed(args.seed)
module, vocab, _, cfg = load_bundle(
args.root, args.variant, args.epoch, args.device, args.steps
)
from morph.inference.generate import _prior_loader, final_mol
from rdkit import Chem
loader = _prior_loader(
module.distributions, args.num_mols, args.batch_size,
uniform_n_atoms_padding=int(cfg.data.uniform_n_atoms_padding),
n_atoms_strategy=str(cfg.data.n_atoms_strategy),
)
records = []
args.output.parent.mkdir(parents=True, exist_ok=True)
with Chem.SDWriter(str(args.output)) as writer, torch.inference_mode():
for batch_index, batch in enumerate(loader):
batch = tuple(graph.to(module.device) for graph in batch)
result = module.predict_step(batch, batch_index)
for trajectory, valid in zip(result["all_mols"], result["is_valid"], strict=True):
frame = trajectory[-1]
if not torch.isfinite(frame.x).all():
raise RuntimeError("Non-finite generated coordinates")
mol = final_mol(trajectory, vocab)
index = len(records)
records.append({
"index": index, "valid": bool(valid),
"smiles": Chem.MolToSmiles(mol) if mol is not None else None,
"represented_atoms": int(frame.x.shape[0]),
})
if valid and mol is not None:
mol.SetProp("_Name", f"sample_{index}")
writer.write(mol)
assert len(records) == args.num_mols
report = {
"variant": args.variant, "completed_epochs": args.epoch,
"seed": args.seed, "integration_steps": args.steps,
"weights": "EMA", "coordinate_units": "angstrom",
"num_attempted": len(records),
"num_valid": sum(row["valid"] for row in records),
"samples": records,
}
args.output.with_suffix(".json").write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps({k: v for k, v in report.items() if k != "samples"}), flush=True)
if __name__ == "__main__":
main()
|