kostyl-toolkit 0.1.21__py3-none-any.whl → 0.1.22__py3-none-any.whl

kostyl/ml/lightning/callbacks/registry_uploading.py

@@ -27,6 +27,8 @@ class ClearMLRegistryUploaderCallback(Callback):
         output_model_tags: list[str] | None = None,
         verbose: bool = True,
         enable_tag_versioning: bool = True,
+        label_enumeration: dict[str, int] | None = None,
+        config_dict: dict[str, str] | None = None,
         uploading_frequency: Literal[
             "after-every-eval", "on-train-end"
         ] = "on-train-end",
@@ -40,6 +42,8 @@ class ClearMLRegistryUploaderCallback(Callback):
             output_model_name: Name for the ClearML output model.
             output_model_tags: Tags for the output model.
             verbose: Whether to log messages.
+            label_enumeration: Optional mapping of label names to integer IDs.
+            config_dict: Optional configuration dictionary to associate with the model.
             enable_tag_versioning: Whether to enable versioning in tags. If True,
                 the version tag (e.g., "v1.0") will be automatically incremented or if not present, added as "v1.0".
             uploading_frequency: When to upload:
@@ -55,6 +59,8 @@ class ClearMLRegistryUploaderCallback(Callback):
         self.ckpt_callback = ckpt_callback
         self.output_model_name = output_model_name
         self.output_model_tags = output_model_tags
+        self.config_dict = config_dict
+        self.label_enumeration = label_enumeration
         self.verbose = verbose
         self.uploading_frequency = uploading_frequency
         self.enable_tag_versioning = enable_tag_versioning
@@ -75,16 +81,21 @@ class ClearMLRegistryUploaderCallback(Callback):
 
         if "LightningCheckpoint" not in self.output_model_tags:
             self.output_model_tags.append("LightningCheckpoint")
-        config = pl_module.model_config
-        if config is not None:
-            config = config.to_dict()
+
+        if self.config_dict is None:
+            config = pl_module.model_config
+            if config is not None:
+                config = config.to_dict()
+        else:
+            config = self.config_dict
 
         return OutputModel(
             task=self.task,
             name=self.output_model_name,
             framework="PyTorch",
             tags=self.output_model_tags,
-            config_dict=config,
+            config_dict=None,
+            label_enumeration=self.label_enumeration,
         )
 
     def _upload_best_checkpoint(self, pl_module: "KostylLightningModule") -> None:
@@ -111,6 +122,13 @@ class ClearMLRegistryUploaderCallback(Callback):
             auto_delete_file=False,
             async_enable=False,
         )
+        if self.config_dict is None:
+            config = pl_module.model_config
+            if config is not None:
+                config = config.to_dict()
+        else:
+            config = self.config_dict
+        self._output_model.update_design(config_dict=config)
 
         self._last_best_model_path = current_best
         return

kostyl/ml/lightning/extenstions/pretrained_model.py

@@ -14,6 +14,7 @@ except ImportError:
 
 from kostyl.utils.logging import log_incompatible_keys
 from kostyl.utils.logging import setup_logger
+from torch import nn
 
 
 logger = setup_logger("LightningPretrainedModelMixin", fmt="only_message")
@@ -67,7 +68,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             raise ValueError(f"{checkpoint_path} is a directory")
         if not checkpoint_path.exists():
             raise FileNotFoundError(f"{checkpoint_path} does not exist")
-        if not checkpoint_path.suffix == ".ckpt":
+        if checkpoint_path.suffix != ".ckpt":
             raise ValueError(f"{checkpoint_path} is not a .ckpt file")
 
         checkpoint_dict = torch.load(
@@ -77,19 +78,21 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             mmap=True,
         )
 
-        config_cls = cast(PretrainedConfig, type(cls.config_class))
+        # 1. Восстанавливаем конфиг
+        config_cls = cast(type[PretrainedConfig], cls.config_class)
         config_dict = checkpoint_dict[config_key]
         config_dict.update(kwargs)
         config = config_cls.from_dict(config_dict)
 
-        kwargs_for_model = {}
-        for key in kwargs:
+        kwargs_for_model: dict[str, Any] = {}
+        for key, value in kwargs.items():
             if not hasattr(config, key):
-                kwargs_for_model[key] = kwargs[key]
+                kwargs_for_model[key] = value
 
         with torch.device("meta"):
             model = cls(config, **kwargs_for_model)
 
+            # PEFT-адаптеры (оставляю твою логику как есть)
             if "peft_config" in checkpoint_dict:
                 if PeftConfig is None:
                     raise ImportError(
@@ -100,26 +103,100 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
                     model.add_adapter(peft_cfg, adapter_name=name)
 
         incompatible_keys: dict[str, list[str]] = {}
-        if weights_prefix != "":
-            if weights_prefix[-1] != ".":
-                weights_prefix += "."
-            model_state_dict = {}
-            mismatched_keys = []
-            for key, value in checkpoint_dict["state_dict"].items():
+
+        raw_state_dict: dict[str, torch.Tensor] = checkpoint_dict["state_dict"]
+
+        if weights_prefix:
+            if not weights_prefix.endswith("."):
+                weights_prefix = weights_prefix + "."
+            state_dict: dict[str, torch.Tensor] = {}
+            mismatched_keys: list[str] = []
+
+            for key, value in raw_state_dict.items():
                 if key.startswith(weights_prefix):
                     new_key = key[len(weights_prefix) :]
-                    model_state_dict[new_key] = value
+                    state_dict[new_key] = value
                 else:
                     mismatched_keys.append(key)
+
+            if mismatched_keys:
                 incompatible_keys["mismatched_keys"] = mismatched_keys
         else:
-            model_state_dict = checkpoint_dict["state_dict"]
-
-        missing_keys, unexpected_keys = model.load_state_dict(
-            model_state_dict, strict=False, assign=True
+            state_dict = raw_state_dict
+
+        # 5. Логика base_model_prefix как в HF:
+        #    поддержка загрузки базовой модели <-> модели с головой
+        #
+        # cls.base_model_prefix обычно "model" / "bert" / "encoder" и т.п.
+        base_prefix: str = getattr(cls, "base_model_prefix", "") or ""
+        model_to_load: nn.Module = model
+
+        if base_prefix:
+            prefix_with_dot = base_prefix + "."
+            loaded_keys = list(state_dict.keys())
+            full_model_state = model.state_dict()
+            expected_keys = list(full_model_state.keys())
+
+            has_prefix_module = any(k.startswith(prefix_with_dot) for k in loaded_keys)
+            expects_prefix_module = any(
+                k.startswith(prefix_with_dot) for k in expected_keys
+            )
+
+            # Кейc 1: загружаем базовую модель в модель с головой.
+            # Пример: StaticEmbeddingsForSequenceClassification (имеет .model)
+            #         state_dict с ключами "embeddings.weight", "token_pos_weights", ...
+            if (
+                hasattr(model, base_prefix)
+                and not has_prefix_module
+                and expects_prefix_module
+            ):
+                # Веса без префикса -> грузим только в model.<base_prefix>
+                model_to_load = getattr(model, base_prefix)
+
+            # Кейc 2: загружаем чекпоинт модели с головой в базовую модель.
+            # Пример: BertModel, а в state_dict ключи "bert.encoder.layer.0..."
+            elif (
+                not hasattr(model, base_prefix)
+                and has_prefix_module
+                and not expects_prefix_module
+            ):
+                new_state_dict: dict[str, torch.Tensor] = {}
+                for key, value in state_dict.items():
+                    if key.startswith(prefix_with_dot):
+                        new_key = key[len(prefix_with_dot) :]
+                    else:
+                        new_key = key
+                    new_state_dict[new_key] = value
+                state_dict = new_state_dict
+
+        load_result = model_to_load.load_state_dict(
+            state_dict, strict=False, assign=True
+        )
+        missing_keys, unexpected_keys = (
+            load_result.missing_keys,
+            load_result.unexpected_keys,
         )
+
+        # Если мы грузили только в base-подмодуль, расширим missing_keys
+        # до полного списка (base + голова), как в старых версиях HF.
+        if model_to_load is not model and base_prefix:
+            base_keys = set(model_to_load.state_dict().keys())
+            # Приводим ключи полной модели к "безпрефиксному" виду
+            head_like_keys = set()
+            prefix_with_dot = base_prefix + "."
+            for k in model.state_dict().keys():
+                if k.startswith(prefix_with_dot):
+                    # отрезаем "model."
+                    head_like_keys.add(k[len(prefix_with_dot) :])
+                else:
+                    head_like_keys.add(k)
+            extra_missing = sorted(head_like_keys - base_keys)
+            missing_keys = list(missing_keys) + extra_missing
+
         incompatible_keys["missing_keys"] = missing_keys
         incompatible_keys["unexpected_keys"] = unexpected_keys
+
         if should_log_incompatible_keys:
             log_incompatible_keys(incompatible_keys=incompatible_keys, logger=logger)
+
         return model

{kostyl_toolkit-0.1.21.dist-info → kostyl_toolkit-0.1.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.21
+Version: 0.1.22
 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.21.dist-info → kostyl_toolkit-0.1.22.dist-info}/RECORD

@@ -13,10 +13,10 @@ kostyl/ml/lightning/__init__.py,sha256=-F3JAyq8KU1d-nACWryGu8d1CbvWbQ1rXFdeRwfE2
 kostyl/ml/lightning/callbacks/__init__.py,sha256=Vd-rozY4T9Prr3IMqbliXxj6sC6y9XsovHQqRwzc2HI,297
 kostyl/ml/lightning/callbacks/checkpoint.py,sha256=FooGeeUz6TtoXQglpcK16NWAmSX3fbu6wntRtK3a_Io,1936
 kostyl/ml/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
-kostyl/ml/lightning/callbacks/registry_uploading.py,sha256=jJdSoFIkTcGLnZIKBzbAkt9MAgpZURLCQdd0DwAC5gk,4659
+kostyl/ml/lightning/callbacks/registry_uploading.py,sha256=32vhMNNuThtEcvRdS5jh5s-wf7LwZNsCTwZA3emcObs,5449
 kostyl/ml/lightning/extenstions/__init__.py,sha256=OY6QGv1agYgqqKf1xJBrxgp_i8FunVfPzYezfaRrGXU,182
 kostyl/ml/lightning/extenstions/custom_module.py,sha256=nB5jW7cqRD1tyh-q5LD2EtiFQwFkLXpnS9Yu6c5xMRg,5987
-kostyl/ml/lightning/extenstions/pretrained_model.py,sha256=ZOKtrVl095cwvI43wAz-Xdzu4l0v0lHH2mfh4WXwxKQ,5059
+kostyl/ml/lightning/extenstions/pretrained_model.py,sha256=x8D2nMDDW8J913qFRSEGKXfQO8ipPJM5SLo4Y5kc3YA,8638
 kostyl/ml/lightning/loggers/__init__.py,sha256=e51dszaoJbuzwBkbdugmuDsPldoSO4yaRgmZUg1Bdy0,71
 kostyl/ml/lightning/loggers/tb_logger.py,sha256=j02HK5ue8yzXXV8FWKmmXyHkFlIxgHx-ahHWk_rFCZs,893
 kostyl/ml/lightning/steps_estimation.py,sha256=fTZ0IrUEZV3H6VYlx4GYn56oco56mMiB7FO9F0Z7qc4,1511
@@ -30,6 +30,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.21.dist-info/WHEEL,sha256=z-mOpxbJHqy3cq6SvUThBZdaLGFZzdZPtgWLcP2NKjQ,79
-kostyl_toolkit-0.1.21.dist-info/METADATA,sha256=nd_zUIEvb2Bda7h6tbJHnUonjK9iBj4nMNYGMivhU2o,4269
-kostyl_toolkit-0.1.21.dist-info/RECORD,,
+kostyl_toolkit-0.1.22.dist-info/WHEEL,sha256=z-mOpxbJHqy3cq6SvUThBZdaLGFZzdZPtgWLcP2NKjQ,79
+kostyl_toolkit-0.1.22.dist-info/METADATA,sha256=GweBJ42Dhbl4Y5PNu-jnffXj1CaJ34DTPUcoFEndJ1M,4269
+kostyl_toolkit-0.1.22.dist-info/RECORD,,