REAL-SWIT

REAL-SWIT is a target-guided molecular generation workflow for exploring synthesizable chemical space. It combines a generative model trained on REAL Space-derived molecules with a target-specific scoring model trained to approximate docking scores, enabling iterative reinforcement-learning-based generation of molecules prioritized for a given protein target.

This repository contains the code, example data, and pretrained components associated with the manuscript.

Access to Synthesizable Chemical Space Through Generative Models Enables Ultra-Large Virtual Screening
Kaiyue Zhang et al.

Pretrained models

Pretrained model weights are hosted on the Hugging Face Hub: REAL‑SWIT checkpoints on Hugging Face. Please download the pretrained file QBL_model.ckpt from that repository and place it in this project's checkpoints/ directory before running the generation workflows.

Requirements

REAL-SWIT was developed and tested on a Linux-based GPU cluster. Most Python scripts should also run on other platforms if the required dependencies are properly installed.

• Micromamba (recommended) or Conda/Miniconda
• Python environment specified in environment.yml
• CUDA-enabled GPU recommended for model training and generation

Installation

Clone the repository and enter the project directory:

git clone https://github.com/JingHuangLab/REAL_SWIT.git
cd REAL_SWIT

Create the real_swit environment from environment.yml. We recommend using micromamba for faster dependency resolution, although Conda is also supported.

Recommended: micromamba

micromamba create -f environment.yml
micromamba activate real_swit

Alternative: Conda

conda env create -f environment.yml
conda activate real_swit

Repository structure

A typical repository layout is:

REAL_SWIT/
├── README.md
├── environment.yml
├── data/
│   └── demo_data/
│   └── surrogate_initial_data/
│   └── surrogate_training_data/
├── examples/
├── checkpoints/
│   └── QBL_model.ckpt
├── gen_models/
├── mpn_models/
├── train_generative_model.sh
├── train_scoring_model.py
├── evaluate_scoring_model.py
├── prepare_rl_config.py
└── submit_real_swit.sh

Main folders and files:

• data/: input datasets and small demo datasets for training and evaluation.
• examples/: runtime outputs organized by task, including trained scoring models, predictions, RL configurations, and generated molecules. This directory is populated when the workflows are run.
• checkpoints/: directory for pretrained generative model checkpoints, including the default QBL_model.ckpt.
• gen_models/: scripts for training generative models and performing molecular generation.
• mpn_models/: implementation of the target-specific molecular scoring model.
• train_generative_model.sh: example SLURM workflow for training a generative model.
• train_scoring_model.py: trains a target-specific scoring model using molecules labeled with docking scores.
• evaluate_scoring_model.py: generates target-specific score predictions and optionally evaluates model performance when labels are available.
• prepare_rl_config.py: creates the reinforcement learning configuration file for target-guided generation.
• submit_real_swit.sh: SLURM-compatible entry point for training or evaluating a target-specific scoring model and running target-guided generation.

Input data format

Training data for the target-specific scoring model

The training file should be a CSV file with at least two columns:

SMILES,score
CCOc1ccc(...),-8.7
CCN1CC(...),-10.2

• The first column should contain SMILES strings.
• The second column should contain docking scores.
• Docking scores are expected to follow the usual convention where more negative values indicate better predicted binding.

Usage

REAL-SWIT contains three main workflows:

1. Optional: train a molecular generative model from custom molecules.
2. Required: train a target-specific scoring model for the target of interest.
3. Required: create an RL configuration file and run target-guided molecular generation.

If you use the pretrained generative model provided with this repository, you can skip Step 1 and start from Step 2.

Step 1. Optional: train a molecular generative model

This step is only needed if you want to train a new generative model using your own molecular dataset. If you use the pretrained model from the manuscript, the default checkpoint is:

checkpoints/QBL_model.ckpt

The generative model training workflow consists of three stages:

1. Create randomized SMILES for the training and validation sets.
2. Create an empty model and prepare the vocabulary from the training data.
3. Train the generative model.

The complete workflow is provided in train_generative_model.sh. Before running it, edit the User-editable settings section to specify the input files, output directory, model hyperparameters, and local computing environment.

Run the workflow on a SLURM cluster with:

sbatch train_generative_model.sh

By default, the script uses the demo training and validation datasets under data/demo_data/ and writes randomized SMILES, model checkpoints, and TensorBoard logs to gm_training_demo/. Individual stages can be enabled or disabled using CREATE_RANDOMIZED_SMILES, CREATE_EMPTY_MODEL, TRAIN_MODEL, and SAMPLE_AFTER_TRAINING.

After training, select the desired trained checkpoint and use it as the prior and initial agent model in the RL configuration.

Step 2. Train a target-specific scoring model

For each new target, train a target-specific scoring model using molecules labeled by docking scores.

Before running submit_real_swit.sh, edit its User-editable settings section as needed:

• CONDA_ENV_NAME: name of the micromamba or Conda environment.
• CUDA_MODULE: optional CUDA module used on the local cluster. Leave it empty if environment modules are not used.
• NCPU: number of CPU cores passed to the scoring-model scripts.
• N_EPOCHS: number of training epochs.

Also update the SLURM directives at the beginning of the script, particularly the partition, GPU, CPU, memory, and log settings, according to the local cluster environment. The task name and input dataset paths are supplied directly in the submission command.

The recommended SLURM command is:

sbatch submit_real_swit.sh \
    ROCK1_demo train \
    data/demo_data/rock1_train_demo.csv \
    data/demo_data/rock1_test_demo.csv

The equivalent Python command is:

python train_scoring_model.py \
    data/demo_data/rock1_train_demo.csv \
    ROCK1_demo \
    --testing_dataset_path data/demo_data/rock1_test_demo.csv \
    --ncpu 6 \
    --epochs 100

Arguments:

• data/demo_data/rock1_train_demo.csv: training dataset containing SMILES and docking scores.
• ROCK1_demo: task name. Outputs will be saved under examples/ROCK1_demo/.
• --testing_dataset_path: optional test dataset for model evaluation.
• --ncpu: number of CPU cores available to each prediction worker.
• --epochs: number of training epochs.

The trained checkpoint is saved under:

examples/<task_name>/lightning_logs/<version>/checkpoints/

If a test dataset with docking scores is provided, the script also evaluates the model and outputs predictions and plots under:

examples/<task_name>/preds/

Step 3. Create an RL configuration file

After training the target-specific scoring model, create an RL configuration file for target-guided generation.

Before running prepare_rl_config.py, check the user-editable settings at the beginning of the script:

• gen_model_name: pretrained or custom generative model checkpoint name.
• n_steps: number of RL optimization steps.
• n_mols: maximum number of molecules generated during the RL run.
• low_mw and high_mw: molecular weight bounds.
• max_inverted_score: upper bound used for the sign-inverted target-specific score. The docking score is multiplied by -1 before transformation, so a raw docking score of -50 corresponds to 50 here.
• target_score_weight: weight of the target-specific scoring component.
• sigma: sigma value used in the augmented likelihood calculation.

Then provide the task name and RL run name in the command:

python prepare_rl_config.py ROCK1_demo run_001

This command creates:

examples/ROCK1_demo/RL_practice/run_001/RL_config.json

By default, the configuration uses the pretrained generative model:

checkpoints/QBL_model.ckpt

and automatically locates the target-specific scoring model checkpoint under:

examples/ROCK1_demo/lightning_logs/<version>/checkpoints/

Step 4. Run target-guided molecular generation

Run reinforcement-learning-based molecular generation using the configuration file generated in Step 3.

Before running submit_real_swit.sh, check the same environment and SLURM settings described in Step 2. For generation, also check DEFAULT_RL_RUN_NAME, which is used when no RL run name is supplied in the command. Ensure that the task name and RL run name match those used to create the configuration file in Step 3.

On a SLURM cluster, use:

sbatch submit_real_swit.sh ROCK1_demo generate run_001

The underlying Python command is:

python gen_models/input.py examples/ROCK1_demo/RL_practice/run_001/RL_config.json

Generated molecules and the associated scores are saved under:

examples/ROCK1_demo/RL_practice/run_001/results/

The main output file is usually:

scaffold_memory.csv

Notes for SLURM users

Example SLURM scripts are provided for running training and generation on GPU clusters. Before submitting a job, update the following settings according to your local cluster environment:

• partition name
• GPU type
• conda environment path
• repository path
• log directory

Avoid hard-coded absolute paths when preparing a public GitHub release. Prefer paths relative to the repository root whenever possible.

Citation

If you use REAL-SWIT in your work, please cite the associated manuscript:

Kaiyue Zhang et al. Access to Synthesizable Chemical Space Through Generative Models Enables Ultra-Large Virtual Screening.

A BibTeX entry will be added after publication.

Contact

For questions or issues, please contact:

• Kaiyue Zhang
• Jing Huang

Alternatively, please open an issue on GitHub.

Acknowledgments

REAL-SWIT builds on ideas and code components from several open-source projects, including:

• REINVENT
• reinvent-randomized
• Chemprop
• MolPAL