Deep learning for molecular discovery with a simple sklearn-style interface

---

torch-molecule is a package that facilitates molecular discovery through deep learning, featuring a user-friendly, sklearn-style interface. It includes model checkpoints for efficient deployment and benchmarking across a range of molecular tasks. The package focuses on three main components: Predictive Models, Generative Models, and Representation Models, which make molecular AI models easy to implement and deploy.

See the List of Supported Models section for all available models.

Installation

1. Create a Conda environment:

   conda create --name torch_molecule python=3.11.7
   conda activate torch_molecule
   

2. Install using pip:

   pip install torch-molecule
   

3. Install from source for the latest version:

   Clone the repository:

   git clone https://github.com/liugangcode/torch-molecule
   cd torch-molecule
   

   Install:

   pip install .
   

Additional Packages

For models that require torch-scatter: Install using the following command: pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html, e.g.,

pip install torch-scatter -f https://data.pyg.org/whl/torch-2.7.1+cu128.html

For models that require transformers: pip install transformers

For MolGen (selfies): pip install "selfies>=2.1". Source: aspuru-guzik-group/selfies.

For Molexar: pip install fragment-selfies loguru (Fragment-SELFIES) and pip install git+https://github.com/fairydance/Molexar.git (Molexar). Molexar itself requires transformers>=5.8.

For SAFE-GPT: pip install safe-mol (SAFE).

from torch_molecule import HFPretrainedMolecularGenerator

model = HFPretrainedMolecularGenerator(
    repo_id="datamol-io/safe-gpt",
)
model.fit()
print(model.generate(n_samples=5))
print(model.generate(n_samples=5, scaffold="c1ccccc1"))

Usage

More examples can be found in the examples and tests folders.

torch-molecule supports applications in broad domains from chemistry, biology, to materials science. To get started, you can load prepared datasets from torch_molecule.datasets (updated after v0.1.3):

from torch_molecule.datasets import load_qm9

# local_dir is the local path where the dataset will be saved
molecular_data = load_qm9(local_dir='torchmol_data')
smiles_list, property_np_array = molecular_data.data, molecular_data.target

# len(smiles_list): 133885
# Property array shape: (133885, 1)

# load_qm9 returns the target "gap" by default, but you can adjust it by passing new target_cols
target_cols = ['homo', 'lumo', 'gap']
molecular_data = load_qm9(local_dir='torchmol_data', target_cols=target_cols)
smiles_list, property_np_array = molecular_data.data, molecular_data.target

# the target could be None if loading an unlabeled dataset
from torch_molecule.datasets import load_zinc250k
molecular_data = load_zinc250k(local_dir='torchmol_data')
smiles_list = molecular_data.data
assert molecular_data.target is None

(We are actively adding more datasets. We welcome your suggestions and contributions on your datasets!)

Fit a Model

After preparing the dataset, split it, then fit a model with an sklearn-style API (no extra SMILES featurization is required):

from torch_molecule.datasets import load_qm9
from torch_molecule import GREAMolecularPredictor

data = load_qm9(local_dir='torchmol_data')
# "random" | "scaffold" | "butina" | "size"
# scaffold: unseen Bemis-Murcko scaffolds; butina: Tanimoto clusters; size: heavy-atom count
# Split the full dataset. subsample() is only for local debugging / CI — do not
# shrink QM9 (or any benchmark) just to make Butina cheaper.
train, val = data.train_test_split(test_size=0.2, method="scaffold", seed=42)

grea = GREAMolecularPredictor(
    num_task=1,
    task_type="regression",
    evaluate_higher_better=False,
    verbose="progress_bar" #or "print_statement" recommended for jupyter notebooks, or "none"
)

# Fit with automatic hyperparameter tuning with 10 attempts, or implement .fit() with the default/manual hyperparameters
grea.autofit(
    X_train=train.data,
    y_train=train.target,
    X_val=val.data,
    y_val=val.target,
    n_trials=10,
)

Checkpoints

torch-molecule provides checkpoint functions that can be interacted with on Hugging Face:

from torch_molecule import GREAMolecularPredictor

repo_id = "user/repo_id"  # replace with your own Hugging Face username and repo_id

# Save the trained model to Hugging Face
grea.save_to_hf(
    repo_id=repo_id,
    task_id="qm9_grea",
    commit_message="Upload qm9_grea",
    private=False
)

# Load a pretrained checkpoint from Hugging Face
model = GREAMolecularPredictor()
model.load_from_hf(repo_id=repo_id, local_cache=f"{model_dir}/GREA_{task_name}.pt")

# Adjust model parameters and make predictions
model.set_params(verbose='none')
predictions = model.predict(smiles_list)

Or you can save the model to a local path:

grea.save_to_local("qm9_grea.pt")

new_model = GREAMolecularPredictor()
new_model.load_from_local("qm9_grea.pt")

List of Supported Models

Predictive Models

Generative Models

Representation Models

Acknowledgements

The project template was adapted from https://github.com/lwaekfjlk/python-project-template. We thank the authors for their contribution to the open-source community.

This project is initiated and maintained by Gang Liu and Meng Jiang, and would not be possible without the valuable contributions from the open-source community, including the following members: