nshtrainer 0.42.0__py3-none-any.whl → 0.43.0__py3-none-any.whl

nshtrainer/config/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -30,25 +32,47 @@ if TYPE_CHECKING:
     )
     from nshtrainer.callbacks import CallbackConfig as CallbackConfig
     from nshtrainer.callbacks import DebugFlagCallbackConfig as DebugFlagCallbackConfig
-    from nshtrainer.callbacks import DirectorySetupConfig as DirectorySetupConfig
-    from nshtrainer.callbacks import EarlyStoppingConfig as EarlyStoppingConfig
-    from nshtrainer.callbacks import EMAConfig as EMAConfig
-    from nshtrainer.callbacks import EpochTimerConfig as EpochTimerConfig
-    from nshtrainer.callbacks import FiniteChecksConfig as FiniteChecksConfig
-    from nshtrainer.callbacks import GradientSkippingConfig as GradientSkippingConfig
+    from nshtrainer.callbacks import (
+        DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
+    )
+    from nshtrainer.callbacks import EMACallbackConfig as EMACallbackConfig
+    from nshtrainer.callbacks import (
+        EpochTimerCallbackConfig as EpochTimerCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        FiniteChecksCallbackConfig as FiniteChecksCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        GradientSkippingCallbackConfig as GradientSkippingCallbackConfig,
+    )
     from nshtrainer.callbacks import (
         LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
     )
-    from nshtrainer.callbacks import NormLoggingConfig as NormLoggingConfig
+    from nshtrainer.callbacks import (
+        NormLoggingCallbackConfig as NormLoggingCallbackConfig,
+    )
     from nshtrainer.callbacks import (
         OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
     )
-    from nshtrainer.callbacks import PrintTableMetricsConfig as PrintTableMetricsConfig
-    from nshtrainer.callbacks import RLPSanityChecksConfig as RLPSanityChecksConfig
-    from nshtrainer.callbacks import SharedParametersConfig as SharedParametersConfig
+    from nshtrainer.callbacks import (
+        PrintTableMetricsCallbackConfig as PrintTableMetricsCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        SharedParametersCallbackConfig as SharedParametersCallbackConfig,
+    )
     from nshtrainer.callbacks import ThroughputMonitorConfig as ThroughputMonitorConfig
-    from nshtrainer.callbacks import WandbUploadCodeConfig as WandbUploadCodeConfig
-    from nshtrainer.callbacks import WandbWatchConfig as WandbWatchConfig
+    from nshtrainer.callbacks import (
+        WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        WandbWatchCallbackConfig as WandbWatchCallbackConfig,
+    )
     from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig
     from nshtrainer.callbacks.checkpoint._base import (
         BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
@@ -155,26 +179,28 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "BaseConfig":
-            return importlib.import_module("nshtrainer").BaseConfig
         if name == "MetricConfig":
             return importlib.import_module("nshtrainer").MetricConfig
-        if name == "CallbackConfigBase":
-            return importlib.import_module("nshtrainer._hf_hub").CallbackConfigBase
-        if name == "HuggingFaceHubConfig":
-            return importlib.import_module("nshtrainer._hf_hub").HuggingFaceHubConfig
+        if name == "BaseConfig":
+            return importlib.import_module("nshtrainer").BaseConfig
         if name == "HuggingFaceHubAutoCreateConfig":
             return importlib.import_module(
                 "nshtrainer._hf_hub"
             ).HuggingFaceHubAutoCreateConfig
+        if name == "HuggingFaceHubConfig":
+            return importlib.import_module("nshtrainer._hf_hub").HuggingFaceHubConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module("nshtrainer._hf_hub").CallbackConfigBase
         if name == "OptimizerConfigBase":
             return importlib.import_module("nshtrainer.optimizer").OptimizerConfigBase
         if name == "AdamWConfig":
             return importlib.import_module("nshtrainer.optimizer").AdamWConfig
+        if name == "DirectorySetupCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).DirectorySetupCallbackConfig
         if name == "DirectoryConfig":
             return importlib.import_module("nshtrainer.model").DirectoryConfig
-        if name == "DirectorySetupConfig":
-            return importlib.import_module("nshtrainer.callbacks").DirectorySetupConfig
         if name == "TrainerConfig":
             return importlib.import_module("nshtrainer.model").TrainerConfig
         if name == "EnvironmentConfig":
@@ -183,36 +209,36 @@ else:
             return importlib.import_module("nshtrainer.nn").BaseNonlinearityConfig
         if name == "MLPConfig":
             return importlib.import_module("nshtrainer.nn").MLPConfig
+        if name == "PReLUConfig":
+            return importlib.import_module("nshtrainer.nn").PReLUConfig
+        if name == "LeakyReLUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").LeakyReLUNonlinearityConfig
         if name == "SwiGLUNonlinearityConfig":
             return importlib.import_module(
                 "nshtrainer.nn.nonlinearity"
             ).SwiGLUNonlinearityConfig
-        if name == "ReLUNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").ReLUNonlinearityConfig
+        if name == "SoftsignNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SoftsignNonlinearityConfig
         if name == "SiLUNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").SiLUNonlinearityConfig
+        if name == "SigmoidNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SigmoidNonlinearityConfig
+        if name == "SoftplusNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SoftplusNonlinearityConfig
         if name == "ELUNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").ELUNonlinearityConfig
+        if name == "SoftmaxNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").SoftmaxNonlinearityConfig
         if name == "GELUNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").GELUNonlinearityConfig
-        if name == "SoftplusNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SoftplusNonlinearityConfig
-        if name == "SoftsignNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SoftsignNonlinearityConfig
         if name == "SwishNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").SwishNonlinearityConfig
-        if name == "SoftmaxNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SoftmaxNonlinearityConfig
         if name == "MishNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").MishNonlinearityConfig
-        if name == "SigmoidNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").SigmoidNonlinearityConfig
         if name == "TanhNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").TanhNonlinearityConfig
-        if name == "PReLUConfig":
-            return importlib.import_module("nshtrainer.nn").PReLUConfig
-        if name == "LeakyReLUNonlinearityConfig":
-            return importlib.import_module("nshtrainer.nn").LeakyReLUNonlinearityConfig
+        if name == "ReLUNonlinearityConfig":
+            return importlib.import_module("nshtrainer.nn").ReLUNonlinearityConfig
         if name == "LRSchedulerConfigBase":
             return importlib.import_module(
                 "nshtrainer.lr_scheduler"
@@ -231,32 +257,40 @@ else:
             return importlib.import_module("nshtrainer.loggers").TensorboardLoggerConfig
         if name == "WandbLoggerConfig":
             return importlib.import_module("nshtrainer.loggers").WandbLoggerConfig
-        if name == "WandbUploadCodeConfig":
-            return importlib.import_module("nshtrainer.callbacks").WandbUploadCodeConfig
-        if name == "WandbWatchConfig":
-            return importlib.import_module("nshtrainer.callbacks").WandbWatchConfig
+        if name == "WandbUploadCodeCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).WandbUploadCodeCallbackConfig
+        if name == "WandbWatchCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).WandbWatchCallbackConfig
         if name == "CSVLoggerConfig":
             return importlib.import_module("nshtrainer.loggers").CSVLoggerConfig
-        if name == "EnvironmentPackageConfig":
-            return importlib.import_module(
-                "nshtrainer.util._environment_info"
-            ).EnvironmentPackageConfig
-        if name == "EnvironmentSnapshotConfig":
+        if name == "EnvironmentLinuxEnvironmentConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentSnapshotConfig
+            ).EnvironmentLinuxEnvironmentConfig
         if name == "EnvironmentLSFInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).EnvironmentLSFInformationConfig
-        if name == "EnvironmentLinuxEnvironmentConfig":
+        if name == "EnvironmentGPUConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentLinuxEnvironmentConfig
-        if name == "EnvironmentSLURMInformationConfig":
+            ).EnvironmentGPUConfig
+        if name == "EnvironmentPackageConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentSLURMInformationConfig
+            ).EnvironmentPackageConfig
+        if name == "EnvironmentHardwareConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentHardwareConfig
+        if name == "EnvironmentSnapshotConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentSnapshotConfig
         if name == "EnvironmentClassInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
@@ -269,14 +303,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).EnvironmentCUDAConfig
-        if name == "EnvironmentGPUConfig":
-            return importlib.import_module(
-                "nshtrainer.util._environment_info"
-            ).EnvironmentGPUConfig
-        if name == "EnvironmentHardwareConfig":
+        if name == "EnvironmentSLURMInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentHardwareConfig
+            ).EnvironmentSLURMInformationConfig
         if name == "EpochsConfig":
             return importlib.import_module("nshtrainer.util.config").EpochsConfig
         if name == "StepsConfig":
@@ -287,52 +317,56 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).CheckpointLoadingConfig
-        if name == "OptimizationConfig":
+        if name == "SanityCheckingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OptimizationConfig
+            ).SanityCheckingConfig
+        if name == "OnExceptionCheckpointCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).OnExceptionCheckpointCallbackConfig
         if name == "GradientClippingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).GradientClippingConfig
-        if name == "LastCheckpointCallbackConfig":
+        if name == "LoggingConfig":
+            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
+        if name == "RLPSanityChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).LastCheckpointCallbackConfig
-        if name == "OnExceptionCheckpointCallbackConfig":
+            ).RLPSanityChecksCallbackConfig
+        if name == "CheckpointSavingConfig":
             return importlib.import_module(
-                "nshtrainer.callbacks"
-            ).OnExceptionCheckpointCallbackConfig
-        if name == "RLPSanityChecksConfig":
-            return importlib.import_module("nshtrainer.callbacks").RLPSanityChecksConfig
-        if name == "EarlyStoppingConfig":
-            return importlib.import_module("nshtrainer.callbacks").EarlyStoppingConfig
+                "nshtrainer.trainer._config"
+            ).CheckpointSavingConfig
         if name == "DebugFlagCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).DebugFlagCallbackConfig
-        if name == "CheckpointSavingConfig":
+        if name == "LastCheckpointCallbackConfig":
             return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).CheckpointSavingConfig
-        if name == "BestCheckpointCallbackConfig":
+                "nshtrainer.callbacks"
+            ).LastCheckpointCallbackConfig
+        if name == "SharedParametersCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).BestCheckpointCallbackConfig
-        if name == "LoggingConfig":
-            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
-        if name == "SanityCheckingConfig":
+            ).SharedParametersCallbackConfig
+        if name == "ReproducibilityConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).SanityCheckingConfig
-        if name == "SharedParametersConfig":
+            ).ReproducibilityConfig
+        if name == "EarlyStoppingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).SharedParametersConfig
-        if name == "ReproducibilityConfig":
+            ).EarlyStoppingCallbackConfig
+        if name == "OptimizationConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).ReproducibilityConfig
+            ).OptimizationConfig
+        if name == "BestCheckpointCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).BestCheckpointCallbackConfig
         if name == "CheckpointMetadata":
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"
@@ -349,28 +383,34 @@ else:
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"
             ).UserProvidedPathCheckpointStrategyConfig
-        if name == "PrintTableMetricsConfig":
+        if name == "PrintTableMetricsCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).PrintTableMetricsConfig
+            ).PrintTableMetricsCallbackConfig
         if name == "ThroughputMonitorConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).ThroughputMonitorConfig
-        if name == "GradientSkippingConfig":
+        if name == "GradientSkippingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).GradientSkippingConfig
-        if name == "EMAConfig":
-            return importlib.import_module("nshtrainer.callbacks").EMAConfig
+            ).GradientSkippingCallbackConfig
+        if name == "EMACallbackConfig":
+            return importlib.import_module("nshtrainer.callbacks").EMACallbackConfig
         if name == "ActSaveConfig":
             return importlib.import_module("nshtrainer.callbacks.actsave").ActSaveConfig
-        if name == "FiniteChecksConfig":
-            return importlib.import_module("nshtrainer.callbacks").FiniteChecksConfig
-        if name == "NormLoggingConfig":
-            return importlib.import_module("nshtrainer.callbacks").NormLoggingConfig
-        if name == "EpochTimerConfig":
-            return importlib.import_module("nshtrainer.callbacks").EpochTimerConfig
+        if name == "FiniteChecksCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).FiniteChecksCallbackConfig
+        if name == "NormLoggingCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).NormLoggingCallbackConfig
+        if name == "EpochTimerCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).EpochTimerCallbackConfig
         if name == "BaseCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint._base"

nshtrainer/config/_checkpoint/loader/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -29,24 +31,24 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CheckpointLoadingConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).CheckpointLoadingConfig
-        if name == "CheckpointMetadata":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).CheckpointMetadata
         if name == "MetricConfig":
             return importlib.import_module("nshtrainer._checkpoint.loader").MetricConfig
         if name == "BestCheckpointStrategyConfig":
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"
             ).BestCheckpointStrategyConfig
+        if name == "CheckpointMetadata":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).CheckpointMetadata
         if name == "LastCheckpointStrategyConfig":
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"
             ).LastCheckpointStrategyConfig
+        if name == "CheckpointLoadingConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.loader"
+            ).CheckpointLoadingConfig
         if name == "UserProvidedPathCheckpointStrategyConfig":
             return importlib.import_module(
                 "nshtrainer._checkpoint.loader"

nshtrainer/config/_checkpoint/metadata/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -14,14 +16,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "EnvironmentConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.metadata"
-            ).EnvironmentConfig
         if name == "CheckpointMetadata":
             return importlib.import_module(
                 "nshtrainer._checkpoint.metadata"
             ).CheckpointMetadata
+        if name == "EnvironmentConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.metadata"
+            ).EnvironmentConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/_directory/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -6,7 +8,9 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from nshtrainer._directory import DirectoryConfig as DirectoryConfig
-    from nshtrainer._directory import DirectorySetupConfig as DirectorySetupConfig
+    from nshtrainer._directory import (
+        DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
+    )
     from nshtrainer._directory import LoggerConfig as LoggerConfig
 else:
 
@@ -15,10 +19,12 @@ else:
 
         if name in globals():
             return globals()[name]
+        if name == "DirectorySetupCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer._directory"
+            ).DirectorySetupCallbackConfig
         if name == "DirectoryConfig":
             return importlib.import_module("nshtrainer._directory").DirectoryConfig
-        if name == "DirectorySetupConfig":
-            return importlib.import_module("nshtrainer._directory").DirectorySetupConfig
         if name == "LoggerConfig":
             return importlib.import_module("nshtrainer._directory").LoggerConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

nshtrainer/config/_hf_hub/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -17,14 +19,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
-            return importlib.import_module("nshtrainer._hf_hub").CallbackConfigBase
-        if name == "HuggingFaceHubConfig":
-            return importlib.import_module("nshtrainer._hf_hub").HuggingFaceHubConfig
         if name == "HuggingFaceHubAutoCreateConfig":
             return importlib.import_module(
                 "nshtrainer._hf_hub"
             ).HuggingFaceHubAutoCreateConfig
+        if name == "HuggingFaceHubConfig":
+            return importlib.import_module("nshtrainer._hf_hub").HuggingFaceHubConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module("nshtrainer._hf_hub").CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -11,25 +13,47 @@ if TYPE_CHECKING:
     from nshtrainer.callbacks import CallbackConfig as CallbackConfig
     from nshtrainer.callbacks import CallbackConfigBase as CallbackConfigBase
     from nshtrainer.callbacks import DebugFlagCallbackConfig as DebugFlagCallbackConfig
-    from nshtrainer.callbacks import DirectorySetupConfig as DirectorySetupConfig
-    from nshtrainer.callbacks import EarlyStoppingConfig as EarlyStoppingConfig
-    from nshtrainer.callbacks import EMAConfig as EMAConfig
-    from nshtrainer.callbacks import EpochTimerConfig as EpochTimerConfig
-    from nshtrainer.callbacks import FiniteChecksConfig as FiniteChecksConfig
-    from nshtrainer.callbacks import GradientSkippingConfig as GradientSkippingConfig
+    from nshtrainer.callbacks import (
+        DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
+    )
+    from nshtrainer.callbacks import EMACallbackConfig as EMACallbackConfig
+    from nshtrainer.callbacks import (
+        EpochTimerCallbackConfig as EpochTimerCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        FiniteChecksCallbackConfig as FiniteChecksCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        GradientSkippingCallbackConfig as GradientSkippingCallbackConfig,
+    )
     from nshtrainer.callbacks import (
         LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
     )
-    from nshtrainer.callbacks import NormLoggingConfig as NormLoggingConfig
+    from nshtrainer.callbacks import (
+        NormLoggingCallbackConfig as NormLoggingCallbackConfig,
+    )
     from nshtrainer.callbacks import (
         OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
     )
-    from nshtrainer.callbacks import PrintTableMetricsConfig as PrintTableMetricsConfig
-    from nshtrainer.callbacks import RLPSanityChecksConfig as RLPSanityChecksConfig
-    from nshtrainer.callbacks import SharedParametersConfig as SharedParametersConfig
+    from nshtrainer.callbacks import (
+        PrintTableMetricsCallbackConfig as PrintTableMetricsCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        SharedParametersCallbackConfig as SharedParametersCallbackConfig,
+    )
     from nshtrainer.callbacks import ThroughputMonitorConfig as ThroughputMonitorConfig
-    from nshtrainer.callbacks import WandbUploadCodeConfig as WandbUploadCodeConfig
-    from nshtrainer.callbacks import WandbWatchConfig as WandbWatchConfig
+    from nshtrainer.callbacks import (
+        WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
+    )
+    from nshtrainer.callbacks import (
+        WandbWatchCallbackConfig as WandbWatchCallbackConfig,
+    )
     from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig
     from nshtrainer.callbacks.checkpoint._base import (
         BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
@@ -45,12 +69,12 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
-            return importlib.import_module("nshtrainer.callbacks").CallbackConfigBase
-        if name == "PrintTableMetricsConfig":
+        if name == "PrintTableMetricsCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).PrintTableMetricsConfig
+            ).PrintTableMetricsCallbackConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module("nshtrainer.callbacks").CallbackConfigBase
         if name == "DebugFlagCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
@@ -59,50 +83,66 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).ThroughputMonitorConfig
-        if name == "GradientSkippingConfig":
+        if name == "GradientSkippingCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).GradientSkippingCallbackConfig
+        if name == "RLPSanityChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).GradientSkippingConfig
-        if name == "RLPSanityChecksConfig":
-            return importlib.import_module("nshtrainer.callbacks").RLPSanityChecksConfig
-        if name == "WandbUploadCodeConfig":
-            return importlib.import_module("nshtrainer.callbacks").WandbUploadCodeConfig
+            ).RLPSanityChecksCallbackConfig
+        if name == "WandbUploadCodeCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).WandbUploadCodeCallbackConfig
         if name == "MetricConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.early_stopping"
             ).MetricConfig
-        if name == "EarlyStoppingConfig":
-            return importlib.import_module("nshtrainer.callbacks").EarlyStoppingConfig
-        if name == "WandbWatchConfig":
-            return importlib.import_module("nshtrainer.callbacks").WandbWatchConfig
-        if name == "EMAConfig":
-            return importlib.import_module("nshtrainer.callbacks").EMAConfig
-        if name == "DirectorySetupConfig":
-            return importlib.import_module("nshtrainer.callbacks").DirectorySetupConfig
+        if name == "EarlyStoppingCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).EarlyStoppingCallbackConfig
+        if name == "WandbWatchCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).WandbWatchCallbackConfig
+        if name == "EMACallbackConfig":
+            return importlib.import_module("nshtrainer.callbacks").EMACallbackConfig
+        if name == "DirectorySetupCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).DirectorySetupCallbackConfig
         if name == "ActSaveConfig":
             return importlib.import_module("nshtrainer.callbacks.actsave").ActSaveConfig
-        if name == "FiniteChecksConfig":
-            return importlib.import_module("nshtrainer.callbacks").FiniteChecksConfig
-        if name == "NormLoggingConfig":
-            return importlib.import_module("nshtrainer.callbacks").NormLoggingConfig
-        if name == "LastCheckpointCallbackConfig":
+        if name == "FiniteChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).LastCheckpointCallbackConfig
-        if name == "BestCheckpointCallbackConfig":
+            ).FiniteChecksCallbackConfig
+        if name == "NormLoggingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).BestCheckpointCallbackConfig
+            ).NormLoggingCallbackConfig
+        if name == "EpochTimerCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).EpochTimerCallbackConfig
         if name == "OnExceptionCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).OnExceptionCheckpointCallbackConfig
-        if name == "EpochTimerConfig":
-            return importlib.import_module("nshtrainer.callbacks").EpochTimerConfig
-        if name == "SharedParametersConfig":
+        if name == "LastCheckpointCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).LastCheckpointCallbackConfig
+        if name == "SharedParametersCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks"
+            ).SharedParametersCallbackConfig
+        if name == "BestCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
-            ).SharedParametersConfig
+            ).BestCheckpointCallbackConfig
         if name == "CheckpointMetadata":
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint._base"