A small GPT that writes molecules.

MolGPT learns from 250,000 drug like molecules and writes new ones as SMILES strings. Ask for a range of logP and QED and it writes to match. Then explore what it invented as a network of similar molecules, in 2D and in 3D.

Open the interface View on GitHub

Version 0.1, an early release. More coming.

The MolGPT interface showing a network of generated molecules with a 2D and 3D view of the selected one
Ninety molecules generated by the model, linked by Tanimoto similarity.

Watch it write, one piece at a time

The model writes a molecule token by token. Each piece carries a bar with the confidence the model had in that choice, and hovering shows the other options it considered. The recording comes straight from the trained model, and only the playback speed is artificial.

Try it in the interface
The GPT writing a SMILES string token by token, with a confidence bar under each token and the resulting 2D structure

What the small model learned

A model with 0.8 million parameters, trained for about 28 minutes on a CPU. Numbers come from 4000 generated molecules compared with a real ZINC250k sample, measured with RDKit.

56%valid molecules, the weak spot for a model this small
100%unique among the valid ones, and almost all new
0.88 and 0.88diversity of the model against real ZINC
0.73 and 0.74mean QED of the model against real ZINC
0.75 to 4.17mean logP from the lowest to the highest requested bin
Training results page with a comparison table, the logP control chart and the learning curve

How it works

The pipeline is short on purpose, so every step is easy to read in the code.

  1. Tokenize every moleculeEach SMILES string is split into atoms, bonds and ring markers.
  2. Add two control tokensOne marks the logP bin and one marks the QED bin, so the model can be asked for a range.
  3. Train a GPT on the next tokenThe model sees a molecule so far and learns to predict the next piece.
  4. Generate and checkRDKit reads each result and keeps only the ones that are real molecules.
  5. Filter by drug likenessLipinski rules and PAINS alerts remove poor candidates before the visual step.

Run it yourself

The small preset trains on a laptop CPU in about half an hour. The full preset is meant for a GPU, for example a free Colab T4.

pip install rdkit torch pandas numpy matplotlib
python molgpt.py cpu     # small model, about 28 minutes on 2 CPU cores
python molgpt.py         # full model, meant for a GPU

MolGPT only knows the shape of drug like molecules from ZINC250k. It knows nothing about protein targets or binding affinity, so what it writes are starting candidates and not drugs. The next step is an affinity model, trained on public benchmarks, acting as a judge for the generated molecules.