nshtrainer 0.10.8__tar.gz → 0.10.9__tar.gz

{nshtrainer-0.10.8 → nshtrainer-0.10.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.10.8
+Version: 0.10.9
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.10.8 → nshtrainer-0.10.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.10.8"
+version = "0.10.9"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -11,14 +11,14 @@ nshrunner = "*"
 nshconfig = "*"
 nshutils = "*"
 psutil = "*"
+numpy = "*"
 torch = "*"
 typing-extensions = "*"
 lightning = "*"
 pytorch-lightning = "*"
-torchmetrics = "*"
-numpy = "*"
-wrapt = "*"
-GitPython = "*"
+torchmetrics = { version = "*", optional = true }
+wrapt = { version = "*", optional = true }
+GitPython = { version = "*", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 pyright = "^1.1.372"

{nshtrainer-0.10.8 → nshtrainer-0.10.9}/src/nshtrainer/callbacks/gradient_skipping.py

@@ -1,8 +1,8 @@
-from logging import getLogger
-from typing import Literal, Protocol, runtime_checkable
+import importlib.util
+import logging
+from typing import Any, Literal, Protocol, runtime_checkable
 
 import torch
-import torchmetrics
 from lightning.pytorch import Callback, LightningModule, Trainer
 from torch.optim import Optimizer
 from typing_extensions import override
@@ -10,23 +10,29 @@ from typing_extensions import override
 from .base import CallbackConfigBase
 from .norm_logging import compute_norm
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 @runtime_checkable
 class HasGradSkippedSteps(Protocol):
-    grad_skipped_steps: torchmetrics.SumMetric
+    grad_skipped_steps: Any
 
 
 class GradientSkipping(Callback):
     def __init__(self, config: "GradientSkippingConfig"):
-        super().__init__()
+        if importlib.util.find_spec("torchmetrics") is not None:
+            raise ImportError(
+                "To use the GradientSkipping callback, please install torchmetrics: pip install torchmetrics"
+            )
 
+        super().__init__()
         self.config = config
 
     @override
     def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
         if not isinstance(pl_module, HasGradSkippedSteps):
+            import torchmetrics  # type: ignore
+
             pl_module.grad_skipped_steps = torchmetrics.SumMetric()
 
     @override
@@ -47,12 +53,7 @@ class GradientSkipping(Callback):
         ):
             return
 
-        norm = compute_norm(
-            pl_module,
-            optimizer,
-            self.config.norm_type,
-            grad=True,
-        )
+        norm = compute_norm(pl_module, optimizer, self.config.norm_type, grad=True)
 
         # If the norm is NaN/Inf, we don't want to skip the step
         # beacuse AMP checks for NaN/Inf grads to adjust the loss scale.

{nshtrainer-0.10.8 → nshtrainer-0.10.9}/src/nshtrainer/data/transform.py

@@ -22,7 +22,13 @@ def transform(
         deepcopy: Whether to deep copy each item before applying the transform.
     """
 
-    import wrapt
+    try:
+        import wrapt
+    except ImportError:
+        raise ImportError(
+            "wrapt is not installed. wrapt is required for the transform function."
+            "Please install it using 'pip install wrapt'"
+        )
 
     class _TransformedDataset(wrapt.ObjectProxy):
         def __getitem__(self, idx):
@@ -52,7 +58,13 @@ def transform_with_index(
         deepcopy: Whether to deep copy each item before applying the transform.
     """
 
-    import wrapt
+    try:
+        import wrapt
+    except ImportError:
+        raise ImportError(
+            "wrapt is not installed. wrapt is required for the transform function."
+            "Please install it using 'pip install wrapt'"
+        )
 
     class _TransformedWithIndexDataset(wrapt.ObjectProxy):
         def __getitem__(self, idx: int):

{nshtrainer-0.10.8 → nshtrainer-0.10.9}/src/nshtrainer/model/_environment.py

@@ -9,7 +9,6 @@ from datetime import timedelta
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
-import git
 import nshconfig as C
 import psutil
 import torch
@@ -618,6 +617,11 @@ class GitRepositoryConfig(C.Config):
 
     @classmethod
     def from_current_directory(cls):
+        try:
+            import git
+        except ImportError:
+            return cls()
+
         draft = cls.draft()
         try:
             repo = git.Repo(os.getcwd(), search_parent_directories=True)

nshtrainer-0.10.8/src/nshtrainer/scripts/check_env.py

@@ -1,41 +0,0 @@
-REQUIRED_PACKAGES = [
-    "beartype",
-    "cloudpickle",
-    "jaxtyping",
-    "lightning",
-    "lightning_fabric",
-    "lightning_utilities",
-    "lovely_numpy",
-    "lovely_tensors",
-    "numpy",
-    "psutil",
-    "pydantic",
-    "pydantic_core",
-    "pysnooper",
-    "rich",
-    "tabulate",
-    "torch",
-    "torchmetrics",
-    "tqdm",
-    "typing_extensions",
-    "wrapt",
-    "yaml",
-]
-
-
-def main():
-    import importlib.util
-    import sys
-
-    missing_packages: list[str] = []
-    for package_name in REQUIRED_PACKAGES:
-        spec = importlib.util.find_spec(package_name)
-        if spec is None:
-            missing_packages.append(package_name)
-
-    if missing_packages:
-        sys.exit(f"Error: Missing required packages: {', '.join(missing_packages)}")
-
-
-if __name__ == "__main__":
-    main()