GenMolMolecularGenerator

Masked diffusion over SAFE fragments, optionally retaining a supplied scaffold.

Usage

GenMol trains on SMILES strings and generates molecules using masked diffusion over SAFE fragments. Pass scaffold="c1ccccc1" to generate() to retain a benzene substructure in generated molecules. Install its optional packages with pip install ".[genmol]" from the source checkout.

from torch_molecule import GenMolMolecularGenerator

Methods

Constructor

class GenMolMolecularGenerator(num_layer=12, num_head=12, hidden_size=768, intermediate_size=3072, max_len=256, dropout=0.1, attention_dropout=0.1, hidden_act='gelu', layer_norm_eps=1e-12, initializer_range=0.02, type_vocab_size=2, ema_decay=0.9999, antithetic_sampling=True, sampling_eps=0.001, global_mean_loss=True, batch_size=64, epochs=100, learning_rate=0.0003, adamw_betas=(0.9, 0.999), adamw_eps=1e-08, weight_decay=0.0, warmup_steps=2500, grad_clip_value=1.0, verbose='none', *, device=None, model_name='GenMolMolecularGenerator')[source]

Generate molecules using GenMol V1 masked diffusion over SAFE sequences.

Parameters:

  • num_layer (int, default=12) – Number of bidirectional transformer layers.

  • num_head (int, default=12) – Number of attention heads; must divide hidden_size.

  • hidden_size (int, default=768) – Transformer hidden dimension.

  • intermediate_size (int, default=3072) – Feed-forward hidden dimension.

  • max_len (int, default=256) – Maximum token length including beginning and end tokens.

  • dropout (float, default=0.1) – Hidden and attention dropout probabilities.

  • attention_dropout (float, default=0.1) – Hidden and attention dropout probabilities.

  • hidden_act (str, default="gelu") – BERT feed-forward activation: gelu, relu, silu, or gelu_new.

  • layer_norm_eps (float, default=1e-12) – Layer normalization epsilon.

  • initializer_range (float, default=0.02) – Standard deviation of the weight initializer.

  • type_vocab_size (int, default=2) – BERT token-type vocabulary size.

  • ema_decay (float, default=0.9999) – Exponential moving average decay. Zero disables EMA.

  • antithetic_sampling (bool, default=True) – Stratify sampled masking times within each batch.

  • sampling_eps (float, default=0.001) – Lower time bound for antithetic sampling.

  • global_mean_loss (bool, default=True) – Average loss over tokens; False gives each molecule equal weight.

  • batch_size (int, default=64) – Training and generation batch size.

  • epochs (int, default=100) – Number of training epochs.

  • learning_rate (float, default=0.0003) – AdamW learning rate after warmup.

  • adamw_betas (tuple, default=(0.9, 0.999)) – AdamW moment decay rates.

  • adamw_eps (float, default=1e-8) – AdamW denominator epsilon.

  • weight_decay (float, default=0.0) – AdamW weight decay.

  • warmup_steps (int, default=2500) – Linear learning-rate warmup steps; zero disables warmup.

  • grad_clip_value (float or None, default=1.0) – Maximum gradient norm; None disables clipping.

  • verbose (str, default="none") – One of none, progress_bar, or print_statement.

  • device (torch.device or str, optional) – CPU or CUDA device; defaults to CUDA when available.

  • model_name (str, default="GenMolMolecularGenerator") – Estimator name.

References

Paper: https://proceedings.mlr.press/v267/lee25o.html Code: https://github.com/NVIDIA-BioNeMo/genmol

Methods

fit

Train a fresh model on SMILES strings.

generate

Generate SMILES, optionally containing a supplied scaffold.

load_from_pretrained

Load GenMol V1 weights with matching constructor parameters.

For shared checkpoint and parameter operations, see Base Modules.

fit

GenMolMolecularGenerator.fit(X_train, y_train=None)[source]

Train a fresh model on SMILES strings.

Parameters:

  • X_train (List[str]) – Training molecules as SMILES strings. Polymer attachment atoms (*) are not preserved in generated molecules.

  • y_train (None) – Property labels are unsupported.

Returns:

The fitted estimator.

Return type:

GenMolMolecularGenerator

generate

GenMolMolecularGenerator.generate(n_samples=10, *, scaffold=None, temperature=0.8, randomness=0.5, min_new_tokens=40, max_attempts=10, guidance_scale=2.0, context_mask_ratio=0.0, fix=True)[source]

Generate SMILES, optionally containing a supplied scaffold.

n_samples is the requested number of valid outputs, which may include duplicates. min_new_tokens controls the minimum masked extension length. max_attempts bounds rejection sampling per output batch; exhaustion warns and returns the available molecules, which may be an empty list. temperature and randomness control token and confidence sampling. context_mask_ratio enables molecular context guidance for scaffold generation. fix enables the original SAFE fragment repair during decoding.

Parameters:

  • n_samples (int)

  • scaffold (str | None)

  • temperature (float)

  • randomness (float)

  • min_new_tokens (int)

  • max_attempts (int)

  • guidance_scale (float)

  • context_mask_ratio (float)

  • fix (bool)

Return type:

List[str]

load_from_pretrained

GenMolMolecularGenerator.load_from_pretrained(path, *, sequence_lengths)[source]

Load GenMol V1 weights with matching constructor parameters.

sequence_lengths supplies the token-length distribution used for generation. The checkpoint must contain tensors and plain Python values. Load saved torch-molecule estimators with load_from_local.

Parameters:

  • path (str)

  • sequence_lengths (Sequence[int])

Return type:

GenMolMolecularGenerator