BatteryGen — a conditional SELFIES-VAE for de-novo molecule & battery-electrolyte design

BatteryGen is a conditional variational autoencoder over SELFIES representations of molecules, with about 42 million parameters (41,966,682), trained on 7,116,053 molecules curated from five public chemistry databases (Molport, ChEMBL, and ZINC for broad chemical coverage, plus electrolyte data from OEDB and CALiSol-23). It learns a smooth 256-dimensional latent space you can sample, traverse, and optimize, and because it decodes SELFIES, essentially 100% of generated strings are valid molecules (measured validity 1.000). It is purpose-built for de-novo battery-electrolyte design — generating candidate solvents and additives across chemistries (Li / Na / K / Mg / Zn / …) and ranking them with a paired electrolyte property model grounded in real electrolyte data.

Code / library: https://github.com/NealKapadia/BatteryGen Weights (this repo): checkpoints/best.pt

Architecture: embedding 512 → bidirectional GRU encoder (1024 × 2 layers) → latent 256 → GRU decoder (1024 × 2 layers), conditioned on 11 RDKit descriptors, with an auxiliary latent→property head. Decoder word-dropout 0.25 (Bowman et al.) for a meaningful latent. SELFIES robust alphabet, 79 tokens, max length 120.

Training data — 7,116,053 molecules from five public databases

Filtered to 3–60 heavy atoms and an organic element set, then de-duplicated:

Database	Molecules	Role
Molport "All Stock"	6,088,143	core corpus of purchasable molecules
ChEMBL-37 (sample)	800,000	bioactive chemical diversity
ZINC	227,902	additional lead-like diversity
OEDB + CALiSol-23 (solvents)	8	electrolyte solvents in the generator
Total	7,116,053	generative training set

OEDB and CALiSol-23 additionally provide the electrolyte solvents and 18,918 electrolyte formulations (conductivity, coordination, viscosity) that train the separate property model. Trained with the default SELFIES constraints (S=6 / P=5 allowed) so sulfonyl/phosphate electrolyte motifs round-trip.

Selected checkpoint: best.pt, selected by val_token_acc + 0.25·valid_rate.

Conditioning properties (fixed order)

MolWt, MolLogP, TPSA, QED, NumHDonors, NumHAcceptors, NumRotatableBonds, NumAromaticRings, NumRings, FractionCSP3, HeavyAtomCount

Evaluation (best.pt, 5,000 samples @ temperature 0.9)

Metric	Value
Validity	1.000
Uniqueness	0.998
Novelty (vs. training set)	0.995
Internal diversity	0.894
Reconstruction (exact)	0.945
Reconstruction (token acc)	0.998

Latent→property head R² (held-out): MolWt 0.994, TPSA 0.977, MolLogP 0.962, NumHAcceptors / NumRotatableBonds 0.969, QED 0.926, NumHDonors 0.922.

On the standard generative benchmark columns (validity / uniqueness / novelty / diversity) this model is competitive with — and on several columns exceeds — the autoregressive ElectrolyteGPT (Kim et al., JACS Au, 2026, 6, 2288–2302). The structural advantage is the latent space: smooth interpolation and gradient-based property optimization, which a left-to-right token model does not offer.

How BatteryGen differs from existing models

• A latent space, not left-to-right text generation. Autoregressive models (ElectrolyteGPT, MolGPT) emit one token at a time. BatteryGen's VAE provides a continuous latent space you can interpolate and optimize with gradients (e.g. "increase molecular weight by 10 while keeping everything else") — a token model cannot.
• Validity by construction. Decoding SELFIES yields essentially 100% valid molecules (measured 1.000), versus SMILES models that emit invalid strings.
• A full inverse-design system, not just a generator. The generator is paired with a predictive model (Optuna-tuned), an electrolyte property model, optional LLM guidance, and literature-grounded retrieval — an end-to-end loop from a plain-English request to a ranked, scored candidate list.
• Electrolyte-formulation awareness. Conductivity, coordination, and viscosity are system properties; BatteryGen models them at the formulation level (multi-cation), grounded in OEDB and CALiSol-23 data — most molecule generators ignore this.
• Multi-database breadth. Trained across five public databases, not a single catalog.

Files

checkpoints/best.pt              # the SELFIES-VAE generator weights
checkpoints/electrolyte_model.pt # optional: formulation property model (conductivity etc.)
processed/vocab.json             # SELFIES token vocabulary
processed/descriptor_stats.json  # descriptor normalization (mean/std)
processed/meta.json              # vocab size, max length, property order, constraints

This is exactly the layout the BatteryGen library expects under BATTERYGEN_ART_DIR.

Usage

pip install "git+https://github.com/NealKapadia/BatteryGen.git"

from huggingface_hub import snapshot_download
from batterygen import BatteryGen

art = snapshot_download("NealKapadia/BatteryGen")   # downloads checkpoints/ + processed/
bg = BatteryGen(device="cpu", artifacts_dir=art)    # or device="cuda"

bg.generate(10)                                   # 10 valid, novel SMILES
bg.generate(5, spec={"MolWt": 250, "QED": 0.8})   # property-targeted
z = bg.encode("OCCN(CCO)CCO"); bg.decode(z)       # latent round-trip
bg.properties("CCO")                              # RDKit descriptors

Or set the path manually instead of artifacts_dir=: export BATTERYGEN_ART_DIR=/path/to/download (Windows PowerShell: $env:BATTERYGEN_ART_DIR="...").

Limitations & intended use

• Research / educational use for molecular design and screening — not a substitute for experimental validation, synthesis feasibility, or safety assessment.
• Soft conditioning: spec targets nudge generation toward a value; they are not exact. For hard constraints, over-generate and filter by RDKit-computed properties.
• The generator covers a broad space of small-to-medium organic molecules; very small electrolyte molecules (EC/DEC/MeCN) sit at the edge of that distribution, so for tight electrolyte focus, specialize via fine-tuning plus the electrolyte property model.
• The electrolyte property model has labeled data for Li / Na / K only; the generator proposes candidates for any chemistry, but quantitative ranking beyond Li/Na/K needs additional labeled data.