kostyl-toolkit 0.1.20__tar.gz → 0.1.22__tar.gz

{kostyl_toolkit-0.1.20 → kostyl_toolkit-0.1.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.20
+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.20 → kostyl_toolkit-0.1.22}/kostyl/ml/configs/hyperparams.py

@@ -4,8 +4,6 @@ from pydantic import model_validator
 
 from kostyl.utils.logging import setup_logger
 
-from .base_model import KostylBaseModel
-
 
 logger = setup_logger(fmt="only_message")
 
@@ -75,7 +73,7 @@ class WeightDecay(BaseModel):
         return self
 
 
-class HyperparamsConfig(KostylBaseModel):
+class HyperparamsConfig(BaseModel):
     """Model training hyperparameters configuration."""
 
     grad_clip_val: float | None = Field(default=None, gt=0, validate_default=False)

{kostyl_toolkit-0.1.20 → kostyl_toolkit-0.1.22}/kostyl/ml/configs/training_settings.py

@@ -5,8 +5,6 @@ from pydantic import Field
 
 from kostyl.utils.logging import setup_logger
 
-from .base_model import KostylBaseModel
-
 
 logger = setup_logger(fmt="only_message")
 
@@ -95,7 +93,7 @@ class DataConfig(BaseModel):
     data_columns: list[str]
 
 
-class TrainingSettings(KostylBaseModel):
+class TrainingSettings(BaseModel):
     """Training parameters configuration."""
 
     trainer: LightningTrainerParameters

{kostyl_toolkit-0.1.20 → kostyl_toolkit-0.1.22}/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_toolkit-0.1.20 → kostyl_toolkit-0.1.22}/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.20 → kostyl_toolkit-0.1.22}/kostyl/utils/logging.py

@@ -118,7 +118,7 @@ def setup_logger(
     else:
         base = name
 
-    if add_rank is None:
+    if (add_rank is None) or add_rank:
         try:
             add_rank = dist.is_available() and dist.is_initialized()
         except Exception:

{kostyl_toolkit-0.1.20 → kostyl_toolkit-0.1.22}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kostyl-toolkit"
-version = "0.1.20"
+version = "0.1.22"
 description = "Kickass Orchestration System for Training, Yielding & Logging "
 readme = "README.md"
 requires-python = ">=3.12"