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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.10
+Version: 0.1.11
 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.11/kostyl/ml}/clearml/config_mixin.py

@@ -5,7 +5,7 @@ import clearml
 from caseconverter import pascalcase
 from caseconverter import snakecase
 
-from kostyl.ml_core.configs.base_model import KostylBaseModel
+from kostyl.ml.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

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.11/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.11/kostyl/ml}/configs/hyperparams.py

@@ -2,7 +2,7 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import model_validator
 
-from kostyl.ml_core.clearml.config_mixin import ClearMLConfigMixin
+from kostyl.ml.clearml.config_mixin import ClearMLConfigMixin
 from kostyl.utils.logging import setup_logger
 
 

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

@@ -3,7 +3,7 @@ from typing import Literal
 from pydantic import BaseModel
 from pydantic import Field
 
-from kostyl.ml_core.clearml.config_mixin import ClearMLConfigMixin
+from kostyl.ml.clearml.config_mixin import ClearMLConfigMixin
 from kostyl.utils.logging import setup_logger
 
 from .base_model import KostylBaseModel

{kostyl_toolkit-0.1.10/kostyl/ml_core → kostyl_toolkit-0.1.11/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.11/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.11/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.11/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.11/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.11/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.11/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.11}/pyproject.toml

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