Update import_fixes.py

unsloth/import_fixes.py

@@ -19,7 +19,16 @@ from importlib.metadata import version as importlib_version
 from packaging.version import Version as TrueVersion
 import re
 import logging
-
+# We cannot do from unsloth_zoo.log import logger since FBGEMM might cause seg faults.
+UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING",  "0") in ("1", "True", "true",)
+logger = logging.getLogger(__name__)
+if UNSLOTH_ENABLE_LOGGING:
+    logging.basicConfig(level = logging.INFO, format = '[%(name)s|%(levelname)s]%(message)s')
+    logger.setLevel(logging.INFO)
+else:
+    logging.basicConfig(level = logging.WARNING, format = '[%(name)s|%(levelname)s]%(message)s')
+    logger.setLevel(logging.WARNING) 
+pass
 
 def Version(version):
     try:
@@ -68,8 +77,6 @@ def fix_message_factory_issue():
             def GetPrototype(self, *args, **kwargs):
                 return
 
-        from unsloth_zoo.log import logger
-
         if not hasattr(google.protobuf.message_factory, "MessageFactory"):
             logger.info("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
             google.protobuf.message_factory.MessageFactory = MessageFactory
@@ -107,8 +114,6 @@ def fix_xformers_performance_issue():
         return
     xformers_version = importlib_version("xformers")
     if Version(xformers_version) < Version("0.0.29"):
-        from unsloth_zoo.log import logger
-
         xformers_location = importlib.util.find_spec("xformers").origin
         xformers_location = os.path.split(xformers_location)[0]
         cutlass = Path(xformers_location) / "ops" / "fmha" / "cutlass.py"
@@ -138,8 +143,6 @@ def fix_vllm_aimv2_issue():
         return
     vllm_version = importlib_version("vllm")
     if Version(vllm_version) < Version("0.10.1"):
-        from unsloth_zoo.log import logger
-
         vllm_version = importlib.util.find_spec("vllm").origin
         vllm_version = os.path.split(vllm_version)[0]
         ovis_config = Path(vllm_version) / "transformers_utils" / "configs" / "ovis.py"
@@ -271,7 +274,6 @@ def check_fbgemm_gpu_version():
         raise ImportError(
             f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected. It might cause unexpected issues like segmentation faults. Please uninstall the current one by doing `pip uninstall fbgemm-gpu` && `pip install fbgemm-gpu` to install fbgemm-gpu 1.4.0 or newer!"
         )
-    from unsloth_zoo.log import logger
 
     logger.info(f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected.")
 
@@ -334,7 +336,6 @@ def patch_enable_input_require_grads():
             self._require_grads_hook = hooks[0]
 
     PreTrainedModel.enable_input_require_grads = _patched_enable_input_require_grads
-    from unsloth_zoo.log import logger
 
     logger.info(
         "Unsloth: Patched enable_input_require_grads for vision model compatibility"
@@ -376,7 +377,6 @@ def torchvision_compatibility_check():
             f"but found torchvision=={torchvision_version}. "
             f"Please refer to https://pytorch.org/get-started/previous-versions/ for more information."
         )
-    from unsloth_zoo.log import logger
 
     logger.info(
         f"Unsloth: torch=={torch_version} and torchvision=={torchvision_version} are compatible."
@@ -392,7 +392,6 @@ def fix_openenv_no_vllm():
     openenv = Path(trl_location) / "experimental" / "openenv" / "utils.py"
     if not openenv.exists():
         return
-    from unsloth_zoo.log import logger
 
     try:
         with open(openenv, "r+", encoding = "utf-8") as f: