From ca998cc6f2b76ccefb168a7c9a92132fe6f70826 Mon Sep 17 00:00:00 2001
From: Zhao Changmin <changmin.zhao@intel.com>
Date: Wed, 2 Aug 2023 16:46:22 +0800
Subject: [PATCH] LLM: Mute shape mismatch output (#8601)

* LLM: Mute shape mismatch output
---
 python/llm/src/bigdl/llm/transformers/model.py     |  9 ++++-----
 python/llm/src/bigdl/llm/utils/common/__init__.py  |  2 +-
 python/llm/src/bigdl/llm/utils/common/log4Error.py | 12 ++++++++++++
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/python/llm/src/bigdl/llm/transformers/model.py b/python/llm/src/bigdl/llm/transformers/model.py
index 06dc88a4..d417203b 100644
--- a/python/llm/src/bigdl/llm/transformers/model.py
+++ b/python/llm/src/bigdl/llm/transformers/model.py
@@ -23,7 +23,7 @@ from .utils import extract_local_archive_file, \
     get_local_shard_files, \
     fix_key
 from bigdl.llm.ggml.quantize import ggml_tensor_qtype
-from bigdl.llm.utils.common import invalidInputError
+from bigdl.llm.utils.common import invalidInputError, MuteHFLogger
 
 
 def save_low_bit(self, *args, **kwargs):
@@ -166,10 +166,9 @@ class _BaseAutoModelClass:
         variant = kwargs.get("variant", None)
 
         from .convert import ggml_convert_quant
-        model = cls.HF_Model.from_pretrained(*args, **kwargs)
-        print("Note: If there are warnings during the model loading process, "
-              "they can be safely ignored; "
-              "the model will be loaded with INT4 optimizations applied.")
+
+        with MuteHFLogger(logger=transformers.modeling_utils.logger):
+            model = cls.HF_Model.from_pretrained(*args, **kwargs)
 
         # add save_low_bit to pretrained model dynamically
         import types
diff --git a/python/llm/src/bigdl/llm/utils/common/__init__.py b/python/llm/src/bigdl/llm/utils/common/__init__.py
index 22fc3510..feaf24bf 100644
--- a/python/llm/src/bigdl/llm/utils/common/__init__.py
+++ b/python/llm/src/bigdl/llm/utils/common/__init__.py
@@ -19,5 +19,5 @@
 # Otherwise there would be module not found error in non-pip's setting as Python would
 # only search the first bigdl package and end up finding only one sub-package.
 
-from .log4Error import invalidInputError, invalidOperationError
+from .log4Error import invalidInputError, invalidOperationError, MuteHFLogger
 from .lazyimport import LazyImport
diff --git a/python/llm/src/bigdl/llm/utils/common/log4Error.py b/python/llm/src/bigdl/llm/utils/common/log4Error.py
index 9db756f0..229fd2f8 100644
--- a/python/llm/src/bigdl/llm/utils/common/log4Error.py
+++ b/python/llm/src/bigdl/llm/utils/common/log4Error.py
@@ -39,3 +39,15 @@ def invalidOperationError(condition, errMsg, fixMsg=None, cause=None):
             raise cause
         else:
             raise RuntimeError(errMsg)
+
+class MuteHFLogger():
+    def __init__(self, logger, speak_level=logging.ERROR) -> None:
+        self.logger = logger
+        self.speak_level = speak_level
+        self.old_level = logger.getEffectiveLevel()
+
+    def __enter__(self):
+        self.logger.setLevel(self.speak_level)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.logger.setLevel(self.old_level)