Rxn-SMARTS-Embeddings

Self-supervised transformer embeddings for reaction SMARTS, pretrained via contiguous span-masked language modeling directly on 361,751 validated enzymatic reaction templates from RetroRules v3.0. No labeled supervision is used at any stage — the model learns chemical patterns (atom environments, bond contexts, reaction centers) purely from SMARTS syntax.

Code: https://github.com/jcorreia11/Rxn-SMARTS-Embeddings Paper: (citation to follow upon publication)

Models

Three sizes are released, all trained with the same rule-based SMARTS tokenizer (100% round-trip fidelity, 99.8% corpus coverage at max_len=512):

Size	d_model	Layers	max_len	Params (encoder)	File prefix
small	256	6	256	5.95M	small/smarts_transformer_small
medium (recommended)	256	6	512	6.01M	medium/smarts_transformer_medium
large	512	8	512	27.77M	large/smarts_transformer_large

The medium model is used for all main-text results in the paper and is the default downloaded by the rxn-smarts-embeddings package. After 100 pretraining epochs it reaches 96.6% masked-token top-1 accuracy, and mean-pooled embeddings from it achieve 81.9% / 70.3% / 64.1% accuracy on EC-number classification at label depths 1–3 (outperforming TF-IDF and random-init baselines by 19.8–34.4 and 14.0–18.8 percentage points respectively), and correlate with structural Tanimoto similarity at Pearson r = 0.620.

Each <size>/ directory contains:

• smarts_transformer_<size>.pt — model weights (PyTorch state_dict)
• smarts_transformer_<size>.json — architecture config + training metadata

A shared vocab.json (rule-based SMARTS tokenizer vocabulary) is at the repo root and is common to all three sizes.

Usage

pip install rxn-smarts-embeddings

from rxn_smarts_embeddings.predict import predict

# Downloads and caches the medium model automatically on first use
embedding = predict("[C:1]-[O:2]>>[C:1]=[O:2]")   # shape (256,), mean-pooled

Or via the CLI:

smarts-embed "[C:1]-[O:2]>>[C:1]=[O:2]" --size large

See the GitHub repo for the Python API, batch usage, and pooling options.

Training data

The full deduplicated/validated SMARTS corpus, SentencePiece tokenizer, rule-based vocabulary, and pre-extracted embeddings are archived on Zenodo: https://doi.org/10.5281/zenodo.22645328

Full experimental results and logs (all pipeline phases) are archived separately: https://doi.org/10.5281/zenodo.22646105

The experiment pipeline and paper-statistics code are in a companion repo: https://github.com/jcorreia11/Rxn-SMARTS-Embeddings-paper

License

MIT