Merge pull request #24 from grok-ai/develop

Version 0.1.1

Author: lucmos

src/nn_core/callbacks.py

@@ -53,4 +53,6 @@ def on_save_checkpoint(
     ) -> None:
         if self._is_nnlogger(trainer):
             trainer.logger.on_save_checkpoint(trainer=trainer, pl_module=pl_module, checkpoint=checkpoint)
-            checkpoint[METADATA_KEY] = trainer.datamodule.metadata
+            metadata = getattr(trainer.datamodule, "metadata", None)
+            if metadata is not None:
+                checkpoint[METADATA_KEY] = metadata

src/nn_core/common/utils.py

@@ -51,6 +51,8 @@ def load_envs(env_file: Optional[str] = None) -> None:
     :param env_file: the file that defines the environment variables to use. If None
                      it searches for a `.env` file in the project.
     """
+    if env_file is None:
+        env_file = dotenv.find_dotenv(usecwd=True)
     dotenv.load_dotenv(dotenv_path=env_file, override=True)
 
 

src/nn_core/serialization.py

@@ -1,3 +1,4 @@
+import collections
 import importlib
 import inspect
 import logging
@@ -17,6 +18,7 @@
 
 pylogger = logging.getLogger(__name__)
 
+from typing import Mapping
 
 _METADATA_MODULE_KEY = f"{METADATA_KEY}_module"
 _METADATA_CLASS_KEY = f"{METADATA_KEY}_class"
@@ -124,14 +126,41 @@ def extract_checkpoint(ckpt_file: Path) -> Path:
         yield Path(tmp_dir)
 
 
+def _substistute(dictionary, substitute_values: Dict[str, str], substitute_keys: Dict[str, str] = {}):
+    if not isinstance(dictionary, Mapping):
+        if isinstance(dictionary, collections.Hashable):
+            if substitute_values is not None and dictionary in substitute_values:
+                return substitute_values[dictionary]
+            elif substitute_keys is not None and dictionary in substitute_keys:
+                return substitute_keys[dictionary]
+            else:
+                return dictionary
+        return dictionary
+
+    return {
+        _substistute(key, substitute_values=substitute_values, substitute_keys=substitute_keys,): _substistute(
+            value,
+            substitute_values=substitute_values,
+            substitute_keys=substitute_keys,
+        )
+        for key, value in dictionary.items()
+    }
+
+
 def load_model(
     module_class: Type[pl.LightningModule],
     checkpoint_path: Path,
     map_location: Optional[Union[Dict[str, str], str, torch.device, int, Callable]] = None,
+    substitute_keys: Optional[Dict[str, str]] = None,
+    substitute_values: Optional[Dict[str, str]] = None,
 ):
     # Lightning checkpoints end with .ckpt, ours with .ckpt.zip
     if checkpoint_path.name.endswith(".ckpt.zip"):
         checkpoint = NNCheckpointIO.load(path=checkpoint_path, map_location=map_location)
+
+        if substitute_values is not None:
+            checkpoint = _substistute(checkpoint, substitute_values=substitute_values, substitute_keys=substitute_keys)
+
         return module_class._load_model_state(checkpoint=checkpoint, metadata=checkpoint.get("metadata", None))
     else:
         pylogger.warning(f"Loading a legacy checkpoint (from vanilla PyTorch Lightning): '{checkpoint_path}'")