kostyl-toolkit 0.1.40__tar.gz → 0.1.43__tar.gz

{kostyl_toolkit-0.1.40 → kostyl_toolkit-0.1.43}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.3
 Name: kostyl-toolkit
-Version: 0.1.40
+Version: 0.1.43
 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.40 → kostyl_toolkit-0.1.43}/kostyl/ml/configs/training_settings.py

@@ -38,18 +38,18 @@ class FSDP1StrategyConfig(BaseModel):
     """Fully Sharded Data Parallel (FSDP) strategy configuration."""
 
     type: Literal["fsdp1"]
-    param_dtype: DTYPE | None
-    reduce_dtype: DTYPE | None
-    buffer_dtype: DTYPE | None
+    param_dtype: DTYPE | None = None
+    reduce_dtype: DTYPE | None = None
+    buffer_dtype: DTYPE | None = None
 
 
 class FSDP2StrategyConfig(BaseModel):
     """Fully Sharded Data Parallel (FSDP) strategy configuration."""
 
     type: Literal["fsdp2"]
-    param_dtype: DTYPE | None
-    reduce_dtype: DTYPE | None
-    buffer_dtype: DTYPE | None
+    param_dtype: DTYPE | None = None
+    reduce_dtype: DTYPE | None = None
+    buffer_dtype: DTYPE | None = None
 
 
 class DDPStrategyConfig(BaseModel):
@@ -59,12 +59,20 @@ class DDPStrategyConfig(BaseModel):
     find_unused_parameters: bool = False
 
 
+SUPPORTED_STRATEGIES = (
+    FSDP1StrategyConfig
+    | FSDP2StrategyConfig
+    | SingleDeviceStrategyConfig
+    | DDPStrategyConfig
+)
+
+
 class LightningTrainerParameters(BaseModel):
     """Lightning Trainer parameters configuration."""
 
     accelerator: str
     max_epochs: int
-    strategy: FSDP1StrategyConfig | SingleDeviceStrategyConfig | DDPStrategyConfig
+    strategy: SUPPORTED_STRATEGIES
     val_check_interval: int | float
     devices: list[int] | int
     precision: PRECISION

{kostyl_toolkit-0.1.40 → kostyl_toolkit-0.1.43}/kostyl/ml/data_collator.py

@@ -93,10 +93,10 @@ class BatchCollatorWithKeyAlignment:
                 if new_key is None:
                     continue
                 value = item[k]
-                if self.max_length is not None and new_key in (
+                if self.max_length is not None and new_key in {
                     "input_ids",
                     "attention_mask",
-                ):
+                }:
                     value = self._truncate_data(new_key, value)
                 new_item[new_key] = value
             aligned_batch.append(new_item)

{kostyl_toolkit-0.1.40 → kostyl_toolkit-0.1.43}/kostyl/ml/optim/factory.py

@@ -1,8 +1,11 @@
-from typing import Any
+from typing import TypedDict
+from typing import Unpack
 
 from torch.optim import Optimizer
+from torch.optim.optimizer import ParamsT
 
 from kostyl.ml.configs import OPTIMIZER_CONFIG
+from kostyl.ml.configs import SCHEDULER
 from kostyl.ml.configs import AdamConfig
 from kostyl.ml.configs import AdamWithPrecisionConfig
 from kostyl.ml.configs import MuonConfig
@@ -18,6 +21,17 @@ from .schedulers import PlateauWithAnnealingScheduler
 logger = setup_logger(fmt="only_message")
 
 
+class OVERRIDABLE_CONFIG_KWARGS(TypedDict, total=False):  # noqa: D101, N801
+    scheduler_type: SCHEDULER | None
+
+    freeze_ratio: float | None
+    warmup_ratio: float | None
+    warmup_value: float | None
+    base_value: float | None
+    final_value: float | None
+    plateau_ratio: float | None
+
+
 def create_scheduler(
     config: ScheduledParamConfig,
     param_group_field: str,
@@ -27,6 +41,7 @@ def create_scheduler(
     skip_if_zero: bool = False,
     apply_if_field: str | None = None,
     ignore_if_field: str | None = None,
+    **kwargs: Unpack[OVERRIDABLE_CONFIG_KWARGS],
 ) -> LinearScheduler | CosineScheduler | PlateauWithAnnealingScheduler:
     """
     Converts a ScheduledParamConfig to a scheduler instance.
@@ -41,37 +56,48 @@ def create_scheduler(
             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.
+        **kwargs: Optional overrides for scheduler configuration parameters (e.g., base_value,
+            final_value, warmup_ratio, etc.). These overrides take precedence over the values
+            provided in the `config` object.
 
     Returns:
         A scheduler instance based on the configuration.
 
     """
-    if config.scheduler_type is None:
+    scheduler_type = kwargs.get("scheduler_type", config.scheduler_type)
+    base_value = kwargs.get("base_value", config.base_value)
+    final_value = kwargs.get("final_value", config.final_value)
+    warmup_ratio = kwargs.get("warmup_ratio", config.warmup_ratio)
+    warmup_value = kwargs.get("warmup_value", config.warmup_value)
+    freeze_ratio = kwargs.get("freeze_ratio", config.freeze_ratio)
+    plateau_ratio = kwargs.get("plateau_ratio", config.plateau_ratio)
+
+    if scheduler_type is None:
         raise ValueError("scheduler_type must be specified in the config.")
 
-    if "plateau" in config.scheduler_type:
-        scheduler_type = "plateau"
+    if "plateau" in scheduler_type:
+        lookup_scheduler_type = "plateau"
     else:
-        scheduler_type = config.scheduler_type
-    scheduler_cls = SCHEDULER_MAPPING[scheduler_type]  # type: ignore
+        lookup_scheduler_type = scheduler_type
+    scheduler_cls = SCHEDULER_MAPPING[lookup_scheduler_type]  # type: ignore
 
     if issubclass(scheduler_cls, PlateauWithAnnealingScheduler):
-        if "cosine" in config.scheduler_type:
+        if "cosine" in scheduler_type:
             annealing_type = "cosine"
-        elif "linear" in config.scheduler_type:
+        elif "linear" in scheduler_type:
             annealing_type = "linear"
         else:
-            raise ValueError(f"Unknown annealing_type: {config.scheduler_type}")
+            raise ValueError(f"Unknown annealing_type: {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
+            plateau_value=base_value,
+            final_value=final_value,  # type: ignore
+            warmup_ratio=warmup_ratio,
+            warmup_value=warmup_value,
+            freeze_ratio=freeze_ratio,
+            plateau_ratio=plateau_ratio,  # type: ignore
             annealing_type=annealing_type,
             multiplier_field=multiplier_field,
             skip_if_zero=skip_if_zero,
@@ -83,8 +109,8 @@ def create_scheduler(
             optimizer=optim,
             param_group_field=param_group_field,
             num_iters=num_iters,
-            initial_value=config.base_value,
-            final_value=config.final_value,  # type: ignore
+            initial_value=base_value,
+            final_value=final_value,  # type: ignore
             multiplier_field=multiplier_field,
             skip_if_zero=skip_if_zero,
             apply_if_field=apply_if_field,
@@ -95,23 +121,23 @@ def create_scheduler(
             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,
+            base_value=base_value,
+            final_value=final_value,  # type: ignore
+            warmup_ratio=warmup_ratio,
+            warmup_value=warmup_value,
+            freeze_ratio=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}")
+        raise ValueError(f"Unsupported scheduler type: {scheduler_type}")
     return scheduler
 
 
 def create_optimizer(  # noqa: C901
-    parameters_groups: dict[str, Any],
+    parameters_groups: ParamsT,
     optimizer_config: OPTIMIZER_CONFIG,
     lr: float,
     weight_decay: float,
@@ -120,8 +146,7 @@ def create_optimizer(  # noqa: C901
     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.).
+        parameters_groups: Parameter groups for the optimizer.
         optimizer_config: Configuration for the optimizer.
         lr: Learning rate.
         weight_decay: Weight decay.

{kostyl_toolkit-0.1.40 → kostyl_toolkit-0.1.43}/kostyl/ml/params_groups.py

@@ -28,8 +28,8 @@ def create_params_groups(
             Defaults to None, which uses an empty set.
         no_decay_keywords (set[str] | None, optional): A set of string keywords. If a parameter's
             name contains any of these keywords, its weight decay is set to 0.0.
-            If additional keywords are provided, they will be added to the default set.
-            Defaults to None, which uses a standard set of exclusion keywords:
+            If keywords are provided, they will be added to the default set, otherwise the default set is used.
+            Default set of keywords:
             {"norm", "bias", "embedding", "tokenizer", "ln", "scale"}.
 
     Returns:

{kostyl_toolkit-0.1.40 → kostyl_toolkit-0.1.43}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kostyl-toolkit"
-version = "0.1.40"
+version = "0.1.43"
 description = "Kickass Orchestration System for Training, Yielding & Logging "
 readme = "README.md"
 requires-python = ">=3.12"
@@ -13,8 +13,8 @@ dependencies = [
 ml = [
     "case-converter>=1.2.0",
     "pydantic>=2.12.4",
-    "torch>=2.9.1",
-    "transformers>=4.57.1",
+    "torch",
+    "transformers",
 ]
 
 [dependency-groups]
@@ -57,17 +57,7 @@ publish-url = "https://test.pypi.org/legacy/"
 explicit = true
 
 [tool.pyright]
-typeCheckingMode = "standard" # "off" | "basic" | "standard" | "strict"
-reportMissingTypeStubs = false
-reportIncompatibleMethodOverride = true
-exclude = [
-    ".venv",
-    "build",
-    "dist",
-    "__pycache__",
-    "migrations",
-    "notebooks",
-]
+typeCheckingMode = "off" # "off" | "basic" | "standard" | "strict"
 
 [tool.ruff.lint]
 exclude = [