Verbosity control for export

Summary: Verbosity control for export

Reviewed By: kimishpatel

Differential Revision: D55374748

Author: Michael Gschwind

examples/models/llama2/export_llama_lib.py

@@ -46,6 +46,7 @@
 logging.basicConfig(level=logging.INFO, format=FORMAT)
 
 pkg_name = __name__
+verbosity_setting = None
 
 
 def set_pkg_name(name: str) -> None:
@@ -57,6 +58,15 @@ def get_resource_path(resource_name) -> str:
     return pkg_resources.resource_filename(pkg_name, resource_name)
 
 
+def set_verbosity(val):
+    global verbosity_setting
+    verbosity_setting = val
+
+
+def verbose_export():
+    return verbosity_setting
+
+
 @dataclass
 class EmbeddingQuantOptions:
     is_per_channel: bool = True
@@ -240,7 +250,8 @@ def quantize(
         from torchao.quantization.quant_api import Int8DynActInt4WeightQuantizer
 
         model = Int8DynActInt4WeightQuantizer(precision=torch_dtype).quantize(model)
-        print("quantized model:", model)
+        if verbose_export():
+            print("quantized model:", model)
         return model
     elif qmode == "8da4w-gptq":
         from torchao.quantization.quant_api import Int8DynActInt4WeightGPTQQuantizer
@@ -448,7 +459,9 @@ def build_args_parser() -> argparse.ArgumentParser:
 
 def canonical_path(path: str, *, dir: bool = False) -> str:
 
-    print(f"creating canonical path for {path}")
+    if verbose_export():
+        print(f"creating canonical path for {path}")
+
     if not path.startswith("par:"):
         return path
 
@@ -457,7 +470,8 @@ def canonical_path(path: str, *, dir: bool = False) -> str:
         return path[4:]
     else:
         return_val = pkg_resources.resource_filename(pkg_name, path[4:])
-        print(f"canonical name is: {return_val}")
+        if verbose_export():
+            print(f"canonical name is: {return_val}")
         return return_val