MIST: Molecular Insight SMILES Transformer

MIST is a family of molecular foundation models for molecular property prediction. The models were pre-trained on Smirk 😏 tokenized SMILES strings from the Enamine REAL Space dataset using the Masked Language Modeling (MLM) objective, then fine-tuned for downstream prediction tasks.

Installation

The following provides installation instructions for the top-level package (electrolyte_fm), optional add-ons for our various additional analysis and downstream applications (See ./opt may require additional configuration.

1. Install uv and julia (only needed for /opt tasks)
2. Instantiate the environment: uv sync
3. Use submit/submit.py to submit a training job or checkout one of our applications in ./opt

You may need to install rust if pre-built wheels for smirk are not available on PyPI. Feel free to open an issue to request additional pre-built wheels.

Installation took under a minute on a MacBook Pro 2023. However, installation times depend heavily on internet connection and local uv cache. Pulling the NVIDIA packages can easily add 5-10 minutes without a uv cache. Installing the Julia environments (for example in /opt) can takes ~5 mins but may take longer on limited hardware (e.g. GitHub runners take ~18 mins).

Local

1. Install rust and uv
2. Run uv sync

Not all configurations will work locally. For example, pre-training configuration files typically use DeepSpeed, which requires NVIDIA GPUs. Some adaptation may be required to run outside of an NVIDIA GPU cluster.

Polaris

1. Install rust and uv

2. Load conda

module purge
module use /soft/modulefiles/
module --ignore_cache load conda/2024-04-29
conda activate base

3. Install the environment

uv sync

Artemis

Same as above except:

1. Skip loading conda (just use uv)
2. Ensure a module for CUDA@12.2 exists, may need to install with spack (make sure buildable: True)

Apptainer

0. Install or load from a module Apptainer
1. Build the image bash container/build.sh, once build relocate the image mv /tmp/mist.sif ./mist.sif
2. Run training within the image apptainer run --nv mist.sif python train.py ...

See submit/dgx.j2 or submit/delta.j2 for a more complete example of using the container

System Requirements

Hardware

Generally, running the code here requires access to GPUs and ideally a dedicated NVIDIA GPU cluster. However, much of the (non-training) code can be run on a single NVIDIA A40 GPU. Notably, for the MIST-28M model, CPU or MPS inference is viable.

Software Dependencies

All software dependencies can be found in the pyproject.toml (python) or Project.toml (julia) files. Specific versions are detailed in the lock files.

MIST has been tested on the following primary dependencies:

Dependency	Versions
Ubuntu	24.04
Python	3.12, 3.13
NVIDIA CUDA	12.8.0.038
NVIDIA cuBLAS	12.8.3.14
NVIDIA cuDNN	9.7.0.66
NVIDIA NCCL	2.25.1
NVIDIA GPU	A40, A100, H100
PyTorch	2.6

Model Weights

Model weights can be retrieved from Zenodo (fine-tuned only) or Hugging Face (pre-trained and fine-tuned).

Pre-training Dataset

MIST was pre-trained on Enamine REAL Space as provided by Enamine. We are currently working to secure permission to publish a subset of that dataset; however, MIST can be trained on a collection of text files of newline-delimited SMILES. We have uploaded an example dataset to Zenodo.

Fine-Tuning Datasets

All fine-tuning datasets are documented in Supplementary Section C.1. Datasets that have not already been publicly released elsewhere can be found on Zenodo.

Demonstration Code

Examples of using or training the MIST models can be found:

• Fine-tuning MIST (42 minutes)
• MIST for Molecular Property Prediction (1 minute)
• Predicting Molecular Scent with MIST (5 minutes)

Submitting Jobs

We use a python script (submit/submit.py) to template training jobs for submission on HPC systems across multiple sites. Templates may need to be modified for your particular HPC cluster, but should provide a starting point.

source ./activate # Activate Environment
./submit/submit.py ./submit/polaris.j2 --data ./submit/pretrain.yaml | qsub

See submit/submit.py --help for more info

Note: ./activate is used to activate the python virtual environment and set various environment variables.

Pre-training and fine-tuning parameters are documented in our paper's methods section. The configuration for the MIST-1.8B is provided for reference. Configurations for all other models, including our fine-tuning configs, are documented in training logs.

Development

Pre-commit

We use pre-commit to preform various linting checks on the code. To enable:

1. Install poetry (See above)
2. Run pre-commit: uv run pre-commit
3. Run before committing: uv run pre-commit install --allow-missing-config