QM9 ChemicalVAE
This repository contains a ChemicalVAE model trained on the QM9 molecular dataset, following the SMILES-based variational autoencoder approach introduced in Automatic Chemical Design Using a Data-Driven Continuous Representation of Molecules (Gomez-Bombarelli, 2017)
Model summary
This model is a SMILES variational autoencoder (VAE) trained on QM9 molecules. It learns:
- an encoder from SMILES to a continuous latent representation
- a decoder from latent vectors back to SMILES
- optionally, a property prediction head trained jointly on QM9 properties
This checkpoint is intended for:
- latent-space molecular generation
- encoding QM9-like molecules into latent vectors
- decoding latent vectors into candidate SMILES
- downstream latent-space optimization methods such as COWBOYS-style Bayesian optimization
Dataset
Training data: QM9
QM9 is a dataset of about 134k small organic molecules with up to 9 heavy atoms and associated quantum-chemical properties. In the original paper, the QM9 VAE was trained on 108,000 molecules using canonicalized SMILES, with:
- maximum SMILES length: 34
- character set size: 22
- latent dimension: 156
The QM9 property-prediction setup in the paper used:
- HOMO
- LUMO
- R2 (electronic spatial extent) :contentReference[oaicite:1]{index=1}
Files in this repository
Required for inference
qm9_encoder.h5 — trained encoder weightsqm9_decoder.h5 — trained decoder weightsexp.json — model/training configurationqm9_chars.json — SMILES character vocabulary
Optional but useful
qm9_prop_pred.h5 — joint property prediction headqm9_prop_norm.csv — property normalization metadatalogs/qm9_history.csv — training history- additional notes/scripts for data conversion or evaluation
Model architecture
This model follows the QM9 configuration reported in the original ChemicalVAE paper:
- input representation: canonicalized SMILES
- maximum sequence length: 34
- encoder: 3 1D convolutional layers
- latent dimension: 156
- decoder: 3 recurrent layers with hidden size 500
- optional property predictor: 2 fully connected layers of 1000 units with dropout 0.2 :contentReference[oaicite:2]{index=2}
Training details
This model was trained with the chemical_vae codebase on QM9-style data using:
- canonicalized SMILES
- fixed maximum SMILES length
- a character vocabulary stored in
qm9_chars.json
- optionally joint supervision on HOMO / LUMO / R2
If you want to reproduce training, the key ingredients are:
- a CSV with first column
smiles
- optional property columns such as
HOMO, LUMO, R2
- a matching
qm9_chars.json
- an
exp.json specifying architecture and training settings
Loading example
Below is an example of loading the model with the upstream chemical_vae utilities.
1from chemvae.vae_utils import VAEUtils
2
3vae = VAEUtils(
4 exp_file="exp.json",
5 directory=".",
6 encoder_file="qm9_encoder.h5",
7 decoder_file="qm9_decoder.h5",
8)
9
10# Encode molecules
11smiles = ["CCO", "CCN"]
12hot = vae.smiles_to_hot(smiles, canonize_smiles=True, check_smiles=False)
13z = vae.encode(hot)
14
15# Decode latent vectors
16decoded_hot = vae.decode(z)
17decoded_smiles = vae.hot_to_smiles(decoded_hot, strip=True)
18
19print(decoded_smiles)
