kostyl-toolkit 0.1.22__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/callbacks/__init__.py

@@ -1,6 +1,6 @@
 from .checkpoint import setup_checkpoint_callback
 from .early_stopping import setup_early_stopping_callback
-from .registry_uploading import ClearMLRegistryUploaderCallback
+from .registry_uploader import ClearMLRegistryUploaderCallback
 
 
 __all__ = [

kostyl/ml/lightning/callbacks/checkpoint.py

@@ -1,21 +1,289 @@
+from datetime import timedelta
 from pathlib import Path
 from shutil import rmtree
+from typing import Literal
+from typing import cast
 
+from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.callbacks import ModelCheckpoint
 
 from kostyl.ml.configs import CheckpointConfig
 from kostyl.ml.dist_utils import is_main_process
 from kostyl.utils import setup_logger
 
+from .registry_uploader import RegistryUploaderCallback
+
 
 logger = setup_logger("callbacks/checkpoint.py")
 
 
+class CustomModelCheckpoint(ModelCheckpoint):
+    r"""
+    Save the model after every epoch by monitoring a quantity. Every logged metrics are passed to the
+    :class:`~lightning.pytorch.loggers.logger.Logger` for the version it gets saved in the same directory as the
+    checkpoint.
+
+    After training finishes, use :attr:`best_model_path` to retrieve the path to the
+    best checkpoint file and :attr:`best_model_score` to get its score.
+
+    .. note::
+        When using manual optimization with ``every_n_train_steps``, you should save the model state
+        in your ``training_step`` before the optimizer step if you want the checkpoint to reflect
+        the pre-optimization state. Example:
+
+        .. code-block:: python
+
+            def training_step(self, batch, batch_idx):
+                # ... forward pass, loss calculation, backward pass ...
+
+                # Save model state before optimization
+                if not hasattr(self, 'saved_models'):
+                    self.saved_models = {}
+                self.saved_models[batch_idx] = {
+                    k: v.detach().clone()
+                    for k, v in self.layer.state_dict().items()
+                }
+
+                # Then perform optimization
+                optimizer.zero_grad()
+                self.manual_backward(loss)
+                optimizer.step()
+
+                # Optional: Clean up old states to save memory
+                if batch_idx > 10:  # Keep last 10 states
+                    del self.saved_models[batch_idx - 10]
+
+    Args:
+        dirpath: Directory to save the model file.
+            Example: ``dirpath='my/path/'``.
+
+            .. warning::
+                In a distributed environment like DDP, it's recommended to provide a `dirpath` to avoid race conditions.
+                When using manual optimization with ``every_n_train_steps``, make sure to save the model state
+                in your training loop as shown in the example above.
+
+            Can be remote file paths such as `s3://mybucket/path/` or 'hdfs://path/'
+            (default: ``None``). If dirpath is ``None``, we only keep the ``k`` best checkpoints
+            in memory, and do not save anything to disk.
+
+        filename: Checkpoint filename. Can contain named formatting options to be auto-filled.
+            If no name is provided, it will be ``None`` and the checkpoint will be saved to
+            ``{epoch}``.and if the Trainer uses a logger, the path will also contain logger name and version.
+
+        filename: checkpoint filename. Can contain named formatting options to be auto-filled.
+
+            Example::
+
+                # save any arbitrary metrics like `val_loss`, etc. in name
+                # saves a file like: my/path/epoch=2-val_loss=0.02-other_metric=0.03.ckpt
+                >>> checkpoint_callback = ModelCheckpoint(
+                ...     dirpath='my/path',
+                ...     filename='{epoch}-{val_loss:.2f}-{other_metric:.2f}'
+                ... )
+
+            By default, filename is ``None`` and will be set to ``'{epoch}-{step}'``, where "epoch" and "step" match
+            the number of finished epoch and optimizer steps respectively.
+        monitor: quantity to monitor. By default it is ``None`` which saves a checkpoint only for the last epoch.
+        verbose: verbosity mode. Default: ``False``.
+        save_last: When ``True``, saves a `last.ckpt` copy whenever a checkpoint file gets saved. Can be set to
+            ``'link'`` on a local filesystem to create a symbolic link. This allows accessing the latest checkpoint
+            in a deterministic manner. Default: ``None``.
+        save_top_k: if ``save_top_k == k``,
+            the best k models according to the quantity monitored will be saved.
+            If ``save_top_k == 0``, no models are saved.
+            If ``save_top_k == -1``, all models are saved.
+            Please note that the monitors are checked every ``every_n_epochs`` epochs.
+            If ``save_top_k >= 2`` and the callback is called multiple times inside an epoch, and the filename remains
+            unchanged, the name of the saved file will be appended with a version count starting with ``v1`` to avoid
+            collisions unless ``enable_version_counter`` is set to False. The version counter is unrelated to the top-k
+            ranking of the checkpoint, and we recommend formatting the filename to include the monitored metric to avoid
+            collisions.
+        save_on_exception: Whether to save a checkpoint when an exception is raised. Default: ``False``.
+        mode: one of {min, max}.
+            If ``save_top_k != 0``, the decision to overwrite the current save file is made
+            based on either the maximization or the minimization of the monitored quantity.
+            For ``'val_acc'``, this should be ``'max'``, for ``'val_loss'`` this should be ``'min'``, etc.
+        auto_insert_metric_name: When ``True``, the checkpoints filenames will contain the metric name.
+            For example, ``filename='checkpoint_{epoch:02d}-{acc:02.0f}`` with epoch ``1`` and acc ``1.12`` will resolve
+            to ``checkpoint_epoch=01-acc=01.ckpt``. Is useful to set it to ``False`` when metric names contain ``/``
+            as this will result in extra folders.
+            For example, ``filename='epoch={epoch}-step={step}-val_acc={val/acc:.2f}', auto_insert_metric_name=False``
+        save_weights_only: if ``True``, then only the model's weights will be
+            saved. Otherwise, the optimizer states, lr-scheduler states, etc are added in the checkpoint too.
+        every_n_train_steps: How many training steps to wait before saving a checkpoint. This does not take into account
+            the steps of the current epoch. If ``every_n_train_steps == None or every_n_train_steps == 0``,
+            no checkpoints
+            will be saved during training. Mutually exclusive with ``train_time_interval`` and ``every_n_epochs``.
+
+            .. note::
+                When using with manual optimization, the checkpoint will be saved after the optimizer step by default.
+                To save the model state before the optimizer step, you need to save the model state in your
+                ``training_step`` before calling ``optimizer.step()``. See the class docstring for an example.
+        train_time_interval: Checkpoints are monitored at the specified time interval.
+            For all practical purposes, this cannot be smaller than the amount
+            of time it takes to process a single training batch. This is not
+            guaranteed to execute at the exact time specified, but should be close.
+            This must be mutually exclusive with ``every_n_train_steps`` and ``every_n_epochs``.
+        every_n_epochs: Number of epochs between checkpoints.
+            This value must be ``None`` or non-negative.
+            To disable saving top-k checkpoints, set ``every_n_epochs = 0``.
+            This argument does not impact the saving of ``save_last=True`` checkpoints.
+            If all of ``every_n_epochs``, ``every_n_train_steps`` and
+            ``train_time_interval`` are ``None``, we save a checkpoint at the end of every epoch
+            (equivalent to ``every_n_epochs = 1``).
+            If ``every_n_epochs == None`` and either ``every_n_train_steps != None`` or ``train_time_interval != None``,
+            saving at the end of each epoch is disabled
+            (equivalent to ``every_n_epochs = 0``).
+            This must be mutually exclusive with ``every_n_train_steps`` and ``train_time_interval``.
+            Setting both ``ModelCheckpoint(..., every_n_epochs=V, save_on_train_epoch_end=False)`` and
+            ``Trainer(max_epochs=N, check_val_every_n_epoch=M)``
+            will only save checkpoints at epochs 0 < E <= N
+            where both values for ``every_n_epochs`` and ``check_val_every_n_epoch`` evenly divide E.
+        save_on_train_epoch_end: Whether to run checkpointing at the end of the training epoch.
+            If ``True``, checkpoints are saved at the end of every training epoch.
+            If ``False``, checkpoints are saved at the end of validation.
+            If ``None`` (default), checkpointing behavior is determined based on training configuration.
+            If ``val_check_interval`` is a str, dict, or `timedelta` (time-based), checkpointing is performed after
+            validation.
+            If ``check_val_every_n_epoch != 1``, checkpointing will not be performed at the end of
+            every training epoch. If there are no validation batches of data, checkpointing will occur at the
+            end of the training epoch. If there is a non-default number of validation runs per training epoch
+            (``val_check_interval != 1``), checkpointing is performed after validation.
+        enable_version_counter: Whether to append a version to the existing file name.
+            If ``False``, then the checkpoint files will be overwritten.
+
+    Note:
+        For extra customization, ModelCheckpoint includes the following attributes:
+
+        - ``CHECKPOINT_JOIN_CHAR = "-"``
+        - ``CHECKPOINT_EQUALS_CHAR = "="``
+        - ``CHECKPOINT_NAME_LAST = "last"``
+        - ``FILE_EXTENSION = ".ckpt"``
+        - ``STARTING_VERSION = 1``
+
+        For example, you can change the default last checkpoint name by doing
+        ``checkpoint_callback.CHECKPOINT_NAME_LAST = "{epoch}-last"``
+
+        If you want to checkpoint every N hours, every M train batches, and/or every K val epochs,
+        then you should create multiple ``ModelCheckpoint`` callbacks.
+
+        If the checkpoint's ``dirpath`` changed from what it was before while resuming the training,
+        only ``best_model_path`` will be reloaded and a warning will be issued.
+
+        If you provide a ``filename`` on a mounted device where changing permissions is not allowed (causing ``chmod``
+        to raise a ``PermissionError``), install `fsspec>=2025.5.0`. Then the error is caught, the file's permissions
+        remain unchanged, and the checkpoint is still saved. Otherwise, no checkpoint will be saved and training stops.
+
+    Raises:
+        MisconfigurationException:
+            If ``save_top_k`` is smaller than ``-1``,
+            if ``monitor`` is ``None`` and ``save_top_k`` is none of ``None``, ``-1``, and ``0``, or
+            if ``mode`` is none of ``"min"`` or ``"max"``.
+        ValueError:
+            If ``trainer.save_checkpoint`` is ``None``.
+
+    Example::
+
+        >>> from lightning.pytorch import Trainer
+        >>> from lightning.pytorch.callbacks import ModelCheckpoint
+
+        # saves checkpoints to 'my/path/' at every epoch
+        >>> checkpoint_callback = ModelCheckpoint(dirpath='my/path/')
+        >>> trainer = Trainer(callbacks=[checkpoint_callback])
+
+        # save epoch and val_loss in name
+        # saves a file like: my/path/sample-mnist-epoch=02-val_loss=0.32.ckpt
+        >>> checkpoint_callback = ModelCheckpoint(
+        ...     monitor='val_loss',
+        ...     dirpath='my/path/',
+        ...     filename='sample-mnist-{epoch:02d}-{val_loss:.2f}'
+        ... )
+
+        # save epoch and val_loss in name, but specify the formatting yourself (e.g. to avoid problems with Tensorboard
+        # or Neptune, due to the presence of characters like '=' or '/')
+        # saves a file like: my/path/sample-mnist-epoch02-val_loss0.32.ckpt
+        >>> checkpoint_callback = ModelCheckpoint(
+        ...     monitor='val/loss',
+        ...     dirpath='my/path/',
+        ...     filename='sample-mnist-epoch{epoch:02d}-val_loss{val/loss:.2f}',
+        ...     auto_insert_metric_name=False
+        ... )
+
+        # retrieve the best checkpoint after training
+        >>> checkpoint_callback = ModelCheckpoint(dirpath='my/path/')
+        >>> trainer = Trainer(callbacks=[checkpoint_callback])
+        >>> model = ...  # doctest: +SKIP
+        >>> trainer.fit(model)  # doctest: +SKIP
+        >>> print(checkpoint_callback.best_model_path)  # doctest: +SKIP
+
+    .. tip:: Saving and restoring multiple checkpoint callbacks at the same time is supported under variation in the
+        following arguments:
+
+        *monitor, mode, every_n_train_steps, every_n_epochs, train_time_interval*
+
+        Read more: :ref:`Persisting Callback State <extensions/callbacks_state:save callback state>`
+
+    """  # noqa: D205
+
+    def __init__(  # noqa: D107
+        self,
+        dirpath: _PATH | None = None,
+        filename: str | None = None,
+        monitor: str | None = None,
+        verbose: bool = False,
+        save_last: bool | Literal["link"] | None = None,
+        save_top_k: int = 1,
+        save_on_exception: bool = False,
+        save_weights_only: bool = False,
+        mode: str = "min",
+        auto_insert_metric_name: bool = True,
+        every_n_train_steps: int | None = None,
+        train_time_interval: timedelta | None = None,
+        every_n_epochs: int | None = None,
+        save_on_train_epoch_end: bool | None = None,
+        enable_version_counter: bool = True,
+        registry_uploader_callback: RegistryUploaderCallback | None = None,
+    ) -> None:
+        super().__init__(
+            dirpath=dirpath,
+            filename=filename,
+            monitor=monitor,
+            verbose=verbose,
+            save_last=save_last,
+            save_top_k=save_top_k,
+            save_on_exception=save_on_exception,
+            save_weights_only=save_weights_only,
+            mode=mode,
+            auto_insert_metric_name=auto_insert_metric_name,
+            every_n_train_steps=every_n_train_steps,
+            train_time_interval=train_time_interval,
+            every_n_epochs=every_n_epochs,
+            save_on_train_epoch_end=save_on_train_epoch_end,
+            enable_version_counter=enable_version_counter,
+        )
+        self.registry_uploader_callback = registry_uploader_callback
+        self._custom_best_model_path = cast(str, self.best_model_path)
+        return
+
+    @property
+    def best_model_path(self) -> str:
+        """Best model path."""
+        return self._custom_best_model_path
+
+    @best_model_path.setter
+    def best_model_path(self, value: str) -> None:
+        self._custom_best_model_path = value
+        if self.registry_uploader_callback is not None:
+            self.registry_uploader_callback.best_model_path = value
+        return
+
+
 def setup_checkpoint_callback(
     dirpath: Path,
     ckpt_cfg: CheckpointConfig,
     save_weights_only: bool = True,
-) -> ModelCheckpoint:
+    registry_uploader_callback: RegistryUploaderCallback | None = None,
+) -> CustomModelCheckpoint:
     """
     Sets up a ModelCheckpoint callback for PyTorch Lightning.
 
@@ -30,6 +298,8 @@ def setup_checkpoint_callback(
             settings such as filename, save_top_k, monitor, and mode.
         save_weights_only (bool, optional): Whether to save only the model weights
             or the full model. Defaults to True.
+        registry_uploader_callback (RegistryUploaderCallback | None, optional):
+            An optional callback for uploading checkpoints to a registry. Defaults to None.
 
     Returns:
         ModelCheckpoint: The configured ModelCheckpoint callback instance.
@@ -44,7 +314,7 @@ def setup_checkpoint_callback(
         logger.info(f"Creating checkpoint directory {dirpath}.")
         dirpath.mkdir(parents=True, exist_ok=True)
 
-    checkpoint_callback = ModelCheckpoint(
+    checkpoint_callback = CustomModelCheckpoint(
         dirpath=dirpath,
         filename=ckpt_cfg.filename,
         save_top_k=ckpt_cfg.save_top_k,
@@ -52,5 +322,6 @@ def setup_checkpoint_callback(
         mode=ckpt_cfg.mode,
         verbose=True,
         save_weights_only=save_weights_only,
+        registry_uploader_callback=registry_uploader_callback,
     )
     return checkpoint_callback

kostyl/ml/lightning/callbacks/{registry_uploading.py → registry_uploader.py}

@@ -1,3 +1,7 @@
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import Callable
+from functools import partial
 from typing import Literal
 from typing import override
 
@@ -5,7 +9,6 @@ from clearml import OutputModel
 from clearml import Task
 from lightning import Trainer
 from lightning.pytorch.callbacks import Callback
-from lightning.pytorch.callbacks import ModelCheckpoint
 
 from kostyl.ml.clearml.logging_utils import find_version_in_tags
 from kostyl.ml.clearml.logging_utils import increment_version
@@ -16,13 +19,31 @@ from kostyl.utils.logging import setup_logger
 logger = setup_logger()
 
 
-class ClearMLRegistryUploaderCallback(Callback):
+class RegistryUploaderCallback(Callback, ABC):
+    """Abstract Lightning callback responsible for tracking and uploading the best-performing model checkpoint."""
+
+    @property
+    @abstractmethod
+    def best_model_path(self) -> str:
+        """Return the file system path pointing to the best model artifact produced during training."""
+        raise NotImplementedError
+
+    @best_model_path.setter
+    @abstractmethod
+    def best_model_path(self, value: str) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _upload_best_checkpoint(self, pl_module: "KostylLightningModule") -> None:
+        raise NotImplementedError
+
+
+class ClearMLRegistryUploaderCallback(RegistryUploaderCallback):
     """PyTorch Lightning callback to upload the best model checkpoint to ClearML."""
 
     def __init__(
         self,
         task: Task,
-        ckpt_callback: ModelCheckpoint,
         output_model_name: str,
         output_model_tags: list[str] | None = None,
         verbose: bool = True,
@@ -56,7 +77,6 @@ class ClearMLRegistryUploaderCallback(Callback):
             output_model_tags = []
 
         self.task = task
-        self.ckpt_callback = ckpt_callback
         self.output_model_name = output_model_name
         self.output_model_tags = output_model_tags
         self.config_dict = config_dict
@@ -66,7 +86,22 @@ class ClearMLRegistryUploaderCallback(Callback):
         self.enable_tag_versioning = enable_tag_versioning
 
         self._output_model: OutputModel | None = None
-        self._last_best_model_path: str = ""
+        self._last_uploaded_model_path: str = ""
+        self._best_model_path: str = ""
+        self._upload_callback: Callable | None = None
+        return
+
+    @property
+    @override
+    def best_model_path(self) -> str:
+        return self._best_model_path
+
+    @best_model_path.setter
+    @override
+    def best_model_path(self, value: str) -> None:
+        self._best_model_path = value
+        if self._upload_callback is not None:
+            self._upload_callback()
         return
 
     def _create_output_model(self, pl_module: "KostylLightningModule") -> OutputModel:
@@ -98,27 +133,29 @@ class ClearMLRegistryUploaderCallback(Callback):
             label_enumeration=self.label_enumeration,
         )
 
+    @override
     def _upload_best_checkpoint(self, pl_module: "KostylLightningModule") -> None:
-        current_best = self.ckpt_callback.best_model_path
-
-        if not current_best:
-            if self.verbose:
-                logger.info("No best model found yet to upload")
-            return
-
-        if current_best == self._last_best_model_path:
-            if self.verbose:
-                logger.info("Best model unchanged since last upload")
+        if not self._best_model_path or (
+            self._best_model_path == self._last_uploaded_model_path
+        ):
+            if not self._best_model_path:
+                if self.verbose:
+                    logger.info("No best model found yet to upload")
+            elif self._best_model_path == self._last_uploaded_model_path:
+                if self.verbose:
+                    logger.info("Best model unchanged since last upload")
+            self._upload_callback = partial(self._upload_best_checkpoint, pl_module)
             return
+        self._upload_callback = None
 
         if self._output_model is None:
             self._output_model = self._create_output_model(pl_module)
 
         if self.verbose:
-            logger.info(f"Uploading best model from {current_best}")
+            logger.info(f"Uploading best model from {self._best_model_path}")
 
         self._output_model.update_weights(
-            current_best,
+            self._best_model_path,
             auto_delete_file=False,
             async_enable=False,
         )
@@ -130,7 +167,7 @@ class ClearMLRegistryUploaderCallback(Callback):
             config = self.config_dict
         self._output_model.update_design(config_dict=config)
 
-        self._last_best_model_path = current_best
+        self._last_uploaded_model_path = self._best_model_path
         return
 
     @override

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.22.dist-info → kostyl_toolkit-0.1.24.dist-info}/METADATA

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

@@ -3,20 +3,20 @@ 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
 kostyl/ml/configs/training_settings.py,sha256=Sq2tiRuwkbmi9zKDG2JghZLXo5DDt_eQqN_KYJSdcTY,2509
 kostyl/ml/dist_utils.py,sha256=Onf0KHVLA8oeUgZTcTdmR9qiM22f2uYLoNwgLbMGJWk,3495
 kostyl/ml/lightning/__init__.py,sha256=-F3JAyq8KU1d-nACWryGu8d1CbvWbQ1rXFdeRwfE2X8,175
-kostyl/ml/lightning/callbacks/__init__.py,sha256=Vd-rozY4T9Prr3IMqbliXxj6sC6y9XsovHQqRwzc2HI,297
-kostyl/ml/lightning/callbacks/checkpoint.py,sha256=FooGeeUz6TtoXQglpcK16NWAmSX3fbu6wntRtK3a_Io,1936
+kostyl/ml/lightning/callbacks/__init__.py,sha256=enexQt3octktsTiEYHltSF_24CM-NeFEVFimXiavGiY,296
+kostyl/ml/lightning/callbacks/checkpoint.py,sha256=rJ05S7BnPopvQjV5b9CI29-S7-ySllyhG7PIzz34VyY,16793
 kostyl/ml/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
-kostyl/ml/lightning/callbacks/registry_uploading.py,sha256=32vhMNNuThtEcvRdS5jh5s-wf7LwZNsCTwZA3emcObs,5449
+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.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,,
+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,,