SPMM: Structure-Property Multi-Modal learning for molecules

The official GitHub for SPMM, a multi-modal molecular pre-trained model for a synergistic comprehension of molecular structure and properties. The details can be found in the following paper: Bidirectional Generation of Structure and Properties Through a Single Molecular Foundation Model. (Nature Communications 2024)

Molecule structure will be given in SMILES, and we used 53 simple chemical properties to build a property vector(PV) of a molecule.

The model checkpoint and data are too heavy to be included in this repo, and they can be found here.

Files

• data/: Contains the data used for the experiments in the paper. (you have to make this folder and put the data that you downloaded from the link above.)
• Pretrain/: Contains the checkpoint of the pre-trained SPMM. (you have to make this folder and put the checkpoint that you downloaded from the link above.)
• vocab_bpe_300.txt: Contains the SMILES tokens for the SMILES tokenizer.
• property_name.txt: Contains the name of the 53 chemical properties.
• normalize.pkl: Contains the mean and standard deviation of the 53 chemical properties that we used for PV.
• calc_property.py: Contains the code to calculate the 53 chemical properties and build a PV for a given SMILES. Modify this code accordingly to utilize SPMM pre-training for your custom PVs.
• SPMM_models.py: Contains the code for the SPMM model and its pre-training codes.
• SPMM_pretrain.py: runs SPMM pre-training.
• d_*.py: Codes for the downstream tasks.

Requirements

Run pip install -r requirements.txt to install the required packages.

Code running

Arguments can be passed with commands, or be edited manually in the running code.

1. Pre-training

   python SPMM_pretrain.py --data_path './data/pretrain.txt'

2. PV-to-SMILES generation

   • batched: The model takes PVs from the molecules in input_file, and generates molecules with those PVs using k-beam search. The generated molecules will be written in generated_molecules.txt.

     python d_pv2smiles_batched.py --checkpoint './Pretrain/checkpoint_SPMM.ckpt' --input_file './data/pubchem_1k_unseen.txt' --k 2

   • single: The model takes one query PV and generates n_generate molecules with that PV using k-beam search. The generated molecules will be written in generated_molecules.txt. Here, you need to build your input PV in the code. Check the four examples that we included.

     python d_pv2smiles_single.py --checkpoint './Pretrain/checkpoint_SPMM.ckpt' --n_generate 1000 --stochastic True --k 2

3. SMILES-to-PV generation

   The model takes the query molecules in input_file, and generates their PV.

   python d_smiles2pv.py --checkpoint './Pretrain/checkpoint_SPMM.ckpt' --input_file './data/pubchem_1k_unseen.txt'

4. MoleculeNet + DILI prediction task

   d_regression.py, d_classification.py, and d_classification_multilabel.py, perform regression, binary classification, and multi-label classification tasks, respectively.

   python d_regression.py --checkpoint './Pretrain/checkpoint_SPMM.ckpt' --name 'bace'
   python d_classification.py --checkpoint './Pretrain/checkpoint_SPMM.ckpt' --name 'bbbp'
   python d_classification_multilabel.py --checkpoint './Pretrain/checkpoint_SPMM.ckpt' --name 'clintox'

5. Forward/retro-reaction prediction tasks

   d_rxn_prediction.py performs both forward/reverse reaction prediction tasks on USPTO-480k and USPTO-50k datasets.

   e.g. forward reaction prediction, no beam search

   python d_rxn_prediction.py --checkpoint './Pretrain/checkpoint_SPMM.ckpt' --mode 'forward' --n_beam 1 

   e.g. retro reaction prediction, beam search with k=3

   python d_rxn_prediction.py --checkpoint './Pretrain/checkpoint_SPMM.ckpt' --mode 'retro' --n_beam 3 

Acknowledgement

• The code for BERT with cross-attention layers xbert.py and schedulers are modified from the one in ALBEF.
• The code for SMILES augmentation is taken from pysmilesutils.