minor changes for gni

Author: aditya0by0

chebai_graph/models/dynamic_gni.py

@@ -33,7 +33,11 @@ def __init__(self, config: dict[str, Any], **kwargs: Any):
         )
         self.distribution = distribution
 
-        self.complete_randomness = config.get("complete_randomness", True)
+        self.complete_randomness = (
+            str(config.get("complete_randomness", "True")).lower() == "true"
+        )
+
+        print("Using complete randomness: ", self.complete_randomness)
 
         if not self.complete_randomness:
             assert (
@@ -44,11 +48,25 @@ def __init__(self, config: dict[str, Any], **kwargs: Any):
                 if config.get("random_pad_node") is not None
                 else None
             )
+            if self.random_pad_node is not None:
+                print(
+                    f"[Info] Node features will be padded with {self.random_pad_node} "
+                    f"new set of random features from distribution {self.distribution} "
+                    f"in each forward pass."
+                )
+
             self.random_pad_edge = (
                 int(config["random_pad_edge"])
                 if config.get("random_pad_edge") is not None
                 else None
             )
+            if self.random_pad_edge is not None:
+                print(
+                    f"[Info] Edge features will be padded with {self.random_pad_edge} "
+                    f"new set of random features from distribution {self.distribution} "
+                    f"in each forward pass."
+                )
+
             assert (
                 self.random_pad_node > 0 or self.random_pad_edge > 0
             ), "'random_pad_node' or 'random_pad_edge' must be positive integers"

chebai_graph/preprocessing/datasets/chebi.py

@@ -388,32 +388,6 @@ def load_processed_data_from_file(self, filename: str) -> list[dict]:
 
         return base_df[base_data[0].keys()].to_dict("records")
 
-    @property
-    def processed_file_names_dict(self) -> dict:
-        """
-        Returns a dictionary for the processed and tokenized data files.
-
-        Returns:
-            dict: A dictionary mapping dataset keys to their respective file names.
-                  For example, {"data": "data.pt"}.
-        """
-        if self.n_token_limit is not None:
-            return {"data": f"data_maxlen{self.n_token_limit}.pt"}
-
-        data_pt_filename = "data"
-        if self.zero_pad_node:
-            data_pt_filename += f"_zpn{self.zero_pad_node}"
-        if self.zero_pad_edge:
-            data_pt_filename += f"_zpe{self.zero_pad_edge}"
-        if self.random_pad_node:
-            data_pt_filename += f"_rpn{self.random_pad_node}"
-        if self.random_pad_edge:
-            data_pt_filename += f"_rpe{self.random_pad_edge}"
-        if self.random_pad_node or self.random_pad_edge:
-            data_pt_filename += f"_D{self.distribution}"
-
-        return {"data": data_pt_filename + ".pt"}
-
 
 class GraphPropAsPerNodeType(DataPropertiesSetter, ABC):
     def __init__(self, properties=None, transform=None, **kwargs):