HFPretrainedMolecularGenerator

NovoMolGen, MolGen, Molexar, or SAFE-GPT checkpoints; constraints depend on the model family.

from torch_molecule import HFPretrainedMolecularGenerator

Examples · Methods

Constructor

class HFPretrainedMolecularGenerator(repo_id, max_length=128, revision=None, trust_remote_code=False, tokenizer_repo_id=None, generate_max_length=64, batch_size=8, epochs=1, learning_rate=5e-05, weight_decay=0.01, grad_norm_clip=1.0, *, device=None, model_name='HFPretrainedMolecularGenerator', verbose='none')[source]

Hugging Face pretrained models as molecular generators.

This class loads pretrained generative models from Hugging Face and exposes a sklearn-style fit / generate interface consistent with other generators in torch-molecule.

Supported generation modes depend on the model family:

  • NovoMolGen: de novo SMILES generation from BOS.

  • MolGen-large / MolGen-large-opt: SELFIES seq2seq generation via prefix_selfies= or scaffold= (SMILES converted internally).

  • Molexar: Fragment-SELFIES de novo and fragment-constrained generation via start_smiles / start_string / conditions (omni).

  • SAFE-GPT: GPT-2 causal LM on SAFE strings; de novo and scaffold= prefix.

Supported model repositories include:

Parameters:

  • repo_id (str) – Hugging Face repository id of the pretrained generator.

  • max_length (int, default=128) – Maximum sequence length used when loading the tokenizer.

  • revision (Optional[str], default=None) – Model revision on the Hugging Face Hub. NovoMolGen defaults to "hf-checkpoint" so standard model.generate works out of the box.

  • trust_remote_code (bool, default=False) – Whether to trust remote code when loading from Hugging Face. Automatically enabled for Molexar.

  • tokenizer_repo_id (Optional[str], default=None) – Optional Hugging Face repo for the tokenizer.

  • generate_max_length (int, default=64) – Default generation length. For causal LMs (NovoMolGen, SAFE-GPT) this is max_new_tokens so a scaffold= prefix does not consume the budget. For MolGen this is still decoder max_length.

  • batch_size (int, default=8) – Batch size used when fine-tuning on SMILES data.

  • epochs (int, default=1) – Number of fine-tuning epochs when fit(X) is called.

  • learning_rate (float, default=5e-5) – Learning rate for fine-tuning.

  • weight_decay (float, default=0.01) – Weight decay for fine-tuning.

  • grad_norm_clip (Optional[float], default=1.0) – Maximum gradient norm during fine-tuning. Set to None to disable clipping.

  • device (Optional[Union[torch.device, str]], default=None) – Device to run the model on.

  • model_name (str, default="HFPretrainedMolecularGenerator") – Name identifier for the model instance.

  • verbose (str, default="none") – Progress display mode: "none", "progress_bar", or "print_statement".

Example

examples/generation/pretrained_safe_gpt.py
import tempfile
import shutil

from torch_molecule import HFPretrainedMolecularGenerator

REPO_ID = "datamol-io/safe-gpt"
N_SAMPLES = 5
SHORT_SCAFFOLD = "c1ccccc1"
LONG_SCAFFOLD = "CC1=CC=C(C=C1)C2=CC(=NN2C3=CC=C(C=C3)S(=O)(=O)N)C(F)(F)F"


def test_safe_gpt_generator():
    print("\n=== Testing SAFE-GPT initialization ===")
    model = HFPretrainedMolecularGenerator(
        repo_id=REPO_ID,
        verbose="progress_bar",
    )
    print("SAFE-GPT initialized successfully")

    print("\n=== Testing SAFE-GPT loading from Hugging Face ===")
    model.fit()
    print("SAFE-GPT loaded successfully")

    print("\n=== Testing SAFE-GPT de novo generation ===")
    generated_smiles = model.generate(n_samples=N_SAMPLES)
    print(f"Generated {len(generated_smiles)} molecules")
    print("Example generated SMILES:", generated_smiles[:2])

    print("\n=== Testing SAFE-GPT short scaffold generation ===")
    generated_smiles = model.generate(n_samples=N_SAMPLES, scaffold=SHORT_SCAFFOLD)
    print(f"Generated {len(generated_smiles)} molecules from scaffold {SHORT_SCAFFOLD}")
    print("Example generated SMILES:", generated_smiles[:2])

    print("\n=== Testing SAFE-GPT long scaffold generation ===")
    generated_smiles = model.generate(n_samples=2, scaffold=LONG_SCAFFOLD)
    print(f"Generated {len(generated_smiles)} molecules from long scaffold")
    print("Example generated SMILES:", generated_smiles[:2])

    print("\n=== Testing SAFE-GPT saving and loading ===")
    save_path = tempfile.mkdtemp(prefix="torchmol-safe-gpt-")
    model.save_to_local(save_path)
    print(f"SAFE-GPT saved to {save_path}")

    loaded_model = HFPretrainedMolecularGenerator(repo_id=REPO_ID)
    loaded_model.load_from_local(save_path)
    print("SAFE-GPT loaded from local directory")

    generated_smiles = loaded_model.generate(n_samples=2, scaffold=SHORT_SCAFFOLD)
    print(f"Generated {len(generated_smiles)} molecules with loaded model")
    print("Example generated SMILES:", generated_smiles[:2])

    shutil.rmtree(save_path)
    print(f"Cleaned up {save_path}")


if __name__ == "__main__":
    test_safe_gpt_generator()

Methods

fit

Load the pretrained model from Hugging Face.

generate

Generate molecules as SMILES strings.

save_to_local

Save the model and tokenizer to a local directory.

load_from_local

Load a model and tokenizer saved by save_to_local().

For shared checkpoint and parameter operations, see Base Modules.

fit

HFPretrainedMolecularGenerator.fit(X=None, y=None)[source]

Load the pretrained model from Hugging Face.

Parameters:

  • X (Optional[List[str]], default=None) – Optional SMILES strings for fine-tuning. When provided, the pretrained weights are adapted on the encoded family-specific representation.

  • y (Optional[np.ndarray], default=None) – Reserved for future conditional fine-tuning. Currently ignored with a warning.

Returns:

The fitted generator instance.

Return type:

HFPretrainedMolecularGenerator

generate

HFPretrainedMolecularGenerator.generate(n_samples=10, **kwargs)[source]

Generate molecules as SMILES strings.

Parameters:

  • n_samples (int, default=10) – Number of molecules to generate.

  • **kwargs – Additional arguments forwarded to the family-specific generator. For causal LMs, common options include max_new_tokens, temperature, do_sample, and scaffold. max_length is still accepted as a Hugging Face total-length override. For MolGen, use prefix_selfies or scaffold plus optional num_beams, min_length, and max_length. For Molexar, use start_smiles, start_string, generation_task, or conditions for omni models. For SAFE-GPT, use scaffold= with a SMILES prefix (converted to SAFE internally). A prefix does not guarantee that the final molecule retains the supplied substructure.

Returns:

Generated SMILES strings. For MolGen, Molexar, and SAFE-GPT, invalid decodes are dropped, so the list may be shorter than n_samples.

Return type:

List[str]

save_to_local

HFPretrainedMolecularGenerator.save_to_local(path)[source]

Save the model and tokenizer to a local directory.

Parameters:

path (str)

Return type:

None

load_from_local

HFPretrainedMolecularGenerator.load_from_local(path)[source]

Load a model and tokenizer saved by save_to_local().

Parameters:

path (str)

Return type:

None