kostyl-toolkit 0.1.23__py3-none-any.whl → 0.1.24__py3-none-any.whl

kostyl/ml/clearml/pulling_utils.py

@@ -95,7 +95,7 @@ def get_model_from_clearml[
             raise ValueError(
                 f"Model class {model.__name__} is not compatible with Lightning checkpoints."
             )
-        model_instance = model.from_lighting_checkpoint(local_path, **kwargs)
+        model_instance = model.from_lightning_checkpoint(local_path, **kwargs)
     else:
         raise ValueError(
             f"Unsupported model format for path: {local_path}. "

kostyl/ml/lightning/extenstions/pretrained_model.py

@@ -6,15 +6,7 @@ import torch
 from transformers import PretrainedConfig
 from transformers import PreTrainedModel
 
-
-try:
-    from peft import PeftConfig
-except ImportError:
-    PeftConfig = None  # ty: ignore
-
-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")
@@ -24,12 +16,11 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
     """A mixin class for loading pretrained models from PyTorch Lightning checkpoints."""
 
     @classmethod
-    def from_lighting_checkpoint[TModelInstance: LightningCheckpointLoaderMixin](  # noqa: C901
+    def from_lightning_checkpoint[TModelInstance: LightningCheckpointLoaderMixin](  # noqa: C901
         cls: type[TModelInstance],
         checkpoint_path: str | Path,
         config_key: str = "config",
         weights_prefix: str = "model.",
-        should_log_incompatible_keys: bool = True,
         **kwargs: Any,
     ) -> TModelInstance:
         """
@@ -50,8 +41,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
                 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.
-            should_log_incompatible_keys (bool, optional): Whether to log incompatible keys. Defaults to True.
-            **kwargs: Additional keyword arguments to pass to the model loading method.
+            kwargs: Additional keyword arguments to pass to the model's `from_pretrained` method.
 
         Returns:
             TModelInstance: The loaded model instance.
@@ -89,114 +79,27 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             if not hasattr(config, 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(
-                        "peft is not installed. Please install it to load PEFT models."
-                    )
-                for name, adapter_dict in checkpoint_dict["peft_config"].items():
-                    peft_cfg = PeftConfig.from_peft_type(**adapter_dict)
-                    model.add_adapter(peft_cfg, adapter_name=name)
-
-        incompatible_keys: dict[str, list[str]] = {}
-
         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) :]
                     state_dict[new_key] = value
                 else:
-                    mismatched_keys.append(key)
-
-            if mismatched_keys:
-                incompatible_keys["mismatched_keys"] = mismatched_keys
+                    state_dict[key] = value
         else:
             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
+        model = cls.from_pretrained(
+            pretrained_model_name_or_path=None,
+            config=config,
+            state_dict=state_dict,
+            **kwargs_for_model,
         )
-        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.23.dist-info → kostyl_toolkit-0.1.24.dist-info}/METADATA

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

@@ -3,7 +3,7 @@ kostyl/ml/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kostyl/ml/clearml/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kostyl/ml/clearml/dataset_utils.py,sha256=eij_sr2KDhm8GxEbVbK8aBjPsuVvLl9-PIGGaKVgXLA,1729
 kostyl/ml/clearml/logging_utils.py,sha256=GBjIIZbH_itd5sj7XpvxjkyZwxxGOpEcQ3BiWaJTyq8,1210
-kostyl/ml/clearml/pulling_utils.py,sha256=07bb7ZYlZy-qoZLn7uWZCtz02eX2idgk3JA-PPooS9E,4077
+kostyl/ml/clearml/pulling_utils.py,sha256=cNa_-_5LHjNVYi9btXBrfl5sPvI6BAAlIFidtpKu310,4078
 kostyl/ml/configs/__init__.py,sha256=IetcivbqYGutowLqxdKp7QR4tkXKBr4m8t4Zkk9jHZU,911
 kostyl/ml/configs/base_model.py,sha256=Eofn14J9RsjpVx_J4rp6C19pDDCANU4hr3JtX-d0FpQ,4820
 kostyl/ml/configs/hyperparams.py,sha256=OqN7mEj3zc5MTqBPCZL3Lcd2VCTDLo_K0yvhRWGfhCs,2924
@@ -16,7 +16,7 @@ kostyl/ml/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jOR
 kostyl/ml/lightning/callbacks/registry_uploader.py,sha256=ksoh02dzIde4E_GaZykfiOgfSjZti-IJt_i61enem3s,6779
 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=x8D2nMDDW8J913qFRSEGKXfQO8ipPJM5SLo4Y5kc3YA,8638
+kostyl/ml/lightning/extenstions/pretrained_model.py,sha256=QJGr2UvYJcU2Gy2w8z_cEvTodjv7hGdd2PPPfdOI-Mw,4017
 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.23.dist-info/WHEEL,sha256=z-mOpxbJHqy3cq6SvUThBZdaLGFZzdZPtgWLcP2NKjQ,79
-kostyl_toolkit-0.1.23.dist-info/METADATA,sha256=8af_sRkZy9w8chOp4NLvercyB57df6FXAvpLsWKPqro,4269
-kostyl_toolkit-0.1.23.dist-info/RECORD,,
+kostyl_toolkit-0.1.24.dist-info/WHEEL,sha256=z-mOpxbJHqy3cq6SvUThBZdaLGFZzdZPtgWLcP2NKjQ,79
+kostyl_toolkit-0.1.24.dist-info/METADATA,sha256=uq8MPJ9vJgWsp9Z2c7C9tcbaH29QM9ux7_SyahPSlHE,4269
+kostyl_toolkit-0.1.24.dist-info/RECORD,,