update doc

Author: liugangcode

docs/source/api/generator.rst

@@ -3,6 +3,25 @@ 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.
 
+The following models support conditional generation (click model name to jump to details):
+
+.. list-table:: Conditional Generation Models
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Model
+     - Description
+   * - :class:`GraphDITMolecularGenerator <torch_molecule.generator.graph_dit.modeling_graph_dit.GraphDITMolecularGenerator>`
+     - `Graph Diffusion Transformers for Multi-Conditional Molecular Generation <https://arxiv.org/abs/2401.13858>`_
+   * - :class:`DeFoGMolecularGenerator <torch_molecule.generator.defog.modeling_defog.DeFoGMolecularGenerator>`
+     - `Discrete Flow Matching for Graph Generation <https://openreview.net/forum?id=KPRIwWhqAZ>`_
+   * - :class:`GraphGAMolecularGenerator <torch_molecule.generator.graph_ga.modeling_graph_ga.GraphGAMolecularGenerator>`
+     - Graph Genetic Algorithm with Random Forests
+   * - :class:`MolGPTMolecularGenerator <torch_molecule.generator.molgpt.modeling_molgpt.MolGPTMolecularGenerator>`
+     - `MolGPT: Molecular Generation Using a Transformer-Decoder Model <https://pubs.acs.org/doi/10.1021/acs.jcim.1c00600>`_
+   * - :class:`LSTMMolecularGenerator <torch_molecule.generator.lstm.modeling_lstm.LSTMMolecularGenerator>`
+     - LSTM
+
 .. rubric:: Training and Generation
 
 - ``fit(X, **kwargs)``: Train the model on given data, where X contains SMILES strings (y should be provided for conditional generation)
@@ -30,21 +49,40 @@ inherited from :class:`torch_molecule.base.base.BaseModel`
 Modeling Molecules as Graphs
 ---------------------------------------------------------------------
 
-.. rubric:: GraphDiT for Un/Multi-conditional Molecular Generation
+.. rubric:: GraphDiT for Unconditional/Multi-Conditional Molecular Generation
 .. autoclass:: torch_molecule.generator.graph_dit.modeling_graph_dit.GraphDITMolecularGenerator
    :exclude-members: fitting_epoch, fitting_loss, model_name, model_class
    :members: fit, generate
    :undoc-members:
    :show-inheritance:
 
+.. rubric:: Discrete Flow Matching for Graph Generation for Unconditional/Multi-Conditional Molecular Generation
+.. autoclass:: torch_molecule.generator.defog.modeling_defog.DeFoGMolecularGenerator
+   :exclude-members: fitting_epoch, fitting_loss, model_name, model_class
+   :members: fit, generate
+   :undoc-members:
+   :show-inheritance:
+
+.. rubric:: Graph Genetic Algorithm for Unconditional/Multi-Conditional Molecular Generation
+.. autoclass:: torch_molecule.generator.graph_ga.modeling_graph_ga.GraphGAMolecularGenerator
+   :exclude-members: fitting_epoch, fitting_loss, save_to_hf, load_from_hf
+   :members: fit, generate
+   :undoc-members:
+   :show-inheritance:
+   
+.. automodule:: torch_molecule.generator.graph_ga.oracle
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 .. rubric:: DiGress for Unconditional Molecular Generation
 .. autoclass:: torch_molecule.generator.digress.modeling_digress.DigressMolecularGenerator
    :exclude-members: fitting_epoch, fitting_loss, model_class, dataset_info, model_name
    :members: fit, generate
    :undoc-members:
    :show-inheritance:
 
-.. rubric:: GDSS for score-based molecular generation
+.. rubric:: GDSS for Unconditional Molecular Generation
 .. autoclass:: torch_molecule.generator.gdss.modeling_gdss.GDSSMolecularGenerator
    :exclude-members: fitting_epoch, fitting_loss, model_name, model_class
    :members: fit, generate
@@ -58,29 +96,17 @@ Modeling Molecules as Graphs
    :undoc-members:
    :show-inheritance:
 
-.. rubric:: Graph Genetic Algorithm for Un/Multi-conditional Molecular Generation
-.. autoclass:: torch_molecule.generator.graph_ga.modeling_graph_ga.GraphGAMolecularGenerator
-   :exclude-members: fitting_epoch, fitting_loss, save_to_hf, load_from_hf
-   :members: fit, generate
-   :undoc-members:
-   :show-inheritance:
-   
-.. automodule:: torch_molecule.generator.graph_ga.oracle
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
 Modeling Molecules as Sequences
 --------------------------------
 
-.. rubric:: MolGPT for Unconditional Molecular Generation
+.. rubric:: MolGPT for Unconditional/Multi-Conditional Molecular Generation
 .. autoclass:: torch_molecule.generator.molgpt.modeling_molgpt.MolGPTMolecularGenerator
    :exclude-members: fitting_epoch, fitting_loss, model_name, model_class
    :members: fit, generate
    :undoc-members:
    :show-inheritance:
 
-.. rubric:: LSTM for Unconditional/Conditional Molecular Generation
+.. rubric:: LSTM for Unconditional/Multi-Conditional Molecular Generation
 .. autoclass:: torch_molecule.generator.lstm.modeling_lstm.LSTMMolecularGenerator
    :exclude-members: fitting_epoch, fitting_loss, model_name, model_class
    :members: fit, generate

torch_molecule/generator/defog/modeling_defog.py

@@ -32,9 +32,9 @@ class DeFoGMolecularGenerator(BaseMolecularGenerator):
     ----------
     num_layer : int, default=6
         Number of transformer layers
-    hidden_mlp_dims : Dict[str, int], default={'X': 256, 'E': 128, 'y': 128}
+    hidden_mlp_dims : Dict[str, int], default={'X': 256, 'E': 128, 'y': 128} if None
         Hidden dimensions for MLP layers in X (node dim), E (edge dim), and y (property dim) components
-    hidden_dims : Dict[str, Any], default={'dx': 256, 'de': 64, 'dy': 64, 'n_head': 8, 'dim_ffX': 256, 'dim_ffE': 128, 'dim_ffy': 128}
+    hidden_dims : Dict[str, Any], default={'dx': 256, 'de': 64, 'dy': 64, 'n_head': 8, 'dim_ffX': 256, 'dim_ffE': 128, 'dim_ffy': 128} if None
         Hidden dimensions for transformer components including attention heads and feed-forward layers
         Keys: 'dx' (node dim), 'de' (edge dim), 'dy' (property dim), 'n_head' (number of attention heads), 'dim_ffX' (feed-forward dim for node features), 'dim_ffE' (feed-forward dim for edge features), 'dim_ffy' (feed-forward dim for property features)
     transition : str, default='marginal'
@@ -43,7 +43,7 @@ class DeFoGMolecularGenerator(BaseMolecularGenerator):
     time_distortion : str, default="polydec"
         Time distortion schedule used during training/sampling.
         Options: 'identity', 'cos', 'revcos', 'polyinc', 'polydec'
-    lambda_train : List[float], default=[5.0, 1.0]
+    lambda_train : List[float], default=[5.0, 1.0] if None
         Loss weights: [edge_loss_weight, property_loss_weight]
     extra_features_type : str, default='rrwp'
         Extra feature type.

torch_molecule/generator/graph_ga/oracle.py

@@ -1,10 +1,8 @@
 import numpy as np
 from sklearn.ensemble import RandomForestRegressor
 from rdkit import Chem
-from typing import List, Any
 from ...utils.graph.features import getmorganfingerprint
 
-
 class Oracle:
     """The default Oracle class for scoring molecules in GraphGA.
     

torch_molecule/generator/molgpt/modeling_molgpt.py

@@ -38,7 +38,7 @@ class MolGPTMolecularGenerator(BaseMolecularGenerator):
     max_len : int, default=128
         Maximum length of SMILES strings.
     num_task : int, default=0
-        Number of property prediction tasks for conditional generation. O for unconditional generation.
+        Number of property prediction tasks for conditional generation. 0 for unconditional generation.
     use_scaffold : bool, default=False
         Whether to use scaffold conditioning.
     use_lstm : bool, default=False