torch-molecule 0.1.0.dev2__py3-none-any.whl

torch_molecule/base/old_predictor2.py

@@ -0,0 +1,436 @@
+import warnings
+import torch
+import numpy as np
+from ..utils import (
+    roc_auc_score,
+    accuracy_score,
+    mean_absolute_error,
+    mean_squared_error,
+    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 ..base.base import BaseModel
+
+@dataclass
+class BaseMolecularPredictor(BaseModel, ABC):
+    """Base class for molecular discovery estimators."""
+    
+    num_task: int
+    task_type: str
+    model_name: str = field(default="BaseMolecularPredictor")
+    
+    DEFAULT_METRICS: ClassVar[Dict] = {
+        "classification": {"default": ("roc_auc", roc_auc_score, True)},
+        "regression": {"default": ("mae", mean_absolute_error, False)},
+    }
+
+    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}")
+
+    def _get_param_names(self) -> List[str]:
+        return super()._get_param_names() + ["num_task", "task_type"]
+
+    @abstractmethod
+    def fit(self, X_train, y_train, X_val=None, y_val=None, X_unlbl=None) -> "BaseMolecularPredictor":
+        pass
+
+    @abstractmethod
+    def autofit(self, X_train, y_train, X_val=None, y_val=None, search_parameters: Optional[dict] = None, n_trials: int = 10) -> "BaseMolecularPredictor": 
+        pass
+
+    @abstractmethod
+    def predict(self, X: np.ndarray) -> np.ndarray:
+        pass
+
+    @abstractmethod
+    def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
+        pass
+
+    @abstractmethod
+    def _train_epoch(self, train_loader, optimizer):
+        pass
+
+    @abstractmethod
+    def _evaluation_epoch(self, evaluate_loader):
+        pass
+        
+    def _setup_evaluation(
+        self,
+        evaluate_criterion: Optional[Union[str, Callable]],
+        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]
+        else:
+            if isinstance(evaluate_criterion, str):
+                metric_map = {
+                    "accuracy": (accuracy_score, True),
+                    "roc_auc": (roc_auc_score, True),
+                    "rmse": (root_mean_squared_error, False),
+                    "mse": (mean_squared_error, False),
+                    "mae": (mean_absolute_error, False),
+                    "r2": (r2_score, True),
+                }
+                if evaluate_criterion not in metric_map:
+                    raise ValueError(
+                        f"Unknown metric: {evaluate_criterion}. "
+                        f"Available metrics: {list(metric_map.keys())}"
+                    )
+                self.evaluate_name = evaluate_criterion
+                self.evaluate_criterion = metric_map[evaluate_criterion][0]
+                self.evaluate_higher_better = (
+                    metric_map[evaluate_criterion][1]
+                    if evaluate_higher_better is None
+                    else evaluate_higher_better
+                )
+            else:
+                if evaluate_higher_better is None:
+                    raise ValueError(
+                        "evaluate_higher_better must be specified for a custom function."
+                    )
+                self.evaluate_name = "custom"
+                self.evaluate_criterion = evaluate_criterion
+                self.evaluate_higher_better = evaluate_higher_better
+    
+    def _load_default_criterion(self):
+        if self.task_type == "regression":
+            return torch.nn.L1Loss()
+        elif self.task_type == "classification":
+            return torch.nn.BCEWithLogitsLoss()
+        else:
+            warnings.warn(
+                "Unknown task type. Using L1 Loss as default. "
+                "Please specify 'regression' or 'classification' for better results."
+            )
+            return torch.nn.L1Loss()
+
+# import os
+# import json
+# import datetime
+# import tempfile
+# import warnings
+# import torch
+# import numpy as np
+# from rdkit import Chem
+# from ..utils import (
+#     roc_auc_score,
+#     accuracy_score,
+#     mean_absolute_error,
+#     mean_squared_error,
+#     root_mean_squared_error,
+#     r2_score,
+# )
+# from ..utils.checkpoint import LocalCheckpointManager, HuggingFaceCheckpointManager
+# from ..utils.checker import MolecularInputChecker
+# from dataclasses import dataclass, field
+# from abc import ABC, abstractmethod
+# from typing import Optional, ClassVar, Union, List, Dict, Any, Tuple, Callable, Type
+
+# @dataclass
+# class BaseMolecularPredictor(ABC):
+#     """Base class for molecular discovery estimators.
+    
+#     This class provides a structure for implementing custom molecular discovery
+#     models. It includes basic functionality for parameter management, fitting,
+#     prediction, and validation.
+#     """
+
+#     num_task: int
+#     task_type: str
+#     model_name: str = field(default="BaseMolecularPredictor")
+#     device: Optional[torch.device] = field(default=None)
+#     model: Optional[torch.nn.Module] = field(default=None, init=False)
+
+#     DEFAULT_METRICS: ClassVar[Dict] = {
+#         "classification": {
+#             "default": ("roc_auc", roc_auc_score, True),
+#         },
+#         "regression": {"default": ("mae", mean_absolute_error, False)},
+#     }
+
+#     is_fitted_: bool = field(default=False, init=False)
+#     model_class: Optional[Type[torch.nn.Module]] = field(default=None, init=False)
+
+#     def __post_init__(self):
+#         if self.task_type not in ["classification", "regression"]:
+#             raise ValueError(
+#                 f"task_type must be one of ['classification', 'regression'], got {self.task_type}"
+#             )
+
+#         if not isinstance(self.num_task, (int, np.integer)) or self.num_task <= 0:
+#             raise ValueError(f"num_task must be a positive integer, got {self.num_task}")
+
+#         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)
+
+#     def get_params(self, deep: bool = True) -> Dict[str, Any]:
+#         out = dict()
+#         for key in self._get_param_names():
+#             value = getattr(self, key)
+#             if deep and hasattr(value, "get_params"):
+#                 deep_items = value.get_params().items()
+#                 out.update((key + "__" + k, val) for k, val in deep_items)
+#             out[key] = value
+#         return out
+
+#     def set_params(self, **params) -> "BaseMolecularPredictor":
+#         if not params:
+#             return self
+
+#         valid_params = self.get_params(deep=True)
+#         for key, value in params.items():
+#             if key not in valid_params:
+#                 raise ValueError(f"Invalid parameter {key} for predictor {self}")
+#             setattr(self, key, value)
+#         return self
+
+#     @abstractmethod
+#     def fit(self, X_train, y_train, X_val=None, y_val=None, X_unlbl=None) -> "BaseMolecularPredictor":
+#         pass
+
+#     @abstractmethod
+#     def autofit(
+#         self,
+#         X_train,
+#         y_train,
+#         X_val=None,
+#         y_val=None,
+#         search_parameters: Optional[dict] = None,
+#         n_trials: int = 10,
+#     ) -> "BaseMolecularPredictor":
+#         pass
+
+#     @abstractmethod
+#     def predict(self, X: np.ndarray) -> np.ndarray:
+#         pass
+
+#     def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
+#         raise NotImplementedError("Child class must implement optimizer and scheduler setup.")
+
+#     def _setup_evaluation(
+#         self,
+#         evaluate_criterion: Optional[Union[str, Callable]],
+#         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]
+#         else:
+#             if isinstance(evaluate_criterion, str):
+#                 metric_map = {
+#                     "accuracy": (accuracy_score, True),
+#                     "roc_auc": (roc_auc_score, True),
+#                     "rmse": (root_mean_squared_error, False),
+#                     "mse": (mean_squared_error, False),
+#                     "mae": (mean_absolute_error, False),
+#                     "r2": (r2_score, True),
+#                 }
+#                 if evaluate_criterion not in metric_map:
+#                     raise ValueError(
+#                         f"Unknown metric: {evaluate_criterion}. "
+#                         f"Available metrics: {list(metric_map.keys())}"
+#                     )
+#                 self.evaluate_name = evaluate_criterion
+#                 self.evaluate_criterion = metric_map[evaluate_criterion][0]
+#                 self.evaluate_higher_better = (
+#                     metric_map[evaluate_criterion][1]
+#                     if evaluate_higher_better is None
+#                     else evaluate_higher_better
+#                 )
+#             else:
+#                 if evaluate_higher_better is None:
+#                     raise ValueError(
+#                         "evaluate_higher_better must be specified for a custom function."
+#                     )
+#                 self.evaluate_name = "custom"
+#                 self.evaluate_criterion = evaluate_criterion
+#                 self.evaluate_higher_better = evaluate_higher_better
+
+#     def _train_epoch(self, train_loader, optimizer):
+#         raise NotImplementedError("Child class must implement training logic.")
+
+#     def _evaluation_epoch(self, evaluate_loader):
+#         raise NotImplementedError("Child class must implement evaluation logic.")
+
+#     def _check_is_fitted(self) -> None:
+#         if not self.is_fitted_:
+#             raise AttributeError(
+#                 "This instance is not fitted yet. "
+#                 "Call 'fit' before using this estimator."
+#             )
+
+#     def __repr__(self, N_CHAR_MAX=700):
+#         attributes = {
+#             name: value
+#             for name, value in sorted(self.__dict__.items())
+#             if not name.startswith("_") and not callable(value)
+#         }
+
+#         def format_value(v):
+#             if isinstance(v, (float, np.float32, np.float64)):
+#                 return f"{v:.3g}"
+#             elif isinstance(v, (list, tuple, np.ndarray)) and len(v) > 6:
+#                 return f"{str(v[:3])[:-1]}, ..., {str(v[-3:])[1:]}"
+#             elif isinstance(v, str) and len(v) > 50:
+#                 return f"'{v[:25]}...{v[-22:]}'"
+#             elif isinstance(v, dict) and len(v) > 6:
+#                 items = list(v.items())
+#                 return f"{dict(items[:3])}...{dict(items[-3:])}"
+#             elif isinstance(v, torch.nn.Module):
+#                 return f"{v.__class__.__name__}(...)"
+#             elif v is None:
+#                 return "None"
+#             else:
+#                 return repr(v)
+
+#         class_name = self.__class__.__name__
+#         attributes_str = []
+
+#         important_attrs = ["num_task", "task_type", "model_name", "is_fitted_"]
+#         for attr in important_attrs:
+#             if attr in attributes:
+#                 value = attributes.pop(attr)
+#                 attributes_str.append(f"{attr}={format_value(value)}")
+
+#         sorted_attrs = sorted(attributes.items())
+#         attributes_str.extend(f"{name}={format_value(value)}" for name, value in sorted_attrs)
+
+#         content = ",\n    ".join(attributes_str)
+#         repr_ = f"{class_name}(\n    {content}\n)"
+
+#         if len(repr_) > N_CHAR_MAX:
+#             lines = repr_.split("\n")
+#             result = [lines[0]]
+#             current_length = len(lines[0])
+#             for line in lines[1:-1]:
+#                 if current_length + len(line) + 5 > N_CHAR_MAX:
+#                     result.append("    ...")
+#                     break
+#                 result.append(line)
+#                 current_length += len(line) + 1
+#             result.append(lines[-1])
+#             repr_ = "\n".join(result)
+
+#         return repr_
+
+#     @abstractmethod
+#     def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
+#         raise NotImplementedError("Child class must implement _get_model_params.")
+
+#     def _initialize_model(
+#         self, model_class: Type[torch.nn.Module], device, checkpoint: Optional[Dict] = None
+#     ) -> None:
+#         """Initialize the model with parameters or a checkpoint."""
+#         try:
+#             model_params = self._get_model_params(checkpoint)
+#             self.model = model_class(**model_params)
+#             try:
+#                 self.model = self.model.to(device)
+#             except Exception as e:
+#                 raise RuntimeError(f"Failed to move model to device {device}: {str(e)}")
+
+#             if checkpoint is not None:
+#                 try:
+#                     self.model.load_state_dict(checkpoint["model_state_dict"])
+#                 except Exception as e:
+#                     raise ValueError(f"Failed to load model state dictionary: {str(e)}")
+
+#         except Exception as e:
+#             raise RuntimeError(f"Model initialization failed: {str(e)}")
+
+#     def _load_default_criterion(self):
+#         if self.task_type == "regression":
+#             return torch.nn.L1Loss()
+#         elif self.task_type == "classification":
+#             return torch.nn.BCEWithLogitsLoss()
+#         else:
+#             warnings.warn(
+#                 "Unknown task type. Using L1 Loss as default. "
+#                 "Please specify 'regression' or 'classification' for better results."
+#             )
+#             return torch.nn.L1Loss()
+
+#     def _validate_inputs(
+#         self, X: List[str], y: Optional[Union[List, np.ndarray]] = None, return_rdkit_mol: bool = True
+#     ) -> Tuple[Union[List[str], List["Chem.Mol"]], Optional[np.ndarray]]:
+#         return MolecularInputChecker.validate_inputs(X, y, self.num_task, 0, return_rdkit_mol)
+
+#     # def _validate_smiles(self, smiles: str, idx: int) -> Tuple[bool, Optional[str], Optional[Chem.Mol]]:
+#     #     return MolecularInputChecker.validate_smiles(smiles, idx)
+
+#     @staticmethod
+#     def _get_param_names() -> List[str]:
+#         return ["num_task", "task_type", "is_fitted_"]
+
+#     def save_to_local(self, path: str) -> None:
+#         """Save to local disk using LocalCheckpointManager."""
+#         LocalCheckpointManager.save_model_to_local(self, path)
+
+#     def load_from_local(self, path: str) -> None:
+#         """Load from local disk using LocalCheckpointManager."""
+#         LocalCheckpointManager.load_model_from_local(self, path)
+
+#     def save_to_hf(
+#         self,
+#         repo_id: str,
+#         task_id: str = "default",
+#         metadata_dict: Optional[Dict[str, Any]] = None,
+#         metrics: Optional[Dict[str, float]] = None,
+#         commit_message: str = "Update model",
+#         token: Optional[str] = None,
+#         private: bool = False,
+#     ) -> None:
+#         HuggingFaceCheckpointManager.push_to_huggingface(
+#             model_instance=self,
+#             repo_id=repo_id,
+#             task_id=task_id,
+#             metadata_dict=metadata_dict,
+#             metrics=metrics,
+#             commit_message=commit_message,
+#             token=token,
+#             private=private,
+#         )
+
+#     def load_from_hf(self, repo_id: str, path: str) -> None:
+#         """Download and load model from Hugging Face Hub."""
+#         HuggingFaceCheckpointManager.load_model_from_hf(self, repo_id, path)
+
+#     def save(self, path: Optional[str] = None, repo_id: Optional[str] = None, **kwargs) -> None:
+#         """
+#         Automatic save. If `repo_id` is given, push to Hugging Face Hub.
+#         Otherwise, save to local path.
+#         """
+#         if repo_id is not None:
+#             self.save_to_hf(repo_id=repo_id, **kwargs)
+#         else:
+#             if path is None:
+#                 raise ValueError("path must be provided if repo_id is not given.")
+#             self.save_to_local(path)
+
+#     def load(self, path: str, repo_id: Optional[str] = None) -> None:
+#         """
+#         Automatic load. If local `path` exists, load from there. 
+#         Otherwise, try loading from Hugging Face Hub if `repo_id` is given.
+#         """
+#         if os.path.exists(path):
+#             self.load_from_local(path)
+#         else:
+#             if repo_id is None:
+#                 raise FileNotFoundError(
+#                     f"No local file found at '{path}' and no repo_id provided."
+#                 )
+#             self.load_from_hf(repo_id, path)

torch_molecule/base/predictor.py

@@ -0,0 +1,112 @@
+import warnings
+import torch
+import numpy as np
+from ..utils import (
+    roc_auc_score,
+    accuracy_score,
+    mean_absolute_error,
+    mean_squared_error,
+    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 ..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 __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"]
+
+    @abstractmethod
+    def autofit(self, X_train, y_train, X_val=None, y_val=None, search_parameters: Optional[dict] = None, n_trials: int = 10) -> "BaseMolecularPredictor": 
+        pass
+    
+    @abstractmethod
+    def fit(self, X_train, y_train, X_val=None, y_val=None, search_parameters: Optional[dict] = None, n_trials: int = 10) -> "BaseMolecularPredictor": 
+        pass
+
+    @abstractmethod
+    def predict(self, X: np.ndarray) -> np.ndarray:
+        pass
+
+    @abstractmethod
+    def _evaluation_epoch(self, evaluate_loader):
+        pass
+        
+    def _setup_evaluation(
+        self,
+        evaluate_criterion: Optional[Union[str, Callable]],
+        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]
+        else:
+            if isinstance(evaluate_criterion, str):
+                metric_map = {
+                    "accuracy": (accuracy_score, True),
+                    "roc_auc": (roc_auc_score, True),
+                    "rmse": (root_mean_squared_error, False),
+                    "mse": (mean_squared_error, False),
+                    "mae": (mean_absolute_error, False),
+                    "r2": (r2_score, True),
+                }
+                if evaluate_criterion not in metric_map:
+                    raise ValueError(
+                        f"Unknown metric: {evaluate_criterion}. "
+                        f"Available metrics: {list(metric_map.keys())}"
+                    )
+                self.evaluate_name = evaluate_criterion
+                self.evaluate_criterion = metric_map[evaluate_criterion][0]
+                self.evaluate_higher_better = (
+                    metric_map[evaluate_criterion][1]
+                    if evaluate_higher_better is None
+                    else evaluate_higher_better
+                )
+            else:
+                if evaluate_higher_better is None:
+                    raise ValueError(
+                        "evaluate_higher_better must be specified for a custom function."
+                    )
+                self.evaluate_name = "custom"
+                self.evaluate_criterion = evaluate_criterion
+                self.evaluate_higher_better = evaluate_higher_better
+    
+    def _load_default_criterion(self):
+        if self.task_type == "regression":
+            return torch.nn.L1Loss(reduction='none')
+        elif self.task_type == "classification":
+            return torch.nn.BCEWithLogitsLoss(reduction='none')
+        else:
+            warnings.warn(
+                "Unknown task type. Using L1 Loss as default. "
+                "Please specify 'regression' or 'classification' for better results."
+            )
+            return torch.nn.L1Loss(reduction='none')
+    
+    def _validate_inputs(
+        self, X: List[str], y: Optional[Union[List, np.ndarray]] = None, num_task: int = 0, num_pretask: int = 0, return_rdkit_mol: bool = True
+    ) -> Tuple[Union[List[str], List["Chem.Mol"]], Optional[np.ndarray]]:
+        return super()._validate_inputs(X, y, self.num_task, 0, return_rdkit_mol)

torch_molecule/encoder/attrmask/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_attrmask import AttrMaskMolecularEncoder
+
+__all__ = ['AttrMaskMolecularEncoder']

torch_molecule/encoder/attrmask/model.py

@@ -0,0 +1,113 @@
+import torch
+import torch.nn as nn
+from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool
+
+from ...nn import GNN_node, GNN_node_Virtualnode, MLP
+from ...utils import init_weights
+
+import random
+
+class_criterion = torch.nn.CrossEntropyLoss()
+
+class GNN(nn.Module):
+    def __init__(
+        self,
+        num_layer,
+        hidden_size,
+        drop_ratio=0.5,
+        norm_layer="batch_norm",
+        encoder_type="gin-virtual",
+        readout="max",
+        mask_num=0,
+        mask_rate=0.15
+    ):
+        super(GNN, self).__init__()
+        gnn_name = encoder_type.split("-")[0]
+        self.num_atom_type = 119
+        self.hidden_size = hidden_size
+        self.mask_num = mask_num
+        self.mask_rate = mask_rate
+
+        encoder_params = {
+            "num_layer": num_layer,
+            "hidden_size": hidden_size,
+            "JK": "last", 
+            "drop_ratio": drop_ratio,
+            "residual": True,
+            "gnn_name": gnn_name,
+            "norm_layer": norm_layer
+        }
+        
+        # Choose encoder type based on encoder_model
+        encoder_class = GNN_node_Virtualnode if "virtual" in encoder_type else GNN_node
+        self.graph_encoder = encoder_class(**encoder_params)
+        pooling_funcs = {
+            "sum": global_add_pool,
+            "mean": global_mean_pool,
+            "max": global_max_pool
+        }
+        self.pool = pooling_funcs.get(readout)
+        if self.pool is None:
+            raise ValueError(f"Invalid graph pooling type {readout}.")
+
+        self.predictor = MLP(hidden_size, hidden_features=2 * hidden_size, out_features=self.num_atom_type)
+    
+    def initialize_parameters(self, seed=None):
+        """
+        Randomly initialize all model parameters using the init_weights function.
+        
+        Args:
+            seed (int, optional): Random seed for reproducibility. Defaults to None.
+        """
+        if seed is not None:
+            torch.manual_seed(seed)
+        
+        # Initialize the main components
+        init_weights(self.graph_encoder)
+        init_weights(self.predictor)
+        
+        # Reset all parameters using PyTorch Geometric's reset function
+        def reset_parameters(module):
+            if hasattr(module, 'reset_parameters'):
+                module.reset_parameters()
+            elif hasattr(module, 'weight') and hasattr(module.weight, 'data'):
+                init_weights(module)
+        
+        self.apply(reset_parameters)
+
+    def compute_loss(self, batched_data):
+        masked_node_indices = list()
+
+        # select indices of masked nodes
+        for i in range(batched_data.batch[-1] + 1):
+            idx = torch.nonzero((batched_data.batch == i).float()).squeeze(-1)
+            num_node = idx.shape[0]
+            if self.mask_num == 0:
+                sample_size = int(num_node * self.mask_rate + 1)
+            else:
+                sample_size = min(self.mask_num, int(num_node * 0.5))
+            masked_node_idx = random.sample(idx.tolist(), sample_size)
+            masked_node_idx.sort()
+            masked_node_indices += masked_node_idx
+
+        batched_data.masked_node_indices = torch.tensor(masked_node_indices)
+
+        batched_data.y = batched_data.x[masked_node_indices][:, 0]
+
+        # mask nodes' features
+        for node_idx in masked_node_indices:
+            batched_data.x[node_idx] = torch.tensor([self.num_atom_type - 1] + [0] * (batched_data.x.shape[1] - 1))
+    
+        # generate predictions
+        h_node, _ = self.graph_encoder(batched_data)
+        prediction_class = self.predictor(h_node[masked_node_indices])
+        
+        target_class = batched_data.y.to(torch.float32)      
+        loss_class = class_criterion(prediction_class.to(torch.float32), target_class.long())
+
+        return loss_class
+
+    def forward(self, batched_data):
+        h_node, _ = self.graph_encoder(batched_data)
+        h_rep = self.pool(h_node, batched_data.batch)
+        return {"graph": h_rep, "node": h_node}