update readme

Author: liugangcode

docs/source/index.rst

@@ -21,19 +21,58 @@ Welcome to the torch-molecule documentation. torch-molecule is an actively devel
 .. toctree::
    :maxdepth: 3
    :caption: Main Components:
+
+.. toggle:: Predictor
+   :hidden:
+
+   .. toctree::
+      :maxdepth: 1
+
+      api/predictor
+
+.. toggle:: Generator
+   :hidden:
+
+   .. toctree::
+      :maxdepth: 1
+
+      api/generator
+
+.. toggle:: Encoder
    :hidden:
 
-   api/predictor
-   api/generator
-   api/encoder
+   .. toctree::
+      :maxdepth: 1
+
+      api/encoder
 
 .. toctree::
    :maxdepth: 2
    :caption: API for utilities:
 
+   .. toggle:: Utils
+      :hidden:
+
+      .. toctree::
+         :maxdepth: 1
+
    api/utils
-   api/nn
-   api/base
+
+   .. toggle:: NN
+      :hidden:
+
+      .. toctree::
+         :maxdepth: 1
+
+         api/nn
+
+   .. toggle:: Base
+      :hidden:
+
+      .. toctree::
+         :maxdepth: 1
+
+         api/base
 
 .. toctree::
    :maxdepth: 2

torch_molecule/generator/graph_dit/modeling_graph_dit.py

@@ -86,11 +86,10 @@ class GraphDITMolecularGenerator(BaseMolecularGenerator):
     scheduler_factor: float = 0.5
     scheduler_patience: int = 5
 
-    # Other parameters
     verbose: bool = False
-    model_name: str = "GraphDITMolecularGenerator"
 
-    # Non-init fields
+    # Attributes
+    model_name: str = "GraphDITMolecularGenerator"
     fitting_loss: List[float] = field(default_factory=list, init=False)
     fitting_epoch: int = field(default=0, init=False)
     dataset_info: Dict[str, Any] = field(default_factory=dict, init=False)