schrodinger.application.matsci.flywheel.ldmol.autoencoder module

Variational autoencoder for SMILES encoding and decoding.

Maps SMILES strings to 64-dim latent vectors and back via pre-trained BERT encoder/decoder loaded from checkpoint.

class schrodinger.application.matsci.flywheel.ldmol.autoencoder.VAEAutoencoder(device='cpu', logger=None)

Bases: BaseModel

VAE autoencoder for SMILES to latent-space conversion.

All model weights are loaded lazily on first use. encode() maps SMILES → 64-dim latents. decode() maps 64-dim latents → SMILES autoregressively.

__init__(device='cpu', logger=None)
Parameters:

  • device (str) – Compute device (‘cpu’, ‘cuda’, ‘mps’)

  • logger – Logger instance for debug output

ensureModelLoaded()

Load encoder, decoder, and prefix layers on first call.

loadCheckpoint()

Download (if needed) and load the VAE checkpoint.

Delegates weight distribution to loadEncoder, loadDecoder, and loadPrefixLayers.

loadModule(module, partial_state=None, strict=True)

Load weights into a module then move it to self.device in eval mode.

Parameters:

  • module (torch.nn.Module) – Module to configure

  • partial_state (dict) – State dict to load, or None to skip loading

  • strict (bool) – Passed to load_state_dict

Return type:

torch.nn.Module

Returns:

The module, on self.device, in eval mode

loadBertModel(model_cls, config_path, key_prefix, state_dict)

Build a BERT model and load weights from a checkpoint state dict.

Parameters:

  • model_cls – BertModel or BertForMaskedLM class to instantiate

  • config_path (str) – Path to JSON config file

  • key_prefix (str) – Checkpoint key prefix to strip when filtering

  • state_dict (dict) – Full checkpoint state dict

Return type:

torch.nn.Module

Returns:

Loaded model in eval mode on self.device

loadEncoder(state_dict)

Build and load the BERT SMILES encoder.

Parameters:

state_dict (dict) – Full checkpoint state dict

loadDecoder(state_dict)

Build and load the BERT SMILES decoder.

Parameters:

state_dict (dict) – Full checkpoint state dict

loadPrefixLayers(state_dict)

Build and load the latent projection layers.

encode_prefix: Linear(1024 → 64) compresses BERT output. decode_prefix: Linear(64 → 1024) expands latent for decoder.

Parameters:

state_dict (dict) – Full checkpoint state dict

encode(smiles_list)

Encode SMILES strings to 64-dim latent vectors.

Parameters:

smiles_list (list(str) or str) – SMILES strings to encode

Return type:

torch.Tensor

Returns:

Latents (batch, seq_len, latent_dim)

decode(latents, max_length=None)

Decode 64-dim latent vectors to SMILES strings.

Uses autoregressive greedy decoding: at each step the highest-logit token is appended and fed back as input.

Parameters:

  • latents (torch.Tensor) – Latents (batch, seq_len, latent_dim)

  • max_length (int) – Maximum tokens to generate

Return type:

list(str)

Returns:

Decoded SMILES strings, one per batch item