mrothroc/molgpt-moses-smiles-mixlab

MolGPT (MOSES), reproduced in mixlab, SMILES

A small SMILES transformer that generates drug-like molecules. It is a faithful reproduction of MolGPT, trained in mixlab on a single Apple-silicon Mac.

This is a model to learn from and tinker with, not a state-of-the-art result. Architecture is something you discover by intuition and iteration, and the point of mixlab is to make that loop fast: reproduce a known model quickly, then experiment from there. This checkpoint is a starting point for that.

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:

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()             # GPT2LMHeadModel
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))                  # SMILES strings

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