kostyl-toolkit 0.1.10__tar.gz → 0.1.12__tar.gz

{kostyl_toolkit-0.1.10 → kostyl_toolkit-0.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.10
+Version: 0.1.12
 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.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/clearml/pulling_utils.py

@@ -3,12 +3,13 @@ from typing import cast
 
 from clearml import InputModel
 from clearml import Task
+from transformers import Any
 from transformers import AutoModel
 from transformers import AutoTokenizer
 from transformers import PreTrainedModel
 from transformers import PreTrainedTokenizerBase
 
-from kostyl.ml_core.lightning.extenstions.pretrained_model import (
+from kostyl.ml.lightning.extenstions.pretrained_model import (
     LightningCheckpointLoaderMixin,
 )
 
@@ -48,6 +49,7 @@ def get_model_from_clearml[
     model: type[TModel],
     task: Task | None = None,
     ignore_remote_overrides: bool = True,
+    **kwargs: Any,
 ) -> TModel:
     """
     Retrieve a pretrained model from ClearML and instantiate it using the appropriate loader.
@@ -59,6 +61,7 @@ def get_model_from_clearml[
             will be connected to this task, with remote overrides optionally ignored.
         ignore_remote_overrides: When connecting the input model to the provided task, determines whether
             remote configuration overrides should be ignored.
+        **kwargs: Additional keyword arguments to pass to the model loading method.
 
     Returns:
         An instantiated model loaded either from a ClearML package directory or a Lightning checkpoint.
@@ -72,13 +75,13 @@ def get_model_from_clearml[
     local_path = Path(input_model.get_local_copy(raise_on_error=True))
 
     if local_path.is_dir() and input_model._is_package():
-        model_instance = model.from_pretrained(local_path)
+        model_instance = model.from_pretrained(local_path, **kwargs)
     elif local_path.suffix == ".ckpt":
         if not issubclass(model, LightningCheckpointLoaderMixin):
             raise ValueError(
                 f"Model class {model.__name__} is not compatible with Lightning checkpoints."
             )
-        model_instance = model.from_lighting_checkpoint(local_path)
+        model_instance = model.from_lighting_checkpoint(local_path, **kwargs)
     else:
         raise ValueError(
             f"Unsupported model format for path: {local_path}. "

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/configs/__init__.py

@@ -4,7 +4,6 @@ from .hyperparams import Lr
 from .hyperparams import Optimizer
 from .hyperparams import WeightDecay
 from .training_settings import CheckpointConfig
-from .training_settings import ClearMLTrainingSettings
 from .training_settings import DataConfig
 from .training_settings import DDPStrategyConfig
 from .training_settings import EarlyStoppingConfig
@@ -16,7 +15,6 @@ from .training_settings import TrainingSettings
 
 __all__ = [
     "CheckpointConfig",
-    "ClearMLTrainingSettings",
     "DDPStrategyConfig",
     "DataConfig",
     "EarlyStoppingConfig",

kostyl_toolkit-0.1.10/kostyl/ml_core/clearml/config_mixin.py → kostyl_toolkit-0.1.12/kostyl/ml/configs/base_model.py

@@ -1,29 +1,75 @@
 from pathlib import Path
+from typing import Self
 from typing import TypeVar
 
-import clearml
 from caseconverter import pascalcase
 from caseconverter import snakecase
+from clearml import Task
+from pydantic import BaseModel as PydanticBaseModel
 
-from kostyl.ml_core.configs.base_model import KostylBaseModel
 from kostyl.utils.dict_manipulations import convert_to_flat_dict
 from kostyl.utils.dict_manipulations import flattened_dict_to_nested
 from kostyl.utils.fs import load_config
 
 
-TModel = TypeVar("TModel", bound="KostylBaseModel")
+TConfig = TypeVar("TConfig", bound=PydanticBaseModel)
 
 
-class ClearMLConfigMixin:
-    """Pydantic mixin class providing ClearML configuration loading and syncing functionality."""
+class BaseModelWithConfigLoading(PydanticBaseModel):
+    """Pydantic class providing basic configuration loading functionality."""
+
+    @classmethod
+    def from_file(
+        cls: type[Self],  # pyright: ignore
+        path: str | Path,
+    ) -> Self:
+        """
+        Create an instance of the class from a configuration file.
+
+        Args:
+            cls_: The class type to instantiate.
+            path (str | Path): Path to the configuration file.
+
+        Returns:
+            An instance of the class created from the configuration file.
+
+        """
+        config = load_config(path)
+        instance = cls.model_validate(config)
+        return instance
+
+    @classmethod
+    def from_dict(
+        cls: type[Self],  # pyright: ignore
+        state_dict: dict,
+    ) -> Self:
+        """
+        Creates an instance from a dictionary.
+
+        Args:
+            cls_: The class type to instantiate.
+            state_dict (dict): A dictionary representing the state of the
+                class that must be validated and used for initialization.
+
+        Returns:
+            An initialized instance of the class based on the
+                provided state dictionary.
+
+        """
+        instance = cls.model_validate(state_dict)
+        return instance
+
+
+class BaseModelWithClearmlSyncing(BaseModelWithConfigLoading):
+    """Pydantic class providing ClearML configuration loading and syncing functionality."""
 
     @classmethod
     def connect_as_file(
-        cls: type[TModel],  # pyright: ignore
-        task: clearml.Task,
+        cls: type[Self],  # pyright: ignore
+        task: Task,
         path: str | Path,
         alias: str | None = None,
-    ) -> TModel:
+    ) -> Self:
         """
         Connects the configuration file to a ClearML task and creates an instance of the class from it.
 
@@ -58,11 +104,11 @@ class ClearMLConfigMixin:
 
     @classmethod
     def connect_as_dict(
-        cls: type[TModel],  # pyright: ignore
-        task: clearml.Task,
+        cls: type[Self],  # pyright: ignore
+        task: Task,
         path: str | Path,
         alias: str | None = None,
-    ) -> TModel:
+    ) -> Self:
         """
         Connects configuration from a file as a dictionary to a ClearML task and creates an instance of the class.
 
@@ -89,3 +135,9 @@ class ClearMLConfigMixin:
 
         model = cls.from_dict(state_dict=config)
         return model
+
+
+class KostylBaseModel(BaseModelWithClearmlSyncing):
+    """A Pydantic model class with basic configuration loading functionality."""
+
+    pass

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/configs/hyperparams.py

@@ -2,9 +2,10 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import model_validator
 
-from kostyl.ml_core.clearml.config_mixin import ClearMLConfigMixin
 from kostyl.utils.logging import setup_logger
 
+from .base_model import KostylBaseModel
+
 
 logger = setup_logger(fmt="only_message")
 
@@ -74,7 +75,7 @@ class WeightDecay(BaseModel):
         return self
 
 
-class HyperparamsConfig(BaseModel, ClearMLConfigMixin):
+class HyperparamsConfig(KostylBaseModel):
     """Model training hyperparameters configuration."""
 
     grad_clip_val: float | None = Field(default=None, gt=0, validate_default=False)

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/configs/training_settings.py

@@ -3,7 +3,6 @@ from typing import Literal
 from pydantic import BaseModel
 from pydantic import Field
 
-from kostyl.ml_core.clearml.config_mixin import ClearMLConfigMixin
 from kostyl.utils.logging import setup_logger
 
 from .base_model import KostylBaseModel
@@ -103,13 +102,3 @@ class TrainingSettings(KostylBaseModel):
     early_stopping: EarlyStoppingConfig | None = None
     checkpoint: CheckpointConfig
     data: DataConfig
-
-
-class ClearMLTrainingSettings(
-    TrainingSettings,
-    ClearMLConfigMixin,
-):
-    """Training parameters configuration with ClearML features support (config syncing, model identifiers tracking and etc)."""
-
-    model_id: str
-    tokenizer_id: str

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/lightning/callbacks/checkpoint.py

@@ -3,8 +3,8 @@ from shutil import rmtree
 
 from lightning.pytorch.callbacks import ModelCheckpoint
 
-from kostyl.ml_core.configs import CheckpointConfig
-from kostyl.ml_core.dist_utils import is_main_process
+from kostyl.ml.configs import CheckpointConfig
+from kostyl.ml.dist_utils import is_main_process
 from kostyl.utils import setup_logger
 
 

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/lightning/callbacks/early_stopping.py

@@ -1,6 +1,6 @@
 from lightning.pytorch.callbacks import EarlyStopping
 
-from kostyl.ml_core.configs import EarlyStoppingConfig
+from kostyl.ml.configs import EarlyStoppingConfig
 
 
 def setup_early_stopping_callback(

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/lightning/callbacks/registry_uploading.py

@@ -7,9 +7,9 @@ from lightning import Trainer
 from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.callbacks import ModelCheckpoint
 
-from kostyl.ml_core.clearml.logging_utils import find_version_in_tags
-from kostyl.ml_core.clearml.logging_utils import increment_version
-from kostyl.ml_core.lightning import KostylLightningModule
+from kostyl.ml.clearml.logging_utils import find_version_in_tags
+from kostyl.ml.clearml.logging_utils import increment_version
+from kostyl.ml.lightning import KostylLightningModule
 from kostyl.utils.logging import setup_logger
 
 

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/lightning/extenstions/custom_module.py

@@ -15,9 +15,8 @@ from torchmetrics import MetricCollection
 from transformers import PretrainedConfig
 from transformers import PreTrainedModel
 
-from kostyl.ml_core.configs import HyperparamsConfig
-from kostyl.ml_core.metrics_formatting import apply_suffix
-from kostyl.ml_core.schedulers.base import BaseScheduler
+from kostyl.ml.metrics_formatting import apply_suffix
+from kostyl.ml.schedulers.base import BaseScheduler
 from kostyl.utils import setup_logger
 
 
@@ -28,7 +27,6 @@ class KostylLightningModule(L.LightningModule):
     """Custom PyTorch Lightning Module with logging, checkpointing, and distributed training utilities."""
 
     model: PreTrainedModel | nn.Module | None
-    hyperparams: HyperparamsConfig
 
     def get_process_group(self) -> ProcessGroup | None:
         """
@@ -70,6 +68,11 @@ class KostylLightningModule(L.LightningModule):
             return model.config  # type: ignore
         return None
 
+    @property
+    def grad_clip_val(self) -> float | None:
+        """Returns the gradient clipping value from hyperparameters if set."""
+        raise NotImplementedError
+
     @override
     def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None:
         model = self.get_model()
@@ -93,20 +96,16 @@ class KostylLightningModule(L.LightningModule):
     def on_before_optimizer_step(self, optimizer) -> None:
         if self.model is None:
             raise ValueError("Model must be configured before optimizer step.")
-        if not hasattr(self, "hyperparams"):
-            logger.warning_once("cannot clip gradients, hyperparams attr missing")
-            return
-        if self.hyperparams.grad_clip_val is None:
+
+        grad_clip_val = self.grad_clip_val
+        if grad_clip_val is None:
             return
 
         if not isinstance(self.trainer.strategy, FSDPStrategy):
-            norm = torch.nn.utils.clip_grad_norm_(
-                self.parameters(), self.hyperparams.grad_clip_val
-            )
+            norm = torch.nn.utils.clip_grad_norm_(self.parameters(), grad_clip_val)
         else:
             module: FSDP = self.trainer.strategy.model  # type: ignore
-            norm = module.clip_grad_norm_(self.hyperparams.grad_clip_val)
-
+            norm = module.clip_grad_norm_(grad_clip_val)
         self.log(
             "grad_norm",
             norm,

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/lightning/extenstions/pretrained_model.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from typing import Any
 from typing import cast
 
 import torch
@@ -28,6 +29,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
         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.
@@ -48,6 +50,7 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
             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.
@@ -75,10 +78,17 @@ class LightningCheckpointLoaderMixin(PreTrainedModel):
         )
 
         config_cls = cast(PretrainedConfig, type(cls.config_class))
-        config = config_cls.from_dict(checkpoint_dict[config_key])
+        config_dict = checkpoint_dict[config_key]
+        config_dict.update(kwargs)
+        config = config_cls.from_dict(config_dict)
+
+        kwargs_for_model = {}
+        for key in kwargs:
+            if not hasattr(config, key):
+                kwargs_for_model[key] = kwargs[key]
 
         with torch.device("meta"):
-            model = cls(config)
+            model = cls(config, **kwargs_for_model)
 
             if "peft_config" in checkpoint_dict:
                 if PeftConfig is None:

kostyl_toolkit-0.1.12/kostyl/ml/lightning/loggers/__init__.py

@@ -0,0 +1,4 @@
+from .tb_logger import setup_tb_logger
+
+
+__all__ = ["setup_tb_logger"]

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.12/kostyl/ml}/lightning/loggers/tb_logger.py

@@ -3,7 +3,7 @@ from shutil import rmtree
 
 from lightning.pytorch.loggers import TensorBoardLogger
 
-from kostyl.ml_core.dist_utils import is_main_process
+from kostyl.ml.dist_utils import is_main_process
 from kostyl.utils.logging import setup_logger
 
 

{kostyl_toolkit-0.1.10 → kostyl_toolkit-0.1.12}/pyproject.toml

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

kostyl_toolkit-0.1.10/kostyl/ml_core/configs/base_model.py

@@ -1,60 +0,0 @@
-from pathlib import Path
-from typing import TypeVar
-
-from pydantic import BaseModel as PydanticBaseModel
-
-from kostyl.utils.fs import load_config
-
-
-TConfig = TypeVar("TConfig", bound=PydanticBaseModel)
-
-
-class ConfigLoadingMixin:
-    """Pydantic mixin class providing basic configuration loading functionality."""
-
-    @classmethod
-    def from_file(
-        cls: type[TConfig],  # pyright: ignore
-        path: str | Path,
-    ) -> TConfig:
-        """
-        Create an instance of the class from a configuration file.
-
-        Args:
-            cls_: The class type to instantiate.
-            path (str | Path): Path to the configuration file.
-
-        Returns:
-            An instance of the class created from the configuration file.
-
-        """
-        config = load_config(path)
-        instance = cls.model_validate(config)
-        return instance
-
-    @classmethod
-    def from_dict(
-        cls: type[TConfig],  # pyright: ignore
-        state_dict: dict,
-    ) -> TConfig:
-        """
-        Creates an instance from a dictionary.
-
-        Args:
-            cls_: The class type to instantiate.
-            state_dict (dict): A dictionary representing the state of the
-                class that must be validated and used for initialization.
-
-        Returns:
-            An initialized instance of the class based on the
-                provided state dictionary.
-
-        """
-        instance = cls.model_validate(state_dict)
-        return instance
-
-
-class KostylBaseModel(PydanticBaseModel, ConfigLoadingMixin):
-    """A Pydantic model class with basic configuration loading functionality."""
-
-    pass