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 (0.1.2):
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 |
| BFGNNMolecularPredictor | torch-scatter |
| GRINMolecularPredictor | torch-scatter |
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
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 |
from torch_molecule.datasets import load_qm9
# local_dir is the local path where the dataset will be saved
smiles_list, property_np_array = load_qm9(local_dir='torchmol_data')
# 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']
smiles_list, property_np_array = load_qm9(local_dir='torchmol_data', target_cols=target_cols)(We welcome your suggestions and contributions on your datasets!)
After preparing the dataset, we can easily fit a model similar to how we use sklearn (actually, the coding is even simpler than sklearn, as we still need to do feature engineering in sklearn to convert molecule SMILES into vectors):
from torch_molecule import GREAMolecularPredictor
split = int(0.8 * len(smiles_list))
grea = GREAMolecularPredictor(
num_task=num_task,
task_type="regression",
evaluate_higher_better=False,
verbose=True
)
# Fit with automatic hyperparameter tuning with 10 attempts, or implement .fit() with the default/manual hyperparameters
grea.autofit(
X_train=smiles_list[:split],
y_train=property_np_array[:split],
X_val=smiles_list[split:],
y_val=property_np_array[split:],
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=False)
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.
