kostyl-toolkit 0.1.34__py3-none-any.whl → 0.1.35__py3-none-any.whl

kostyl/ml/lightning/extensions/pretrained_model.py

@@ -20,7 +20,8 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
         cls: type[TModelInstance],
         checkpoint_path: str | Path,
         config_key: str = "config",
-        weights_prefix: str = "model.",
+        weights_prefix: str | None = "model.",
+        strict_prefix: bool = False,
         **kwargs: Any,
     ) -> TModelInstance:
         """
@@ -39,8 +40,10 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             checkpoint_path (str | Path): Path to the checkpoint file. Must be a .ckpt file.
             config_key (str, optional): Key in the checkpoint dictionary where the config is stored.
                 Defaults to "config".
-            weights_prefix (str, optional): Prefix to strip from state dict keys. Defaults to "model.".
-                If not empty and doesn't end with ".", a "." is appended.
+            weights_prefix (str | None, optional): Prefix to strip from state dict keys. Defaults to "model.".
+                If not empty and doesn't end with ".", a "." is appended. If empty or None, no prefix stripping will be skipped.
+            strict_prefix (bool, optional): If True, drop tensors those keys that do not start with the
+                specified prefix. Defaults to False.
             kwargs: Additional keyword arguments to pass to the model's `from_pretrained` method.
 
         Returns:
@@ -53,6 +56,13 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
         """
         if isinstance(checkpoint_path, str):
             checkpoint_path = Path(checkpoint_path)
+        if weights_prefix is None:
+            weights_prefix = ""
+        weights_prefix = cast(str, weights_prefix)
+        if weights_prefix == "" and strict_prefix:
+            logger.warning(
+                "strict_prefix=True has no effect when weights_prefix is empty or None."
+            )
 
         if checkpoint_path.is_dir():
             raise ValueError(f"{checkpoint_path} is a directory")
@@ -85,13 +95,25 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             if not weights_prefix.endswith("."):
                 weights_prefix = weights_prefix + "."
             state_dict: dict[str, torch.Tensor] = {}
-
+            matched_keys_counter = 0
             for key, value in raw_state_dict.items():
                 if key.startswith(weights_prefix):
                     new_key = key[len(weights_prefix) :]
                     state_dict[new_key] = value
-                else:
+                    matched_keys_counter += 1
+                elif not strict_prefix:
                     state_dict[key] = value
+
+            if matched_keys_counter == 0:
+                if strict_prefix:
+                    raise ValueError(
+                        f"No keys in the checkpoint start with the specified prefix '{weights_prefix}'. "
+                        "Try to load with `strict_prefix=False` or verify the prefix."
+                    )
+                else:
+                    logger.warning(
+                        f"No keys in the checkpoint start with the specified prefix '{weights_prefix}'. "
+                    )
         else:
             state_dict = raw_state_dict
 

{kostyl_toolkit-0.1.34.dist-info → kostyl_toolkit-0.1.35.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.34
+Version: 0.1.35
 Summary: Kickass Orchestration System for Training, Yielding & Logging 
 Requires-Dist: case-converter>=1.2.0
 Requires-Dist: loguru>=0.7.3

{kostyl_toolkit-0.1.34.dist-info → kostyl_toolkit-0.1.35.dist-info}/RECORD

@@ -16,7 +16,7 @@ kostyl/ml/lightning/callbacks/checkpoint.py,sha256=sZ9OqudO-gXp7FqtWaOH46TXVpeCJ
 kostyl/ml/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
 kostyl/ml/lightning/extensions/__init__.py,sha256=OY6QGv1agYgqqKf1xJBrxgp_i8FunVfPzYezfaRrGXU,182
 kostyl/ml/lightning/extensions/custom_module.py,sha256=iQrnPz-WTmRfvLo94C5fQc2Qwa1IpHtUh1sCpVwTSFM,6602
-kostyl/ml/lightning/extensions/pretrained_model.py,sha256=QJGr2UvYJcU2Gy2w8z_cEvTodjv7hGdd2PPPfdOI-Mw,4017
+kostyl/ml/lightning/extensions/pretrained_model.py,sha256=eRfQBzAjVernHl9A4PP5uTLvjjmcNKPdTu7ABFLq7HI,5196
 kostyl/ml/lightning/loggers/__init__.py,sha256=e51dszaoJbuzwBkbdugmuDsPldoSO4yaRgmZUg1Bdy0,71
 kostyl/ml/lightning/loggers/tb_logger.py,sha256=j02HK5ue8yzXXV8FWKmmXyHkFlIxgHx-ahHWk_rFCZs,893
 kostyl/ml/lightning/training_utils.py,sha256=u7X9ysF9Gqy8CdwacdcDlNQNsbagYAhslbv-1WLJ45k,9052
@@ -32,6 +32,6 @@ kostyl/utils/__init__.py,sha256=hkpmB6c5pr4Ti5BshOROebb7cvjDZfNCw83qZ_FFKMM,240
 kostyl/utils/dict_manipulations.py,sha256=e3vBicID74nYP8lHkVTQc4-IQwoJimrbFELy5uSF6Gk,1073
 kostyl/utils/fs.py,sha256=gAQNIU4R_2DhwjgzOS8BOMe0gZymtY1eZwmdgOdDgqo,510
 kostyl/utils/logging.py,sha256=Vye0u4-yeOSUc-f03gpQbxSktTbFiilTWLEVr00ZHvc,5796
-kostyl_toolkit-0.1.34.dist-info/WHEEL,sha256=ZyFSCYkV2BrxH6-HRVRg3R9Fo7MALzer9KiPYqNxSbo,79
-kostyl_toolkit-0.1.34.dist-info/METADATA,sha256=AduVYS1jJViZoxxpO7KYHO0_YNZ23f3P9eKgZAfyteE,4269
-kostyl_toolkit-0.1.34.dist-info/RECORD,,
+kostyl_toolkit-0.1.35.dist-info/WHEEL,sha256=ZyFSCYkV2BrxH6-HRVRg3R9Fo7MALzer9KiPYqNxSbo,79
+kostyl_toolkit-0.1.35.dist-info/METADATA,sha256=KL4-Z421DpchI6KUZ6tVATy99urk1OP2OY4Uf5r9R3U,4269
+kostyl_toolkit-0.1.35.dist-info/RECORD,,