Skip to content

Reproducing this work end to end

Every number in reports/ comes from the commands below. Nothing here depends on the specific cluster used during development: any machine with a CUDA GPU works for training, and the evaluation half runs on a laptop.

0. Environment

Two environments are needed, because MolScribe (the external baseline) pins torch<2 and cannot share an interpreter with a current PyTorch.

Training / evaluation (Python 3.11+):

python -m venv .venv && . .venv/bin/activate
pip install -e ".[dev]"
python -m pic2smiles.cli env          # confirms rdkit, torch and the accelerator
python scripts/capture_environment.py reports/environment.json

On a headless Linux node, RDKit's Cairo drawing also needs X libraries that a pip install does not bring:

conda install -c conda-forge xorg-libxrender xorg-libxext

Without them, from rdkit.Chem.Draw import rdMolDraw2D fails with ImportError: libXrender.so.1.

MolScribe baseline (optional, Python 3.10, Linux):

conda create -y -n molscribe python=3.10 && conda activate molscribe
pip install torch==1.13.1 torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cpu
pip install molscribe
pip install "opencv-python-headless==4.10.0.84"   # not opencv-python: compute nodes lack libGL.so.1
pip install --no-deps "numpy==1.26.4"             # rdkit 2022.09 predates numpy 2
curl -sSL -o swin_base_char_aux_1m.pth \
  https://huggingface.co/yujieq/MolScribe/resolve/main/swin_base_char_aux_1m.pth

This environment does not work on macOS 26+; see docs/BASELINES.md.

1. Molecule corpus

curl -sSL -o chembl_37_chemreps.txt.gz \
  https://ftp.ebi.ac.uk/pub/databases/chembl/ChEMBLdb/latest/chembl_37_chemreps.txt.gz
zcat chembl_37_chemreps.txt.gz | tail -n +2 | cut -f2 > chembl_37_smiles.txt

# Hand-drawn evaluation molecules, to be excluded from training
cut -f2 DECIMER_HDM_Dataset_SMILES.tsv | tail -n +2 > decimer_holdout.smi

python -m pic2smiles.cli corpus \
  --source chembl_37_smiles.txt \
  --output data/corpus/chembl \
  --holdout decimer_holdout.smi \
  --min-heavy-atoms 4 --max-heavy-atoms 48 \
  --val-size 5000 --test-size 5000 --min-token-count 20 \
  --seed 20260819

Writes train.smi, val.smi, test.smi, vocab.json and corpus_meta.json. The metadata file records the curation counts and how many molecules the leakage guard removed; the report embeds it verbatim as the declared scope.

2. Frozen evaluation sets

Evaluation images are written once and reused, so runs stay comparable.

# Rendered ChEMBL test molecules, three difficulty tiers
python -m pic2smiles.cli eval-set \
  --smiles data/corpus/chembl/test.smi \
  --output data/eval/chembl_test_linux \
  --count 1000 --image-size 384 --seed 20260819

# Memorization-floor gallery: training molecules, clean style
head -3000 data/corpus/chembl/train.smi > gallery.smi
python -m pic2smiles.cli eval-set --smiles gallery.smi \
  --output data/eval/gallery_train --count 3000 --styles clean --image-size 384

Build the eval set on the same host that trains, then copy the images to wherever they are scored. Fonts differ between operating systems, and a set rendered on another machine measures font generalization rather than accuracy (see D14 in docs/DECISIONS.md). To measure that axis deliberately, build a second set on the other host and score both.

The hand-drawn set comes from Zenodo 10.5281/zenodo.6456306; convert its TSV into the manifest schema with columns sample_id,source,split,style,canonical_smiles,image_path.

3. Train

Size the run from a measurement, not a guess:

python -m pic2smiles.cli train --corpus data/corpus/chembl --output runs/probe \
  --preset base --steps 300 --batch-size 64 --amp bf16 --device cuda \
  --eval-every 0 --checkpoint-every 0

Read img/s from the log; wall clock is steps * batch_size / img_per_second. Then launch the real run:

python -m pic2smiles.cli train \
  --corpus data/corpus/chembl --output runs/molmini_base \
  --preset base --steps 240000 --batch-size 64 \
  --lr 3e-4 --warmup 2000 --num-workers 14 \
  --amp bf16 --device cuda \
  --log-every 100 --eval-every 5000 --eval-samples 384 --checkpoint-every 5000 \
  --seed 20260819

Or, under Slurm: sbatch scripts/slurm/train_a100.sbatch <corpus> <output> 240000 64 base.

Watch enc_cos in the log. It is the cosine similarity between the encoder's memory for two different images. It should sit well below 1. Values approaching 1.0 mean the encoder has stopped distinguishing images and the decoder is learning an unconditional SMILES prior — a failure a falling loss curve does not reveal. The run aborts by itself above 0.995 after step 500.

Running from a git worktree

finalize.sh and the CLI resolve dataset paths under the repository root. A fresh clone has a real data/ directory and needs nothing extra, but a git worktree does not share the main checkout's ignored files, so link it:

ln -s /path/to/main/checkout/data data

4. Score and report

scripts/finalize.sh <training_host> <remote_root>

which is equivalent to, for each checkpoint and evaluation set:

python -m pic2smiles.cli benchmark \
  --manifest data/eval/chembl_test_linux/manifest.csv --dataset chembl_test_linux \
  --models molmini --checkpoint checkpoints/molmini_base.pt \
  --batch-size 24 --output reports/results --run-id molmini_base_chembl_test_linux

python -m pic2smiles.cli benchmark \
  --manifest data/eval/chembl_test_linux/manifest.csv --dataset chembl_test_linux \
  --models oracle image_hash_nn --gallery data/eval/gallery_train/manifest.csv \
  --output reports/results --run-id reference_chembl_test_linux

python -m pic2smiles.cli report --results reports/results \
  --output reports/benchmark_report.md --scope reports/corpus_meta.json

oracle must score exactly 1.000. If it does not, the scoring path is broken and no other number in the table is trustworthy.

Both best.pt and last.pt are scored. best.pt is chosen during training on a 384-sample evaluation whose molecules differ between evaluations -- at a ~0.3 rate that carries roughly 2.3 points of standard error, so "best" can be a lucky draw, and under a WSD schedule the genuinely best weights usually come last, after the decay phase. The 3000-image frozen set is the estimate that can carry a selection decision, so whichever checkpoint wins there is the one reported.

Comparing against MolScribe fairly

Only numbers produced by this harness on the same frozen manifest are comparable. In particular, the [eval] lines in a training log are not comparable to a benchmark table: they draw val-split molecules from an infinite stream, render them fresh each time, and sample 384 of them. Different molecules, different pixels, different sample each call. Use them to read a trajectory, never to compare against another system.

5. Predict on your own image

python -m pic2smiles.cli predict --checkpoint checkpoints/molmini_base.pt structure.png

Seeds and determinism

what seed effect
corpus split 20260819 which molecules land in train/val/test
frozen eval rendering 20260819 the exact evaluation images
training data stream --seed order and styling of training samples
in-training eval 999983 (fixed) comparability across evals within a run

Corpus curation and evaluation-set construction are deterministic: the same inputs and seeds give byte-identical images, which tests/test_data.py::test_eval_manifest_is_frozen_and_reproducible asserts. Training is not bit-deterministic — GPU kernel scheduling is not — so a repeat run lands near, not exactly on, the reported numbers.

Test suite

python -m pytest tests -q
python -m ruff check src tests scripts

The tests worth knowing about are in tests/test_init_health.py: they assert that encoder and decoder gradients are within 50x of each other at step 0, that the encoder distinguishes real rendered depictions, and that zeroing the encoder memory changes the decoder's logits. They fail on the pre-fix code and are the guard against the failure documented as D18 in docs/DECISIONS.md.

Gotchas worth knowing

Do not edit finalize.sh while it is running. Bash reads a script lazily by byte offset, so editing it mid-execution shifts the remaining bytes and the shell misparses from wherever it had reached — the symptom is a syntax error on a line that is perfectly valid. If you need to change it during a long run, copy it first and run the copy.

Checkpoint sizes. A training checkpoint holds the live weights, the EMA shadow and two AdamW moments — about four copies of the model, so 381 MB for the 23.8 M parameter base preset. pic2smiles export strips it to 95 MB, or 48 MB with --half, and is verified to produce identical predictions.

Long rsync pulls. The checkpoints are the slow part of finalize.sh. Four files at ~380 MB each takes a few minutes over a normal link.