update readme

Author: liugangcode

README.md

@@ -2,6 +2,15 @@
   <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/generator/graphga/modeling_graph_ga.py

@@ -25,7 +25,7 @@ class GraphGAMolecularGenerator(BaseMolecularGenerator):
     Paper: `A Graph-Based Genetic Algorithm and Its Application to the Multiobjective Evolution of Median Molecules 
     <https://pubs.acs.org/doi/full/10.1021/ci034290p>`_
 
-    Reference code: `https://github.com/wenhao-gao/mol_opt/blob/main/main/graph_ga/run.py`_
+    Reference code: https://github.com/wenhao-gao/mol_opt/blob/main/main/graph_ga/run.py
 
     :param num_task: Number of properties to condition on. Set to 0 for unconditional generation. Default is 0.
     :type num_task: int

torch_molecule/predictor/grea/modeling_grea.py

@@ -23,7 +23,7 @@ class GREAMolecularPredictor(GNNMolecularPredictor):
     Paper: `Graph Rationalization with Environment-based Augmentations 
     <https://dl.acm.org/doi/10.1145/3534678.3539347>`_
 
-    Reference Code: `https://github.com/liugangcode/GREA`_
+    Reference Code: https://github.com/liugangcode/GREA
 
     :param gamma: GREA-specific parameter. Default is 0.4.
     :type gamma: float

torch_molecule/predictor/sgir/modeling_sgir.py

@@ -20,26 +20,26 @@
 
 @dataclass
 class SGIRMolecularPredictor(GREAMolecularPredictor):
-    """This predictor trains the GREA model based on pseudo-labeling and data augmentation.
-    Paper: Semi-Supervised Graph Imbalanced Regression (https://dl.acm.org/doi/10.1145/3580305.3599497)
-    Reference Code: https://github.com/liugangcode/SGIR
-
-    Parameters
-    ----------
-    num_anchor : int, default=10
-        Number of anchor points used to split the label space during pseudo-labeling
-    warmup_epoch : int, default=20 
-        Number of epochs to train before starting pseudo-labeling and data augmentation
-    labeling_interval : int, default=5
-        Interval (in epochs) between pseudo-labeling steps
-    augmentation_interval : int, default=5
-        Interval (in epochs) between data augmentation steps
-    top_quantile : float, default=0.1
-        Quantile threshold for selecting high confidence predictions during pseudo-labeling
-    label_logscale : bool, default=False
-        Whether to use log scale for the label space during pseudo-labeling and data augmentation
-    lw_aug : float, default=1
-        Weight for the data augmentation loss
+    """
+    This predictor trains the GREA model based on pseudo-labeling and data augmentation.
+
+    Paper: `Semi-Supervised Graph Imbalanced Regression <https://dl.acm.org/doi/10.1145/3580305.3599497>`_
+    Reference Code: `SGIR GitHub <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
+    :param warmup_epoch: Number of epochs to train before starting pseudo-labeling and data augmentation
+    :type warmup_epoch: int, default=20
+    :param labeling_interval: Interval (in epochs) between pseudo-labeling steps
+    :type labeling_interval: int, default=5
+    :param augmentation_interval: Interval (in epochs) between data augmentation steps
+    :type augmentation_interval: int, default=5
+    :param top_quantile: Quantile threshold for selecting high confidence predictions during pseudo-labeling
+    :type top_quantile: float, default=0.1
+    :param label_logscale: Whether to use log scale for the label space during pseudo-labeling and data augmentation
+    :type label_logscale: bool, default=False
+    :param lw_aug: Weight for the data augmentation loss
+    :type lw_aug: float, default=1
     """
     # SGIR-specific parameters
     num_anchor: int = 10

torch_molecule/utils/search.py

@@ -13,16 +13,7 @@ class ParameterType(Enum):
     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.
-
-    Attributes
-    ----------
-    param_type : ParameterType
-        The type of parameter (categorical, integer, float, log_float)
-    value_range : Union[Tuple[Any, Any], List[Any]]
-        For numerical parameters: (min_value, max_value)
-        For categorical parameters: list of possible values
-    """
+    """Specification for a hyperparameter including its type and valid range/options."""
 
     param_type: ParameterType
     value_range: Union[Tuple[Any, Any], List[Any]]