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…
torch-molecule is a deep learning package for molecular discovery, designed with an sklearn-style interface for property prediction, inverse design and representation learning.
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.
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 folder.
Datasets
Explore small-molecule and polymer datasets in the molecular dataset hub, with original sources, licenses, units, source splits, and quality tags.
See the dataset documentation
for dataset selection, caching, target columns, and train/test splits.
Current counts follow the hub’s published metadata. Hugging Face entries include public and gated datasets and files without SMILES tables. Source rows may repeat structures; each dataset retains its own license.
Load published SMILES tables by their original dataset name:
from torch_molecule.datasets import list_datasets, load_dataset
list_datasets(category="adme", grade="ready") # browse by category, quality grade, or license
data = load_dataset("qm9", targets=["homo", "lumo", "gap"], local_dir="torchmol_data")
smiles_list, property_np_array = data.data, data.target # 133,885 molecules, shape (133885, 3)
data.info["targets"]["gap"]["units"] # 'Hartree'
train = load_dataset("moleculeace_chembl204_ki", split="train") # source splits where defined;
test = load_dataset("moleculeace_chembl204_ki", split="test") # test drops molecules seen in train
zinc = load_dataset("zinc250k", targets=[]) # structures only: target is None
load_qm9, load_chembl2k, load_broad6k, load_toxcast, load_admet, load_gasperm, and load_zinc250k still work but are deprecated aliases of load_dataset and will be removed.
(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_dataset
from torch_molecule import GREAMolecularPredictor
data = load_dataset("qm9", targets=["gap"], local_dir="torchmol_data")
# "random" | "scaffold" | "butina" | "size"
# scaffold: unseen Bemis-Murcko scaffolds; butina: Tanimoto clusters; size: heavy-atom count
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,
)
Random splitting assigns individual rows, so repeated structures can occur
in both splits. The splitters preserve data.groups metadata but do not use it
to assign rows. Use scaffold or butina, or split by group yourself, when
repeated measurements must stay together.
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="GREAMolecularPredictor.pt")
# Adjust model parameters and make predictions
model.set_params(verbose='none')
predictions = model.predict(smiles_list)
NovoMolGen: Causal LM pretrained on ZINC-22 for de novo SMILES generation. MolGen-large: Seq2Seq SELFIES generator with high chemical validity. MolGen-large-opt: MolGen-large fine-tuned for QED / p-logP optimization. Molexar-10M-base: Fragment-SELFIES de novo and fragment-constrained generation. Molexar-10M-omni: Multi-condition Molexar model for property-guided generation. SAFE-GPT: GPT-2 causal LM pretrained on SAFE strings for de novo generation and scaffold-prefix completion.
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
评论