refine save and load funcs

Author: liugangcode

docs/source/api/encoder.rst

@@ -18,10 +18,10 @@ inherited from :class:`torch_molecule.base.base.BaseModel`
 
 - ``save_to_local(path)``: Save the trained model to a local file
 - ``load_from_local(path)``: Load a trained model from a local file
-- ``push_to_huggingface(repo_id)``: Push the model to Hugging Face Hub
-- ``load_from_huggingface(repo_id)``: Load a model from Hugging Face Hub
-- ``save``: Save the model to either local storage or Hugging Face
-- ``load``: Load a model from either local storage or Hugging Face
+- ``save_to_hf(repo_id)``: Push the model to Hugging Face Hub
+- ``load_from_hf(repo_id, local_cache)``: Load a model from Hugging Face Hub and save it to a local file
+- ``save(path, repo_id)``: Save the model to either local storage or Hugging Face
+- ``load(path, repo_id)``: Load a model from either local storage or Hugging Face
 
 
 Self-supervised Molecular Representation Learning

docs/source/api/generator.rst

@@ -14,18 +14,18 @@ inherited from :class:`torch_molecule.base.base.BaseModel`
 
 - ``save_to_local(path)``: Save the trained model to a local file
 - ``load_from_local(path)``: Load a trained model from a local file
-- ``push_to_huggingface(repo_id)``: Push the model to Hugging Face Hub
+- ``save_to_hf(repo_id)``: Push the model to Hugging Face Hub
 
   Not implemented for:
   - :class:`torch_molecule.generator.graph_ga.modeling_graph_ga.GraphGAMolecularGenerator`
 
-- ``load_from_huggingface(repo_id)``: Load a model from Hugging Face Hub 
+- ``load_from_hf(repo_id, local_cache)``: Load a model from Hugging Face Hub and save it to a local file
 
   Not implemented for:
   - :class:`torch_molecule.generator.graph_ga.modeling_graph_ga.GraphGAMolecularGenerator`
 
-- ``save``: Save the model to either local storage or Hugging Face
-- ``load``: Load a model from either local storage or Hugging Face
+- ``save(path, repo_id)``: Save the model to either local storage or Hugging Face
+- ``load(path, repo_id)``: Load a model from either local storage or Hugging Face
 
 Modeling Molecules as Graphs with GNN / Transformer-based Generators
 ---------------------------------------------------------------------

docs/source/api/predictor.rst

@@ -15,10 +15,10 @@ inherited from :class:`torch_molecule.base.base.BaseModel`
 
 - ``save_to_local(path)``: Save the trained model to a local file
 - ``load_from_local(path)``: Load a trained model from a local file
-- ``push_to_huggingface(repo_id)``: Push the model to Hugging Face Hub
-- ``load_from_huggingface(repo_id)``: Load a model from Hugging Face Hub
-- ``save``: Save the model to either local storage or Hugging Face
-- ``load``: Load a model from either local storage or Hugging Face
+- ``save_to_hf(repo_id)``: Push the model to Hugging Face Hub
+- ``load_from_hf(repo_id, local_cache)``: Load a model from Hugging Face Hub and save it to a local file
+- ``save(path, repo_id)``: Save the model to either local storage or Hugging Face
+- ``load(path, repo_id)``: Load a model from either local storage or Hugging Face
 
 Modeling Molecules as Graphs with Graph Neural Networks 
 -------------------------------------------------------

torch_molecule/base/base.py

@@ -225,6 +225,7 @@ def save_to_hf(
         commit_message: str = "Update model",
         hf_token: Optional[str] = None,
         private: bool = False,
+        config_filename: Optional[str] = 'config.json',
     ) -> None:
         """Save model to Hugging Face Hub.
         
@@ -244,7 +245,8 @@ def save_to_hf(
             Hugging Face authentication token
         private : bool, default=False
             Whether the repository should be private
-            
+        config_filename : Optional[str], default='config.json'
+            Name of the configuration file to save to the repository
         Raises
         ------
         ValueError
@@ -261,19 +263,24 @@ def save_to_hf(
             commit_message=commit_message,
             token=hf_token,
             private=private,
+            config_filename=config_filename,
         )
 
-    def load_from_hf(self, repo_id: str, path: str) -> None:
+    def load_from_hf(self, repo_id: str, local_cache: Optional[str] = None, config_filename: Optional[str] = 'config.json') -> None:
         """Load model from Hugging Face Hub.
         
         Parameters
         ----------
         repo_id : str
             Hugging Face repository ID
-        path : str
-            Path within the repository to load the model from
+        local_cache : str, default=None
+            Local path to save the model
+        config_filename : str, default='config.json'
+            Name of the configuration file to load from the repository
         """
-        HuggingFaceCheckpointManager.load_model_from_hf(self, repo_id, path)
+        if local_cache is None:
+            local_cache = 'model.pt'
+        HuggingFaceCheckpointManager.load_model_from_hf(self, repo_id, local_cache, config_filename)
 
     def save(self, path: Optional[str] = None, repo_id: Optional[str] = None, **kwargs) -> None:
         """Automatic save to either local disk or Hugging Face Hub.
@@ -292,34 +299,42 @@ def save(self, path: Optional[str] = None, repo_id: Optional[str] = None, **kwar
         ValueError
             If path is None when repo_id is None
         """
+        # if both path and repo_id are None, raise an error
+        if path is None and repo_id is None:
+            raise ValueError("path must be provided if repo_id is not given.")
+        
         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.")
+
+        if path is not None:
             self.save_to_local(path)
 
-    def load(self, path: str, repo_id: Optional[str] = None) -> None:
+    def load(self, path: Optional[str] = None, repo_id: Optional[str] = None, **kwargs) -> None:
         """Automatic load from either local disk or Hugging Face Hub.
         
         Parameters
         ----------
-        path : str
-            File path for local loading or path within the repository
+        path : Optional[str], default=None
+            File path for local loading.
         repo_id : Optional[str], default=None
-            Hugging Face repository ID for remote loading
+            Hugging Face repository ID for remote loading. If path is provided, repo_id is ignored.
+        **kwargs
+            Additional arguments passed to load_from_hf
             
         Raises
         ------
         FileNotFoundError
             If no local file is found and no repo_id is provided
         """
-        if os.path.exists(path):
-            self.load_from_local(path)
+        if path is not None:
+            if os.path.exists(path):
+                self.load_from_local(path)
+            else:
+                raise FileNotFoundError(f"No local file found at '{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)
+                raise ValueError("repo_id must be provided if path is not given.")
+            self.load_from_hf(repo_id, **kwargs)
 
     def _check_is_fitted(self) -> None:
         """Check if the model is fitted.

torch_molecule/utils/checkpoint.py

@@ -107,7 +107,7 @@ class HuggingFaceCheckpointManager:
     """Handles saving and loading of models to and from the Hugging Face Hub."""
 
     @staticmethod
-    def load_model_from_hf(model_instance, repo_id: str, path: str) -> None:
+    def load_model_from_hf(model_instance, repo_id: str, path: str, config_filename: str = "config.json") -> None:
         """Load model from Hugging Face Hub, saving locally to `path` first."""
         try:
             from huggingface_hub import hf_hub_download
@@ -129,7 +129,7 @@ def load_model_from_hf(model_instance, repo_id: str, path: str) -> None:
 
             hf_hub_download(
                 repo_id=repo_id,
-                filename="config.json",
+                filename=config_filename,
                 local_dir=os.path.dirname(path),
             )
 
@@ -196,6 +196,7 @@ def push_to_huggingface(
         commit_message: str = "Update model",
         token: Optional[str] = None,
         private: bool = False,
+        config_filename: str = "config.json",
     ) -> None:
         """Push a task-specific model checkpoint to Hugging Face Hub."""
         try:
@@ -252,8 +253,8 @@ def push_to_huggingface(
                     num_params=num_params,
                 )
 
-                # Save config.json
-                config_path = os.path.join(tmp_dir, "config.json")
+                # Save config file
+                config_path = os.path.join(tmp_dir, config_filename)
                 with open(config_path, "w") as f:
                     json.dump(final_config, f, indent=2)