torch-molecule 0.1.0.dev2__py3-none-any.whl

torch_molecule/utils/hf.py

@@ -0,0 +1,232 @@
+import os
+import json
+import datetime
+from typing import Dict, Tuple, Optional, List, Union
+from huggingface_hub import hf_hub_download
+
+def merge_task_configs(
+    task_id: str,
+    existing_config: Dict,
+    new_task_config: Dict,
+    num_params: int
+) -> Dict:
+    """Merge task-specific configuration and maintain version history.
+    
+    Parameters
+    ----------
+    task_id : str
+        Task identifier (e.g., 'O2', 'N2')
+    existing_config : Dict
+        Existing configuration dictionary
+    new_task_config : Dict
+        New task configuration to merge
+    num_params : int
+        Number of model parameters
+        
+    Returns
+    -------
+    Dict
+        Updated configuration with task history
+    """
+    merged = existing_config.copy()
+    tasks_config = merged.get("tasks", {})
+    task_config = tasks_config.get(task_id, {
+        "versions": [],
+        "current_version": "0.0.0",
+        "date_added": datetime.datetime.now().isoformat(),
+    })
+    
+    # Update task version history
+    current_version = task_config["current_version"]
+    version_parts = current_version.split('.')
+    version_parts[-1] = str(int(version_parts[-1]) + 1)
+    new_version = '.'.join(version_parts)
+    
+    # Archive current version if it exists
+    if task_config["versions"]:
+        task_config["versions"].append({
+            "version": current_version,
+            "date": datetime.datetime.now().isoformat(),
+            "config": task_config.copy()
+        })
+    
+    # Update task config
+    task_config.update({
+        "current_version": new_version,
+        "last_updated": datetime.datetime.now().isoformat(),
+        "num_parameters": num_params,
+        **new_task_config
+    })
+    
+    # Update tasks configuration
+    tasks_config[task_id] = task_config
+    merged["tasks"] = tasks_config
+    merged["last_updated"] = datetime.datetime.now().isoformat()
+    
+    return merged
+
+def get_existing_repo_data(repo_id: str, token: Optional[str] = None) -> Tuple[bool, Dict, str]:
+    """Get existing repository data from HuggingFace Hub.
+    
+    Parameters
+    ----------
+    repo_id : str
+        Repository ID
+    token : Optional[str]
+        HuggingFace token
+        
+    Returns
+    -------
+    Tuple[bool, Dict, str]
+        Tuple containing (repo_exists, existing_config, existing_readme)
+    """
+    repo_exists = False
+    existing_config = {}
+    existing_readme = ""
+    
+    try:
+        config_path = hf_hub_download(
+            repo_id=repo_id,
+            filename="config.json",
+            token=token
+        )
+        with open(config_path, 'r') as f:
+            existing_config = json.load(f)
+        repo_exists = True
+        
+        try:
+            readme_path = hf_hub_download(
+                repo_id=repo_id,
+                filename="README.md",
+                token=token
+            )
+            with open(readme_path, 'r') as f:
+                existing_readme = f.read()
+        except Exception:
+            pass
+            
+    except Exception:
+        pass
+        
+    return repo_exists, existing_config, existing_readme
+
+def create_model_card(
+    model_class: str,
+    model_name: str,
+    tasks_config: Dict,
+    model_config: Dict,
+    repo_id: str,
+    existing_readme: str = ""
+) -> str:
+    """Create a model card for multiple tasks.
+    
+    Parameters
+    ----------
+    model_class : str
+        Class name of the model
+    model_name : str
+        Name of the model
+    tasks_config : Dict
+        Configuration for all tasks
+    model_config : Dict
+        General model configuration
+    repo_id : str
+        Repository ID
+    existing_readme : str
+        Existing README content
+        
+    Returns
+    -------
+    str
+        Generated model card content
+    """
+    # Create tasks summary table
+    tasks_table = "| Task | Version | Last Updated | Parameters | Metrics |\n"
+    tasks_table += "|------|---------|--------------|------------|----------|\n"
+    
+    for task_id, task_info in sorted(tasks_config.items()):
+        metrics = task_info.get("metrics", {})
+        metrics_str = ", ".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
+        tasks_table += (
+            f"| {task_id} | {task_info['current_version']} | "
+            f"{task_info['last_updated'][:10]} | "
+            f"{task_info['num_parameters']:,} | {metrics_str} |\n"
+        )
+
+# - molecular-property-prediction
+    readme_content = f"""---
+tags:
+- torch_molecule
+library_name: torch_molecule
+---
+
+# {model_class} Model
+
+## Model Description
+- **Model Type**: {model_class}
+- **Framework**: torch_molecule
+- **Last Updated**: {model_config.get('last_updated', datetime.datetime.now().isoformat())[:10]}
+
+## Task Summary
+{tasks_table}
+
+## Usage
+
+```python
+from torch_molecule import {model_class}
+
+# Load model for specific task
+model = {model_class}()
+model.load(
+    "local_model_dir/{model_name}.pt",
+    repo="{repo_id}"
+)
+
+# For predictor: Make predictions
+# predictions = model.predict(smiles_list)
+# For generator: Make generations
+# generations = model.generate(n_samples)
+# For encoder: Make encodings
+# encodings = model.encode(smiles_list)
+```
+
+## Tasks Details
+
+"""
+    # Add detailed task sections (latest version only, excluding fitting_loss)
+    for task_id, task_info in sorted(tasks_config.items()):
+        # Filter out fitting_loss from config
+        task_config = task_info.get('config', {}).copy()
+        if 'fitting_loss' in task_config:
+            del task_config['fitting_loss']
+        
+        readme_content += f"""
+### {task_id} Task
+- **Current Version**: {task_info['current_version']}
+- **Last Updated**: {task_info['last_updated'][:10]}
+- **Parameters**: {task_info['num_parameters']:,}
+- **Configuration**:
+```python
+{json.dumps(task_config, indent=2)}
+```
+"""
+
+    # Preserve custom sections from existing README
+    # if existing_readme:
+    #     try:
+    #         custom_sections = ""
+    #         split_points = ["## Tasks Details", "## Usage"]
+    #         for split_point in split_points:
+    #             if split_point in existing_readme:
+    #                 parts = existing_readme.split(split_point)
+    #                 if len(parts) > 1 and "##" in parts[1]:
+    #                     custom_sections = parts[1].split("##", 1)[1]
+    #                     custom_sections = "##" + custom_sections
+    #                     break
+            
+    #         if custom_sections:
+    #             readme_content += f"\n{custom_sections}"
+    #     except Exception as e:
+    #         print(f"Warning: Error preserving custom README sections: {e}")
+
+    return readme_content

torch_molecule/utils/search.py

@@ -0,0 +1,64 @@
+from enum import Enum
+from typing import Dict, Any, Union, List, Tuple, NamedTuple
+
+class ParameterType(Enum):
+    """Enum defining types of hyperparameters for optimization.
+
+    Each type corresponds to a specific Optuna suggest method and parameter behavior.
+    """
+
+    CATEGORICAL = "categorical"  # Uses suggest_categorical for discrete choices
+    INTEGER = "integer"  # Uses suggest_int for whole numbers
+    FLOAT = "float"  # Uses suggest_float for continuous values
+    LOG_FLOAT = "log_float"  # Uses suggest_float with log=True for exponential scale
+
+class ParameterSpec(NamedTuple):
+    """Specification for a hyperparameter including its type and valid range/options."""
+
+    param_type: ParameterType
+    value_range: Union[Tuple[Any, Any], List[Any]]
+
+def suggest_parameter(trial: Any, param_name: str, param_spec: ParameterSpec) -> Any:
+    """Suggest a parameter value using the appropriate Optuna suggest method.
+
+    Parameters
+    ----------
+    trial : optuna.Trial
+        The Optuna trial object
+    param_name : str
+        Name of the parameter
+    param_spec : ParameterSpec
+        Specification of the parameter type and range
+
+    Returns
+    -------
+    Any
+        The suggested parameter value
+
+    Raises
+    ------
+    ValueError
+        If the parameter type is not recognized
+    """
+    if param_spec.param_type == ParameterType.CATEGORICAL:
+        return trial.suggest_categorical(param_name, param_spec.value_range)
+
+    elif param_spec.param_type == ParameterType.INTEGER:
+        min_val, max_val = param_spec.value_range
+        return trial.suggest_int(param_name, min_val, max_val)
+
+    elif param_spec.param_type == ParameterType.FLOAT:
+        min_val, max_val = param_spec.value_range
+        return trial.suggest_float(param_name, min_val, max_val)
+
+    elif param_spec.param_type == ParameterType.LOG_FLOAT:
+        min_val, max_val = param_spec.value_range
+        return trial.suggest_float(param_name, min_val, max_val, log=True)
+
+    else:
+        raise ValueError(f"Unknown parameter type: {param_spec.param_type}")
+    
+def parse_list_params(params_str):
+    if params_str is None:
+        return None
+    return params_str.split(',')

torch_molecule-0.1.0.dev2.dist-info/METADATA

@@ -0,0 +1,243 @@
+Metadata-Version: 2.4
+Name: torch_molecule
+Version: 0.1.0.dev2
+Summary: Deep learning packages for molecular discovery with a simple sklearn-style interface
+Author-email: Gang Liu <gliu7@nd.edu>
+License: MIT
+Project-URL: Homepage, https://github.com/liugangcode/torch-molecule
+Project-URL: Bug Tracker, https://github.com/liugangcode/torch-molecule
+Project-URL: Documentation, https://liugangcode.github.io/torch-molecule/
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.2.0
+Requires-Dist: torch-geometric>=2.6.1
+Requires-Dist: numpy
+Requires-Dist: pandas>=2.2.3
+Requires-Dist: click
+Requires-Dist: huggingface_hub>=0.22.2
+Requires-Dist: joblib>=1.3.2
+Requires-Dist: networkx>=3.2.1
+Requires-Dist: PyYAML>=6.0.2
+Requires-Dist: rdkit>=2023.9.5
+Requires-Dist: scikit_learn>=1.4.1.post1
+Requires-Dist: scipy>=1.14.1
+Requires-Dist: tqdm>=4.66.2
+Requires-Dist: optuna>=4.0.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: isort; extra == "dev"
+Requires-Dist: flake8; extra == "dev"
+Dynamic: license-file
+
+<p align="center">
+  <img src="assets/logo.png" alt="torch-molecule logo" width="600"/>
+</p>
+
+<p align="center">
+  <a href="https://github.com/liugangcode/torch-molecule">
+    <img src="https://img.shields.io/badge/GitHub-Repository-blue?logo=github" alt="GitHub Repository">
+  </a>
+  <a href="https://liugangcode.github.io/torch-molecule/">
+    <img src="https://img.shields.io/badge/Documentation-Online-brightgreen?logo=readthedocs" alt="Documentation">
+  </a>
+</p>
+
+<p align="center">
+  <b>Deep learning for molecular discovery with a simple sklearn-style interface</b>
+</p>
+
+---
+
+`torch-molecule` is a package under active development 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. Currently, the package focuses on three main components: **Predictive Models**, **Generative Models**, and **Representation Models**. See the [List of Supported Models](#list-of-supported-models) section for all available models.
+
+> **Note**: This project is in active development, and features may change.
+
+## Installation
+
+1. **Create a Conda environment**:
+   ```bash
+   conda create --name torch_molecule python=3.11.7
+   conda activate torch_molecule
+   ```
+
+2. **Install `torch_molecule` from GitHub**:
+
+   Clone the repository:
+   ```bash
+   git clone https://github.com/liugangcode/torch-molecule
+   ```
+
+   Install the requirements:
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+   Editable install:
+   ```bash
+   pip install -e .
+   ```
+
+3. **Install `torch_molecule` from PyPI** (Legacy):
+   ```bash
+   pip install -i https://test.pypi.org/simple/ torch-molecule
+   ```
+
+### Additional Packages
+
+| Model | Required Packages |
+|-------|-------------------|
+| HFPretrainedMolecularEncoder | transformers |
+| BFGNNMolecularPredictor | torch-scatter |
+| GRINMolecularPredictor | torch-scatter |
+
+## Usage
+
+Refer to the `tests` folder for more use cases.
+
+### Python API Example
+
+The following example demonstrates how to use the `GREAMolecularPredictor` class from `torch_molecule`:
+
+More examples could be found in the folders `examples` and `tests`.
+
+```python
+from torch_molecule import GREAMolecularPredictor
+
+# Train GREA model
+grea_model = GREAMolecularPredictor(
+    num_task=num_task,
+    task_type="regression",
+    model_name="GREA_multitask",
+    evaluate_criterion='r2',
+    evaluate_higher_better=True,
+    verbose=True
+)
+
+# Fit the model
+X_train = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
+y_train = [[0.5], [1.5]]
+X_val = ['C1=CC=CC=C1', 'C1=CC=CC=C1']
+y_val = [[0.5], [1.5]]
+N_trial = 10
+
+grea_model.autofit(
+    X_train=X_train.tolist(),
+    y_train=y_train,
+    X_val=X_val.tolist(),
+    y_val=y_val,
+    n_trials=N_trial,
+)
+```
+
+### Checkpoints
+
+`torch-molecule` provides checkpoint functions that can be interacted with on Hugging Face.
+
+```python
+from torch_molecule import GREAMolecularPredictor
+from sklearn.metrics import mean_absolute_error
+
+# Define the repository ID for Hugging Face
+repo_id = "user/repo_id"
+
+# Initialize the GREAMolecularPredictor model
+model = GREAMolecularPredictor()
+
+# Train the model using autofit
+model.autofit(
+    X_train=X.tolist(),  # List of SMILES strings for training
+    y_train=y_train,     # numpy array [n_samples, n_tasks] for training labels
+    X_val=X_val.tolist(),# List of SMILES strings for validation
+    y_val=y_val,         # numpy array [n_samples, n_tasks] for validation labels
+)
+
+# Make predictions on the test set
+output = model.predict(X_test.tolist()) # (n_sample, n_task)
+
+# Calculate the mean absolute error
+mae = mean_absolute_error(y_test, output['prediction'])
+metrics = {'MAE': mae}
+
+# Save the trained model to Hugging Face
+model.save_to_hf(
+    repo_id=repo_id,
+    task_id=f"{task_name}",
+    metrics=metrics,
+    commit_message=f"Upload GREA_{task_name} model with metrics: {metrics}",
+    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")
+
+# Set model parameters
+model.set_params(verbose=True)
+
+# Make predictions using the loaded model
+predictions = model.predict(smiles_list)
+```
+
+## List of Supported Models
+
+### Predictive Models
+
+| Model                | Reference           |
+|----------------------|---------------------|
+| GRIN                 | [Learning Repetition-Invariant Representations for Polymer Informatics. May 2025](https://arxiv.org/abs/2505.10726) |
+| BFGNN                 | [Graph neural networks extrapolate out-of-distribution for shortest paths. March 2025](https://arxiv.org/abs/2503.19173) |
+| SGIR                 | [Semi-Supervised Graph Imbalanced Regression. KDD 2023](https://dl.acm.org/doi/10.1145/3580305.3599497) |
+| GREA                | [Graph Rationalization with Environment-based Augmentations. KDD 2022](https://dl.acm.org/doi/abs/10.1145/3534678.3539347) |
+| DIR                  | [Discovering Invariant Rationales for Graph Neural Networks. ICLR 2022](https://arxiv.org/abs/2201.12872) |
+| SSR                  | [SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks. NeurIPS 2022](https://arxiv.org/abs/2206.07096) |
+| IRM                  | [Invariant Risk Minimization (2019)](https://arxiv.org/abs/1907.02893) |
+| RPGNN                | [Relational Pooling for Graph Representations. ICML 2019](https://arxiv.org/abs/1903.02541) |
+| GNNs                 | [Graph Convolutional Networks. ICLR 2017](https://arxiv.org/abs/1609.02907) and [Graph Isomorphism Network. ICLR 2019](https://arxiv.org/abs/1810.00826) |
+| Transformer (SMILES) | [Transformer (Attention is All You Need. NeurIPS 2017)](https://arxiv.org/abs/1706.03762) based on SMILES strings |
+| LSTM (SMILES)        | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
+
+### Generative Models
+
+| Model      | Reference           |
+|------------|---------------------|
+| Graph DiT  | [Graph Diffusion Transformers for Multi-Conditional Molecular Generation. NeurIPS 2024](https://openreview.net/forum?id=cfrDLD1wfO) |
+| DiGress    | [DiGress: Discrete Denoising Diffusion for Graph Generation. ICLR 2023](https://openreview.net/forum?id=UaAD-Nu86WX) |
+| GDSS       | [Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations. ICML 2022](https://proceedings.mlr.press/v162/jo22a/jo22a.pdf) |
+| MolGPT     | [MolGPT: Molecular Generation Using a Transformer-Decoder Model. Journal of Chemical Information and Modeling 2021](https://pubs.acs.org/doi/10.1021/acs.jcim.1c00600) |
+| JTVAE     | [Junction Tree Variational Autoencoder for Molecular Graph Generation. ICML 2018.](https://proceedings.mlr.press/v80/jin18a) |
+| GraphGA    | [A Graph-Based Genetic Algorithm and Its Application to the Multiobjective Evolution of Median Molecules. Journal of Chemical Information and Computer Sciences 2004](https://pubs.acs.org/doi/10.1021/ci034290p) |
+| LSTM (SMILES)        | [Long short-term memory (Neural Computation 1997)](https://ieeexplore.ieee.org/abstract/document/6795963) based on SMILES strings |
+
+### Representation Models
+
+| Model        | Reference           |
+|--------------|---------------------|
+| MoAMa        | [Motif-aware Attribute Masking for Molecular Graph Pre-training. LoG 2024](https://arxiv.org/abs/2309.04589) |
+| GraphMAE        | [GraphMAE: Self-Supervised Masked Graph Autoencoders. KDD 2022](https://arxiv.org/abs/2205.10803) |
+| AttrMasking  | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
+| ContextPred  | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
+| EdgePred     | [Strategies for Pre-training Graph Neural Networks. ICLR 2020](https://arxiv.org/abs/1905.12265) |
+| InfoGraph    | [InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization. ICLR 2020](https://arxiv.org/abs/1908.01000) |
+| Supervised   | Supervised pretraining |
+| Pretrained   | More than ten pretrained models from [Hugging Face](https://huggingface.co) |
+
+## Project Structure
+
+See the structure of `torch_molecule` with the command `tree -L 2 torch_molecule -I '__pycache__|*.pyc|*.pyo|.git|old*'`
+
+## Plan
+
+1. **Predictive Models**: Done: GREA, SGIR, IRM, GIN/GCN w/ virtual, DIR. SMILES-based LSTM/Transformers. TODO more
+2. **Generative Models**: Done: Graph DiT, GraphGA, DiGress, GDS, MolGPT TODO: more
+3. **Representation Models**: Done: MoAMa, AttrMasking, ContextPred, EdgePred. Many pretrained models from HF. TODO: checkpoints, more 
+
+## Acknowledgements
+
+This project is under active development, and some features may change over time.
+
+The project template was adapted from [https://github.com/lwaekfjlk/python-project-template](https://github.com/lwaekfjlk/python-project-template). We thank the authors for their contribution to the open-source community.