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

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

@@ -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 → kostyl_toolkit-0.1.24}/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_toolkit-0.1.22 → kostyl_toolkit-0.1.24}/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.24/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.22/kostyl/ml/lightning/callbacks/registry_uploading.py → kostyl_toolkit-0.1.24/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,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_toolkit-0.1.24/kostyl/ml/lightning/extenstions/pretrained_model.py

@@ -0,0 +1,105 @@
+from pathlib import Path
+from typing import Any
+from typing import cast
+
+import torch
+from transformers import PretrainedConfig
+from transformers import PreTrainedModel
+
+from kostyl.utils.logging import setup_logger
+
+
+logger = setup_logger("LightningPretrainedModelMixin", fmt="only_message")
+
+
+class LightningCheckpointLoaderMixin(PreTrainedModel):
+    """A mixin class for loading pretrained models from PyTorch Lightning checkpoints."""
+
+    @classmethod
+    def from_lightning_checkpoint[TModelInstance: LightningCheckpointLoaderMixin](  # noqa: C901
+        cls: type[TModelInstance],
+        checkpoint_path: str | Path,
+        config_key: str = "config",
+        weights_prefix: str = "model.",
+        **kwargs: Any,
+    ) -> TModelInstance:
+        """
+        Load a model from a Lightning checkpoint file.
+
+        This class method loads a pretrained model from a PyTorch Lightning checkpoint file (.ckpt).
+        It extracts the model configuration from the checkpoint, instantiates the model, and loads
+        the state dictionary, handling any incompatible keys.
+
+        Note:
+            The method uses `torch.load` with `weights_only=False` and `mmap=True` for loading.
+            Incompatible keys (missing, unexpected, mismatched) are collected and optionally logged.
+
+        Args:
+            cls (type["LightningPretrainedModelMixin"]): The class of the model to instantiate.
+            checkpoint_path (str | Path): Path to the checkpoint file. Must be a .ckpt file.
+            config_key (str, optional): Key in the checkpoint dictionary where the config is stored.
+                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.
+            kwargs: Additional keyword arguments to pass to the model's `from_pretrained` method.
+
+        Returns:
+            TModelInstance: The loaded model instance.
+
+        Raises:
+            ValueError: If checkpoint_path is a directory, not a .ckpt file, or invalid.
+            FileNotFoundError: If the checkpoint file does not exist.
+
+        """
+        if isinstance(checkpoint_path, str):
+            checkpoint_path = Path(checkpoint_path)
+
+        if checkpoint_path.is_dir():
+            raise ValueError(f"{checkpoint_path} is a directory")
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(f"{checkpoint_path} does not exist")
+        if checkpoint_path.suffix != ".ckpt":
+            raise ValueError(f"{checkpoint_path} is not a .ckpt file")
+
+        checkpoint_dict = torch.load(
+            checkpoint_path,
+            map_location="cpu",
+            weights_only=False,
+            mmap=True,
+        )
+
+        # 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: dict[str, Any] = {}
+        for key, value in kwargs.items():
+            if not hasattr(config, key):
+                kwargs_for_model[key] = value
+
+        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] = {}
+
+            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:
+                    state_dict[key] = value
+        else:
+            state_dict = raw_state_dict
+
+        model = cls.from_pretrained(
+            pretrained_model_name_or_path=None,
+            config=config,
+            state_dict=state_dict,
+            **kwargs_for_model,
+        )
+
+        return model

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

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

kostyl_toolkit-0.1.22/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

kostyl_toolkit-0.1.22/kostyl/ml/lightning/extenstions/pretrained_model.py

@@ -1,202 +0,0 @@
-from pathlib import Path
-from typing import Any
-from typing import cast
-
-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")
-
-
-class LightningCheckpointLoaderMixin(PreTrainedModel):
-    """A mixin class for loading pretrained models from PyTorch Lightning checkpoints."""
-
-    @classmethod
-    def from_lighting_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:
-        """
-        Load a model from a Lightning checkpoint file.
-
-        This class method loads a pretrained model from a PyTorch Lightning checkpoint file (.ckpt).
-        It extracts the model configuration from the checkpoint, instantiates the model, and loads
-        the state dictionary, handling any incompatible keys.
-
-        Note:
-            The method uses `torch.load` with `weights_only=False` and `mmap=True` for loading.
-            Incompatible keys (missing, unexpected, mismatched) are collected and optionally logged.
-
-        Args:
-            cls (type["LightningPretrainedModelMixin"]): The class of the model to instantiate.
-            checkpoint_path (str | Path): Path to the checkpoint file. Must be a .ckpt file.
-            config_key (str, optional): Key in the checkpoint dictionary where the config is stored.
-                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.
-
-        Returns:
-            TModelInstance: The loaded model instance.
-
-        Raises:
-            ValueError: If checkpoint_path is a directory, not a .ckpt file, or invalid.
-            FileNotFoundError: If the checkpoint file does not exist.
-
-        """
-        if isinstance(checkpoint_path, str):
-            checkpoint_path = Path(checkpoint_path)
-
-        if checkpoint_path.is_dir():
-            raise ValueError(f"{checkpoint_path} is a directory")
-        if not checkpoint_path.exists():
-            raise FileNotFoundError(f"{checkpoint_path} does not exist")
-        if checkpoint_path.suffix != ".ckpt":
-            raise ValueError(f"{checkpoint_path} is not a .ckpt file")
-
-        checkpoint_dict = torch.load(
-            checkpoint_path,
-            map_location="cpu",
-            weights_only=False,
-            mmap=True,
-        )
-
-        # 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: dict[str, Any] = {}
-        for key, value in kwargs.items():
-            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
-        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
-        )
-        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