What it is
- 6,342,656 parameters: 8 layers, 256 dim, 8 heads, causal GPT with learned absolute positions and a GELU MLP. Exported as a native Hugging Face
GPT2LMHeadModel.
- Character-level SMILES tokenizer, vocab 30.
- Trained on the MOSES benchmark, one molecule per sequence, 33,000 steps (about 10 epochs), roughly 5 hours on an M1 Max. Newer Apple chips are much faster.
How close is the reproduction
Scored on 30,000 unconditional samples with the official moses.metrics:
Table with columns: Metric, this model, MolGPT (published)| Metric | this model | MolGPT (published) |
|---|
| Valid | 0.988 | 0.994 |
| Unique@10k | 0.991 | ~1.000 |
| Novelty | 0.783 | 0.797 |
| IntDiv | 0.857 | 0.857 |
Within a point on every firm metric. FCD/Test is 2.83; MolGPT never published an FCD at this size, so there is no firm number to match (the SELFIES sibling below reaches 0.43).
Load and generate
It loads with stock transformers as a GPT2LMHeadModel, no trust_remote_code. The framing tokens are BOS=1, EOS=2, PAD=0: prompt with BOS, sample until EOS, and read the SMILES off the tokens.
import torch
from transformers import AutoModelForCausalLM
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download
repo = "mrothroc/molgpt-moses-smiles-mixlab"
model = AutoModelForCausalLM.from_pretrained(repo).eval()
tok = Tokenizer.from_file(hf_hub_download(repo, "tokenizer.json"))
BOS, EOS, PAD = 1, 2, 0
out = model.generate(torch.full((8, 1), BOS), do_sample=True, temperature=1.0,
top_k=0, max_length=65, eos_token_id=EOS, pad_token_id=PAD)
for row in out.tolist():
seq = [t for t in row[1:] if t not in (BOS, PAD)]
if EOS in seq: seq = seq[:seq.index(EOS)]
print("".join(tok.id_to_token(t) for t in seq))
The SELFIES sibling
mrothroc/molgpt-moses-selfies-mixlab retrains the same model on SELFIES, a representation where every string decodes to a valid molecule by construction: validity 1.000, Novelty 0.859, FCD/Test 0.43. Same size, same budget, better numbers, from changing the representation rather than the model.
Learn more
The full recipe (reproduce, three ways to guarantee validity, and goal-directed fine-tuning) is in the mixlab cookbook: https://github.com/mrothroc/mixlab-cookbook. Trainer: mixlab.
For the story behind it, read You Don't Code the Model, You Edit the JSON.
Citation
This is a reproduction, so please cite what it builds on:
- MolGPT: Bagal, V.; Aggarwal, R.; Vinod, P. K.; Priyakumar, U. D. "MolGPT: Molecular Generation Using a Transformer-Decoder Model." J. Chem. Inf. Model. 2022, 62 (9), 2064-2076. doi:10.1021/acs.jcim.1c00600
- MOSES: Polykovskiy, D. et al. "Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models." Front. Pharmacol. 2020, 11, 565644. doi:10.3389/fphar.2020.565644