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.
-
Create a Conda environment:
conda create --name torch_molecule python=3.11.7 conda activate torch_molecule
-
Install using pip:
pip install torch-molecule
-
Install from source for the latest version:
Clone the repository:
git clone https://github.com/liugangcode/torch-molecule cd torch-moleculeInstall:
pip install .
| Model | Required Packages |
|---|---|
| HFPretrainedMolecularEncoder | transformers |
| HFPretrainedMolecularGenerator | transformers |
| HFPretrainedMolecularGenerator (MolGen) | transformers, selfies |
| HFPretrainedMolecularGenerator (Molexar) | transformers, fragment-selfies, molexar |
| HFPretrainedMolecularGenerator (SAFE-GPT) | transformers, safe-mol |
| BFGNNMolecularPredictor | torch-scatter |
| GRINMolecularPredictor | torch-scatter |
GRINMolecularPredictor (if enable repetition_augmentation=True) |
CombineMols |
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"))More examples can be found in the
examplesandtestsfolders.
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):
| Dataset | Description | Function |
|---|---|---|
| qm9 | Quantum chemical properties (DFT level) | load_qm9 |
| chembl2k | Bioactive molecules with drug-like properties | load_chembl2k |
| broad6k | Bioactive molecules with drug-like properties | load_broad6k |
| toxcast | Toxicity of chemical compounds | load_toxcast |
| admet | Chemical absorption, distribution, metabolism, excretion, and toxicity | load_admet |
| gasperm | Six gas permeability properties for polymeric materials | load_gasperm |
| zinc250k | A common subset of ZINC dataset, which does not have labels and could be used for unconditional generation or virtual screening | load_zinc250k |
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!)
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,
)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")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:
 Man Hei Matthew Thom Integrating Pretrained Generator (NovoMolGen, MolGen, Molexar, SAFE-GPT), dataset splitting modules |
 Eric Inae Integrating MoAMa, GraphMAE, AttrMasking, ContextPred, EdgePred |
 Yihan Zhu Integrating DeFoG, GRIN, BFGNN, RPGNN, Transformer (SMILES) |
 Monica C. S Debugging and integrating training logs (#19) |