kostyl-toolkit 0.1.14__tar.gz → 0.1.16__tar.gz

{kostyl_toolkit-0.1.14 → kostyl_toolkit-0.1.16}/PKG-INFO

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

@@ -15,7 +15,10 @@ from kostyl.ml.lightning.extenstions.pretrained_model import (
 
 
 def get_tokenizer_from_clearml(
-    model_id: str, task: Task | None = None, ignore_remote_overrides: bool = True
+    model_id: str,
+    task: Task | None = None,
+    ignore_remote_overrides: bool = True,
+    name: str | None = None,
 ) -> PreTrainedTokenizerBase:
     """
     Retrieve a Hugging Face tokenizer stored in a ClearML.
@@ -26,6 +29,8 @@ def get_tokenizer_from_clearml(
             the model. Defaults to None.
         ignore_remote_overrides (bool, optional): Whether to ignore remote hyperparameter
             overrides when connecting the ClearML task. Defaults to True.
+        name: The model name to be stored on the Task
+            (default to filename of the model weights, without the file extension, or to Input Model if that is not found)
 
     Returns:
         PreTrainedTokenizerBase: The instantiated tokenizer loaded from the local copy
@@ -34,7 +39,9 @@ def get_tokenizer_from_clearml(
     """
     clearml_tokenizer = InputModel(model_id=model_id)
     if task is not None:
-        clearml_tokenizer.connect(task, ignore_remote_overrides=ignore_remote_overrides)
+        clearml_tokenizer.connect(
+            task, ignore_remote_overrides=ignore_remote_overrides, name=name
+        )
 
     tokenizer = AutoTokenizer.from_pretrained(
         clearml_tokenizer.get_local_copy(raise_on_error=True)
@@ -49,6 +56,7 @@ def get_model_from_clearml[
     model: type[TModel],
     task: Task | None = None,
     ignore_remote_overrides: bool = True,
+    name: str | None = None,
     **kwargs: Any,
 ) -> TModel:
     """
@@ -59,8 +67,11 @@ def get_model_from_clearml[
         model: The model class that implements either PreTrainedModel or LightningCheckpointLoaderMixin.
         task: Optional ClearML task used to resolve the input model reference. If provided, the input model
             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
+        ignore_remote_overrides: When connecting the input model to the provided task,
+        determines whether
             remote configuration overrides should be ignored.
+        name: The model name to be stored on the Task
+            (default to filename of the model weights, without the file extension, or to Input Model if that is not found)
         **kwargs: Additional keyword arguments to pass to the model loading method.
 
     Returns:
@@ -70,7 +81,11 @@ def get_model_from_clearml[
     input_model = InputModel(model_id=model_id)
 
     if task is not None:
-        input_model.connect(task, ignore_remote_overrides=ignore_remote_overrides)
+        input_model.connect(
+            task,
+            ignore_remote_overrides=ignore_remote_overrides,
+            name=name,
+        )
 
     local_path = Path(input_model.get_local_copy(raise_on_error=True))
 

{kostyl_toolkit-0.1.14 → kostyl_toolkit-0.1.16}/kostyl/ml/configs/training_settings.py

@@ -100,5 +100,5 @@ class TrainingSettings(KostylBaseModel):
 
     trainer: LightningTrainerParameters
     early_stopping: EarlyStoppingConfig | None = None
-    checkpoint: CheckpointConfig
+    checkpoint: CheckpointConfig = CheckpointConfig()
     data: DataConfig

{kostyl_toolkit-0.1.14 → kostyl_toolkit-0.1.16}/kostyl/ml/lightning/callbacks/registry_uploading.py

@@ -78,49 +78,41 @@ class ClearMLRegistryUploaderCallback(Callback):
         )
         return output_model
 
-    @override
-    def on_validation_epoch_end(
-        self, trainer: Trainer, pl_module: KostylLightningModule
-    ) -> None:
-        if (not trainer.is_global_zero) or (
-            self.uploading_frequency != "after-every-eval"
-        ):
-            return
+    def _upload_best_checkpoint(self, pl_module: KostylLightningModule) -> None:
         if self._output_model is None:
             self._output_model = self._create_output_model(pl_module)
 
-        if self.ckpt_callback.best_model_path != self._last_best_model_path:
+        if self.ckpt_callback.best_model_path == self._last_best_model_path:
+            if self.verbose and (self._last_best_model_path != ""):
+                logger.info("Best model unchanged since last upload")
+            elif self.verbose:
+                logger.info("No best model found yet to upload")
+        else:
             if self.verbose:
-                logger.info(f"Best model path: {self.ckpt_callback.best_model_path}")
+                logger.info(
+                    f"Uploading best model from {self.ckpt_callback.best_model_path}"
+                )
             self._output_model.update_weights(
                 self.ckpt_callback.best_model_path,
                 auto_delete_file=False,
-                async_enable=True,
+                async_enable=False,
             )
-            self._last_best_model_path = self.ckpt_callback.best_model_path
-        elif self.verbose:
-            logger.info("Best model path: unchanged since last upload, skipping...")
+        return
+
+    @override
+    def on_validation_epoch_end(
+        self, trainer: Trainer, pl_module: KostylLightningModule
+    ) -> None:
+        if (not trainer.is_global_zero) or (
+            self.uploading_frequency != "after-every-eval"
+        ):
+            return
+        self._upload_best_checkpoint(pl_module)
         return
 
     @override
     def on_train_end(self, trainer: Trainer, pl_module: KostylLightningModule) -> None:
         if not trainer.is_global_zero:
             return
-        if self._output_model is None:
-            self._output_model = self._create_output_model(pl_module)
-
-        if self.ckpt_callback.best_model_path != self._last_best_model_path:
-            if self.verbose:
-                logger.info(
-                    f"Best model path at the end of training: {self.ckpt_callback.best_model_path}"
-                )
-            self._output_model.update_weights(
-                self.ckpt_callback.best_model_path,
-                auto_delete_file=False,
-                async_enable=False,
-            )
-        elif self.verbose:
-            logger.info(
-                "Best model path at the end of training: unchanged since last upload, skipping..."
-            )
+        self._upload_best_checkpoint(pl_module)
         return

{kostyl_toolkit-0.1.14 → kostyl_toolkit-0.1.16}/pyproject.toml

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