Official implementation of FragFM: Hierarchical Framework for Efficient Molecule Generation via Fragment-Level Discrete Flow Matching by Joongwon Lee^*, Seongwhan Kim^*, Seokhyun Moon^*, Hyeunwoo Kim^†, and Woo Youn Kim^†. [Project Page] [OpenReview] [arXiv]

FragFM Model

FragFM is a hierarchical fragment-level discrete flow matching framework that generates molecular graphs via stochastic fragment bags and lossless coarse-to-fine atom-level reconstruction.

Installation

To train and generate molecules with FragFM, install the required packages using the following script.

Option 1: Conda Environment File (recommended)

conda env create -f environment.yaml
conda activate fragfm

Option 2: Manual Conda + pip

conda create -n fragfm python=3.11 -y
conda activate fragfm
conda install scipy==1.14.1 numpy==1.26.4 pandas==2.2.3 scikit-learn==1.5.2 -y
pip install torch==2.1.0+cu118 -f https://download.pytorch.org/whl/torch_stable.html
pip install torch_geometric
pip install rdkit==2023.9.2
pip install git+https://github.com/bp-kelley/descriptastorus
pip install wandb==0.18.6 lmdb==1.5.1 pyyaml==6.0.1 easydict==1.13 parmap==1.7.0
pip install matplotlib==3.8.1 seaborn==0.13.0
pip install -e .

Option 3: uv

curl -LsSf https://astral.sh/uv/install.sh | sh
uv venv -p 3.11
uv sync
source .venv/bin/activate

Data and Checkpoints

We provide raw data, processed data, and checkpoints for MOSES, GuacaMol, ZINC250k, and NPGen benchmarks. To generate molecules with FragFM, you will need the processed data along with the autoencoder and flow generator checkpoints. Preprocessing only needs to be done once, but can take several hours (12+ hours depending on the dataset). We recommend downloading the preprocessed data instead.

Data	Size	Path
MOSES (Raw)	27MB	data/raw/
GuacaMol (Raw)	21MB	data/raw/
ZINC250k (Raw)	15MB	data/raw/
NPGen (Raw)	12MB	data/raw/
MOSES (Processed)	706MB	data/processed/
GuacaMol (Processed)	980MB	data/processed/
ZINC250k (Processed)	129MB	data/processed/
NPGen (Processed)	576MB	data/processed/

Checkpoint	Size	Path
Coarse-to-fine Autoencoders	94MB	save/ae_model/
Flow Generators	307MB	save/flow_model/

Data Preparation

You can run the script to process the data from scratch. If you have already downloaded the processed data above, you can skip this step. The following is an example for moses:

bash process/process_all_data.sh moses

Training Models

The training consists of two steps: (1) training the coarse-to-fine autoencoder and (2) training the flow generator. Configuration files are stored in cfgs/.

Training the Coarse-to-Fine Autoencoder

To train the coarse-to-fine autoencoder, run the following. The configuration and model checkpoints will be saved to save/ae_model/.

python exe/train_ae.py cfgs/train_ae/moses.yaml

Key options:

Option	Description
tag	Experiment name; checkpoints saved to save/ae_model/<time>_<tag>/
data_dirn	Path to the processed LMDB file (e.g. data/processed/moses_brics_all.lmdb)
latent_z_dim	Dimension of the autoencoder latent space

Training the Flow Generator

To train the flow generator, run the following. The configuration and model checkpoints will be saved to save/flow_model/.

python exe/train_flow.py cfgs/train_flow/moses.yaml

Key options:

Option	Description
tag	Experiment name; checkpoints saved to save/flow_model/<time>_<tag>/
data_dirn	Path to the processed molecule LMDB file (e.g. data/processed/moses_brics_all.lmdb)
frag_data_dirn	Path to the processed fragment LMDB file (e.g. data/processed/moses_brics_all_fragment.lmdb)
frag_smi_to_idx_fn	Path to the fragment SMILES-to-index mapping pickle file (e.g. data/processed/moses_brics_all_fragment_to_idx.pkl)
ae_model_dirn	Path to the pretrained autoencoder directory (e.g. save/ae_model/moses)

Molecule Sampling

To generate molecules with FragFM, run the following. The generated samples will be automatically saved to save/generate unless you set force_save_fn in the .yaml file.

python exe/generate.py cfgs/generate/moses.yaml

Key options:

Option	Description
tag	Experiment name; results saved to save/generate/<time>_<tag>/
force_save_dirn	Override the output directory path (optional)
data_dirn	Path to the processed molecule LMDB file
frag_data_dirn	Path to the processed fragment LMDB file (e.g. data/processed/moses_brics_all_fragment.lmdb)
frag_smi_to_idx_fn	Path to the fragment SMILES-to-index mapping pickle file (e.g. data/processed/moses_brics_all_fragment_to_idx.pkl)
fm_model_dirn	Path to the trained flow generator directory (e.g. save/flow_model/moses)
fragment_bag	Fragment bag to sample from (train, valid, test)
n_sample	Number of molecules to generate

Conditional Generation

FragFM offers a conditional molecule generation scheme with fragment-bag and graph-level predictor guidance. Conditional sampling builds on an unconditional generator: you must first train the lightweight graph-level property predictor described below.

Training the Predictor

To train the predictor, run the following. The configuration and model checkpoints will be saved to save/disc_model/.

python exe/train_disc.py cfgs/train_disc/moses_logp.yaml

Key options:

Option	Description
fm_model_dirn	Path to the pretrained flow generator directory (e.g. save/flow_model/moses)
calc_prop	Molecular property to predict (logp, nring, tpsa, qed)

Molecule Sampling

Conditional sampling follows the same procedure as unconditional generation and can be performed by modifying the configuration in the .yaml file as follows.

python exe/generate.py cfgs/generate/moses_logp.yaml

Key options:

Option	Description
disc_model_dirn	Path to the trained predictor directory (e.g. save/disc_model/moses_logp); set to enable conditional sampling
guide_val	Guidance target value for the predictor
bag_guide_strength	Strength of fragment-bag level guidance
graph_guide_strength	Strength of graph-level guidance

NPGen Benchmark

NPGen is a large-scale natural product generation benchmark designed to evaluate generative models on structurally complex and biologically meaningful chemical space beyond standard drug-like datasets, with molecular structures curated from the COCONUT database.

Models

We evaluate against a diverse set of baseline models: a sequence-based model (SAFE-GPT), an atom-level graph autoregressive model (GraphAF), atom-level graph diffusion models (DiGress, DeFoG), and fragment-level graph models (MolHF, JT-VAE, HierVAE). We provide example molecules generated by each baseline model in npgen-benchmark/baselines/, where each file contains 30,000 SMILES strings of generated molecules.

Installation

Option 1: Conda Environment File

conda env create -f npgen-benchmark/environment.yml
conda activate npgenbenchmark

Option 2: Manual Conda + pip

conda create -n npgenbenchmark python=3.8
conda activate npgenbenchmark
conda install -c conda-forge numpy=1.21 rdkit=2020.03.2 icu=68.2
pip install pandas huggingface_hub scikit-learn scipy tqdm torch fcd_torch
pip install -e .

Option 3: Pixi

A Pixi config is provided in pixi.toml for reproducible setup.

Use pixi.lock for fully reproducible environments:

• If pixi.lock is already in the repo, install exactly that environment.
• Use --frozen to avoid lockfile changes during install.

pixi install --frozen
pixi run pip-deps

When dependencies in pixi.toml change, update the lockfile explicitly:

pixi lock
pixi install
pixi run pip-deps

Run Python commands directly in the Pixi environment:

pixi run python npgen-benchmark/tests/test_benchmark.py

Then start a shell in the environment:

pixi shell

For Apple Silicon Macs, use an x86_64/Rosetta terminal when running this project.

Benchmark Evaluation

Run benchmark experiments using the provided scripts:

import torch
from npgenbenchmark import NPGenBenchmark

# define smiles list to benchmark (usually generated molecules)
# benchmark also includes calculation of validity
smiles_list = ["CCCC", "c1ccccc1", "X", "CCO", "X"]

benchmark = NPGenBenchmark(
    device="cuda:0" if torch.cuda.is_available() else "cpu",
    n_jobs=4,  # Adjust based on your system cores (num_workers in dataloader)
    batch_size=256,  # Example batch size
    verbose=True,
    n_eval_data=30000, # Number of data to evaluate
)

# Run the benchmark
results = benchmark.run_benchmark(smiles_list)
print(results)

For better convenience, we re-package the original NP-Classifier code to use pytorch instead of Keras.