chemtsv3 1.3.1

ChemTSv3:

ChemTSv3

ChemTSv3¹ is an extended and improved version of ChemTSv2². The original implementations are available at https://github.com/molecule-generator-collection/ChemTSv2.

In ChemTSv3,

• Node is modular: Supports any molecular representation (e.g., SMILES, SELFIES, FASTA, or HELM) in either string or tensor format.
• Transition is modular: Allows any molecular transformation strategy, including graph-based editing, sequence generation with RNN or GPT-2, sequence mutation, or LLM-guided modification.
• Filter is modular: Enables flexible constraints such as structural alerts, scaffold preservation, or physicochemical property filters.
• Reward is modular: Anything can be optimized, including QSAR predictions or simulation results, for both single- and multi-objective tasks.

Setup

Minimal installation (Mac, Linux)

Available classes

• Transition: GBGATransition, GPT2Transition, RNNBasedMutation, RNNTransition, SMIRKSTransition
• Reward: GFPReward, SimilarityReward, JScoreReward, LogPReward
• Policy: UCT, PUCT
• The corresponding Node classes and all implemented Filter classes are also available in this environment.

Setup steps

1. Clone the repository
2. Install uv: https://docs.astral.sh/uv/getting-started/installation/
3. Restart the shell
4. Move to the repository root (e.g., cd molgen)
5. Run the following commands:

uv venv --python 3.11.11
source .venv/bin/activate
uv pip install chemtsv3 numpy==1.26.4 pandas==2.3.3 matplotlib==3.10.7 rdkit==2023.09.6 ipykernel==6.30.0 transformers==4.43.4 torch==2.5.1 --torch-backend=auto

To activate the virtual environment, run the following command from the repository root (this process can also be automated through VS Code settings):

source .venv/bin/activate

To deactivate the virtual environment, run:

deactivate

Minimal installation (Windows)

Available classes

• Transition: GBGATransition, GPT2Transition, RNNBasedMutation, RNNTransition, SMIRKSTransition
• Reward: GFPReward, SimilarityReward, JScoreReward, LogPReward
• Policy: UCT, PUCT
• The corresponding Node classes and all implemented Filter classes are also available in this environment.

Setup steps

1. Clone the repository
2. Install uv: https://docs.astral.sh/uv/getting-started/installation/
3. Restart the shell (and VSCode if used)
4. Move to the repository root (e.g., cd molgen)
5. Run the following commands:

uv venv --python 3.11.11
.venv\Scripts\activate
uv pip install chemtsv3 numpy==1.26.4 pandas==2.3.3 matplotlib==3.10.7 rdkit==2023.09.6 ipykernel==6.30.0 transformers==4.43.4 torch==2.5.1 --torch-backend=auto

To activate the virtual environment, run the following command from the repository root (this process can also be automated through VS Code settings):

.venv\Scripts\activate

To deactivate the virtual environment, run:

deactivate

Full installation (Mac, Linux)

Available classes

• Transition: BioT5Transition, ChatGPTTransition, ChatGPTTransitionWithMemory, GBGATransition, GPT2Transition, RNNBasedMutation, RNNTransition, SMIRKSTransition
• Reward: DScoreReward, DyRAMOReward, GFPReward, SimilarityReward, JScoreReward, LogPReward, TDCReward
• The corresponding Node classes, along with all implemented Filter and Policy classes, are also available in this environment.
• ChatGPTTransition and ChatGPTTransitionWithMemory requires openai api key to use.

Setup steps

1. Clone the repository
2. Install uv: https://docs.astral.sh/uv/getting-started/installation/
3. Restart the shell
4. Move to the repository root (e.g., cd molgen)
5. Run the following commands:

uv venv --python 3.11.11
source .venv/bin/activate
uv pip install chemtsv3 pytdc==1.1.14 numpy==1.26.4 pandas==2.3.3 matplotlib==3.10.7 rdkit==2023.09.6 selfies==2.2.0 ipykernel==6.30.0 transformers==4.43.4 setuptools==78.1.1 lightgbm==4.6.0 openai==2.6.0 torch==2.5.1 --torch-backend=auto

To activate the virtual environment, run the following command from the repository root (this process can also be automated through VS Code settings):

source .venv/bin/activate

To deactivate the virtual environment, run:

deactivate

Optional dependencies

The full installation includes the following optional packages:

Package	Required for	Tested version
lightgbm	DScoreReward, DyRAMOReward, PUCTWithPredictor	3.3.5, 4.6.0
selfies	SELFIESStringNode	2.2.0
openai	ChatGPT2Transition, ChatGPT2TransitionWithMemory	2.6.0
pytdc	TDCReward	1.1.14

Troubleshooting

CUDA not available

In some cases (for example, when setting up environments on a control node), it may be necessary to reinstall torch with a different backend to enable CUDA support. However, since major implemented classes (including RNNTransition) are likely to run faster on the CPU, this is not strictly required. After reinstalling torch, you may also need to downgrade numpy to version 1.26.4 if it was upgraded during the process.

Generation via CLI

See config/examples/example.yaml and config/examples/example_lead.yaml for an example YAML configuration.

# Simple generation
chemtsv3 -c config/examples/example.yaml
# Auto switch
chemtsv3 -c config/examples/example_chain_1.yaml
# Load a checkpoint and continue the generation
chemtsv3 -l generation_results/~~~ --max_generations 100 --time_limit 60
# Parallel generation using MPI (requires mpi4py)
mpirun -np 4 chemtsv3-mpi -c config/mpi/mpi_example.yaml

More examples can be found in config/examples directory. Settings that were used in the paper are available in config/paper.

Tutorials and notebooks

• Tutorials: sandbox/tutorial/***.ipynb
• Generation via notebook: sandbox/generation.ipynb

Options

All options for each component (class) are defined as arguments in the __init__() method of the corresponding class. For detailed documentation of each class, see chemtsv3/generator/README.md, chemtsv3/transition/README.md, and chemtsv3/policy/README.md.

Nodes and Transitions

For general usage:

Node class	Transition class	Description
MolSentenceNode	RNNTransition	For de novo generation. Uses the RNN (GRU / LSTM) model specified by model_dir.
MolSentenceNode	GPT2Transition	For de novo generation. Uses the Transformer (GPT-2) model specified by model_dir.
CanonicalSMILESStringNode	GBGATransition	For lead optimization. Uses GB-GA mutation rules.
CanonicalSMILESStringNode	SMIRKSTransition	For lead optimization. Uses the specified SMIRKS rules (e.g. MMP-based ones).
SMILESStringNode	ChatGPTTransition	For lead optimization. Uses the specified prompt(s) as input to the GPT model specified by model (e.g., "gpt-5.4-mini"). Set use_web_search to enable the API web search tool. Requires an OpenAI API key.

For research purposes (did not perform well in our testing):

Node class	Transition class	Description
CanonicalSMILESStringNode	GBGMTransition	For de novo generation. Uses GB-GM rules. Rollouts iteratively apply transitions until the molecule size reaches a sampled value determined by size_mean and size_std.
FASTAStringNode	ProtGPT2Transition	For de novo protein generation. Uses the ProtGPT2 model.
SELFIESStringNode	BioT5Transition	For lead optimization. Uses the specified prompt(s) as input to the BioT5 text2mol model.
SMILESStringNode	ChatGPTTransitionWithMemory	For lead optimization. Unlike ChatGPTTransition, retains conversation history and feedback reward calculation results to the model.

Policies

• UCT: Does not use transition probabilities. Performed better with RNNTransition in our testing.
• PUCT: Incorporates transition probabilities (follows the modification introduced in AlphaGo Zero). Performed better with GBGATransition in our testing.
• PUCTWithPredictor: Trains an optimistic predictor of leaf-node evaluations using the generation history, and uses its output as the score for unvisited nodes when the model’s performance (measured by the normalized pinball loss) exceeds a specified threshold. This option adds a few seconds of overhead per generation (depending on the number of child nodes per transition and the computational cost of each prediction), and is recommended only when the reward calculations are expensive. Inherits all the arguments of UCT and PUCT. For non-molecular nodes, a function that returns a feature vector must be defined (see policy/puct_with_predictor.py for details.)

Basic options

Class	Option	Default	Description
-	max_generations	-	Stops generation after producing the specified number of molecules.
-	time_limit	-	Stops generation once the time limit (in seconds) is reached.
-	root	""	Key (string) for the root node (e.g. In SMILESStringNode, SMILES of the starting molecule(s)). Multiple roots can be specified by list input. If not specified, an empty string "" will be used as the root node's key.
MCTS	n_eval_width	∞	By default (= ∞), evaluates all new leaf nodes after each transition. Setting n_eval_width = 1 often improves sample efficiency and can be beneficial when reward computation is expensive.
MCTS	filter_reward	0	Substitutes the reward with this value when nodes are filtered. Use a list to specify different reward values for each filtering step. Set to "ignore" to skip reward assignment (in this case, other penalty types for filtered nodes, such as failed_parent_reward, needs to be set).
UCT, PUCT, PUCTWithPredictor	c	0.3	A larger value prioritizes exploration over exploitation. Recommended range: [0.01, 1]
UCT, PUCT, PUCTWithPredictor	best_rate	0	A value between 0 and 1. The exploitation term is calculated as: best_rate * {best reward} + (1 - best_rate) * {average reward}. For better sample efficiency, it might be better to set this value to around 0.5 for de novo generations, and around 0.9 for lead optimizations.

Advanced options

For other options and further details, please refer to each class’s __init__() method.

Class	Option	Default	Description
-	seed	-	The seed value for random, np.random and torch.
-	device	-	Torch device specification (e.g., "cpu", "cuda", "cuda:0"). For RNNTransition, using the CPU tends to be faster.
-	debug	False	If True, debug logging are enabled.
-	silent	False	If True, console logging are disabled.
-	next_yaml_path	False	If a path to the YAML config for the next generator is set, the generated molecules will be passed for chain generation.
-	n_keys_to_pass	3	Number of top-k generated molecules (keys) to be used as root nodes for the next generator.
MCTS	precalculated_csv_paths	None	Paths of result CSV files of the previous runs with the same reward can be specified here so that their reward and objective values are reused instead of recalculated. Later files take priority when keys overlap.
MCTS	save​_on​_completion	False	If True, saves a checkpoint when completing the generation.
MCTS	n_eval_iters	1	The number of child node evaluations. This value should not be > 1 unless the evaluations are undeterministic (e.g. involve rollouts).
MCTS	n_tries	1	The number of attempts to obtain an unfiltered node in a single evaluation. This value should not be >1 unless the evaluations are undeterministic (e.g. involve rollouts).
MCTS	allow​_eval​_overlaps	False	Whether to allow overlap nodes when sampling eval candidates (recommended: False)
MCTS	reward_cutoff	None	Child nodes are removed if their reward is lower than this value. This applies only to nodes for which has_reward() = True (i.e., complete molecules).
MCTS	reward​_cutoff​_warmups	None	If specified, reward_cutoff will be inactive until reward_cutoff_warmups generations.
MCTS	cut_failed_child	False	If True, child nodes will be removed when {n_eval_iters * n_tries} evals are filtered.
MCTS	failed​_parent​_reward	"ignore"	Backpropagate this value when {n_eval_width * n_eval_iters * n_tries} evals are filtered from the node.
MCTS	terminal_reward	"ignore"	If a float value is set, that value is backpropagated when a leaf node reaches a terminal state. If set to "ignore", no value is backpropagated.
MCTS	cut_terminal	True	If True, terminal nodes are pruned from the search tree and will not be visited more than once.
MCTS	avoid_duplicates	True	If True, duplicate nodes won't be added to the search tree. Should be True if the transition forms a cyclic graph. Unneeded if the tree structure of the transition graph is guranteed, and can be set to False to reduce memory usage.
MCTS	discard​_unneeded​_states	True	If True, discards node variables that are no longer needed after expansion. Set this to False when using custom classes that utilize these values.
UCT, PUCT, PUCTWithPredictor	pw_c, pw_alpha, pw_beta	None, 0, 0	If pw_c is set, the number of available child nodes is limited to pw_c * ({visit count} ** pw_alpha) + pw_beta.
UCT, PUCT, PUCTWithPredictor	max_prior	None (0)	A lower bound for the best reward. If the actual best reward is lower than this value, this value is used instead.
UCT, PUCT, PUCTWithPredictor	epsilon	0	The probability of randomly selecting a child node while descending the search tree.
PUCTWithPredictor	alpha	0.9	Quantile level for the predictor, representing the target percentile of the response variable to be estimated and used.
PUCTWithPredictor	score_threshold	0.6	If the recent prediction score (1 - {pinball loss} / {baseline pinball loss}) is better than this threshold, the model will be used afterwards.
MolSentenceNode​, MolStringNode	use​_canonical​_smiles​_as​_key	False	Whether to convert generated molecules to canonical SMILES when generating keys. If False, the same molecule may be counted multiple times.
RNNTransition, GPT2Transition	top_p	0.995	Nucleus sampling threshold in (0, 1]; keeps the smallest probability mass ≥ top_p.
RNNTransition, GPT2Transition	temperature	1	Logit temperature > 0 applied before top_p; values < 1.0 sharp, > 1.0 smooth.
RNNTransition	sharpness	1	Probability distribution sharpness > 0 applied after top_p; values < 1.0 smooth, > 1.0 sharp.
RNNTransition	disable​_top​_p​_on​_rollout	False	If True, top_p won't be applied for rollouts.
SMIRKSTransition	limit	None	If the number of generated SMILES exceeded this value, stops applying further SMIRKS patterns. The order of SMIRKS patterns are shuffled with weights before applying transition if this option is enabled.

Filters

Sanity

• ValidityFilter: Excludes invalid molecule objects. Since other filters and rewards typically assume validity and do not recheck it, usually this filter should be applied first in molecular generation.
• RadicalFilter: Excludes molecules whose number of radical electrons is not 0.
• ConnectivityFilter: Excludes molecules whose number of disconnected fragments is not 1.

Topological

• SubstructureFilter: Excludes molecules that do not contain the specified (list of) substructure(s) by smiles or smarts arguments. If preserve is set to False, excludes molecules that do contain the specified (list of) substructure(s) instead. By specifying appropriate SMARTS patterns, it is possible to control where substitutions or structural modifications (i.e., adding a substituent or arm) are allowed to occur.
• MurckoScaffoldFilter: Excludes molecules whose Murcko scaffold does not contain the Murcko scaffold of the reference smiles. If generic is set to True, atom and bond types are generalized before comparison. (Default: generic=True)
• AromaticRingFilter: Excludes molecules whose number of aromatic rings falls outside the range [min, max]. (Default: [1, ∞))
• HeavyAtomCountFilter: Excludes molecules whose number of heavy atoms falls outside the range [min, max]. (Default: [0, 45])
• MaxRingSizeFilter: Excludes molecules whose largest ring size falls outside the range [min, max]. (Default: [0, 6])
• MinRingSizeFilter: Excludes molecules whose smallest ring size falls outside the range [min, max]. (Default: (-∞, ∞))
• RingBondFilter: Excludes molecules containing ring allenes ([R]=[R]=[R]) or double bonds in small rings ([r3,r4]=[r3,r4]).
• RotatableBondsFilter: Excludes molecules whose number of rotatable bonds falls outside the range [min, max]. (Default: [0, 10])

Structural alert

• ROCFilter: Excludes molecules that contain structural alerts defined by Ohta and Cho.
• CatalogFilter: Excludes molecules that contain structural alerts in the specified list of rdkit.Chem.FilterCatalogParams.FilterCatalogs. (e.g. catalogs = ["PAINS_A", "PAINS_B", "PAINS_C", "NIH", "BRENK"])

Drug-likeness

• PubChemFilter: Excludes molecules based on the frequency of occurrence of molecular patterns in the PubChem database. Reported in Ma et al..
• LipinskiFilter: Excludes molecules based on Lipinski’s Rule of Five. Set rule_of to 3 to apply the Rule of Three instead.
• SAScoreFilter: Excludes molecules whose synthetic accessibility score (SA Score) falls outside the range [min, max]. (Default: [1, 3.5])

Physicochemical

• ChargeFilter: Excludes molecules whose formal charge is not 0.
• HBAFilter: Excludes molecules whose number of hydrogen bond acceptors falls outside the range [min, max]. (Default: [0, 10])
• HBDFilter: Excludes molecules whose number of hydrogen bond donors falls outside the range [min, max]. (Default: [0, 5])
• LogPFilter: Excludes molecules whose LogP value falls outside the range [min, max]. (Default: (-∞, 5])
• TPSAFilter: Excludes molecules whose topological polar surface area (TPSA) falls outside the range [min, max]. (Default: [0, 140])
• WeightFilter: Excludes molecules whose molecular weight falls outside the range [min, max]. (Default: [0, 500])

Filters can also be specified using filters argument of transitions that inherit from TemplateTransition (e.g. GBGATransition, SMIRKSTransition, ChatGPTTransition) to directly exclude molecules from child nodes.

Model training

• RNN (GRU) training (example): chemtsv3-train -c config/training/train_rnn_smiles.yaml
• Transformer (GPT-2) training (example): chemtsv3-train -c config/training/train_gpt2.yaml

Change dataset_path in YAML to train on an arbitrary dataset (1 sentence per line).

GPT-2 training also requires accelerate package. Installation example: uv pip install accelerate==1.14.0

License

The source code of ChemTSv3 is licensed under the MIT License.

The files under data/ are provided for reproducibility and are not covered by the MIT License. See data/***/README.md for the corresponding attribution and license information.

How to cite

@article{
doi:10.26434/chemrxiv-2025-kdvrt,
author = {Satoru Fujii  and Yuki Murakami  and Tatsuya Yoshizawa  and Shoichi Ishida  and Nobuo Cho  and Masateru Ohta  and Teruki Honma  and Kazuki Yoshizoe  and Masato Sumita  and Koji Tsuda  and Kei Terayama },
title = {ChemTSv3: Generalizing Molecular Design via Flexible Search Space Control},
journal = {ChemRxiv},
volume = {2025},
number = {1228},
pages = {},
year = {2025},
doi = {10.26434/chemrxiv-2025-kdvrt},
URL = {https://chemrxiv.org/doi/abs/10.26434/chemrxiv-2025-kdvrt},
eprint = {https://chemrxiv.org/doi/pdf/10.26434/chemrxiv-2025-kdvrt}}

1. Fujii, S. and Murakami, Y. and Yoshizawa, T. and Ishida, S. and Cho, N. and Ohta, M. and Honma, T. and Yoshizoe, K. and Sumita, M. and Tsuda, K. and Terayama, K. (2025). ChemTSv3: Generalizing Molecular Design via Flexible Search Space Control. ChemRxiv https://chemrxiv.org/doi/abs/10.26434/chemrxiv-2025-kdvrt ↩

2. Ishida, S. and Aasawat, T. and Sumita, M. and Katouda, M. and Yoshizawa, T. and Yoshizoe, K. and Tsuda, K. and Terayama, K. (2023). ChemTSv2: Functional molecular design using de novo molecule generator. WIREs Computational Molecular Science https://wires.onlinelibrary.wiley.com/doi/10.1002/wcms.1680 ↩