kostyl-toolkit 0.1.39__py3-none-any.whl → 0.1.41__py3-none-any.whl

kostyl/ml/configs/__init__.py

@@ -1,6 +1,10 @@
-from .hyperparams import OPTIMIZER
+from .hyperparams import OPTIMIZER_CONFIG
+from .hyperparams import AdamConfig
+from .hyperparams import AdamWithPrecisionConfig
 from .hyperparams import HyperparamsConfig
 from .hyperparams import Lr
+from .hyperparams import MuonConfig
+from .hyperparams import ScheduledParamConfig
 from .hyperparams import WeightDecay
 from .mixins import ConfigLoadingMixin
 from .training_settings import CheckpointConfig
@@ -14,7 +18,9 @@ from .training_settings import TrainingSettings
 
 
 __all__ = [
-    "OPTIMIZER",
+    "OPTIMIZER_CONFIG",
+    "AdamConfig",
+    "AdamWithPrecisionConfig",
     "CheckpointConfig",
     "ConfigLoadingMixin",
     "DDPStrategyConfig",
@@ -24,6 +30,8 @@ __all__ = [
     "HyperparamsConfig",
     "LightningTrainerParameters",
     "Lr",
+    "MuonConfig",
+    "ScheduledParamConfig",
     "SingleDeviceStrategyConfig",
     "TrainingSettings",
     "WeightDecay",

kostyl/ml/configs/hyperparams.py

@@ -21,6 +21,7 @@ class MuonConfig(BaseModel):
     """Muon optimizer hyperparameters configuration."""
 
     type: Literal["Muon"]
+    momentum: float = 0.95
     nesterov: bool = True
     ns_coefficients: tuple[float, float, float] = (3.4445, -4.7750, 2.0315)
     ns_steps: int = 5
@@ -29,14 +30,15 @@ class MuonConfig(BaseModel):
 class AdamWithPrecisionConfig(BaseModel):
     """Adam optimizer with low-precision hyperparameters configuration."""
 
-    type: Literal["Adam8bit", "Adam4bit", "AdamFp8"]
+    type: Literal[
+        "Adam8bit", "Adam4bit", "AdamFp8", "AdamW8bit", "AdamW4bit", "AdamWFp8"
+    ]
     betas: tuple[float, float] = (0.9, 0.999)
     block_size: int
     bf16_stochastic_round: bool = False
-    is_adamw: bool = True
 
 
-OPTIMIZER = AdamConfig | AdamWithPrecisionConfig | MuonConfig
+OPTIMIZER_CONFIG = AdamConfig | AdamWithPrecisionConfig | MuonConfig
 SCHEDULER = Literal[
     "linear",
     "cosine",
@@ -45,8 +47,8 @@ SCHEDULER = Literal[
 ]
 
 
-class Lr(BaseModel):
-    """Learning rate hyperparameters configuration."""
+class ScheduledParamConfig(BaseModel):
+    """Base configuration for a scheduled hyperparameter."""
 
     scheduler_type: SCHEDULER | None = None
 
@@ -60,14 +62,14 @@ class Lr(BaseModel):
     )
 
     @model_validator(mode="after")
-    def _validate_freeze_ratio(self) -> "Lr":
+    def _validate_freeze_ratio(self) -> "ScheduledParamConfig":
         if self.scheduler_type is None and self.freeze_ratio is not None:
             logger.warning("use_scheduler is False, freeze_ratio will be ignored.")
             self.freeze_ratio = None
         return self
 
     @model_validator(mode="after")
-    def _validate_warmup(self) -> "Lr":
+    def _validate_warmup(self) -> "ScheduledParamConfig":
         if ((self.warmup_value is not None) or (self.warmup_ratio is not None)) and self.scheduler_type is None:  # fmt: skip
             logger.warning(
                 "scheduler_type is None, warmup_value and warmup_ratio will be ignored."
@@ -81,7 +83,7 @@ class Lr(BaseModel):
         return self
 
     @model_validator(mode="after")
-    def _validate_final_value(self) -> "Lr":
+    def _validate_final_value(self) -> "ScheduledParamConfig":
         if (self.scheduler_type in {"linear"}) and (self.final_value is not None):
             raise ValueError("If scheduler_type is 'linear', final_value must be None.")
         if (self.scheduler_type is None) and (self.final_value is not None):
@@ -90,7 +92,7 @@ class Lr(BaseModel):
         return self
 
     @model_validator(mode="after")
-    def _validate_plateau_ratio(self) -> "Lr":
+    def _validate_plateau_ratio(self) -> "ScheduledParamConfig":
         if self.scheduler_type is not None:
             if self.scheduler_type.startswith("plateau") and self.plateau_ratio is None:
                 raise ValueError(
@@ -107,7 +109,11 @@ class Lr(BaseModel):
         return self
 
 
-class WeightDecay(Lr):
+class Lr(ScheduledParamConfig):
+    """Learning rate hyperparameters configuration."""
+
+
+class WeightDecay(ScheduledParamConfig):
     """Weight decay hyperparameters configuration."""
 
 
@@ -115,6 +121,6 @@ class HyperparamsConfig(BaseModel):
     """Model training hyperparameters configuration."""
 
     grad_clip_val: float | None = Field(default=None, gt=0, validate_default=False)
-    optimizer: OPTIMIZER
+    optimizer: OPTIMIZER_CONFIG
     lr: Lr
     weight_decay: WeightDecay

kostyl/ml/integrations/clearml/__init__.py

@@ -5,3 +5,25 @@ except ImportError as e:
         "ClearML integration requires the 'clearml' package. "
         "Please install it via 'pip install clearml'."
     ) from e
+from .checkpoint_uploader import ClearMLCheckpointUploader
+from .config_mixin import ConfigSyncingClearmlMixin
+from .dataset_utils import collect_clearml_datasets
+from .dataset_utils import download_clearml_datasets
+from .dataset_utils import get_datasets_paths
+from .loading_utils import load_model_from_clearml
+from .loading_utils import load_tokenizer_from_clearml
+from .version_utils import find_version_in_tags
+from .version_utils import increment_version
+
+
+__all__ = [
+    "ClearMLCheckpointUploader",
+    "ConfigSyncingClearmlMixin",
+    "collect_clearml_datasets",
+    "download_clearml_datasets",
+    "find_version_in_tags",
+    "get_datasets_paths",
+    "increment_version",
+    "load_model_from_clearml",
+    "load_tokenizer_from_clearml",
+]

kostyl/ml/integrations/clearml/checkpoint_uploader.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from functools import partial
+from datetime import datetime
 from pathlib import Path
 from typing import override
 
@@ -24,7 +24,7 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
         comment: str | None = None,
         framework: str | None = None,
         base_model_id: str | None = None,
-        new_model_per_upload: bool = True,
+        upload_as_new_model: bool = True,
         verbose: bool = True,
     ) -> None:
         """
@@ -38,20 +38,22 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
             comment: A comment / description for the model.
             framework: The framework of the model (e.g., "PyTorch", "TensorFlow").
             base_model_id: Optional ClearML model ID to use as a base for the new model
-            new_model_per_upload: Whether to create a new ClearML model
-                for every upload or update weights of the same model. When updating weights,
-                the last uploaded checkpoint will be replaced (and deleted).
+            upload_as_new_model: Whether to create a new ClearML model
+                for every upload or update weights of the same model. When True,
+                each checkpoint is uploaded as a separate model with timestamp added to the name.
+                When False, weights of the same model are updated.
             verbose: Whether to log messages during upload.
 
         """
         super().__init__()
-        if base_model_id is not None and new_model_per_upload:
+        if base_model_id is not None and upload_as_new_model:
             raise ValueError(
-                "Cannot set base_model_id when new_model_per_upload is True."
+                "Cannot set base_model_id when upload_as_new_model is True."
             )
 
         self.verbose = verbose
-        self.new_model_per_upload = new_model_per_upload
+        self.upload_as_new_model = upload_as_new_model
+        self.model_name = model_name
         self.best_model_path: str = ""
         self.config_dict = config_dict
         self._output_model: OutputModel | None = None
@@ -59,15 +61,13 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
         self._upload_callback: Callable | None = None
 
         self._validate_tags(tags)
-        self.model_fabric = partial(
-            OutputModel,
-            name=model_name,
-            label_enumeration=label_enumeration,
-            tags=tags,
-            comment=comment,
-            framework=framework,
-            base_model_id=base_model_id,
-        )
+        self.model_fabric_kwargs = {
+            "label_enumeration": label_enumeration,
+            "tags": tags,
+            "comment": comment,
+            "framework": framework,
+            "base_model_id": base_model_id,
+        }
         return
 
     @staticmethod
@@ -78,16 +78,22 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
             tags.append("LightningCheckpoint")
         return None
 
-    @property
-    def output_model_(self) -> OutputModel:
-        """Returns the OutputModel instance based on `new_model_per_upload` setting."""
-        if self.new_model_per_upload:
-            model = self.model_fabric()
-            self._output_model = self.model_fabric()
-        else:
-            if self._output_model is None:
-                self._output_model = self.model_fabric()
-            model = self._output_model
+    def _create_new_model(self) -> OutputModel:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        model_name_with_timestamp = f"{self.model_name}_{timestamp}"
+        model = OutputModel(
+            name=model_name_with_timestamp,
+            **self.model_fabric_kwargs,
+        )
+        return model
+
+    def _get_output_model(self) -> OutputModel:
+        if self._output_model is None:
+            self._output_model = OutputModel(
+                name=self.model_name,
+                **self.model_fabric_kwargs,
+            )
+        model = self._output_model
         return model
 
     @override
@@ -105,12 +111,17 @@ class ClearMLCheckpointUploader(ModelCheckpointUploader):
         if self.verbose:
             logger.info(f"Uploading model from {path}")
 
-        self.output_model_.update_weights(
+        if self.upload_as_new_model:
+            output_model = self._create_new_model()
+        else:
+            output_model = self._get_output_model()
+
+        output_model.update_weights(
             path,
             auto_delete_file=False,
             async_enable=False,
         )
-        self.output_model_.update_design(config_dict=self.config_dict)
+        output_model.update_design(config_dict=self.config_dict)
 
         self._last_uploaded_model_path = path
         return

kostyl/ml/integrations/clearml/config_mixin.py

@@ -10,7 +10,7 @@ from kostyl.utils.dict_manipulations import flattened_dict_to_nested
 from kostyl.utils.fs import load_config
 
 
-class BaseModelWithClearmlSyncing[TConfig: ConfigLoadingMixin]:
+class ConfigSyncingClearmlMixin[TConfig: ConfigLoadingMixin]:
     """Mixin providing ClearML task configuration syncing functionality for Pydantic models."""
 
     @classmethod

kostyl/ml/integrations/clearml/loading_utils.py

@@ -31,7 +31,7 @@ except ImportError:
     LIGHTING_MIXIN_AVAILABLE = False
 
 
-def get_tokenizer_from_clearml(
+def load_tokenizer_from_clearml(
     model_id: str,
     task: Task | None = None,
     ignore_remote_overrides: bool = True,
@@ -66,7 +66,7 @@ def get_tokenizer_from_clearml(
     return tokenizer, clearml_tokenizer
 
 
-def get_model_from_clearml[
+def load_model_from_clearml[
     TModel: PreTrainedModel | LightningCheckpointLoaderMixin | AutoModel
 ](
     model_id: str,

kostyl/ml/integrations/lightning/callbacks/checkpoint.py

@@ -15,7 +15,7 @@ from kostyl.ml.dist_utils import is_local_zero_rank
 from kostyl.utils import setup_logger
 
 
-logger = setup_logger("callbacks/checkpoint.py")
+logger = setup_logger()
 
 
 class ModelCheckpointWithCheckpointUploader(ModelCheckpoint):
@@ -278,6 +278,10 @@ class ModelCheckpointWithCheckpointUploader(ModelCheckpoint):
                 case "only-best":
                     if filepath == self.best_model_path:
                         self.registry_uploader.upload_checkpoint(filepath)
+                case _:
+                    logger.warning_once(
+                        "Unknown upload strategy for checkpoint uploader. Skipping upload."
+                    )
         return
 
 

kostyl/ml/integrations/lightning/module.py

@@ -14,7 +14,7 @@ from transformers import PretrainedConfig
 from transformers import PreTrainedModel
 
 from kostyl.ml.integrations.lightning.metrics_formatting import apply_suffix
-from kostyl.ml.schedulers.base import BaseScheduler
+from kostyl.ml.optim.schedulers import BaseScheduler
 from kostyl.utils import setup_logger
 
 

kostyl/ml/optim/__init__.py

@@ -0,0 +1,8 @@
+from .factory import create_optimizer
+from .factory import create_scheduler
+
+
+__all__ = [
+    "create_optimizer",
+    "create_scheduler",
+]

kostyl/ml/optim/factory.py

@@ -0,0 +1,257 @@
+from typing import Any
+
+from torch.optim import Optimizer
+
+from kostyl.ml.configs import OPTIMIZER_CONFIG
+from kostyl.ml.configs import AdamConfig
+from kostyl.ml.configs import AdamWithPrecisionConfig
+from kostyl.ml.configs import MuonConfig
+from kostyl.ml.configs import ScheduledParamConfig
+from kostyl.utils import setup_logger
+
+from .schedulers import SCHEDULER_MAPPING
+from .schedulers import CosineScheduler
+from .schedulers import LinearScheduler
+from .schedulers import PlateauWithAnnealingScheduler
+
+
+logger = setup_logger(fmt="only_message")
+
+
+def create_scheduler(
+    config: ScheduledParamConfig,
+    param_group_field: str,
+    num_iters: int,
+    optim: Optimizer,
+    multiplier_field: str | None = None,
+    skip_if_zero: bool = False,
+    apply_if_field: str | None = None,
+    ignore_if_field: str | None = None,
+) -> LinearScheduler | CosineScheduler | PlateauWithAnnealingScheduler:
+    """
+    Converts a ScheduledParamConfig to a scheduler instance.
+
+    Args:
+        config: Configuration object for the scheduler.
+        param_group_field: The field name in the optimizer's param groups to schedule.
+        num_iters: Total number of iterations.
+        optim: The optimizer instance.
+        multiplier_field: Optional per-group field name that contains a multiplier applied to the scheduled value. If None, no multiplier is applied.
+        skip_if_zero: Leave groups untouched when their target field equals zero.
+            Default is False.
+        apply_if_field: Require this key to be present in a param group before updating.
+        ignore_if_field: Skip groups that declare this key in their dictionaries.
+
+    Returns:
+        A scheduler instance based on the configuration.
+
+    """
+    if config.scheduler_type is None:
+        raise ValueError("scheduler_type must be specified in the config.")
+
+    if "plateau" in config.scheduler_type:
+        scheduler_type = "plateau"
+    else:
+        scheduler_type = config.scheduler_type
+    scheduler_cls = SCHEDULER_MAPPING[scheduler_type]  # type: ignore
+
+    if issubclass(scheduler_cls, PlateauWithAnnealingScheduler):
+        if "cosine" in config.scheduler_type:
+            annealing_type = "cosine"
+        elif "linear" in config.scheduler_type:
+            annealing_type = "linear"
+        else:
+            raise ValueError(f"Unknown annealing_type: {config.scheduler_type}")
+        scheduler = scheduler_cls(
+            optimizer=optim,
+            param_group_field=param_group_field,
+            num_iters=num_iters,
+            plateau_value=config.base_value,
+            final_value=config.final_value,  # type: ignore
+            warmup_ratio=config.warmup_ratio,
+            warmup_value=config.warmup_value,
+            freeze_ratio=config.freeze_ratio,
+            plateau_ratio=config.plateau_ratio,  # type: ignore
+            annealing_type=annealing_type,
+            multiplier_field=multiplier_field,
+            skip_if_zero=skip_if_zero,
+            apply_if_field=apply_if_field,
+            ignore_if_field=ignore_if_field,
+        )
+    elif issubclass(scheduler_cls, LinearScheduler):
+        scheduler = scheduler_cls(
+            optimizer=optim,
+            param_group_field=param_group_field,
+            num_iters=num_iters,
+            initial_value=config.base_value,
+            final_value=config.final_value,  # type: ignore
+            multiplier_field=multiplier_field,
+            skip_if_zero=skip_if_zero,
+            apply_if_field=apply_if_field,
+            ignore_if_field=ignore_if_field,
+        )
+    elif issubclass(scheduler_cls, CosineScheduler):
+        scheduler = scheduler_cls(
+            optimizer=optim,
+            param_group_field=param_group_field,
+            num_iters=num_iters,
+            base_value=config.base_value,
+            final_value=config.final_value,  # type: ignore
+            warmup_ratio=config.warmup_ratio,
+            warmup_value=config.warmup_value,
+            freeze_ratio=config.freeze_ratio,
+            multiplier_field=multiplier_field,
+            skip_if_zero=skip_if_zero,
+            apply_if_field=apply_if_field,
+            ignore_if_field=ignore_if_field,
+        )
+    else:
+        raise ValueError(f"Unsupported scheduler type: {config.scheduler_type}")
+    return scheduler
+
+
+def create_optimizer(  # noqa: C901
+    parameters_groups: dict[str, Any],
+    optimizer_config: OPTIMIZER_CONFIG,
+    lr: float,
+    weight_decay: float,
+) -> Optimizer:
+    """
+    Creates an optimizer based on the configuration.
+
+    Args:
+        parameters_groups: Dictionary containing model parameters
+            (key "params" and per-group options, i.e. "lr", "weight_decay" and etc.).
+        optimizer_config: Configuration for the optimizer.
+        lr: Learning rate.
+        weight_decay: Weight decay.
+
+    Returns:
+        An instantiated optimizer.
+
+    """
+    if isinstance(optimizer_config, AdamConfig):
+        match optimizer_config.type:
+            case "Adam":
+                from torch.optim import Adam
+
+                optimizer = Adam(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                )
+
+            case "AdamW":
+                from torch.optim import AdamW
+
+                optimizer = AdamW(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                )
+                return optimizer
+            case _:
+                raise ValueError(f"Unsupported optimizer type: {optimizer_config.type}")
+    elif isinstance(optimizer_config, MuonConfig):
+        from torch.optim import Muon
+
+        optimizer = Muon(
+            params=parameters_groups["params"],
+            lr=lr,
+            weight_decay=weight_decay,
+            momentum=optimizer_config.momentum,
+            nesterov=optimizer_config.nesterov,
+            ns_coefficients=optimizer_config.ns_coefficients,
+            ns_steps=optimizer_config.ns_steps,
+        )
+    elif isinstance(optimizer_config, AdamWithPrecisionConfig):
+        try:
+            import torchao  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "torchao is required for low-precision Adam optimizers. "
+                "Please install it via 'pip install torchao'."
+            ) from e
+        match optimizer_config.type:
+            case "Adam8bit":
+                from torchao.optim import Adam8bit
+
+                logger.warning(
+                    "Ignoring weight_decay for Adam8bit optimizer as it is not supported."
+                )
+
+                optimizer = Adam8bit(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "Adam4bit":
+                from torchao.optim import Adam4bit
+
+                logger.warning(
+                    "Ignoring weight_decay for Adam4bit optimizer as it is not supported."
+                )
+
+                optimizer = Adam4bit(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "AdamFp8":
+                from torchao.optim import AdamFp8
+
+                logger.warning(
+                    "Ignoring weight_decay for AdamFp8 optimizer as it is not supported."
+                )
+
+                optimizer = AdamFp8(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "AdamW8bit":
+                from torchao.optim import AdamW8bit
+
+                optimizer = AdamW8bit(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "AdamW4bit":
+                from torchao.optim import AdamW4bit
+
+                optimizer = AdamW4bit(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case "AdamWFp8":
+                from torchao.optim import AdamWFp8
+
+                optimizer = AdamWFp8(
+                    params=parameters_groups["params"],
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    betas=optimizer_config.betas,
+                    block_size=optimizer_config.block_size,
+                    bf16_stochastic_round=optimizer_config.bf16_stochastic_round,
+                )
+            case _:
+                raise ValueError(f"Unsupported optimizer type: {optimizer_config.type}")
+    else:
+        raise ValueError("Unsupported optimizer configuration type.")
+    return optimizer

kostyl/ml/optim/schedulers/__init__.py

@@ -0,0 +1,56 @@
+from typing import TypedDict
+
+from .base import BaseScheduler
+from .composite import CompositeScheduler
+from .cosine import CosineParamScheduler
+from .cosine import CosineScheduler
+from .linear import LinearParamScheduler
+from .linear import LinearScheduler
+from .plateau import PlateauWithAnnealingParamScheduler
+from .plateau import PlateauWithAnnealingScheduler
+
+
+class SchedulerMapping(TypedDict):
+    """Map names to scheduler classes."""
+
+    linear: type[LinearScheduler]
+    cosine: type[CosineScheduler]
+    plateau: type[PlateauWithAnnealingScheduler]
+    composite: type[CompositeScheduler]
+
+
+class ParamSchedulerMapping(TypedDict):
+    """Map names to scheduler classes."""
+
+    linear: type[LinearParamScheduler]
+    cosine: type[CosineParamScheduler]
+    plateau: type[PlateauWithAnnealingParamScheduler]
+
+
+SCHEDULER_MAPPING: SchedulerMapping = {
+    "linear": LinearScheduler,
+    "cosine": CosineScheduler,
+    "plateau": PlateauWithAnnealingScheduler,
+    "composite": CompositeScheduler,
+}
+
+
+PARAM_SCHEDULER_MAPPING: ParamSchedulerMapping = {
+    "linear": LinearParamScheduler,
+    "cosine": CosineParamScheduler,
+    "plateau": PlateauWithAnnealingParamScheduler,
+}
+
+
+__all__ = [
+    "PARAM_SCHEDULER_MAPPING",
+    "SCHEDULER_MAPPING",
+    "BaseScheduler",
+    "CompositeScheduler",
+    "CosineParamScheduler",
+    "CosineScheduler",
+    "LinearParamScheduler",
+    "LinearScheduler",
+    "PlateauWithAnnealingParamScheduler",
+    "PlateauWithAnnealingScheduler",
+]

kostyl/ml/{schedulers → optim/schedulers}/cosine.py

@@ -145,8 +145,8 @@ class CosineScheduler(_CosineSchedulerCore):
             freeze_ratio: Optional fraction of iterations to keep the value frozen at zero at the beginning.
             multiplier_field: Optional per-group multiplier applied to the scheduled value.
             skip_if_zero: Leave groups untouched when their target field equals zero.
-            apply_if_field: Require this flag to be present in a param group before updating.
-            ignore_if_field: Skip groups that declare this flag.
+            apply_if_field: Require this key to be present in a param group before updating.
+            ignore_if_field: Skip groups that declare this key in their dictionaries.
 
         """
         self.apply_if_field = apply_if_field

kostyl/ml/{schedulers → optim/schedulers}/linear.py

@@ -13,21 +13,21 @@ class _LinearScheduleBase(BaseScheduler):
         self,
         param_name: str,
         num_iters: int,
-        base_value: float,
+        initial_value: float,
         final_value: float,
     ) -> None:
         self.param_name = param_name
         self.num_iters = num_iters
-        self.base_value = base_value
+        self.initial_value = initial_value
         self.final_value = final_value
 
         self.scheduled_values: npt.NDArray[np.float64] = np.array([], dtype=np.float64)
-        self.current_value_ = self.base_value
+        self.current_value_ = self.initial_value
         return
 
     def _create_scheduler(self) -> None:
         self.scheduled_values = np.linspace(
-            self.base_value, self.final_value, num=self.num_iters, dtype=np.float64
+            self.initial_value, self.final_value, num=self.num_iters, dtype=np.float64
         )
         self._verify()
         return
@@ -68,7 +68,7 @@ class LinearScheduler(_LinearScheduleBase):
         optimizer: torch.optim.Optimizer,
         param_group_field: str,
         num_iters: int,
-        base_value: float,
+        initial_value: float,
         final_value: float,
         multiplier_field: str | None = None,
         skip_if_zero: bool = False,
@@ -82,12 +82,12 @@ class LinearScheduler(_LinearScheduleBase):
             optimizer: Optimizer whose param groups are updated in-place.
             param_group_field: Name of the field that receives the scheduled value.
             num_iters: Number of scheduler iterations before clamping at ``final_value``.
-            base_value: Value used on the first iteration.
+            initial_value: Value used on the first iteration.
             final_value: Value used once ``num_iters`` iterations are consumed.
             multiplier_field: Optional per-group multiplier applied to the scheduled value.
             skip_if_zero: Leave groups untouched when their target field equals zero.
-            apply_if_field: Require this flag to be present in a param group before updating.
-            ignore_if_field: Skip groups that declare this flag.
+            apply_if_field: Require this key to be present in a param group before updating.
+            ignore_if_field: Skip groups that declare this key in their dictionaries.
 
         """
         self.apply_if_field = apply_if_field
@@ -98,7 +98,7 @@ class LinearScheduler(_LinearScheduleBase):
         super().__init__(
             param_name=param_group_field,
             num_iters=num_iters,
-            base_value=base_value,
+            initial_value=initial_value,
             final_value=final_value,
         )
         self.param_group_field = param_group_field

kostyl/ml/{schedulers → optim/schedulers}/plateau.py

@@ -198,8 +198,8 @@ class PlateauWithAnnealingScheduler(_PlateauWithAnnealingCore):
             annealing_type: Type of annealing from plateau to final value ("cosine" or "linear").
             multiplier_field: Optional per-group multiplier applied to the scheduled value.
             skip_if_zero: Leave groups untouched when their target field equals zero.
-            apply_if_field: Require this flag to be present in a param group before updating.
-            ignore_if_field: Skip groups that declare this flag.
+            apply_if_field: Require this key to be present in a param group before updating.
+            ignore_if_field: Skip groups that declare this key in their dictionaries.
 
         """
         self.apply_if_field = apply_if_field

{kostyl_toolkit-0.1.39.dist-info → kostyl_toolkit-0.1.41.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.39
+Version: 0.1.41
 Summary: Kickass Orchestration System for Training, Yielding & Logging 
 Requires-Dist: case-converter>=1.2.0
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: case-converter>=1.2.0 ; extra == 'ml'
 Requires-Dist: pydantic>=2.12.4 ; extra == 'ml'
-Requires-Dist: torch>=2.9.1 ; extra == 'ml'
-Requires-Dist: transformers>=4.57.1 ; extra == 'ml'
+Requires-Dist: torch ; extra == 'ml'
+Requires-Dist: transformers ; extra == 'ml'
 Requires-Python: >=3.12
 Provides-Extra: ml
 Description-Content-Type: text/markdown

{kostyl_toolkit-0.1.39.dist-info → kostyl_toolkit-0.1.41.dist-info}/RECORD

@@ -1,40 +1,42 @@
 kostyl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kostyl/ml/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kostyl/ml/base_uploader.py,sha256=KxHuohCcNK18kTVFBBqDu_IOQefluhSXOzwC56O66wc,484
-kostyl/ml/configs/__init__.py,sha256=9AbM3ugsvc0bRKHA99DTRuUc2N7iSV9BH5FTDOp0cpw,913
-kostyl/ml/configs/hyperparams.py,sha256=wU6NoZoTG9WlnU4z3g4ViUfjjp8fb8YWUS_2gAd5wsY,4131
+kostyl/ml/configs/__init__.py,sha256=ytr2HFtsqMIs1pqQ-Ma0jv8c6Ni1-UxPnDTLu16cL68,1189
+kostyl/ml/configs/hyperparams.py,sha256=poOsR6xbE6jHO02GzxXfbz-6PA0HqULgcTgiuT_4oOM,4401
 kostyl/ml/configs/mixins.py,sha256=xHHAoRoPbzP9ECFP9duzg6SzegHcoLI8Pr9NrLoWNHs,1411
 kostyl/ml/configs/training_settings.py,sha256=wT9CHuLaKrLwonsc87Ee421EyFis_c9fqOgn9bSClm8,2747
 kostyl/ml/data_collator.py,sha256=kxiaMDKwSKXGBtrF8yXxHcypf7t_6syU-NwO1LcX50k,4062
 kostyl/ml/dist_utils.py,sha256=UFNMLEHc0A5F6KvTRG8GQPpRDwG4m5dvM__UvXNc2aQ,4526
 kostyl/ml/integrations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kostyl/ml/integrations/clearml/__init__.py,sha256=3TBVI-3fE9ZzuvOLEohW9TOK0BZTLD5JiYalAVDkocc,217
-kostyl/ml/integrations/clearml/checkpoint_uploader.py,sha256=PupFi7jKROsIddOz7X5DhV7nUNdDZg5kKaaLvzdCHlY,4012
-kostyl/ml/integrations/clearml/config_mixin.py,sha256=70QRicU7etiDzLX-MplqVX8uFm5siuPrM8KbTOriZnQ,3308
+kostyl/ml/integrations/clearml/__init__.py,sha256=cIJMmiTgRrBNlCwejs6x4DrJydv7zYRPqUCWpVgQJYo,970
+kostyl/ml/integrations/clearml/checkpoint_uploader.py,sha256=hwE2pbi3EK0iYw0BCDTK4KjuxS3lfOWDFq5aqnU7c1U,4381
+kostyl/ml/integrations/clearml/config_mixin.py,sha256=F9Ni6k80KyJsVGcgHjAIIsdpmuDcVqO7eogwiNHccqg,3306
 kostyl/ml/integrations/clearml/dataset_utils.py,sha256=eij_sr2KDhm8GxEbVbK8aBjPsuVvLl9-PIGGaKVgXLA,1729
-kostyl/ml/integrations/clearml/loading_utils.py,sha256=NAMmB9NTGCXCHh-bR_nrQZyqImUVZqicNjExDyPM2mU,5224
+kostyl/ml/integrations/clearml/loading_utils.py,sha256=1wo5QX8cCYC9vGkkupHmaQFirJj4ZUSm_K8zp685DNU,5226
 kostyl/ml/integrations/clearml/version_utils.py,sha256=GBjIIZbH_itd5sj7XpvxjkyZwxxGOpEcQ3BiWaJTyq8,1210
 kostyl/ml/integrations/lightning/__init__.py,sha256=r96os8kTuKIAymx3k9Td1JBrO2PH7nQAWUC54NsY5yY,392
 kostyl/ml/integrations/lightning/callbacks/__init__.py,sha256=EnKkNwwNDZnEqKRlpY4FVrqP88ECPF6nlT2bSLUIKRk,194
-kostyl/ml/integrations/lightning/callbacks/checkpoint.py,sha256=iJQKaQUzCMNza7SUKtCVDol_Cy3rZy9KfBKQ6kT6Swg,18434
+kostyl/ml/integrations/lightning/callbacks/checkpoint.py,sha256=Qpp4LrLloXTIwuM2TDqYi3vvYKn1KZOK_VCi2bs26qY,18588
 kostyl/ml/integrations/lightning/callbacks/early_stopping.py,sha256=D5nyjktCJ9XYAf28-kgXG8jORvXLl1N3nbDQnvValPM,615
 kostyl/ml/integrations/lightning/loggers/__init__.py,sha256=e51dszaoJbuzwBkbdugmuDsPldoSO4yaRgmZUg1Bdy0,71
 kostyl/ml/integrations/lightning/loggers/tb_logger.py,sha256=jVzEArHr1bpO-HRdKjlLJ4BJnxomHPAuBKCdYQRHRn4,1023
 kostyl/ml/integrations/lightning/metrics_formatting.py,sha256=U6vdNENZLvp2dT1L3HqFKtXrHwGKoDXN93hvamPGHjM,1341
 kostyl/ml/integrations/lightning/mixins.py,sha256=hVIsIUu6Iryrz6S7GQTqog9vNq8LQyjJd2aoJ5Ws6KU,5253
-kostyl/ml/integrations/lightning/module.py,sha256=39hcVNZSGyj5tLpXyX8IoqMGWt5vf6-Bx5JnNJ2-Wag,5218
+kostyl/ml/integrations/lightning/module.py,sha256=APodFpspSoBYmrtTnCpUFq3WFEvpmHp2hlWUxwwAChM,5219
 kostyl/ml/integrations/lightning/utils.py,sha256=QhbK5iTv07xkRKuyXiK05aoY_ObSGorzGu1WcFvvFtI,1712
+kostyl/ml/optim/__init__.py,sha256=GgCr94WbiM-7MVDf6fPlbeUSxS33ufn_UAmqfKerElc,140
+kostyl/ml/optim/factory.py,sha256=1WYbAOOGepVL55gSVCJs5WBsiiE4k4ayrS5svFh9kbE,9760
+kostyl/ml/optim/schedulers/__init__.py,sha256=lVwWj_mtUDo5OLtr8NNljLQgWGJUKPJWV7bm6-997ag,1486
+kostyl/ml/optim/schedulers/base.py,sha256=bjmwgdZpnSqpCnHPnKC6MEiRO79cwxMJpZq-eQVNs2M,1353
+kostyl/ml/optim/schedulers/composite.py,sha256=ee4xlMDMMtjKPkbTF2ue9GTr9DuGCGjZWf11mHbi6aE,2387
+kostyl/ml/optim/schedulers/cosine.py,sha256=LnFujAKaM8v5soPwiWG899nUMeaeAe8oTiaJqvoQxg8,8705
+kostyl/ml/optim/schedulers/linear.py,sha256=hlHfU88WLwfke8XGmRUzRKmEpqyyaLb4RRqq6EIzIxU,5881
+kostyl/ml/optim/schedulers/plateau.py,sha256=kjHivSp2FIMXExgS0q_Wbxk9pEHnuzhrIu4TQ59kYNw,11255
 kostyl/ml/params_groups.py,sha256=nUyw5d06Pvy9QPiYtZzLYR87xwXqJLxbHthgQH8oSCM,3583
-kostyl/ml/schedulers/__init__.py,sha256=VIo8MOP4w5Ll24XqFb3QGi2rKvys6c0dEFYPIdDoPlw,526
-kostyl/ml/schedulers/base.py,sha256=bjmwgdZpnSqpCnHPnKC6MEiRO79cwxMJpZq-eQVNs2M,1353
-kostyl/ml/schedulers/composite.py,sha256=ee4xlMDMMtjKPkbTF2ue9GTr9DuGCGjZWf11mHbi6aE,2387
-kostyl/ml/schedulers/cosine.py,sha256=y8ylrgVOkVcr2-ExoqqNW--tdDX88TBYPQCOppIf2_M,8685
-kostyl/ml/schedulers/linear.py,sha256=a-Dq5VYoM1Z7vBb4_2Np0MAOfRmO1QVZnOzEoY8nM6k,5834
-kostyl/ml/schedulers/plateau.py,sha256=N-hiostPtTR0W4xnEJYB_1dv0DRx39iufLkGUrSIoWE,11235
 kostyl/utils/__init__.py,sha256=hkpmB6c5pr4Ti5BshOROebb7cvjDZfNCw83qZ_FFKMM,240
 kostyl/utils/dict_manipulations.py,sha256=e3vBicID74nYP8lHkVTQc4-IQwoJimrbFELy5uSF6Gk,1073
 kostyl/utils/fs.py,sha256=gAQNIU4R_2DhwjgzOS8BOMe0gZymtY1eZwmdgOdDgqo,510
 kostyl/utils/logging.py,sha256=CgNFNogcK0hoZmygvBWlTcq5A3m2Pfv9eOAP_gwx0pM,6633
-kostyl_toolkit-0.1.39.dist-info/WHEEL,sha256=e_m4S054HL0hyR3CpOk-b7Q7fDX6BuFkgL5OjAExXas,80
-kostyl_toolkit-0.1.39.dist-info/METADATA,sha256=9AA-uH9-jSciHx3xh4WGxHkzTyjEESi0AFRitTkIM5w,4136
-kostyl_toolkit-0.1.39.dist-info/RECORD,,
+kostyl_toolkit-0.1.41.dist-info/WHEEL,sha256=e_m4S054HL0hyR3CpOk-b7Q7fDX6BuFkgL5OjAExXas,80
+kostyl_toolkit-0.1.41.dist-info/METADATA,sha256=RoZ8dXiU-CwUldPkhjY2achhvcQ2EE4RzhsJKTXe-AM,4121
+kostyl_toolkit-0.1.41.dist-info/RECORD,,

kostyl/ml/schedulers/__init__.py

@@ -1,18 +0,0 @@
-from .composite import CompositeScheduler
-from .cosine import CosineParamScheduler
-from .cosine import CosineScheduler
-from .linear import LinearParamScheduler
-from .linear import LinearScheduler
-from .plateau import PlateauWithAnnealingParamScheduler
-from .plateau import PlateauWithAnnealingScheduler
-
-
-__all__ = [
-    "CompositeScheduler",
-    "CosineParamScheduler",
-    "CosineScheduler",
-    "LinearParamScheduler",
-    "LinearScheduler",
-    "PlateauWithAnnealingParamScheduler",
-    "PlateauWithAnnealingScheduler",
-]