File size: 9,923 Bytes
5539717 55d9b0c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
---
license: other
language:
- en
pipeline_tag: text-generation
tags:
- llamol
- molecular generation
- machine learning
- transformers
- generation
- de novo molecular design
---
# Llamol
<p align="center">
<img src="assets/llamol.png" width="300" height="300" alt="LLamol">
</p>
This is the official repository for the paper ["LLamol: A Dynamic Multi-Conditional Generative Transformer for De Novo Molecular Design"](https://arxiv.org/abs/2311.14407).
In this repository are the weights for LLamol (`out/llama2-M-Full-RSS.pt`) and the dataset OrganiX13.
Image made with [Hotspot.ai](https://hotpot.ai/art-generator)
## Installation
Install using Mamba to be fast: https://mamba.readthedocs.io/en/latest/micromamba-installation.html
```bash
$ "${SHELL}" <(curl -L micro.mamba.pm/install.sh)
$ micromamba env create -f torch2-env.yaml
$ micromamba activate torch2-llamol
$ python sample.py
```
# Download and preprocess the OrganiX13 dataset:
If you want to train with the full 13 Million dataset do the following steps. These are *not* necessary if you just want to use the model for inference:
1. Download and preprocess the OPV dataset by running `/data/opv/prepare_opv.py`
2. Download and preprocess the ZINC dataset by running `/data/zinc/zinc_complete/run_download.py` followed by `/data/zinc/convert_to_parquet.py`
(we recommend at least 16GB RAM for this)
3. Download and preprocess the ZINC dataset by running `/data/qm9_zinc250k_cep/convert_to_parquet.py`
4. Run `data/combine_all.py` to combine the dataset to `data/OrganiX13.parquet` (this can take a while, especially on the zinc dataset. In total it took ~2 hours when using my Laptop, which has 16 GB ram and an Intel i7 10th Gen)
5. Run `preprocess_dataset.py` which should create the file `.cache/processed_dataset_None.pkl`
Now you can use that in the training of the model by specifing the file under the `processed_dataset_ckpt` of the training .yaml files.
# Interactive Demo
After installation you can play around with the model using the `demonstrator.ipynb` file. Just run all and scroll down to the last cell.
After a short time there should be a UI where you can play around with the model.
## Training
First the env needs to be activated so:
```bash
$ conda activate torch2-llamol # When installed with conda instead of micromamba
OR
$ micromamba activate torch2-llamol
``````
To train locally you can run:
```bash
# To set the config that you want to train with
$ python train.py train=llama2-M-Full-RSS
```
Parameters can also be overriden by using the following, for example:
```bash
$ python train.py train=llama2-M-Full-RSS train.model.dim=1024
```
For more information look at [Hydra](https://hydra.cc/docs/1.3/intro/)
To start a job on a SLURM cluster use the following script:
```bash
$ sbatch trainLLamaMol.sh
``````
## Training Multi-GPU on 1 Node with multiple GPUS (nproc_per_node)
```bash
torchrun --standalone --max_restarts=3 --nnodes=1 --nproc_per_node=2 --rdzv-backend=c10d --rdzv-endpoint="$localhost:12345" train.py train=llama2-M-Full-RSS > "train_runs/run_MultiGPU.out"
```
## Training Multi-GPU on 1 Node with multiple GPUS on a Cluster
Currently there is only one script to train with DDP. To change the number of GPUS in that script you have to change the bash script itself.
TODO: Make it more dynamic, with allowing console commands to change the number of GPUS etc.
```bash
sbatch trainLLamaMolDDPSingleNode.sh
```
## Sampling
Sampling can be changed by the OPTIONAL parameters as shown below.
```bash
$ python sample.py --help
$ python sample.py --num_samples 2000 --ckpt_path "out/llama2-M-Full-RSS.pt" --max_new_tokens 256 --cmp_dataset_path="data/OrganiX13.parquet" --seed 4312 --context_cols logp sascore mol_weight --temperature 0.8
```
## Using own dataset
Use the `preprocess_dataset.py` file to tokenize the dataset. The dataset should be either in the parquet or csv format.
The SMILES used for training should be in the `smiles` column in the dataset. All conditions, should be given to the pretokenize function.
After the preprocessing is done a file should be stored in the .cache directory with the name `processed_dataset_{limit}.pkl`.
You could also rename this file to not overwrite it every time you run the preprocessing.
The `.cache/processed_dataset_{limit}.pkl` can then be set in the `config/train/llama2-M-Full-RSS.yaml file` to change the training with the new dataset in the `processed_dataset_ckpt` field in the yaml file.
# Training methods
The training method we used and described in the paper is here called RSS for "Random Smiles Sampling" which was the method then described in the "Stochastic Context Learning" as taking a random subsequence from the current SMILES while training and feeding that into the model as a token sequence condition. So the model we used in the paper was the `out/llama2-M-Full-RSS.pt`.
We also tried other approached for including the token sequence.
One was using murcko scaffolds as they were used in the MolGPT paper, but this approach did not yield great results for our purposes.
The other was using BRICKS decomposition, which also did not yield very good results.
The different methods are implemented in the `fragment_creator.py` file.
Each of the models were trained with their respective configurations in the `config/train` folder.
# Thanks
- [Karpathy](https://github.com/karpathy/llama2.c) for the implementation of the Llama 2 architecture and training code
- [DeepChem](https://github.com/deepchem/deepchem) for the SmilesTokenizer
- [TorchDrug](https://github.com/DeepGraphLearning/torchdrug/) for the downloads scripts for the OPV and CEP datasets
- Zinc 15 dataset (Teague Sterling and John J. Irwin. ZINC 15 – ligand discovery for everyone. Journal of Chemical Information
and Modeling, 55(11):2324–2337, November 2015.)
- QM9 dataset (
Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quantum chemistry
structures and properties of 134 kilo molecules. Scientific Data, 1(1), aug 2014.)
- PC9 dataset (Marta Glavatskikh, Jules Leguy, Gilles Hunault, Thomas Cauchy, and Benoit Da Mota. Dataset’s chemical
diversity limits the generalizability of machine learning predictions. Journal of Cheminformatics, 11(1), nov 2019)
- ZINC 250k (Rafael Gó mez-Bombarelli, Jennifer N. Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín
Sánchez-Lengeling, Dennis Sheberla, Jorge Aguilera-Iparraguirre, Timothy D. Hirzel, Ryan P. Adams, and Alán
Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS
Central Science, 4(2):268–276, jan 2018.)
- RedDB (Elif Sorkun, Qi Zhang, Abhishek Khetan, Murat Cihan Sorkun, and Süleyman Er. RedDB, a computational
database of electroactive molecules for aqueous redox flow batteries. Scientific Data, 9(1), nov 2022.)
- OPV (Peter C. St. John, Caleb Phillips, Travis W. Kemper, A. Nolan Wilson, Yanfei Guan, Michael F. Crowley, Mark R.
Nimlos, and Ross E. Larsen. Message-passing neural networks for high-throughput polymer screening. The
Journal of Chemical Physics, 150(23):234111, jun 2019.)
- PubchemQC 2020 (Maho Nakata, Tomomi Shimazaki, Masatomo Hashimoto, and Toshiyuki Maeda. PubChemQC PM6: Data sets
of 221 million molecules with optimized molecular geometries and electronic properties. Journal of Chemical
Information and Modeling, 60(12):5891–5899, oct 2020.)
- PubchemQC 2017 (Maho Nakata and Tomomi Shimazaki. PubChemQC project: A large-scale first-principles electronic structure
database for data-driven chemistry. Journal of Chemical Information and Modeling, 57(6):1300–1308, may 2017.)
- CEP (Johannes Hachmann, Roberto Olivares-Amaya, Sule Atahan-Evrenk, Carlos Amador-Bedolla, Roel S. Sánchez-
Carrera, Aryeh Gold-Parker, Leslie Vogt, Anna M. Brockway, and Alán Aspuru-Guzik. The Harvard clean energy
project: Large-scale computational screening and design of organic photovoltaics on the world community grid.
The Journal of Physical Chemistry Letters, 2(17):2241–2251, aug 2011.) subset ( David Duvenaud, Dougal Maclaurin, Jorge Aguilera-Iparraguirre, Rafael Gómez-Bombarelli, Timothy Hirzel,
Alán Aspuru-Guzik, and Ryan P. Adams. Convolutional networks on graphs for learning molecular fingerprints,
2015.)
- ChEMBL (James Blackshaw, Anna Gaulton, A. Patrícia Bento, Marleen De Veij, David Mendez Lopez, Nicolas Bosc, Juan
Felipe Mosquera Morales, María Paula Margariños, Andrew Leach, Emma Manners, Barbara Zdrazil, Harris
Ioannidis, Fiona Hunter, Eloy Félix, and Ricardo Arcila Toro. CHEMBL database release 31, September 2009.)
# Funding disclaimer
This project has received funding from the European Union’s Horizon 2020 research and innovation programme under Grant Agreement no. 875489.
This website reflects only the author’s view. The funding agency is not responsible for any use made of the information it contains.
# License
<p xmlns:cc="http://creativecommons.org/ns#" xmlns:dct="http://purl.org/dc/terms/"><span property="dct:title">LLamol is licensed under <a href="http://creativecommons.org/licenses/by-nc-sa/4.0/?ref=chooser-v1" target="_blank" rel="license noopener noreferrer" style="display:inline-block;">CC BY-NC-SA 4.0<img style="height:22px!important;margin-left:3px;vertical-align:text-bottom;" src="https://mirrors.creativecommons.org/presskit/icons/cc.svg?ref=chooser-v1"><img style="height:22px!important;margin-left:3px;vertical-align:text-bottom;" src="https://mirrors.creativecommons.org/presskit/icons/by.svg?ref=chooser-v1"><img style="height:22px!important;margin-left:3px;vertical-align:text-bottom;" src="https://mirrors.creativecommons.org/presskit/icons/nc.svg?ref=chooser-v1"><img style="height:22px!important;margin-left:3px;vertical-align:text-bottom;" src="https://mirrors.creativecommons.org/presskit/icons/sa.svg?ref=chooser-v1"></a></p> |