dataclass to class and global pbar

Author: liugangcode

.gitignore

@@ -167,4 +167,6 @@ scripts/*
 examples/temp*
 examples/prediction/predict_GREA_for_gas.py
 examples/prediction/train_GNN_for_any.py
-examples/prediction/train_GREA_for_gas.py
+examples/prediction/train_GREA_for_gas.py
+tools/
+**/*.bak

tests/generator/digress.py

@@ -6,7 +6,7 @@
 from torch_molecule import DigressMolecularGenerator
 from torch_molecule.utils.search import ParameterType, ParameterSpec
 
-EPOCHS = 5
+EPOCHS = 100
 BATCH_SIZE = 32
 
 def test_digress_generator():

tests/predictor/bfgnn.py

@@ -21,7 +21,8 @@ def test_bfgnn_predictor():
         num_layer=3,
         hidden_size=128,
         batch_size=4,
-        epochs=5,  # Small number for testing
+        epochs=100,  # Small number for testing
+        patience=100,
         verbose=True,
         l1_penalty=1e-3
     )

tests/predictor/lstm.py

@@ -27,7 +27,8 @@ def test_lstm_predictor():
         output_dim=5,  # Output dimension matches number of tasks
         LSTMunits=60,
         batch_size=2,
-        epochs=2,
+        epochs=200,
+        patience=200,
         device="cpu",
         verbose=True
     )

tests/predictor/smilestransformer.py

@@ -26,7 +26,8 @@ def test_transformer_predictor():
         n_heads=4,
         num_layers=2,
         batch_size=2,
-        epochs=2,
+        epochs=200,
+        patience=200,
         device="cpu",
         verbose=True,
         use_lr_scheduler=True,

torch_molecule/base/base.py

@@ -1,37 +1,48 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Optional, Dict, List, Type, Any, ClassVar, Union, Tuple, Callable, Literal
+from typing import Optional, Dict, List, Type, Any, Union, Tuple
 import torch
 import os
 import numpy as np
 from ..utils.checkpoint import LocalCheckpointManager, HuggingFaceCheckpointManager
 from ..utils.checker import MolecularInputChecker
 
-@dataclass
 class BaseModel(ABC):
     """Base class for molecular models with shared functionality.
     
     This abstract class provides common methods and utilities for molecular models,
     including model initialization, saving/loading, and parameter management.
-    """
     
-    device: Optional[torch.device] = field(default=None)
-    model_name: str = field(default="BaseModel")
-    model_class: Optional[Type[torch.nn.Module]] = field(default=None, init=False) # used for model initialization
-    model: Optional[torch.nn.Module] = field(default=None, init=False) # initialized model
-    is_fitted_: bool = field(default=False, init=False)
-
-    def __post_init__(self):
-        """Initialize common device settings after instance creation.
+    Parameters
+    ----------
+    device : torch.device, optional
+        Device to run the model on. If None, automatically selects CUDA if available,
+        otherwise CPU.
+    model_name : str, default="BaseModel"
+        String identifier for the model name which can be specified by the user.
+        
+    Attributes
+    ----------
+    model_class : type or None
+        The class of the model used to initialize the model instance.
+    model : object or None
+        The fitted model instance if the model has been trained, None otherwise.
+    is_fitted_ : bool
+        Whether the model has been fitted/trained. False by default.
+    """
+    def __init__(self, device: Optional[torch.device] = None, model_name: str = "BaseModel"):
+        self.device = device
+        self.model_name = model_name # string of the model name which could be specified by the user
 
-        Sets the device to CUDA if available, otherwise CPU, when no device is specified.
-        Converts string device specifications to torch.device objects.
-        """
         if self.device is None:
             self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         elif isinstance(self.device, str):
             self.device = torch.device(self.device)
 
+        self.is_fitted_ = False # whether the model is fitted
+        self.model = None # the fitted model if not None
+        self.model_class = None # the class of the model used to initialize the model
+
     @abstractmethod
     def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
         """Set up optimizers for model training.
@@ -78,7 +89,7 @@ def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]
         pass
 
     @staticmethod
-    def _get_param_names(self) -> List[str]:
+    def _get_param_names() -> List[str]:
         """Get parameter names in the modeling class.
         
         Returns
@@ -104,7 +115,7 @@ def get_params(self, deep: bool = True) -> Dict[str, Any]:
             Dictionary of parameter names mapped to their values
         """
         out = {}
-        for key in self._get_param_names():
+        for key in self.__class__._get_param_names():
             value = getattr(self, key)
             if deep and hasattr(value, "get_params"):
                 deep_items = value.get_params().items()
@@ -392,5 +403,4 @@ def format_value(v):
         if len(repr_str) > N_CHAR_MAX:
             repr_str = "\n".join([repr_str[:N_CHAR_MAX//2], "...", repr_str[-N_CHAR_MAX//2:]])
 
-        return repr_str
-    
+        return repr_str

torch_molecule/base/encoder.py

@@ -1,17 +1,20 @@
-from dataclasses import dataclass, field
 from abc import ABC, abstractmethod
-from typing import Optional, ClassVar, Union, List, Dict, Any, Tuple, Callable, Type, Literal
+from typing import Optional, Union, List, Literal
 
 import torch
 import numpy as np
 from .base import BaseModel
 
-@dataclass
 class BaseMolecularEncoder(BaseModel, ABC):
     """Base class for molecular representation learning."""
-    
-    model_name: str = field(default="BaseMolecularEncoder")
-
+    def __init__(
+        self,
+        *,
+        device: Optional[Union[torch.device, str]] = None,
+        model_name: str = "BaseMolecularEncoder",
+    ):
+        super().__init__(device=device, model_name=model_name)
+        
     @abstractmethod
     def encode(self, X: List[str], return_type: Literal["np", "pt"] = "pt") -> Union[np.ndarray, torch.Tensor]:
         pass

torch_molecule/base/generator.py

@@ -1,16 +1,18 @@
-from dataclasses import dataclass, field
 from abc import ABC, abstractmethod
-from typing import Optional, ClassVar, Union, List, Dict, Any, Tuple, Callable, Type, Literal
-
+from typing import Optional, List, Union
 import torch
 import numpy as np
 from .base import BaseModel
 
-@dataclass
 class BaseMolecularGenerator(BaseModel, ABC):
     """Base class for molecular generation."""
-    
-    model_name: str = field(default="BaseMolecularGenerator")
+    def __init__(
+        self,
+        *,
+        device: Optional[Union[torch.device, str]] = None,
+        model_name: str = "BaseMolecularGenerator",
+    ):
+        super().__init__(device=device, model_name=model_name)
 
     @abstractmethod
     def fit(self, X: List[str], y: Optional[np.ndarray] = None) -> "BaseMolecularGenerator":
@@ -20,5 +22,4 @@ def fit(self, X: List[str], y: Optional[np.ndarray] = None) -> "BaseMolecularGen
     def generate(self, n_samples: int, **kwargs) -> List[str]:
         """Generate molecular structures.
         """
-        pass
-    
+        pass

torch_molecule/base/predictor.py

@@ -9,33 +9,32 @@
     root_mean_squared_error,
     r2_score,
 )
-from dataclasses import dataclass, field
 from abc import ABC, abstractmethod
-from typing import Optional, ClassVar, Union, List, Dict, Any, Tuple, Callable, Type
+from typing import Optional, Union, List, Tuple, Callable
 from ..base.base import BaseModel
 
-@dataclass
 class BaseMolecularPredictor(BaseModel, ABC):
     """Base class for molecular discovery estimators."""
-    
-    model_name: str = field(default="BaseMolecularPredictor")
-    num_task: int = field(default=0)
-    task_type: str = field(default=None)
-    DEFAULT_METRICS: ClassVar[Dict] = {
-        "classification": {"default": ("roc_auc", roc_auc_score, True)},
-        "regression": {"default": ("mae", mean_absolute_error, False)},
-    }
+    def __init__(
+        self,
+        *,
+        device: Optional[Union[torch.device, str]] = None,
+        model_name: str = "BaseMolecularPredictor",
+        num_task: int = 0,
+        task_type: Optional[str] = None,
+    ):
+        super().__init__(device=device, model_name=model_name)
+        self.num_task = num_task
+        self.task_type = task_type
 
-    def __post_init__(self):
-        super().__post_init__()
         if self.task_type not in ["classification", "regression"]:
             raise ValueError(f"Invalid task_type: {self.task_type}")
         if self.num_task <= 0:
             raise ValueError(f"num_task must be positive, got {self.num_task}")
 
     @staticmethod
-    def _get_param_names(self) -> List[str]:
-        return super()._get_param_names() + ["num_task", "task_type"]
+    def _get_param_names() -> List[str]:
+        return BaseModel._get_param_names() + ["num_task", "task_type"]
 
     @abstractmethod
     def autofit(self, X_train, y_train, X_val=None, y_val=None, search_parameters: Optional[dict] = None, n_trials: int = 10) -> "BaseMolecularPredictor": 
@@ -59,10 +58,16 @@ def _setup_evaluation(
         evaluate_higher_better: Optional[bool],
     ) -> None:
         if evaluate_criterion is None:
-            default_metric = self.DEFAULT_METRICS[self.task_type]["default"]
-            self.evaluate_name = default_metric[0]
-            self.evaluate_criterion = default_metric[1]
-            self.evaluate_higher_better = default_metric[2]
+            if self.task_type == 'classification':
+                self.evaluate_name = 'roc_auc'
+                self.evaluate_criterion = roc_auc_score
+                self.evaluate_higher_better = True
+            elif self.task_type == 'regression':
+                self.evaluate_name = 'mae'
+                self.evaluate_criterion = mean_absolute_error
+                self.evaluate_higher_better = False
+            else:
+                raise ValueError(f"The task type {self.task_type} does not have a default metric.")
         else:
             if isinstance(evaluate_criterion, str):
                 metric_map = {