liugangcode/torch-molecule

torch‑molecule – Deep‑learning for molecular discovery (sklearn‑style API)

What it is – A Python library built on PyTorch that lets researchers and engineers train, evaluate, and deploy molecular AI models (predictive, generative, and representation) using a familiar scikit‑learn‑like interface.

Why it matters – Molecular discovery (drug design, materials science, toxicity prediction, etc.) increasingly relies on graph neural networks and transformer models. torch‑molecule bundles a large collection of state‑of‑the‑art architectures, pretrained checkpoints, and benchmark datasets, while hiding the boiler‑plate of data handling, model construction, and hyper‑parameter tuning.

---

Key components

Component	What it does	Example models
Predictive Models	Supervised regression / classification on molecular graphs or SMILES strings.	GRIN, BFGNN, GREA, DIR, SGIR, GNNs, Transformer (SMILES), LSTM (SMILES)
Generative Models	De‑novo molecule generation, conditional generation, scaffold‑guided sampling.	DeFoG, Graph DiT, DiGress, GDSS, MolGPT, JTVAE, GraphGA, pretrained generators (NovoMolGen, MolGen‑large, Molexar, SAFE‑GPT)
Representation Models	Self‑supervised pre‑training of graph or SMILES encoders for downstream tasks.	MoAMa, GraphMAE, AttrMasking, ContextPred, EdgePred, InfoGraph, a suite of pretrained BERT/GPT‑style models

Datasets bundled

• qm9 – quantum‑chemical properties (∼134 k molecules)
• chembl2k / broad6k – drug‑like bioactivity data
• toxcast – toxicity labels
• admet – absorption/distribution/metabolism/excretion/toxicity
• gasperm – polymer gas‑permeability
• zinc250k – unlabeled subset of ZINC for unconditional generation

All are loaded via torch_molecule.datasets.load_<name>(), returning a simple object with .data (SMILES list) and .target (numpy array or None).

Installation

# create a fresh conda env (Python 3.11 recommended)
conda create -n torch_molecule python=3.11.7 && conda activate torch_molecule

# install the package from PyPI
pip install torch-molecule

For the latest source:

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

Some models need extra dependencies (e.g., torch-scatter, transformers, selfies, fragment-selfies, molexar, safe-mol). The README lists the exact pip commands.

Typical workflow (predictive example)

from torch_molecule.datasets import load_qm9
from torch_molecule import GREAMolecularPredictor

# 1️⃣ Load and split data
qm9 = load_qm9(local_dir='torchmol_data')
train, val = qm9.train_test_split(test_size=0.2, method='scaffold', seed=42)

# 2️⃣ Instantiate a model (sklearn‑style)
model = GREAMolecularPredictor(
    num_task=1,
    task_type='regression',
    evaluate_higher_better=False,
    verbose='progress_bar'
)

# 3️⃣ Auto‑tune hyper‑parameters and fit
model.autofit(
    X_train=train.data, y_train=train.target,
    X_val=val.data,   y_val=val.target,
    n_trials=10
)

# 4️⃣ Predict on new SMILES
preds = model.predict(['C1=CC=CC=C1', 'CCO'])

The same pattern works for generative models (HFPretrainedMolecularGenerator) and representation encoders.

Check‑point handling

• Save to Hugging Face – model.save_to_hf(repo_id='user/repo', task_id='qm9_grea', ...)
• Load from Hugging Face – model.load_from_hf(repo_id='user/repo', local_cache='.../GREA_qm9.pt')
• Local save / load – model.save_to_local('qm9_grea.pt') and model.load_from_local('qm9_grea.pt')

Extensibility

• New datasets can be added by contributing a loader function.
• Community contributors have already integrated many pretrained generators (MolGen, Molexar, SAFE‑GPT) and dataset‑splitting utilities.
• The library is open‑source; pull‑requests are welcomed.

---

In a nutshell – torch‑molecule gives you a one‑stop, sklearn‑style toolbox for modern graph‑ and transformer‑based molecular AI, complete with benchmark data, pretrained checkpoints, and easy model persistence.