MolecularDiffusion.modules.models.tabasco.callbacks.save_molecules

Attributes

log

Classes

SaveGeneratedMolsCallback

Log generated molecules, trajectories, and SDF files to Weights & Biases.

Module Contents

class MolecularDiffusion.modules.models.tabasco.callbacks.save_molecules.SaveGeneratedMolsCallback(num_samples: int = 100, num_trajectories: int = 10, num_display_mols: int = 10, num_sampling_steps: int = 100, compute_every: int = 1000)

Bases: lightning.Callback

Log generated molecules, trajectories, and SDF files to Weights & Biases.

Performs several logging passes:

• Sample preview molecules (_log_mols_to_display).

• Highlight novel molecules not present in the training set.

• Store invalid molecules for debugging.

• Persist full SDF and selected trajectories as W&B artifacts.

Heavy I/O is gated by compute_every and executed only on rank-0.

Args: num_samples: Molecules sampled each evaluation. num_trajectories: Trajectories to save as SDF files. num_display_mols: Maximum molecules displayed inline in W&B. num_sampling_steps: Reverse-diffusion steps in sampling. compute_every: Scheduler interval in global steps.

on_validation_epoch_end(trainer: lightning.Trainer, lightning_module: lightning.LightningModule) → None

Core hook: sample molecules and trigger logging helpers.

compute_every = 1000

mol_converter

next_compute = 0

num_display_mols = 10

num_samples = 100

num_sampling_steps = 100

num_trajectories = 10

MolecularDiffusion.modules.models.tabasco.callbacks.save_molecules.log