torch-molecule: Deep Learning for Molecular Discovery

The torch_molecule library provides high-level functions to load various molecular datasets either locally or directly from the Hugging Face Hub. Depending on the dataset, you can load specific chemical properties or unlabeled SMILES lists.

Available dataset loading functions include:

• load_qm9(): Loads the QM9 dataset. By default, it returns the target gap, but you can specify different target columns using the target_cols parameter.
• load_chembl2k(): Loads a small subset from the ChEMBL database.
• load_broad6k(): Loads the Broad6k dataset.
• load_toxcast(): Loads the ToxCast dataset (US EPA Toxicity Forecaster program).
• load_admet(): Loads the ADMET (absorption-distribution-metabolism-excretion-toxicity) dataset.
• load_gasperm(): Loads the Gas Permeability dataset for polymers.

Note: If loading an unlabeled dataset, the target may be None.

from torch_molecule.datasets import load_qm9, load_chembl2k, load_broad6k, load_toxcast, load_admet, load_gasperm

# Example: Loading QM9 with specific targets
dataset = load_qm9(target_cols=['some_property'])

# Example: Loading ChEMBL2k
dataset = load_chembl2k()

All molecular encoder models in torch_molecule inherit from BaseMolecularEncoder. They follow a consistent interface for training on molecular data and generating representations.

Core Methods

• fit(X, **kwargs): Pretrain the model using a dataset X containing SMILES strings.
• encode(X, **kwargs): Encode new SMILES strings and return a dictionary containing the resulting molecular representations.

Model Persistence

Encoders also inherit model persistence methods from BaseModel:

• save_to_local(path): Save the trained model to a local file.
• load_from_local(path): Load a trained model from a local file.
• save_to_hf(repo_id): Push the model to the Hugging Face Hub.
• load_from_hf(repo_id, local_cache): Load a model from the Hugging Face Hub and save it to a local file.
• save(path, repo_id): Save the model to either local storage or Hugging Face.
• load(path, repo_id): Load a model from either local storage or Hugging Face.

To predict molecular properties, use the GREAMolecularPredictor class. You can perform automatic hyperparameter tuning using the .autofit() method by providing a search space defined with ParameterSpec and ParameterType.

Key steps:

1. Define search parameters using ParameterSpec(ParameterType.TYPE, range_or_list).
2. Initialize GREAMolecularPredictor with the task type (e.g., 'regression'), model name, and evaluation criteria.
3. Call .autofit() with training and validation data (SMILES lists and property arrays) and your search configuration.

from torch_molecule import GREAMolecularPredictor
from torch_molecule.utils.search import ParameterType, ParameterSpec

# Define search parameters
search_GREA = {
    "gamma": ParameterSpec(ParameterType.FLOAT, (0.25, 0.75)),
    "gnn_type": ParameterSpec(ParameterType.CATEGORICAL, ["gin-virtual", "gcn-virtual", "gin", "gcn"]),
    "num_layer": ParameterSpec(ParameterType.INTEGER, (2, 5)),
    "learning_rate": ParameterSpec(ParameterType.LOG_FLOAT, (1e-5, 1e-2)),
}

# Initialize model
grea_model = GREAMolecularPredictor(
    num_task=1,
    task_type="regression",
    model_name="GREA_multitask",
    batch_size=32,
    epochs=10,
    evaluate_criterion='r2',
    evaluate_higher_better=True,
    verbose='progress_bar'
)

# Fit with automatic hyperparameter tuning
grea_model.autofit(
    X_train=['C1=CC=CC=C1'],
    y_train=[[0.5]],
    X_val=['C1=CC=CC=C1'],
    y_val=[[0.5]],
    n_trials=100,
    search_parameters=search_GREA
)

The GraphDITMolecularGenerator can be used to generate new molecules based on target property values.

Workflow:

1. Initialize GraphDITMolecularGenerator specifying the task_type (e.g., a list of 'regression' for each property).
2. Train the model using .fit(train_smiles_list, train_property_array).
3. Generate molecules by passing target property arrays to .generate(test_property_array).

Note: Since generated SMILES may be chemically invalid, it is recommended to implement a retry loop using rdkit.Chem to validate the output and regenerate only the invalid entries.

from torch_molecule import GraphDITMolecularGenerator
from rdkit import Chem

# Initialize
model_cond = GraphDITMolecularGenerator(
    task_type=['regression'],
    batch_size=1024,
    drop_condition=0.1,
    verbose='progress_bar',
    epochs=10000,
)

# Fit
model_cond.fit(['C1=CC=CC=C1'], [[1]])

# Generate
test_properties = [[1.5], [2.5]]
generated_smiles_list = model_cond.generate(test_properties)

# Validation helper
def is_valid_smiles(smiles):
    if smiles is None: return False
    return Chem.MolFromSmiles(smiles) is not None

You can manage model checkpoints locally or via Hugging Face.

Hugging Face Integration

Use save_to_hf to upload a trained model and load_from_hf to retrieve it.

Local Storage

Use save_to_local and load_from_local for standard file-based persistence.

Inference

After loading a model, use .set_params() to adjust settings (like verbose) and .predict() to run inference on SMILES strings.

from torch_molecule import GREAMolecularPredictor

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

# Load from Hugging Face
model = GREAMolecularPredictor()
model.load_from_hf(repo_id=repo_id, local_cache="./cache/model.pt")

# Predict
model.set_params(verbose='none')
predictions = model.predict(smiles_list)

# --- OR Local Storage ---

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