nshtrainer 0.21.0__py3-none-any.whl → 0.22.0__py3-none-any.whl

nshtrainer/_hf_hub.py

@@ -10,11 +10,7 @@ from nshrunner._env import SNAPSHOT_DIR
 from typing_extensions import override
 
 from ._callback import NTCallbackBase
-from .callbacks.base import (
-    CallbackConfigBase,
-    CallbackMetadataConfig,
-    CallbackWithMetadata,
-)
+from .callbacks.base import CallbackConfigBase
 
 if TYPE_CHECKING:
     from huggingface_hub import HfApi  # noqa: F401
@@ -81,10 +77,7 @@ class HuggingFaceHubConfig(CallbackConfigBase):
 
     @override
     def create_callbacks(self, root_config):
-        yield CallbackWithMetadata(
-            HFHubCallback(self),
-            CallbackMetadataConfig(ignore_if_exists=True),
-        )
+        yield self.with_metadata(HFHubCallback(self), ignore_if_exists=True)
 
 
 def _api(token: str | None = None):

nshtrainer/callbacks/base.py

@@ -2,29 +2,24 @@ from abc import ABC, abstractmethod
 from collections import Counter
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, TypeAlias, TypedDict
+from typing import TYPE_CHECKING, ClassVar, TypeAlias
 
 import nshconfig as C
 from lightning.pytorch import Callback
+from typing_extensions import TypedDict, Unpack
 
 if TYPE_CHECKING:
     from ..model.config import BaseConfig
 
 
-class CallbackMetadataDict(TypedDict, total=False):
+class CallbackMetadataConfig(TypedDict, total=False):
     ignore_if_exists: bool
-    """If `True`, the callback will not be added if another callback with the same class already exists."""
+    """If `True`, the callback will not be added if another callback with the same class already exists.
+    Default is `False`."""
 
     priority: int
-    """Priority of the callback. Callbacks with higher priority will be loaded first."""
-
-
-class CallbackMetadataConfig(C.Config):
-    ignore_if_exists: bool = False
-    """If `True`, the callback will not be added if another callback with the same class already exists."""
-
-    priority: int = 0
-    """Priority of the callback. Callbacks with higher priority will be loaded first."""
+    """Priority of the callback. Callbacks with higher priority will be loaded first.
+    Default is `0`."""
 
 
 @dataclass(frozen=True)
@@ -37,13 +32,18 @@ ConstructedCallback: TypeAlias = Callback | CallbackWithMetadata
 
 
 class CallbackConfigBase(C.Config, ABC):
-    metadata: CallbackMetadataConfig = CallbackMetadataConfig()
+    metadata: ClassVar[CallbackMetadataConfig] = CallbackMetadataConfig()
     """Metadata for the callback."""
 
-    def with_metadata(self, callback: Callback, **metadata: CallbackMetadataDict):
-        return CallbackWithMetadata(
-            callback=callback, metadata=self.metadata.model_copy(update=metadata)
-        )
+    @classmethod
+    def with_metadata(
+        cls, callback: Callback, **kwargs: Unpack[CallbackMetadataConfig]
+    ):
+        metadata: CallbackMetadataConfig = {}
+        metadata.update(cls.metadata)
+        metadata.update(kwargs)
+
+        return CallbackWithMetadata(callback=callback, metadata=metadata)
 
     @abstractmethod
     def create_callbacks(
@@ -73,7 +73,7 @@ def _filter_ignore_if_exists(callbacks: list[CallbackWithMetadata]):
     for callback in callbacks:
         # If `ignore_if_exists` is `True` and there is already a callback of the same class, skip this callback
         if (
-            callback.metadata.ignore_if_exists
+            callback.metadata.get("ignore_if_exists", False)
             and callback_classes[callback.callback.__class__] > 1
         ):
             continue
@@ -89,7 +89,10 @@ def _process_and_filter_callbacks(
     callbacks = list(callbacks)
 
     # Sort by priority (higher priority first)
-    callbacks.sort(key=lambda callback: callback.metadata.priority, reverse=True)
+    callbacks.sort(
+        key=lambda callback: callback.metadata.get("priority", 0),
+        reverse=True,
+    )
 
     # Process `ignore_if_exists`
     callbacks = _filter_ignore_if_exists(callbacks)

{nshtrainer-0.21.0.dist-info → nshtrainer-0.22.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.21.0
+Version: 0.22.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -26,7 +26,6 @@ Requires-Dist: torchmetrics ; extra == "extra"
 Requires-Dist: typing-extensions
 Requires-Dist: wandb ; extra == "extra"
 Requires-Dist: wrapt ; extra == "extra"
-Requires-Dist: zstandard ; extra == "extra"
 Description-Content-Type: text/markdown
 
 

{nshtrainer-0.21.0.dist-info → nshtrainer-0.22.0.dist-info}/RECORD

@@ -4,11 +4,11 @@ nshtrainer/_checkpoint/loader.py,sha256=5vjg-OFChXJjgiOVv8vnV8nwTscfdDtEdxQRz6uP
 nshtrainer/_checkpoint/metadata.py,sha256=TLAt7yR3KhSRbXCtomLMxcMvOiAju873A1ZRo8VWNwA,5179
 nshtrainer/_checkpoint/saver.py,sha256=6W-Rbc3QGuhcF_mcwN8v31uEjLQCsZvt8CPuqPs4m5g,1342
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
-nshtrainer/_hf_hub.py,sha256=0bOhJNyIjQGJsMRaW7qQJc1oTnUMHj08auuztzTQvZ0,16906
+nshtrainer/_hf_hub.py,sha256=iqhXH54RhSqmot_K3UCVcHTC_TC81_YY7cwvHGHXXlw,16782
 nshtrainer/callbacks/__init__.py,sha256=4qocBDzQbLLhhbIEfvbA3SQB_Dy9ZJH7keMwPay-ZS8,2359
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
-nshtrainer/callbacks/base.py,sha256=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,3555
+nshtrainer/callbacks/base.py,sha256=NpjeKmonJ1Kaz5_39XSn3LlDwvbGjk6WV8BpHSNCvI4,3508
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=g-3zIthupERKqWZQw-A_busQPaPRkto6iHBV-M7nK1Y,527
 nshtrainer/callbacks/checkpoint/_base.py,sha256=r6IPpl3sGUmxBNv80y9r326lTrPAIVSU3Fu-3LrYH2s,6691
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=DJiLo7NDzd-lp-O3v7Cv8WejyjXPV_6_RmfltKO9fvE,2165
@@ -87,6 +87,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.21.0.dist-info/METADATA,sha256=7QfSX_yXi-Up6uxOVFfDPn4ieGK5b3UgQfO_KFsNzXk,979
-nshtrainer-0.21.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.21.0.dist-info/RECORD,,
+nshtrainer-0.22.0.dist-info/METADATA,sha256=sdjt9S4X3xiIGgD6FNF06yIyC1tJA89B9Qm9mxy29tc,935
+nshtrainer-0.22.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.22.0.dist-info/RECORD,,