diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 2c4dbcffb0..308bd92db7 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -71,6 +71,36 @@ class HideLoggingMessage(logging.Filter): return not (self.text in x.getMessage()) +class HidePrintMessage: + __slots__ = ("_original_stream", "_hidden_texts") + + def __init__(self, original_stream): + self._original_stream = original_stream + self._hidden_texts = [] + + def add_filter(self, text): + self._hidden_texts.append(text) + + def write(self, message): + if not any(text in message for text in self._hidden_texts): + self._original_stream.write(message) + + def flush(self): + self._original_stream.flush() + + def __getattr__(self, name): + return getattr(self._original_stream, name) + + +if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") != "1": + import sys + + # Apply to stderr for FBGEMM + sys.stderr = HidePrintMessage(sys.stderr) + # https://github.com/pytorch/FBGEMM/blob/d99cd96490ec4aabac2ee95b1e76ea4dcfcfa628/fbgemm_gpu/experimental/gemm/triton_gemm/utils.py#L43-L52 + sys.stderr.add_filter("TMA benchmarks will be running") + + # Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype' # MUST do this at the start primarily due to tensorflow causing issues def fix_message_factory_issue(): diff --git a/unsloth/save.py b/unsloth/save.py index 640a7ffe14..f5ea8d7d8f 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -3010,7 +3010,11 @@ def patch_saving_functions(model, vision = False): original_model = model while True: - if original_model.push_to_hub.__name__ != "unsloth_push_to_hub": + # Check if push_to_hub exists before accessing its __name__ + if ( + hasattr(original_model, "push_to_hub") + and original_model.push_to_hub.__name__ != "unsloth_push_to_hub" + ): original_model.original_push_to_hub = original_model.push_to_hub original_model.push_to_hub = types.MethodType( unsloth_push_to_hub, original_model