[zeta Module CLEAN UP OPERATIO]

zeta/__init__.py

@@ -1,36 +1,13 @@
-import logging
-import os
-import warnings
-
-# disable warnings
-
-warnings.filterwarnings("ignore")
-
-# disable tensorflow warnings
-
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
-
-# disable bnb warnings and others
-
-logging.getLogger().setLevel(logging.WARNING)
-
-
-class CustomFilter(logging.Filter):
-    def filter(self, record):
-        msg = "Created a temporary directory at"
-        return msg not in record.getMessage()
-
-
-logger = logging.getLogger()
-f = CustomFilter()
-logger.addFilter(f)
-
-from zeta.nn import *
-from zeta.models import *
-from zeta.utils import *
-from zeta.training import *
-from zeta.tokenizers import *
-from zeta.rl import *
-from zeta.optim import *
-from zeta.ops import *
-from zeta.quant import *
+from zeta.utils.disable_logging import disable_warnings_and_logs
+
+disable_warnings_and_logs()
+
+from zeta.nn import *  # noqa: F403, E402
+from zeta.models import *  # noqa: F403, E402
+from zeta.utils import *  # noqa: F403, E402
+from zeta.training import *  # noqa: F403, E402
+from zeta.tokenizers import *  # noqa: F403, E402
+from zeta.rl import *  # noqa: F403, E402
+from zeta.optim import *  # noqa: F403, E402
+from zeta.ops import *  # noqa: F403, E402
+from zeta.quant import *  # noqa: F403, E402

zeta/ops/__Init__.py

@@ -1,7 +1,4 @@
-from zeta.ops.main import *
-from zeta.ops.softmax import *
 from zeta.ops.unitwise_norm import unitwise_norm
-from zeta.ops.mos import MixtureOfSoftmaxes
 
 from zeta.ops.softmax import (
     standard_softmax,

zeta/utils/__init__.py

@@ -7,11 +7,12 @@
     print_cuda_memory_usage,
     save_memory_snapshot,
 )
-
+from zeta.utils.disable_logging import disable_warnings_and_logs
 
 __all__ = [
     "track_cuda_memory_usage",
     "benchmark",
     "print_cuda_memory_usage",
     "save_memory_snapshot",
+    "disable_warnings_and_logs",
 ]

zeta/utils/disable_logging.py

@@ -0,0 +1,31 @@
+import logging
+import os
+import warnings
+
+
+def disable_warnings_and_logs():
+    """Disable warnings and logs.
+
+    Returns:
+        _type_: _description_
+    """
+    # disable warnings
+    warnings.filterwarnings("ignore")
+
+    # disable tensorflow warnings
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+
+    # disable bnb warnings and others
+    logging.getLogger().setLevel(logging.WARNING)
+
+    class CustomFilter(logging.Filter):
+        def filter(self, record):
+            msg = "Created a temporary directory at"
+            return msg not in record.getMessage()
+
+    logger = logging.getLogger()
+    f = CustomFilter()
+    logger.addFilter(f)
+
+
+disable_warnings_and_logs()