Export datamodule to huggingface

Author: dfulu

pvnet_summation/models/base_model.py

@@ -23,6 +23,11 @@
 )
 
 
+def santize_datamodule(config: dict) -> dict:
+    """Create new datamodule config which only keeps the details required for inference"""
+    return {"pvnet_model": config["pvnet_model"]}
+
+
 def download_from_hf(
     repo_id: str,
     filename: str | list[str],
@@ -114,6 +119,32 @@ def from_pretrained(
         model.eval()  # type: ignore
 
         return model
+    
+    @classmethod
+    def get_datamodule_config(
+        cls,
+        model_id: str,
+        revision: str,
+        cache_dir: str | None = None,
+        force_download: bool = False,
+    ) -> str:
+        """Load data config file."""
+        if os.path.isdir(model_id):
+            print("Loading datamodule config from local directory")
+            datamodule_config_file = os.path.join(model_id, DATAMODULE_CONFIG_NAME)
+        else:
+            print("Loading datamodule config from huggingface repo")
+            datamodule_config_file = download_from_hf(
+                repo_id=model_id,
+                filename=DATAMODULE_CONFIG_NAME,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                max_retries=5,
+                wait_time=10,
+            )
+
+        return datamodule_config_file
 
     def _save_model_weights(self, save_directory: str) -> None:
         """Save weights from a Pytorch model to a local directory."""
@@ -126,7 +157,7 @@ def save_pretrained(
         wandb_repo: str,
         wandb_id: str,
         card_template_path: str,
-        datamodule_config_path: str | None = None,
+        datamodule_config_path,
         experiment_config_path: str | None = None,
         hf_repo_id: str | None = None,
         push_to_hub: bool = False,
@@ -142,16 +173,15 @@ def save_pretrained(
             wandb_id: Identifier of the model on wandb.
             datamodule_config_path:
                 The path to the datamodule config.
+            card_template_path: Path to the HuggingFace model card template. Defaults to card in
+                PVNet library if set to None.
             experiment_config_path:
                 The path to the full experimental config.
             hf_repo_id:
                 ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to
                 the folder name if not provided.
             push_to_hub (`bool`, *optional*, defaults to `False`):
                 Whether or not to push your model to the HuggingFace Hub after saving it.
-
-            card_template_path: Path to the HuggingFace model card template. Defaults to card in
-                PVNet library if set to None.
         """
 
         save_directory = Path(save_directory)
@@ -165,9 +195,14 @@ def save_pretrained(
             with open(save_directory / MODEL_CONFIG_NAME, "w") as outfile:
                 yaml.dump(model_config, outfile, sort_keys=False, default_flow_style=False)
 
-        # Save the datamodule config
-        if datamodule_config_path is not None:
-            shutil.copyfile(datamodule_config_path, save_directory / DATAMODULE_CONFIG_NAME)
+        # Sanitize and save the datamodule config
+        with open(datamodule_config_path) as cfg:
+            datamodule_config = yaml.load(cfg, Loader=yaml.FullLoader)
+
+        datamodule_config = santize_datamodule(datamodule_config)
+
+        with open(save_directory / DATAMODULE_CONFIG_NAME, "w") as outfile:
+            yaml.dump(datamodule_config, outfile, sort_keys=False, default_flow_style=False)
 
         # Save the full experimental config
         if experiment_config_path is not None:

pvnet_summation/training/train.py

@@ -13,7 +13,7 @@
 from tqdm import tqdm
 
 from pvnet_summation.data.datamodule import PresavedDataModule, StreamedDataModule
-from pvnet_summation.utils import MODEL_CONFIG_NAME
+from pvnet_summation.utils import MODEL_CONFIG_NAME, FULL_CONFIG_NAME, DATAMODULE_CONFIG_NAME
 
 log = logging.getLogger(__name__)
 
@@ -163,6 +163,14 @@ def train(config: DictConfig) -> None:
                 os.makedirs(save_dir, exist_ok=True)
                 OmegaConf.save(config.model, f"{save_dir}/{MODEL_CONFIG_NAME}")
 
+                # Save the datamodule config
+                OmegaConf.save(config.datamodule, f"{save_dir}/{DATAMODULE_CONFIG_NAME}")
+
+                # Save the full hydra config to the output directory and to wandb
+                OmegaConf.save(config, f"{save_dir}/{FULL_CONFIG_NAME}")
+                wandb_logger.experiment.save(f"{save_dir}/{FULL_CONFIG_NAME}", base_path=save_dir)
+
+
     # Init lightning model
     model = hydra.utils.instantiate(config.model)