Easy access to state-of-the-art pre-trained models with python

HuggingMolecules

We envision models that are pre-trained on a vast range of domain-relevant tasks to become key for molecule propertyprediction. 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

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

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.

GitHub

https://github.com/gmum/huggingmolecules