update docs for reference

liugangcode · liugangcode · commit 5c0b7dbc6d1b · 2025-03-31T22:54:55.000-04:00
diff --git a/docs/source/additional.rst b/docs/source/additional.rst
@@ -6,11 +6,12 @@ Project Overview
 Main components:
 
 1. **Predictive Models**  
-   - ✔ GREA, SGIR, IRM, GIN/GCN w/ virtual, DIR  
-   - ⏳ SMILES-based LSTM/Transformers and more
+   - ✔ GREA, SGIR, IRM, GIN/GCN w/ virtual, DIR
+   - ✔ SMILES-based LSTM/Transformers
+   - ⏳ More models
 
 2. **Generative Models**  
-   - ✔ Graph DiT, GraphGA, DiGress  
+   - ✔ Graph DiT, GraphGA, DiGress, MolGPT
    - ⏳ GDSS and more
 
 3. **Representation Models**  
diff --git a/docs/source/api/encoder.rst b/docs/source/api/encoder.rst
@@ -3,20 +3,19 @@ Molecular Encoder Models
 
 The encoder models inherit from the :class:`torch_molecule.base.encoder.BaseMolecularEncoder` class and share common methods for model pretraining and encoding, as well as model persistence.
 
-Common Methods
---------------
+.. rubric:: Training and Encoding
 
-Training and Encoding
-^^^^^^^^^^^^^^^^^^^^^^^
 - ``fit(X, **kwargs)``: Pretrain the model on given data, where X contains SMILES strings
 
   Not implemented for:
   - N/A
 
 - ``encode(X, **kwargs)``: Encode new SMILES strings and return a dictionary containing encoded representations
 
-Model Persistence
-^^^^^^^^^^^^^^^^^
+.. rubric:: Model Persistence
+
+inherited from :class:`torch_molecule.base.base.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
 - ``push_to_huggingface(repo_id)``: Push the model to Hugging Face Hub
diff --git a/docs/source/api/generator.rst b/docs/source/api/generator.rst
@@ -3,27 +3,26 @@ Molecular Generation Models
 
 The generator models inherit from the :class:`torch_molecule.base.generator.BaseMolecularGenerator` class and share common methods for model training, generation and persistence.
 
-Common Methods
--------------
+.. rubric:: Training and Generation
 
-Training and Generation
-^^^^^^^^^^^^^^^^^^^^^^^
 - ``fit(X, **kwargs)``: Train the model on given data, where X contains SMILES strings (y should be provided for conditional generation)
 - ``generate(n_samples, **kwargs)``: Generate new molecules and return a list of SMILES strings (y should be provided for conditional generation)
 
-Model Persistence
-^^^^^^^^^^^^^^^^^
+.. rubric:: Model Persistence
+
+inherited from :class:`torch_molecule.base.base.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
 - ``push_to_huggingface(repo_id)``: Push the model to Hugging Face Hub
 
   Not implemented for:
-  - :class:`torch_molecule.generator.graph_ga.GraphGAMolecularGenerator`
+  - :class:`torch_molecule.generator.graph_ga.modeling_graph_ga.GraphGAMolecularGenerator`
 
 - ``load_from_huggingface(repo_id)``: Load a model from Hugging Face Hub 
 
   Not implemented for:
-  - :class:`torch_molecule.generator.graph_ga.GraphGAMolecularGenerator`
+  - :class:`torch_molecule.generator.graph_ga.modeling_graph_ga.GraphGAMolecularGenerator`
 
 - ``save``: Save the model to either local storage or Hugging Face
 - ``load``: Load a model from either local storage or Hugging Face
diff --git a/docs/source/api/predictor.rst b/docs/source/api/predictor.rst
@@ -3,17 +3,16 @@ Molecular Property Prediction Models
 
 The predictor models inherit from the :class:`torch_molecule.base.predictor.BaseMolecularPredictor` class and share common methods for model training, evaluation, prediction and persistence.
 
-Common Methods
---------------
+.. rubric:: Training and Prediction
 
-Training and Prediction
-^^^^^^^^^^^^^^^^^^^^^^^
 - ``fit(X, y, **kwargs)``: Train the model on given data, where X contains SMILES strings and y contains target values
 - ``autofit(X, y, search_parameters, n_trials=50, **kwargs)``: Automatically search for optimal hyperparameters using Optuna and train the model
 - ``predict(X, **kwargs)``: Make predictions on new SMILES strings and return a dictionary containing predictions and optional uncertainty estimates
 
-Model Persistence
-^^^^^^^^^^^^^^^^^
+.. rubric:: Model Persistence
+
+inherited from :class:`torch_molecule.base.base.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
 - ``push_to_huggingface(repo_id)``: Push the model to Hugging Face Hub
diff --git a/torch_molecule/encoder/moama/modeling_moama.py b/torch_molecule/encoder/moama/modeling_moama.py
@@ -21,8 +21,10 @@
 class MoamaMolecularEncoder(BaseMolecularEncoder):
     """This encoder implements Motif-aware Attribute Masking for Molecular Graph Pre-training.
 
-    Paper: https://openreview.net/forum?id=uqPnesiGGi
-    Code: https://github.com/einae-nd/MoAMa-dev
+    References
+    ----------
+    - Paper: https://openreview.net/forum?id=uqPnesiGGi
+    - Code: https://github.com/einae-nd/MoAMa-dev
 
     Parameters
     ----------
diff --git a/torch_molecule/generator/digress/modeling_digress.py b/torch_molecule/generator/digress/modeling_digress.py
@@ -18,13 +18,13 @@
 @dataclass
 class DigressMolecularGenerator(BaseMolecularGenerator):
     """
-    This generator implements the DiGress model for unconditional molecular generation.
+    This generator implements the DiGress for unconditional molecular generation.
 
     References
     ----------
-    .. [1] DiGress: Discrete Denoising diffusion for graph generation. International Conference on 
-           Learning Representations (ICLR) 2023. https://openreview.net/forum?id=UaAD-Nu86WX
-    .. [2] Adapted from: https://github.com/cvignac/DiGress
+    - DiGress: Discrete Denoising diffusion for graph generation. International Conference on 
+      Learning Representations (ICLR) 2023. https://openreview.net/forum?id=UaAD-Nu86WX
+    - Code: https://github.com/cvignac/DiGress
 
     :param hidden_size_X: Hidden dimension size for node features, defaults to 256
     :type hidden_size_X: int, optional
diff --git a/torch_molecule/generator/graph_dit/modeling_graph_dit.py b/torch_molecule/generator/graph_dit/modeling_graph_dit.py
@@ -18,15 +18,14 @@
 @dataclass
 class GraphDITMolecularGenerator(BaseMolecularGenerator):
     """
-    This generator implements the graph diffusion transformer for 
-    (multi-conditional and unconditional) molecular generation.
+    This generator implements the graph diffusion transformer for (multi-conditional and unconditional) molecular generation.
 
     References
     ----------
-    .. [1] Graph Diffusion Transformers for Multi-Conditional Molecular Generation.
-           International Conference on Learning Representations (ICLR) 2024.
-           https://openreview.net/forum?id=cfrDLD1wfO
-    .. [2] Implementation: https://github.com/liugangcode/Graph-DiT
+    - Graph Diffusion Transformers for Multi-Conditional Molecular Generation.
+      International Conference on Learning Representations (ICLR) 2024.
+      https://openreview.net/forum?id=cfrDLD1wfO
+    - Implementation: https://github.com/liugangcode/Graph-DiT
 
     Parameters
     ----------
diff --git a/torch_molecule/generator/graph_ga/modeling_graph_ga.py b/torch_molecule/generator/graph_ga/modeling_graph_ga.py
@@ -24,25 +24,27 @@ class GraphGAMolecularGenerator(BaseMolecularGenerator):
     
     References
     ----------
-    .. [1] A Graph-BasedGenetic Algorithm and Its Application to the Multiobjective Evolution of 
+    - A Graph-Based Genetic Algorithm and Its Application to the Multiobjective Evolution of 
            Median Molecules. Journal of Chemical Information and Computer Sciences.
            https://pubs.acs.org/doi/10.1021/ci034290p
-    .. [2] Implementation adapted from: https://github.com/wenhao-gao/mol_opt
+    - Implementation: https://github.com/wenhao-gao/mol_opt
 
-    :param num_task: Number of properties to condition on. Set to 0 for unconditional generation. Default is 0.
-    :type num_task: int
-    :param population_size: Size of the population in each iteration. Default is 100.
-    :type population_size: int
-    :param offspring_size: Number of offspring molecules to generate in each iteration. Default is 50.
-    :type offspring_size: int
-    :param mutation_rate: Probability of mutation occurring during reproduction. Default is 0.0067.
-    :type mutation_rate: float
-    :param n_jobs: Number of parallel jobs to run. -1 means using all processors. Default is 1.
-    :type n_jobs: int
-    :param iteration: Number of iterations for each target label (or random sample) to run the genetic algorithm. Default is 5.
-    :type iteration: int
-    :param verbose: Whether to display progress bars and logs. Default is False.
-    :type verbose: bool
+    Parameters
+    ----------
+    num_task : int, default=0
+        Number of properties to condition on. Set to 0 for unconditional generation.
+    population_size : int, default=100
+        Size of the population in each iteration.
+    offspring_size : int, default=50
+        Number of offspring molecules to generate in each iteration.
+    mutation_rate : float, default=0.0067
+        Probability of mutation occurring during reproduction.
+    n_jobs : int, default=1
+        Number of parallel jobs to run. -1 means using all processors.
+    iteration : int, default=5
+        Number of iterations for each target label (or random sample) to run the genetic algorithm.
+    verbose : bool, default=False
+        Whether to display progress bars and logs.
     """
 
     # GA parameters
diff --git a/torch_molecule/generator/molgpt/modeling_molgpt.py b/torch_molecule/generator/molgpt/modeling_molgpt.py
@@ -26,9 +26,9 @@ class MolGPTMolecularGenerator(BaseMolecularGenerator):
     
     References
     ----------
-    .. [1] MolGPT:Molecular Generation Using a Transformer-Decoder Model. Journal of Chemical 
-           Information and Modeling. https://pubs.acs.org/doi/10.1021/acs.jcim.1c00600
-    .. [2] Adapted from: https://github.com/devalab/molgpt
+    - MolGPT: Molecular Generation Using a Transformer-Decoder Model. Journal of Chemical 
+      Information and Modeling. https://pubs.acs.org/doi/10.1021/acs.jcim.1c00600
+    - Code: https://github.com/devalab/molgpt
 
     Parameters
     ----------
diff --git a/torch_molecule/predictor/dir/modeling_dir.py b/torch_molecule/predictor/dir/modeling_dir.py
@@ -18,12 +18,16 @@
 
 @dataclass
 class DIRMolecularPredictor(GNNMolecularPredictor):
-    """This predictor implements the DIR model for discovering the invariant rationales.
+    """This predictor implements the DIR for molecular property prediction tasks.
 
-    Paper: `Discovering Invariant Rationales for Graph Neural Networks <https://openreview.net/forum?id=hGXij5rfiHw>`_
-    
-    Reference Code: https://github.com/Wuyxin/DIR-GNN
-    
+    The full name of DIR is Discovering Invariant Rationales.
+
+    References
+    ----------
+    - Discovering Invariant Rationales for Graph Neural Networks.
+      https://openreview.net/forum?id=hGXij5rfiHw
+    - Code: https://github.com/Wuyxin/DIR-GNN
+        
     Parameters
     ----------
     causal_ratio : float, default=0.8
diff --git a/torch_molecule/predictor/grea/modeling_grea.py b/torch_molecule/predictor/grea/modeling_grea.py
@@ -18,12 +18,15 @@
 
 @dataclass
 class GREAMolecularPredictor(GNNMolecularPredictor):
-    """This predictor implements a Graph Rationalization model.
+    """This predictor implements a Graph Rationalization model called GREA.
 
-    Paper: `Graph Rationalization with Environment-based Augmentations 
-    <https://dl.acm.org/doi/10.1145/3534678.3539347>`_
+    The full name of GREA is Graph Rationalization with Environment-based Augmentations. During model training, it learns the rationales (explainable subgraphs) and use them for molecular property prediction tasks.
 
-    Reference Code: https://github.com/liugangcode/GREA
+    References
+    ----------
+    - Graph Rationalization with Environment-based Augmentations.
+      https://dl.acm.org/doi/10.1145/3534678.3539347
+    - Code: https://github.com/liugangcode/GREA
 
     :param gamma: GREA-specific parameter. Default is 0.4.
     :type gamma: float
diff --git a/torch_molecule/predictor/irm/modeling_irm.py b/torch_molecule/predictor/irm/modeling_irm.py
@@ -21,9 +21,14 @@
 class IRMMolecularPredictor(GNNMolecularPredictor):
     """This predictor implements a Invariant Risk Minimization model with the GNN.
     
-    Paper: `Invariant Risk Minimization <https://arxiv.org/abs/1907.02893>`_
+    The full name of IRM is Invariant Risk Minimization.
 
-    Reference Code: https://github.com/facebookresearch/InvariantRiskMinimization
+    References
+    ----------
+    - Invariant Risk Minimization.
+      https://arxiv.org/abs/1907.02893
+
+    - Reference Code: https://github.com/facebookresearch/InvariantRiskMinimization
     
     Parameters
     ----------
diff --git a/torch_molecule/predictor/lstm/modeling_lstm.py b/torch_molecule/predictor/lstm/modeling_lstm.py
@@ -34,6 +34,11 @@
 class LSTMMolecularPredictor(BaseMolecularPredictor):
     """This predictor implements a LSTM model for molecular property prediction tasks.
 
+    References
+    ----------
+    - Predicting Polymers' Glass Transition Temperature by a Chemical Language Processing Model.
+      https://www.semanticscholar.org/reader/f43ed533b2520567be2d8c24f6396f4e63e96430
+
     Parameters
     ----------
     num_task : int, default=1
@@ -73,14 +78,6 @@ class LSTMMolecularPredictor(BaseMolecularPredictor):
     verbose : bool, default=False
         Whether to print progress information during training.
 
-    Notes
-    -----
-    This implementation is based on the paper "Predicting Polymers' Glass Transition 
-    Temperature by a Chemical Language Processing Model".
-
-    References
-    ----------
-    .. [1] https://www.semanticscholar.org/reader/f43ed533b2520567be2d8c24f6396f4e63e96430
     """
     # Model parameters
     num_task: int = 1
diff --git a/torch_molecule/predictor/rpgnn/modeling_rpgnn.py b/torch_molecule/predictor/rpgnn/modeling_rpgnn.py
@@ -19,8 +19,23 @@
 @dataclass
 class RPGNNMolecularPredictor(GNNMolecularPredictor):
     """This predictor implements a GNN model based on Relational pooling.
-    Paper: Relational Pooling for Graph Representations (https://arxiv.org/abs/1903.02541)
-    Reference Code: https://github.com/PurdueMINDS/RelationalPooling/tree/master?tab=readme-ov-file
+
+    The full name of RPGNN is Relational Pooling for Graph Representations.
+
+    References
+    ----------
+    - Relational Pooling for Graph Representations.
+      https://arxiv.org/abs/1903.02541
+    - Reference Code: https://github.com/PurdueMINDS/RelationalPooling/tree/master?tab=readme-ov-file
+
+    Parameters
+    ----------
+    num_perm : int, default=3
+        TODO: Add description
+    fixed_size : int, default=10 
+        TODO: Add description
+    num_node_feature : int, default=9
+        TODO: Add description
     """
     
     # RPGNN-specific parameter
diff --git a/torch_molecule/predictor/sgir/modeling_sgir.py b/torch_molecule/predictor/sgir/modeling_sgir.py
@@ -21,10 +21,15 @@
 @dataclass
 class SGIRMolecularPredictor(GREAMolecularPredictor):
     """
-    This predictor trains the GREA model based on pseudo-labeling and data augmentation.
+    This predictor implements SGIR for semi-supervised graph imbalanced regression.
 
-    Paper: `Semi-Supervised Graph Imbalanced Regression <https://dl.acm.org/doi/10.1145/3580305.3599497>`_
-    Reference Code: `SGIR GitHub <https://github.com/liugangcode/SGIR>`_
+    It trains the GREA model based on pseudo-labeling and data augmentation.
+
+    References
+    ----------
+    - Semi-Supervised Graph Imbalanced Regression.
+      https://dl.acm.org/doi/10.1145/3580305.3599497
+    - Code: https://github.com/liugangcode/SGIR
 
     :param num_anchor: Number of anchor points used to split the label space during pseudo-labeling
     :type num_anchor: int, default=10
diff --git a/torch_molecule/predictor/smiles_transformer/modeling_transformer.py b/torch_molecule/predictor/smiles_transformer/modeling_transformer.py
@@ -36,6 +36,11 @@
 class SMILESTransformerMolecularPredictor(LSTMMolecularPredictor):
     """This predictor implements a Transformer model for SMILES-based molecular property predictions.
     
+    Notes
+    -----
+    This implementation uses a transformer encoder architecture to learn
+    representations of molecular structures from SMILES strings.
+
     Parameters
     ----------
     num_task : int, default=1
@@ -80,11 +85,6 @@ class SMILESTransformerMolecularPredictor(LSTMMolecularPredictor):
         Number of epochs with no improvement after which learning rate will be reduced.
     verbose : bool, default=False
         Whether to print progress information during training.
-
-    Notes
-    -----
-    This implementation uses a transformer encoder architecture to learn
-    representations of molecular structures from SMILES strings.
     """
     # Model parameters
     num_task: int = 1
diff --git a/torch_molecule/predictor/ssr/modeling_ssr.py b/torch_molecule/predictor/ssr/modeling_ssr.py
@@ -21,8 +21,12 @@
 @dataclass
 class SSRMolecularPredictor(GNNMolecularPredictor):
     """This predictor implements a SizeShiftReg model with the GNN.
-    Paper: SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks (https://arxiv.org/abs/2207.07888)
-    Reference Code: https://github.com/DavideBuffelli/SizeShiftReg/tree/main
+    
+    References
+    ----------
+    - Paper: SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks.
+      https://arxiv.org/abs/2207.07888
+    - Reference Code: https://github.com/DavideBuffelli/SizeShiftReg/tree/main
     
     Parameters
     ----------
