Spaces:
Running on Zero
Running on Zero
| """MOSES checkpoint sampling using the release sampler and built-in metadata.""" | |
| from pathlib import Path | |
| import csv | |
| import torch | |
| from hydra import compose, initialize_config_dir | |
| from gem import sampler | |
| from gem.checkpoint_utils import load_model_checkpoint | |
| from gem.datasets import moses_dataset | |
| from gem.metrics_over_time import _molecule_records, _run_mcmc, _set_seed | |
| from gem.models.extra_features import ExtraFeatures | |
| from gem.models.extra_features_molecular import ExtraMolecularFeatures | |
| from gem.models.transformer_model import GraphTransformer | |
| TRANSPORT_STEPS = 225 | |
| CSV_FIELDS = ( | |
| "step", "index", "valid", "connected", "valid_connected", | |
| "num_fragments", "smiles", "error", | |
| ) | |
| def load_sampling_config(config_dir: Path): | |
| with initialize_config_dir(version_base="1.3", config_dir=str(config_dir.resolve())): | |
| cfg = compose(config_name="gem_metrics_over_time_moses") | |
| cfg.release_sampler.chain_warmup.steps = TRANSPORT_STEPS | |
| cfg.metrics_run.compute_fcd = False | |
| cfg.dataset.compute_fcd = False | |
| return cfg | |
| def build_sampling_model(cfg, device): | |
| # Inference needs only the shipped MOSES distributions, not processed data | |
| # or statistics files from the caller's working directory. | |
| infos = moses_dataset.MOSESinfos(None, cfg, load_cached_statistics=False) | |
| extra = ExtraFeatures(cfg.model.extra_features, cfg.model.rrwp_steps, infos) | |
| domain = ExtraMolecularFeatures(infos) | |
| # Derive feature widths on a deterministic two-carbon graph, using the same | |
| # feature functions as training. The +1 global feature is diffusion time; | |
| # GraphTransformer adds its own timestep embedding internally. | |
| atom_classes, bond_classes = infos.num_atom_types, len(infos.edge_types) | |
| x = torch.zeros(1, 2, atom_classes) | |
| x[..., infos.atom_encoder["C"]] = 1 | |
| e = torch.zeros(1, 2, 2, bond_classes) | |
| e[0, 0, 1, 1] = e[0, 1, 0, 1] = 1 | |
| example = dict(X_t=x, E_t=e, y_t=torch.zeros(1, 0), | |
| node_mask=torch.ones(1, 2, dtype=torch.bool)) | |
| extra_dims, domain_dims = extra(example), domain(example) | |
| infos.output_dims = {"X": atom_classes, "E": bond_classes, "y": 0} | |
| infos.input_dims = { | |
| key: base + getattr(extra_dims, key).shape[-1] + getattr(domain_dims, key).shape[-1] | |
| for key, base in {"X": atom_classes, "E": bond_classes, "y": 1}.items() | |
| } | |
| activation = torch.nn.SiLU() if cfg.model.activation == "silu" else torch.nn.ReLU() | |
| torch.set_float32_matmul_precision("medium") | |
| model = GraphTransformer( | |
| n_layers=cfg.model.n_layers, | |
| input_dims=infos.input_dims, | |
| hidden_mlp_dims=cfg.model.hidden_mlp_dims, | |
| hidden_dims=cfg.model.hidden_dims, | |
| output_dims=infos.output_dims, | |
| act_fn_in=activation, act_fn_out=activation, tf_activation=activation, | |
| ).to(device).eval() | |
| return model, infos, extra, domain | |
| def generate_batches( | |
| *, model, dataset_infos, extra_features, domain_features, device, run_cfg, | |
| num_samples, mixing_steps, batch_size, | |
| ): | |
| if num_samples <= 0 or batch_size <= 0 or mixing_steps < 0: | |
| raise ValueError("Sample count and batch size must be positive; mixing steps must be nonnegative.") | |
| for offset in range(0, num_samples, batch_size): | |
| graphs = sampler.initialize_random_graphs( | |
| batch_size=min(batch_size, num_samples - offset), | |
| dataset_info=dataset_infos, device=device, transition="uniform", | |
| ) | |
| nodes, edges = [g[0] for g in graphs], [g[1] for g in graphs] | |
| nodes, edges, _, _, _ = sampler.run_simple_v2_warmup_vectorized( | |
| model=model, dataset_info=dataset_infos, | |
| node_types_list=nodes, edge_types_list=edges, | |
| extra_features=extra_features, domain_features=domain_features, | |
| steps=TRANSPORT_STEPS, device=device, edits_per_step=1, | |
| amp_dtype=None, stop_when_unchanged=True, collect_stats=False, | |
| ) | |
| if mixing_steps: | |
| # Keep the calibrated annealing schedule; only the number of mixing | |
| # steps changes. Graph-input gradients remain enabled for sampling. | |
| nodes, edges, _, _, _ = _run_mcmc( | |
| model=model, dataset_infos=dataset_infos, | |
| node_types=nodes, edge_types=edges, | |
| extra_features=extra_features, domain_features=domain_features, | |
| device=device, run_cfg=run_cfg, steps=mixing_steps, step_offset=0, | |
| ) | |
| yield _molecule_records( | |
| nodes, edges, dataset_infos, step=mixing_steps, sample_offset=offset, | |
| ) | |
| def sample_checkpoint( | |
| checkpoint: Path, output: Path, *, num_samples: int, mixing_steps: int, | |
| batch_size: int, seed: int, config_dir: Path, | |
| ): | |
| cfg = load_sampling_config(config_dir) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Sampling on {device}: {TRANSPORT_STEPS} transport steps + {mixing_steps} mixing steps.", flush=True) | |
| model, infos, extra, domain = build_sampling_model(cfg, device) | |
| load_model_checkpoint(model, str(checkpoint), map_location=device, use_ema=False) | |
| _set_seed(seed) | |
| output.parent.mkdir(parents=True, exist_ok=True) | |
| # Retain all attempts, including invalid molecules, and flush each batch so | |
| # completed batches remain available if generation is interrupted. | |
| with output.open("x", newline="", encoding="utf-8") as handle: | |
| writer = csv.DictWriter(handle, fieldnames=CSV_FIELDS) | |
| writer.writeheader() | |
| count = 0 | |
| for records in generate_batches( | |
| model=model, dataset_infos=infos, extra_features=extra, | |
| domain_features=domain, device=device, run_cfg=cfg.metrics_run, | |
| num_samples=num_samples, mixing_steps=mixing_steps, batch_size=batch_size, | |
| ): | |
| writer.writerows({key: row.get(key) for key in CSV_FIELDS} for row in records) | |
| handle.flush() | |
| count += len(records) | |
| print(f"Generated {count}/{num_samples} samples.", flush=True) | |
| print(f"Saved samples to {output}", flush=True) | |