HuggingMolecules

We envision models that are pre-trained on a vast range of domain-relevant tasks to become key for molecule property prediction. This repository aims to give easy access to state-of-the-art pre-trained models.

Quick tour

To quickly fine-tune a model on a dataset using the pytorch lightning package follow the below example based on the MAT model and the freesolv dataset:

from huggingmolecules import MatModel, MatFeaturizer

# The following import works only from the source code directory:
from experiments.src import TrainingModule, get_data_loaders

from torch.nn import MSELoss
from torch.optim import Adam

from pytorch_lightning import Trainer
from pytorch_lightning.metrics import MeanSquaredError

# Build and load the pre-trained model and the appropriate featurizer:
model = MatModel.from_pretrained('mat_masking_20M')
featurizer = MatFeaturizer.from_pretrained('mat_masking_20M')

# Build the pytorch lightning training module:
pl_module = TrainingModule(model,
                           loss_fn=MSELoss(),
                           metric_cls=MeanSquaredError,
                           optimizer=Adam(model.parameters()))

# Build the data loader for the freesolv dataset:
train_dataloader, _, _ = get_data_loaders(featurizer,
                                          batch_size=32,
                                          task_name='ADME',
                                          dataset_name='hydrationfreeenergy_freesolv')

# Build the pytorch lightning trainer and fine-tune the module on the train dataset:
trainer = Trainer(max_epochs=100)
trainer.fit(pl_module, train_dataloader=train_dataloader)

# Make the prediction for the batch of SMILES strings:
batch = featurizer(['C/C=C/C', '[C]=O'])
output = pl_module.model(batch)

Installation

Create your conda environment and install the rdkit package:

conda create -n huggingmolecules python=3.8.5
conda activate huggingmolecules
conda install -c conda-forge rdkit==2020.09.1

Then install huggingmolecules from the cloned directory:

conda activate huggingmolecules
pip install -e ./src

Huggingmolecules caches weights and configs of the models. To avoid issues with incompatibility of different package versions, it is recommended to clean up the cache directory after every package update:

python -m src.clean_cache --all

Project Structure

The project consists of two main modules: src/ and experiments/ modules:

• The src/ module contains abstract interfaces for pre-trained models along with their implementations based on the pytorch library. This module makes configuring, downloading and running existing models easy and out-of-the-box.
• The experiments/ module makes use of abstract interfaces defined in the src/ module and implements scripts based on the pytorch lightning package for running various experiments. This module makes training, benchmarking and hyper-tuning of models flawless and easily extensible.

Supported models architectures

Huggingmolecules currently provides the following models architectures:

• MAT
• GROVER
• R-MAT (weights were obtained by joint efforts with Nvidia)

For ease of benchmarking, we also include wrappers in the experiments/ module for three other models architectures:

• chemprop
• ChemBERTa
• MolBERT

The src/ module

The implementations of the models in the src/ module are divided into three modules: configuration, featurization and models module. The relation between these modules is shown on the following examples based on the MAT model:

Configuration examples

from huggingmolecules import MatConfig

# Build the config with default parameters values, 
# except 'd_model' parameter, which is set to 1200:
config = MatConfig(d_model=1200)

# Build the pre-defined config:
config = MatConfig.from_pretrained('mat_masking_20M')

# Build the pre-defined config with 'init_type' parameter set to 'normal':
config = MatConfig.from_pretrained('mat_masking_20M', init_type='normal')

# Save the pre-defined config with the previous modification:
config.save_to_cache('mat_masking_20M_normal.json')

# Restore the previously saved config:
config = MatConfig.from_pretrained('mat_masking_20M_normal.json')

Featurization examples

from huggingmolecules import MatConfig, MatFeaturizer

# Build the featurizer with pre-defined config:
config = MatConfig.from_pretrained('mat_masking_20M')
featurizer = MatFeaturizer(config)

# Build the featurizer in one line:
featurizer = MatFeaturizer.from_pretrained('mat_masking_20M')

# Encode (featurize) the batch of two SMILES strings: 
batch = featurizer(['C/C=C/C', '[C]=O'])

Models examples

from huggingmolecules import MatConfig, MatFeaturizer, MatModel

# Build the model with the pre-defined config:
config = MatConfig.from_pretrained('mat_masking_20M')
model = MatModel(config)

# Load the pre-trained weights 
# (which do not include the last layer of the model)
model.load_weights('mat_masking_20M')

# Build the model and load the pre-trained weights in one line:
model = MatModel.from_pretrained('mat_masking_20M')

# Encode (featurize) the batch of two SMILES strings: 
featurizer = MatFeaturizer.from_pretrained('mat_masking_20M')
batch = featurizer(['C/C=C/C', '[C]=O'])

# Feed the model with the encoded batch:
output = model(batch)

# Save the weights of the model (usually after the fine-tuning process):
model.save_weights('tuned_mat_masking_20M.pt')

# Load the previously saved weights
# (which now includes all layers of the model):
model.load_weights('tuned_mat_masking_20M.pt')

# Load the previously saved weights, but without 
# the last layer of the model ('generator' in the case of the 'MatModel')
model.load_weights('tuned_mat_masking_20M.pt', excluded=['generator'])

# Build the model and load the previously saved weights:
config = MatConfig.from_pretrained('mat_masking_20M')
model = MatModel.from_pretrained('tuned_mat_masking_20M.pt',
                                 excluded=['generator'],
                                 config=config)

Running tests

To run base tests for src/ module, type:

pytest src/ --ignore=src/tests/downloading/

To additionally run tests for downloading module (which will download all models to your local computer and therefore may be slow), type:

pytest src/tests/downloading

The experiments/ module

Requirements

In addition to dependencies defined in the src/ module, the experiments/ module goes along with few others. To install them, run:

pip install -r experiments/requirements.txt

The following packages are crucial for functioning of the experiments/ module:

• pytorch lightning
• optuna
• gin-config
• TDC

Neptune.ai

In addition, we recommend installing the neptune.ai package:

1. Sign up to neptune.ai at https://neptune.ai/.

2. Get your Neptune API token (see getting-started for help).

3. Export your Neptune API token to NEPTUNE_API_TOKEN environment variable.

4. Install neptune-client: pip install neptune-client.

5. Enable neptune.ai in the experiments/configs/setup.gin file.

6. Update neptune.project_name parameters in experiments/configs/bases/*.gin files.

Running scripts:

We recommend running experiments scripts from the source code. For the moment there are three scripts implemented:

• experiments/scripts/train.py - for training with the pytorch lightning package
• experiments/scripts/tune_hyper.py - for hyper-parameters tuning with the optuna package
• experiments/scripts/benchmark.py - for benchmarking based on the hyper-parameters tuning (grid-search)

In general running scripts can be done with the following syntax:

python -m experiments.scripts.<script_name> /
       -d <dataset_name> / 
       -m <model_name> /
       -b <parameters_bindings>

Then the script <script_name>.py runs with functions/methods parameters values defined in the following gin-config files:

1. experiments/configs/bases/<script_name>.gin
2. experiments/configs/datasets/<dataset_name>.gin
3. experiments/configs/models/<model_name>.gin

If the binding flag -b is used, then bindings defined in <parameters_binding> overrides corresponding bindings defined in above gin-config files.

So for instance, to fine-tune the MAT model (pre-trained on masking_20M task) on the freesolv dataset using GPU 1, simply run:

python -m experiments.scripts.train /
       -d freesolv / 
       -m mat /
       -b model.pretrained_name=\"mat_masking_20M\"#train.gpus=[1]

or equivalently:

python -m experiments.scripts.train /
       -d freesolv / 
       -m mat /
       --model.pretrained_name mat_masking_20M /
       --train.gpus [1]

Local dataset

To use a local dataset, create an appropriate gin-config file in the experiments/configs/datasets directory and specify the data.data_path parameter within. For details see the get_data_split implementation.

Benchmarking

For the moment there is one benchmark available. It works as follows:

• experiments/scripts/benchmark.py: on the given dataset we fine-tune the given model on 10 learning rates and 6 seeded data splits (60 fine-tunings in total). Then we choose that learning rate that minimizes an averaged (on 6 data splits) validation metric (metric computed on the validation dataset, e.g. RMSE). The result is the averaged value of test metric for the chosen learning rate.

Running a benchmark is essentially the same as running any other script from the experiments/ module. So for instance to benchmark the vanilla MAT model (without pre-training) on the Caco-2 dataset using GPU 0, simply run:

python -m experiments.scripts.benchmark /
       -d caco2 / 
       -m mat /
       --model.pretrained_name None /
       --train.gpus [0]

However, the above script will only perform 60 fine-tunings. It won't compute the final benchmark result. To do that wee need to run:

python -m experiments.scripts.benchmark --results_only /
       -d caco2 / 
       -m mat

The above script won't perform any fine-tuning, but will only compute the benchmark result. If we had neptune enabled in experiments/configs/setup.gin, all data necessary to compute the result will be fetched from the neptune server.

Benchmark results

We performed the benchmark described in Benchmarking as experiments/scripts/benchmark.py for various models architectures and pre-training tasks.

Summary

We report mean/median ranks of tested models across all datasets (both regression and classification ones). For detailed results see Regression and Classification sections.

Regression

As the metric we used MAE for QM7 and RMSE for the rest of datasets.

¹ chemprop with additional rdkit_2d_normalized features generator
² chemprop with additional morgan_count features generator

Classification

We used ROC AUC as the metric.