Gate debug messages in dump_gguf.py behind -v flag (#1105)

I need this because I've been using `dump_gguf.py` for some integration
tests (see
https://github.com/nod-ai/shark-ai/blob/41a09a836c4558745ec1ca01e80a819cd761e608/app_tests/integration_tests/llm/model_management.py#L223-L239)

Author: renxida

sharktank/sharktank/tools/dump_gguf.py

@@ -6,6 +6,7 @@
 
 from pathlib import Path
 import re
+import logging
 
 import numpy as np
 import torch
@@ -17,6 +18,9 @@
 def main():
     from ..utils import cli
 
+    # Set up logging
+    logger = logging.getLogger(__name__)
+
     parser = cli.create_parser()
     cli.add_input_dataset_options(parser)
     parser.add_argument(
@@ -28,63 +32,79 @@ def main():
     parser.add_argument(
         "--save", type=Path, help="Save the GGUF dataset to an IRPA file"
     )
+    parser.add_argument(
+        "--verbose", "-v", action="store_true", help="Enable verbose output"
+    )
     args = cli.parse(parser)
+
+    # Configure logging based on verbosity
+    if args.verbose:
+        logging.basicConfig(level=logging.DEBUG)
+    else:
+        logging.basicConfig(level=logging.INFO)
+
     config = cli.get_input_dataset(args)
 
     if args.save is not None:
 
         def report(s):
-            print(f"Save: {s}")
+            logger.info(f"Save: {s}")
 
-        print(f"Saving to: {args.save}")
+        logger.info(f"Saving to: {args.save}")
         config.save(args.save, io_report_callback=report)
         return
 
-    print(f"Properties:")
+    logger.debug("Properties:")
     for key, value in config.properties.items():
-        print(f"  {key} = {value} (of {type(value)})")
-    print("Tensors:")
+        logger.debug(f"  {key} = {value} (of {type(value)})")
+
+    logger.debug("Tensors:")
     for tensor in config.root_theta.flatten().values():
         if args.tensor_regex is not None:
             if not re.search(args.tensor_regex, tensor.name):
                 continue
-        print(f"  {tensor}")
+
+        logger.debug(f"  {tensor}")
         if isinstance(tensor, PrimitiveTensor):
             torch_tensor = tensor.as_torch()
-            print(
+            logger.debug(
                 f"    : torch.Tensor({list(torch_tensor.shape)}, "
                 f"dtype={torch_tensor.dtype}) = {tensor.as_torch()}"
             )
         elif isinstance(tensor, QuantizedTensor):
-            print(f"    : QuantizedTensor({tensor.layout_type.__name__})")
+            logger.debug(f"    : QuantizedTensor({tensor.layout_type.__name__})")
             try:
                 unpacked = tensor.unpack()
-                print(f"    {unpacked}")
+                logger.debug(f"    {unpacked}")
             except NotImplementedError:
-                print(f"     NOT IMPLEMENTED")
+                logger.warning(f"    Unpacking NOT IMPLEMENTED for {tensor.name}")
         elif isinstance(tensor, ShardedTensor):
             for i, pt in enumerate(tensor.shards):
-                print(f"    {i}: {pt}")
+                logger.debug(f"    {i}: {pt}")
 
-        _maybe_dump_tensor(args, tensor)
+        _maybe_dump_tensor(args, tensor, logger)
 
 
-def _maybe_dump_tensor(args, t: InferenceTensor):
+def _maybe_dump_tensor(args, t: InferenceTensor, logger: logging.Logger):
     if not args.dump_tensor_dir:
         return
     dir: Path = args.dump_tensor_dir
     dir.mkdir(parents=True, exist_ok=True)
-    print(f"    (Dumping to {dir})")
-
-    if isinstance(t, PrimitiveTensor):
-        torch_tensor = t.as_torch()
-        np.save(dir / f"{t.name}.npy", torch_tensor.detach().numpy())
-    elif isinstance(t, QuantizedTensor):
-        layout: QuantizedLayout = t.unpack()
-        dq = layout.dequant()
-        np.save(dir / f"{t.name}.dequant.npy", dq.detach().numpy())
-    else:
-        raise AssertionError(f"Unexpected tensor type: {type(t)}")
+    logger.info(f"Dumping tensor {t.name} to {dir}")
+
+    try:
+        if isinstance(t, PrimitiveTensor):
+            torch_tensor = t.as_torch()
+            np.save(dir / f"{t.name}.npy", torch_tensor.detach().numpy())
+        elif isinstance(t, QuantizedTensor):
+            layout: QuantizedLayout = t.unpack()
+            dq = layout.dequant()
+            np.save(dir / f"{t.name}.dequant.npy", dq.detach().numpy())
+        else:
+            logger.error(f"Unexpected tensor type: {type(t)}")
+            raise AssertionError(f"Unexpected tensor type: {type(t)}")
+    except Exception as e:
+        logger.error(f"Failed to dump tensor {t.name}: {str(e)}")
 
 
 if __name__ == "__main__":