Synthemol

Note

1. Before starting training and evaluation, please download the dataset used in the experiment Data.zip, and modify data_dir in the yaml configuration file to the path of the decompressed dataset. For example: "./data/Data/..."; download resources.zip, and unzip it to examples/synthemol/synthemol/.
2. If you need to use a pre-trained model for evaluation, please download the pre-trained model pretrained.zip and unzip it, for example to the path ./pretrained/pretrained_chemprop.pdparams, and specify the path in PRE_COMPUTE.model_path of the yaml configuration file.
3. Before starting training and generation, please install rdkit etc. For related dependencies, please execute pip install requirements.txt to install.

Property Predictor Model Training CommandProperty Predictor Model Evaluation CommandPre-compute building blocks score commandUse synthemol to generate molecules command

# Use antibiotics and other data to train chemprop model to implement Property Predict
# Configuration can be modified in conf/synthemol.yaml
python main.py

# Download pre-trained model (optional, or specify your own trained model in configuration file)
mkdir -p ./pretrained && wget -O ./pretrained/pretrained_chemprop.pdparams https://paddle-org.bj.bcebos.com/paddlescience/models/synthemol/pretrained_chemprop.pdparams
# Use antibiotics and other data to evaluate chemprop model to implement Property Predict
# Configuration can be modified in conf/synthemol.yaml
python main.py mode=eval

# Download pre-trained model (optional, or specify your own trained model in configuration file)
mkdir -p ./pretrained && wget -O ./pretrained/pretrained_chemprop.pdparams https://paddle-org.bj.bcebos.com/paddlescience/models/synthemol/pretrained_chemprop.pdparams
# Use trained model to score and compute building blocks to accelerate the next generation phase
# Configuration can be modified in conf/synthemol.yaml
python main.py mode=pre-compute

# Use pre-computed building blocks score guidance, combined with synthemol using Monte Carlo Tree Search to generate molecules
# Configuration can be modified in conf/synthemol.yaml
python main.py mode=generate

1. Background Introduction

The rapid emergence of pan-drug-resistant bacteria makes the development of structurally novel antibiotics urgent. Although artificial intelligence can discover new antibiotics, existing methods still have obvious flaws: property prediction models can only evaluate molecules one by one, which has extremely poor scalability when facing huge chemical spaces; while generative models can quickly explore huge chemical spaces, they often output molecules that are difficult to synthesize. To this end, the authors proposed SyntheMol, a generative model that can design new compounds that are easy to synthesize from a chemical space of nearly 30 billion molecules. The authors used SyntheMol to inhibit the growth of Acinetobacter baumannii (a tricky Gram-negative pathogen), synthesized 58 generated molecules and experimentally verified them, of which 6 structurally novel molecules showed antibacterial activity against Acinetobacter baumannii and other bacteria with significant phylogenetic differences. This study demonstrates the potential of generative AI to design structurally novel, synthesizable, and effective small-molecule antibiotic candidates in a huge chemical space, and provides experimental validation.

2. Synthemol Principle

This chapter only briefly introduces the model principle of Synthemol. For detailed theoretical derivation, please read Generative AI for designing and validating easily synthesizable and structurally novel antibiotics.

2.1 Property Predictor

Chemprop is a molecular property prediction model that uses directed message passing neural networks to process molecules and predict their properties. Chemprop first extracts simple atom and bond features (such as atom type and bond type) from the molecular graph to construct feature vectors for each atom and bond. Then, the model performs three rounds of message passing: in each round, the neural network layer iteratively fuses information from neighboring atoms and bonds. After message passing is completed, Chemprop sums all fused feature vectors to generate a single feature vector representing the entire molecule. This vector is then input into a two-layer feedforward neural network to predict molecular properties; in this study, it predicts the probability of inhibiting the growth of Acinetobacter baumannii. We use Chemprop v1.5.2, migrated from PyTorch v1.12.0.post2. For two other predictors, please refer to the original text.

2.2 Synthemol

SyntheMol is a generative model that explores a combinatorial chemical space composed of molecules generated by chemical reactions of molecular building blocks to find molecules with target properties. SyntheMol uses a Monte Carlo Tree Search (MCTS) algorithm similar to AlphaGo to efficiently search for ideal molecules in this chemical space. SyntheMol can not only quickly identify promising molecules, but also give their synthesis routes (that is, the complete steps of combining molecular building blocks through a series of one-step or multi-step chemical reactions). Below, we give the mathematical symbols required to describe the SyntheMol MCTS algorithm and provide the corresponding pseudocode.

SyntheMol MCTS Algorithm

Requires:

• Synthesis tree T
• Property prediction model M
• Maximum number of rollouts n_rollout
• Maximum number of reactions n_reaction

---

function MCTS():     for i = 1 to n_rollout do:         rollout(T.root)     end for     return all visited nodes in T with:         1 molecule and ≥ 1 reaction

---

function rollout(N):     if node N has undergone ≥ n_reaction reactions then         return property prediction score of M applied to molecules in N     end if     E ← expand_node(N)     S ← select child node in E with largest MCTS score     return rollout(S)

---

function expand_node(N):     E ← empty set of nodes     foreach reaction R do         if R is compatible with molecules in N then             Add new node to E with each product of R applied to molecules in N         end if     end for     foreach building block B do         if any reaction is compatible with B and molecules in N then             Add new node to E with B and molecules in N         end if     end for     return E

3. Synthemol Model Implementation

Next, we will explain how to implement the training, pre-calculation score and generation of the Synthemol model based on PaddleScience code. For other details in this case, please refer to API Documentation.

3.1 Dataset Introduction

The dataset uses the Data.zip dataset from the author's repository Synthemol.

The training set consists of 3 compound libraries:

• Library 1 contains 2371 molecules from the Pharmakon-1760 library (containing 1360 FDA-approved drugs and 400 internationally approved drugs) and 800 natural products isolated from plant, animal and microbial sources.
• Library 2 is the Broad Drug Repurposing Hub, containing 6680 molecules, most of which are FDA-approved drugs or clinical candidate compounds.
• Library 3 is a small molecule synthesis screening library containing 5376 molecules, randomly sampled from a larger compound library of the Broad Institute.

All 3 libraries were screened for growth inhibition activity against Acinetobacter baumannii ATCC 17978 in duplicate biological replicates. The experimental process is as follows:

1. The strain was cultured overnight in 2 ml LB medium at 37 °C, and then diluted 1:10 000 in fresh LB.
2. Take 49.5 µl (384-well plate) or 99 µl (96-well plate) bacterial solution and add it to Corning flat-bottom microplate using manual or Agilent Bravo pipetting system.
3. Add the test compound to each well, final concentration 50 µM, final volume 50 µl (384-well plate) or 100 µl (96-well plate).
4. Incubate at 37 °C for 16 h.
5. Read absorbance at 600 nm using SpectraMax M3 microplate reader (Molecular Devices), normalize data by intra-plate quartile mean, and then aggregate and determine positive hits.

For more details, including hyperparameter adjustment space for each model, please refer to the author's original paper. Specific hyperparameters used in this repository are preset in the yaml configuration file and can be adjusted according to the situation.

3.2 Chemprop Model Training

3.2.1 Constraint Construction

This case solves the problem based on data-driven methods, so it is necessary to use SupervisedConstraint built in PaddleScience to construct supervised constraints. Before defining constraints, you need to first specify various parameters used for data loading in supervised constraints.

The code for data loading is as follows:

# set dataloader config
train_dataloader_cfg = {
    "dataset": {
        "name": cfg.DATA.dataset_name,  # "MoleculeDatasetIter",
        "input_keys": tuple(cfg.MODEL.input_keys),
        "args": args,
        "smiles": train_smiles,
        "fingerprints": train_fingerprints,
        "properties": train_properties,
        "label_keys": tuple(cfg.MODEL.label_keys),
    },
    "num_workers": cfg.TRAIN.num_workers,
}

Among them, the "dataset" field defines the used Dataset class name as MoleculeDatasetIter, and num_works is 1.

The code for defining supervised constraints is as follows:

# set constraint
sup_constraint = ppsci.constraint.SupervisedConstraint(
    train_dataloader_cfg,
    output_expr={"pred": lambda out: out["pred"]},
    loss=ppsci.loss.FunctionalLoss(get_train_loss_func(args)),
    name="Sup",
)

# wrap constraints together
constraint = {sup_constraint.name: sup_constraint}

The first parameter of SupervisedConstraint is the data loading method, here train_dataloader_cfg defined above is used;

The second parameter is the definition of loss function, here a custom loss function is used; the author controls loss function selection by passing parameters through get_loss_func function: the Chemprop model in the paper uses CrossEntropyLoss;

The third parameter is the name of the constraint condition, which is convenient for subsequent indexing. Here it is named Sup.

3.2.2 Model Construction

In this case, the molecular property prediction model is implemented based on the Chemprop network model, expressed in PaddleScience code as follows:

# set model
model = ppsci.arch.chemprop_molecule.MoleculeModel(cfg=cfg)

The parameters of the network model are set through the configuration file as follows:

MODEL:
  input_keys: ["mol_batch", "features_batch", "atom_descriptors_batch",
                "atom_features_batch", "bond_features_batch"]
  output_keys: ["pred"]
  label_keys: ["targets", "data_weights", "mask", "target_weights"]

Among them, input_keys and output_keys represent the names of the input and output variables of the network model respectively.

3.2.3 Learning Rate and Optimizer Construction

The learning rate size used in this case is set to 0.0001. The optimizer uses Adam, and parameters are grouped, expressed in PaddleScience code as follows:

# set optimizer
optimizer = ppsci.optimizer.Adam(
    learning_rate=cfg.TRAIN.learning_rate,
    weight_decay=None,  # 0.001
)(model)

3.2.4 Model Training

After completing the above settings, you only need to pass the instantiated objects to ppsci.solver.Solver in order, and then start training.

# initialize solver
solver = ppsci.solver.Solver(
    model,
    constraint,
    cfg.output_dir,
    optimizer,
    None,
    cfg.TRAIN.epochs,
    cfg.TRAIN.iters_per_epoch,
    save_freq=cfg.TRAIN.save_freq,
    eval_during_train=cfg.TRAIN.eval_during_train,
    eval_freq=cfg.TRAIN.eval_freq,
    eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
    checkpoint_path=cfg.TRAIN.checkpoint_path,
)

# train model
solver.train()

3.3 Pre-compute building blocks score

The code for constructing the model is:

model = ppsci.arch.chemprop_molecule.MoleculeModel(cfg=cfg)

3.4 Synthemol Generate Molecules

The code for constructing Generator is:

print("Setting up generator...")
generator = Generator(
    building_block_smiles_to_id=building_block_smiles_to_id,
    max_reactions=max_reactions,
    scoring_fn=model_scoring_fn,
    explore_weight=explore_weight,
    num_expand_nodes=num_expand_nodes,
    optimization=optimization,
    reactions=reactions,
    rng_seed=rng_seed,
    no_building_block_diversity=no_building_block_diversity,
    store_nodes=store_nodes,
    verbose=verbose,
    replicate=replicate,
)

4. Complete Code

# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from datetime import datetime
from pathlib import Path
from random import Random

import hydra
import numpy as np
import paddle
import pandas as pd
from chemprop_models import chemprop_predict
from chemprop_models import my_chemprop_load
from evaluation import evaluate_auto
from loss_functions import get_loss_func
from omegaconf import DictConfig
from synthemol.generate.generator import Generator
from synthemol.generate.utils import create_model_scoring_fn
from synthemol.generate.utils import save_generated_molecules
from synthemol.reactions import REACTIONS
from synthemol.reactions import load_and_set_allowed_reaction_building_blocks
from synthemol.reactions import set_all_building_blocks
from tqdm import tqdm

import ppsci
import ppsci.arch.chemprop_molecule
from ppsci.arch.chemprop_molecule_utils import TrainArgs


def get_train_loss_func(args):  #:paddle.Tensor=None):
    def train_loss_func(output_dict, label_dict, weight_dict):
        preds = output_dict["pred"]

        targets = label_dict["targets"]
        target_weights = label_dict["target_weights"]
        data_weights = label_dict["data_weights"]
        mask = label_dict["mask"]

        loss_func = get_loss_func(args)
        if args.loss_function == "bounded_mse":
            # lt_target_batch = lt_target_batch
            # gt_target_batch = gt_target_batch
            pass
        if args.loss_function == "mcc" and args.dataset_type == "classification":
            loss = loss_func(
                preds, targets, data_weights, mask
            ) * target_weights.squeeze(axis=0)
        elif args.loss_function == "mcc":
            targets = targets.astype(dtype="int64")
            target_losses = []
            for target_index in range(preds.shape[1]):
                target_loss = loss_func(
                    preds[:, target_index, :],
                    targets[:, target_index],
                    data_weights,
                    mask[:, target_index],
                ).unsqueeze(axis=0)
                target_losses.append(target_loss)
            loss = paddle.concat(x=target_losses) * target_weights.squeeze(axis=0)
        elif args.dataset_type == "multiclass":
            targets = targets.astype(dtype="int64")
            if args.loss_function == "dirichlet":
                loss = (
                    loss_func(preds, targets, args.evidential_regularization)
                    * target_weights
                    * data_weights
                    * mask
                )
            else:
                target_losses = []
                for target_index in range(preds.shape[1]):
                    target_loss = loss_func(
                        preds[:, target_index, :], targets[:, target_index]
                    ).unsqueeze(axis=1)
                    target_losses.append(target_loss)
                loss = (
                    paddle.concat(x=target_losses, axis=1)
                    * target_weights
                    * data_weights
                    * mask
                )
        elif args.dataset_type == "spectra":
            loss = (
                loss_func(preds, targets, mask) * target_weights * data_weights * mask
            )
        elif args.loss_function == "bounded_mse":
            pass
            """
            loss = (
                loss_func(preds, targets, lt_target_batch, gt_target_batch)
                * target_weights
                * data_weights
                * mask
            )
            """
        elif args.loss_function == "evidential":
            loss = (
                loss_func(preds, targets, args.evidential_regularization)
                * target_weights
                * data_weights
                * mask
            )
        elif args.loss_function == "dirichlet":
            loss = (
                loss_func(preds, targets, args.evidential_regularization)
                * target_weights
                * data_weights
                * mask
            )
        else:
            loss = loss_func(preds, targets) * target_weights * data_weights * mask
        loss = loss.sum() / mask.sum()

        return {"pred": loss.astype("float32")}

    return train_loss_func


def make_args(
    dataset_type,
    epochs,
    use_gpu,
    fingerprint_type,
    property_name,
    train_smiles,
    train_fingerprints,
):
    # Create args
    arg_list = [
        "--data_path",
        "foo.csv",
        "--dataset_type",
        dataset_type,
        "--save_dir",
        "foo",
        "--epochs",
        str(epochs),
        "--quiet",
    ] + ([] if use_gpu else ["--no_cuda"])

    if fingerprint_type == "morgan":
        arg_list += ["--features_generator", "morgan"]
    elif fingerprint_type == "rdkit":
        arg_list += [
            "--features_generator",
            "rdkit_2d_normalized",
            "--no_features_scaling",
        ]
    elif fingerprint_type is None:
        pass
    else:
        raise ValueError(f'Fingerprint type "{fingerprint_type}" is not supported.')

    args = TrainArgs().parse_args(arg_list)
    args.task_names = [property_name]
    if train_smiles is not None:
        args.train_data_size = len(train_smiles)

    if fingerprint_type is not None:
        args.features_size = train_fingerprints.shape[1]
    return args


def load_raw_data(cfg):
    data_path = cfg.DATA.data_path
    data = pd.read_csv(data_path)
    print(f"Data size = {len(data):,}")
    num_models = cfg.DATA.num_models  # 10
    num_folds = cfg.DATA.num_folds  # 10
    indices = np.tile(np.arange(num_folds), 1 + len(data) // num_folds)[: len(data)]
    random = Random(0)
    random.shuffle(indices)
    assert 1 <= num_models <= num_folds
    smiles_column = cfg.DATA.smiles_column  #'smiles'
    property_column = cfg.DATA.property_column  #'antibiotic_activity'

    model_num = 1
    test_index = model_num
    val_index = (model_num + 1) % num_folds
    test_mask = indices == test_index
    val_mask = indices == val_index
    train_mask = ~(test_mask | val_mask)
    test_data = data[test_mask]
    val_data = data[val_mask]
    train_data = data[train_mask]
    print(
        "test_data:",
        len(test_data),
        "train_data:",
        len(train_data),
        "val_data:",
        len(val_data),
    )
    train_smiles = train_data[smiles_column]
    train_fingerprints = None
    train_properties = train_data[property_column]
    return train_smiles, train_fingerprints, train_properties


def train(cfg: DictConfig):
    train_smiles, train_fingerprints, train_properties = load_raw_data(cfg)

    args = make_args(
        dataset_type=cfg.DATA.dataset_type,  # "classification",
        epochs=cfg.TRAIN.epochs,  # 1,
        use_gpu=cfg.TRAIN.use_gpu,
        fingerprint_type=cfg.DATA.fingerprint_type,  # None,
        property_name=cfg.DATA.property_column,  # "antibiotic_activity"
        train_smiles=train_smiles,
        train_fingerprints=train_fingerprints,
    )

    # set dataloader config
    train_dataloader_cfg = {
        "dataset": {
            "name": cfg.DATA.dataset_name,  # "MoleculeDatasetIter",
            "input_keys": tuple(cfg.MODEL.input_keys),
            "args": args,
            "smiles": train_smiles,
            "fingerprints": train_fingerprints,
            "properties": train_properties,
            "label_keys": tuple(cfg.MODEL.label_keys),
        },
        "num_workers": cfg.TRAIN.num_workers,
    }

    # set constraint
    sup_constraint = ppsci.constraint.SupervisedConstraint(
        train_dataloader_cfg,
        output_expr={"pred": lambda out: out["pred"]},
        loss=ppsci.loss.FunctionalLoss(get_train_loss_func(args)),
        name="Sup",
    )

    # wrap constraints together
    constraint = {sup_constraint.name: sup_constraint}

    # set model
    model = ppsci.arch.chemprop_molecule.MoleculeModel(cfg=cfg)

    # set optimizer
    optimizer = ppsci.optimizer.Adam(
        learning_rate=cfg.TRAIN.learning_rate,
        weight_decay=None,  # 0.001
    )(model)

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        constraint,
        cfg.output_dir,
        optimizer,
        None,
        cfg.TRAIN.epochs,
        cfg.TRAIN.iters_per_epoch,
        save_freq=cfg.TRAIN.save_freq,
        eval_during_train=cfg.TRAIN.eval_during_train,
        eval_freq=cfg.TRAIN.eval_freq,
        eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
        checkpoint_path=cfg.TRAIN.checkpoint_path,
    )

    # train model
    solver.train()


def evaluate(cfg):
    data_path = cfg.DATA.data_path
    data = pd.read_csv(data_path)
    print(f"Data size = {len(data):,}")
    num_models = cfg.DATA.num_models
    num_folds = cfg.DATA.num_folds
    indices = np.tile(np.arange(num_folds), 1 + len(data) // num_folds)[: len(data)]
    random = Random(0)
    random.shuffle(indices)
    assert 1 <= num_models <= num_folds
    smiles_column = cfg.DATA.smiles_column  #'smiles'
    property_column = cfg.DATA.property_column  #'antibiotic_activity'

    model_num = 1
    test_index = model_num
    val_index = (model_num + 1) % num_folds
    test_mask = indices == test_index
    val_mask = indices == val_index
    train_mask = ~(test_mask | val_mask)
    test_data = data[test_mask]
    val_data = data[val_mask]
    train_data = data[train_mask]
    print(
        "test_data:",
        len(test_data),
        "train_data:",
        len(train_data),
        "val_data:",
        len(val_data),
    )

    # load model
    model_path = Path(cfg.PRE_COMPUTE.model_path)
    use_gpu = cfg.PRE_COMPUTE.use_gpu
    model_type = cfg.PRE_COMPUTE.model_type  #'chemprop'

    model = ppsci.arch.chemprop_molecule.MoleculeModel(cfg=cfg)

    if model_type == "chemprop":
        if use_gpu:
            device = str("cuda").replace("cuda", "gpu")
        else:
            device = paddle.CPUPlace()
        paddle.seed(seed=0)
    m = my_chemprop_load(model, model_path=model_path, device=device)
    test_preds = chemprop_predict(
        model=m, smiles=test_data[smiles_column], fingerprints=None, num_workers=1
    )

    scores = evaluate_auto(
        true=test_data[property_column],
        preds=test_preds,
        dataset_type=cfg.DATA.dataset_type,
    )
    for score_name, score_value in scores.items():
        print(f"Test {score_name} = {score_value:.3f}")


def pre_compute(cfg):
    data_path = Path(cfg.PRE_COMPUTE.data_path)
    model_path = Path(cfg.PRE_COMPUTE.model_path)
    smiles_column = cfg.PRE_COMPUTE.smiles_column
    model_type = cfg.PRE_COMPUTE.model_type  #'chemprop'
    fingerprint_type = cfg.PRE_COMPUTE.fingerprint_type
    use_gpu = cfg.PRE_COMPUTE.use_gpu
    average_preds = cfg.PRE_COMPUTE.average_preds
    num_workers = cfg.PRE_COMPUTE.num_workers
    preds_column_prefix = cfg.PRE_COMPUTE.preds_column_prefix
    save_path = Path(cfg.PRE_COMPUTE.save_path)

    model = ppsci.arch.chemprop_molecule.MoleculeModel(cfg=cfg)

    data = pd.read_csv(data_path)
    smiles = list(data[smiles_column])
    if model_type != "chemprop" and fingerprint_type is None:
        raise ValueError("Must define fingerprint_type if using sklearn model.")
    if fingerprint_type is not None:
        # fingerprints = compute_fingerprints(smiles, fingerprint_type=
        #    fingerprint_type)
        pass
    else:
        fingerprints = None
    if model_path.is_dir():
        model_paths = list(
            model_path.glob("**/*.pt" if model_type == "chemprop" else "**/*.pkl")
        )
        if len(model_paths) == 0:
            raise ValueError(f"Could not find any models in directory {model_path}.")
    else:
        model_paths = [model_path]
    if model_type == "chemprop":
        if use_gpu:
            device = str("cuda").replace("cuda", "gpu")
        else:
            device = paddle.CPUPlace()
        paddle.seed(seed=0)

        models = [
            my_chemprop_load(model, model_path=model_path, device=device)
            for model_path in model_paths
        ]

    print(model_paths, models)

    if model_type == "chemprop":
        preds = np.array(
            [
                chemprop_predict(
                    model=m,
                    smiles=smiles,
                    fingerprints=fingerprints,
                    num_workers=num_workers,
                )
                for m in tqdm(models, desc="models")
            ]
        )

    if average_preds:
        preds = np.mean(preds, axis=0)
    model_string = (
        f"{model_type}{f'_{fingerprint_type}' if fingerprint_type is not None else ''}"
    )
    preds_string = f"{f'{preds_column_prefix}_' if preds_column_prefix is not None else ''}{model_string}"
    if average_preds:
        data[f"{preds_string}_ensemble_preds"] = preds
    else:
        for model_num, model_preds in enumerate(preds):
            data[f"{preds_string}_model_{model_num}_preds"] = model_preds
    if save_path is None:
        save_path = data_path
    save_path.parent.mkdir(parents=True, exist_ok=True)
    data.to_csv(save_path, index=False)


def generate(cfg):
    model_path = cfg.GENERATE.model_path
    model_type = cfg.GENERATE.model_type  #'chemprop'
    save_dir = Path(cfg.GENERATE.save_dir)

    building_blocks_path = cfg.GENERATE.building_blocks_path
    fingerprint_type = cfg.GENERATE.fingerprint_type
    reaction_to_building_blocks_path = cfg.GENERATE.reaction_to_building_blocks_path
    building_blocks_id_column = cfg.GENERATE.building_blocks_id_column
    building_blocks_score_column = cfg.GENERATE.building_blocks_score_column

    building_blocks_smiles_column = cfg.GENERATE.building_blocks_smiles_column
    reactions = REACTIONS
    max_reactions = cfg.GENERATE.max_reactions
    n_rollout = cfg.GENERATE.n_rollout

    explore_weight = cfg.GENERATE.explore_weight
    num_expand_nodes = cfg.GENERATE.num_expand_nodes

    optimization = cfg.GENERATE.optimization
    rng_seed = cfg.GENERATE.rng_seed

    no_building_block_diversity = cfg.GENERATE.no_building_block_diversity
    store_nodes = cfg.GENERATE.store_nodes

    verbose = cfg.GENERATE.verbose
    replicate = cfg.GENERATE.replicate

    save_dir.mkdir(parents=True, exist_ok=True)
    print("Loading building blocks...")
    if replicate:
        building_block_data = pd.read_csv(
            building_blocks_path, dtype={building_blocks_score_column: str}
        )
        building_block_data[building_blocks_score_column] = building_block_data[
            building_blocks_score_column
        ].astype(float)
        old_reactions_order = [
            275592,
            22,
            11,
            527,
            2430,
            2708,
            240690,
            2230,
            2718,
            40,
            1458,
            271948,
            27,
        ]
        reactions = tuple(
            sorted(
                reactions, key=lambda reaction: old_reactions_order.index(reaction.id)
            )
        )
        building_block_data.drop_duplicates(
            subset=building_blocks_smiles_column, inplace=True
        )
    else:
        building_block_data = pd.read_csv(building_blocks_path)
    print(f"Loaded {len(building_block_data):,} building blocks")
    if building_block_data[building_blocks_id_column].nunique() != len(
        building_block_data
    ):
        raise ValueError("Building block IDs are not unique.")
    building_block_smiles_to_id = dict(
        zip(
            building_block_data[building_blocks_smiles_column],
            building_block_data[building_blocks_id_column],
        )
    )
    building_block_id_to_smiles = dict(
        zip(
            building_block_data[building_blocks_id_column],
            building_block_data[building_blocks_smiles_column],
        )
    )
    building_block_smiles_to_score = dict(
        zip(
            building_block_data[building_blocks_smiles_column],
            building_block_data[building_blocks_score_column],
        )
    )
    print(f"Found {len(building_block_smiles_to_id):,} unique building blocks")
    set_all_building_blocks(
        reactions=reactions, building_blocks=set(building_block_smiles_to_id)
    )
    if reaction_to_building_blocks_path is not None:
        print("Loading and setting allowed building blocks for each reaction...")
        load_and_set_allowed_reaction_building_blocks(
            reactions=reactions,
            reaction_to_reactant_to_building_blocks_path=reaction_to_building_blocks_path,
        )
    print("Loading models and creating model scoring function...")
    model_scoring_fn = create_model_scoring_fn(
        model_path=model_path,
        model_type=model_type,
        fingerprint_type=fingerprint_type,
        smiles_to_score=building_block_smiles_to_score,
    )
    print("Setting up generator...")
    generator = Generator(
        building_block_smiles_to_id=building_block_smiles_to_id,
        max_reactions=max_reactions,
        scoring_fn=model_scoring_fn,
        explore_weight=explore_weight,
        num_expand_nodes=num_expand_nodes,
        optimization=optimization,
        reactions=reactions,
        rng_seed=rng_seed,
        no_building_block_diversity=no_building_block_diversity,
        store_nodes=store_nodes,
        verbose=verbose,
        replicate=replicate,
    )
    print("Generating molecules...")
    start_time = datetime.now()
    nodes = generator.generate(n_rollout=n_rollout)
    stats = {
        "mcts_time": datetime.now() - start_time,
        "num_nonzero_reaction_molecules": len(nodes),
        "approx_num_nodes_searched": generator.approx_num_nodes_searched,
    }
    print(f"MCTS time = {stats['mcts_time']}")
    print(
        f"Number of full molecule, nonzero reaction nodes = {stats['num_nonzero_reaction_molecules']:,}"
    )
    print(
        f"Approximate total number of nodes searched = {stats['approx_num_nodes_searched']:,}"
    )
    if store_nodes:
        stats["num_nodes_searched"] = generator.num_nodes_searched
        print(f"Total number of nodes searched = {stats['num_nodes_searched']:,}")
    pd.DataFrame(data=[stats]).to_csv(save_dir / "mcts_stats.csv", index=False)
    print("Saving molecules...")
    save_generated_molecules(
        nodes=nodes,
        building_block_id_to_smiles=building_block_id_to_smiles,
        save_path=save_dir / "molecules.csv",
    )


@hydra.main(version_base=None, config_path="./conf", config_name="synthemol.yaml")
def main(cfg: DictConfig):
    if cfg.mode == "train":
        train(cfg)
    elif cfg.mode == "eval":
        evaluate(cfg)
    elif cfg.mode == "pre-compute":
        pre_compute(cfg)
    elif cfg.mode == "generate":
        generate(cfg)
    else:
        raise ValueError(
            f"cfg.mode should in ['train', 'eval', 'pre-compute', 'generate'], but got '{cfg.mode}'"
        )


if __name__ == "__main__":
    main()

5. Result Display

Evaluate the training effect of the first step Chemprop model. By loading the pre-trained model and executing the evaluation command, the results can be obtained:

	roc_auc	prc_auc
chemprop	0.797	0.332

Checking the generated molecules.csv, you can see the generated molecular information similar to the table below:

smiles	node_id	num_expansions	rollout_num	score	Q_value	num_reactions	reaction_1_id	building_block_1_1_id	building_block_1_1_smiles	building_block_1_2_id	building_block_1_2_smiles
C#CCN(C(=O)C(C)(C)C#C)C1CCN(C(=O)OC(C)(C)C)CC1	91431	20	1			1	22	4349560	C#CCNC1CCN(C(=O)OC(C)(C)C)CC1	2998277	C#CC(C)(C)C(=O)O

It can be seen that molecular information meeting the requirements is generated, which is consistent with the author's design purpose.

6. References

• Generative AI for designing and validating easily synthesizable and structurally novel antibiotics
• Author's original repository