README.md

---
license: cc-by-nc-4.0
library_name: pytorch
tags:
  - chemistry
  - molecule-generation
  - generative-model
  - vae
  - selfies
  - rdkit
  - drug-discovery
  - electrolyte
  - batteries
  - cheminformatics
---

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

MolForge is a **conditional variational autoencoder** over
[SELFIES](https://github.com/aspuru-guzik-group/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/molforge
- **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 MolForge differs from existing models

- **A latent space, not left-to-right text generation.** Autoregressive models (ElectrolyteGPT,
  MolGPT) emit one token at a time. MolForge'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; MolForge 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 `molforge` library expects under `MOLVAE_ART_DIR`.

## Usage

```bash
pip install "git+https://github.com/NealKapadia/molforge.git"
```

```python
from huggingface_hub import snapshot_download
from molforge import MolForge

art = snapshot_download("NealKapadia/Molforge")   # downloads checkpoints/ + processed/
mf = MolForge(device="cpu", artifacts_dir=art)    # or device="cuda"

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

Or set the path manually instead of `artifacts_dir=`:
`export MOLVAE_ART_DIR=/path/to/download` (Windows: `$env:MOLVAE_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.