kostyl-toolkit 0.1.21__tar.gz → 0.1.23__tar.gz

{kostyl_toolkit-0.1.21 → kostyl_toolkit-0.1.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.21
+Version: 0.1.23
 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 → kostyl_toolkit-0.1.23}/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_toolkit-0.1.23/kostyl/ml/lightning/callbacks/checkpoint.py

@@ -0,0 +1,327 @@
+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,
+    registry_uploader_callback: RegistryUploaderCallback | None = None,
+) -> CustomModelCheckpoint:
+    """
+    Sets up a ModelCheckpoint callback for PyTorch Lightning.
+
+    This function prepares a checkpoint directory and configures a ModelCheckpoint
+    callback based on the provided configuration. If the directory already exists,
+    it is removed (only by the main process) to ensure a clean start. Otherwise,
+    the directory is created.
+
+    Args:
+        dirpath (Path): The path to the directory where checkpoints will be saved.
+        ckpt_cfg (CheckpointConfig): Configuration object containing checkpoint
+            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.
+
+    """
+    if dirpath.exists():
+        if is_main_process():
+            logger.warning(f"Checkpoint directory {dirpath} already exists.")
+            rmtree(dirpath)
+            logger.warning(f"Removed existing checkpoint directory {dirpath}.")
+    else:
+        logger.info(f"Creating checkpoint directory {dirpath}.")
+        dirpath.mkdir(parents=True, exist_ok=True)
+
+    checkpoint_callback = CustomModelCheckpoint(
+        dirpath=dirpath,
+        filename=ckpt_cfg.filename,
+        save_top_k=ckpt_cfg.save_top_k,
+        monitor=ckpt_cfg.monitor,
+        mode=ckpt_cfg.mode,
+        verbose=True,
+        save_weights_only=save_weights_only,
+        registry_uploader_callback=registry_uploader_callback,
+    )
+    return checkpoint_callback

kostyl_toolkit-0.1.21/kostyl/ml/lightning/callbacks/registry_uploading.py → kostyl_toolkit-0.1.23/kostyl/ml/lightning/callbacks/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,17 +19,37 @@ 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,
         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 +63,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:
@@ -52,15 +77,31 @@ 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
+        self.label_enumeration = label_enumeration
         self.verbose = verbose
         self.uploading_frequency = uploading_frequency
         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:
@@ -75,44 +116,58 @@ 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,
         )
 
+    @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,
         )
-
-        self._last_best_model_path = current_best
+        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_uploaded_model_path = self._best_model_path
         return
 
     @override

{kostyl_toolkit-0.1.21 → kostyl_toolkit-0.1.23}/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 → kostyl_toolkit-0.1.23}/pyproject.toml

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

kostyl_toolkit-0.1.21/kostyl/ml/lightning/callbacks/checkpoint.py

@@ -1,56 +0,0 @@
-from pathlib import Path
-from shutil import rmtree
-
-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
-
-
-logger = setup_logger("callbacks/checkpoint.py")
-
-
-def setup_checkpoint_callback(
-    dirpath: Path,
-    ckpt_cfg: CheckpointConfig,
-    save_weights_only: bool = True,
-) -> ModelCheckpoint:
-    """
-    Sets up a ModelCheckpoint callback for PyTorch Lightning.
-
-    This function prepares a checkpoint directory and configures a ModelCheckpoint
-    callback based on the provided configuration. If the directory already exists,
-    it is removed (only by the main process) to ensure a clean start. Otherwise,
-    the directory is created.
-
-    Args:
-        dirpath (Path): The path to the directory where checkpoints will be saved.
-        ckpt_cfg (CheckpointConfig): Configuration object containing checkpoint
-            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.
-
-    Returns:
-        ModelCheckpoint: The configured ModelCheckpoint callback instance.
-
-    """
-    if dirpath.exists():
-        if is_main_process():
-            logger.warning(f"Checkpoint directory {dirpath} already exists.")
-            rmtree(dirpath)
-            logger.warning(f"Removed existing checkpoint directory {dirpath}.")
-    else:
-        logger.info(f"Creating checkpoint directory {dirpath}.")
-        dirpath.mkdir(parents=True, exist_ok=True)
-
-    checkpoint_callback = ModelCheckpoint(
-        dirpath=dirpath,
-        filename=ckpt_cfg.filename,
-        save_top_k=ckpt_cfg.save_top_k,
-        monitor=ckpt_cfg.monitor,
-        mode=ckpt_cfg.mode,
-        verbose=True,
-        save_weights_only=save_weights_only,
-    )
-    return checkpoint_callback