new: added --quiet/-Q flag to inspect action

dreadnode · Oct 28, 2024 · 3b2e374 · 3b2e374
1 parent c0dd0c4
commit 3b2e374
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 70 deletions.
diff --git a/src/cli/inspect.rs b/src/cli/inspect.rs
@@ -3,15 +3,18 @@ use crate::core::FileType;
 use super::InspectArgs;
 
 pub(crate) fn inspect(args: InspectArgs) -> anyhow::Result<()> {
-    println!(
-        "Inspecting {:?} (detail={:?}{}):\n",
-        args.file_path,
-        args.detail,
-        args.filter
-            .as_ref()
-            .map(|f| format!(" filter_by={:?}", f))
-            .unwrap_or("".to_string())
-    );
+    let quiet = args.quiet.unwrap_or(false);
+    if !quiet {
+        println!(
+            "Inspecting {:?} (detail={:?}{}):\n",
+            args.file_path,
+            args.detail,
+            args.filter
+                .as_ref()
+                .map(|f| format!(" filter_by={:?}", f))
+                .unwrap_or("".to_string())
+        );
+    }
 
     let forced_format = args.format.unwrap_or(FileType::Unknown);
     let inspection = if forced_format.is_safetensors()
@@ -28,84 +31,88 @@ pub(crate) fn inspect(args: InspectArgs) -> anyhow::Result<()> {
         anyhow::bail!("unsupported file format")
     };
 
-    println!("file type:     {}", inspection.file_type);
-    println!("version:       {}", inspection.version);
-    println!(
-        "file size:     {} ({})",
-        humansize::format_size(inspection.file_size, humansize::DECIMAL),
-        inspection.file_size
-    );
-    if inspection.header_size > 0 {
+    if !quiet {
+        println!("file type:     {}", inspection.file_type);
+        println!("version:       {}", inspection.version);
         println!(
-            "header size:   {} ({})",
-            humansize::format_size(inspection.header_size, humansize::DECIMAL),
-            inspection.header_size
+            "file size:     {} ({})",
+            humansize::format_size(inspection.file_size, humansize::DECIMAL),
+            inspection.file_size
+        );
+        if inspection.header_size > 0 {
+            println!(
+                "header size:   {} ({})",
+                humansize::format_size(inspection.header_size, humansize::DECIMAL),
+                inspection.header_size
+            );
+        }
+        println!("total tensors: {}", inspection.num_tensors);
+        println!(
+            "data size:     {} ({})",
+            humansize::format_size(inspection.data_size, humansize::DECIMAL),
+            inspection.data_size
+        );
+        println!(
+            "average size:  {}",
+            humansize::format_size(inspection.average_tensor_size(), humansize::DECIMAL)
         );
-    }
-    println!("total tensors: {}", inspection.num_tensors);
-    println!(
-        "data size:     {} ({})",
-        humansize::format_size(inspection.data_size, humansize::DECIMAL),
-        inspection.data_size
-    );
-    println!(
-        "average size:  {}",
-        humansize::format_size(inspection.average_tensor_size(), humansize::DECIMAL)
-    );
 
-    println!("data types:    {}", inspection.unique_dtypes.join(", "));
-    println!(
-        "shapes:        {}",
-        inspection
-            .unique_shapes
-            .iter()
-            .map(|s| format!("{:?}", s))
-            .collect::<Vec<_>>()
-            .join(", ")
-    );
+        println!("data types:    {}", inspection.unique_dtypes.join(", "));
+        println!(
+            "shapes:        {}",
+            inspection
+                .unique_shapes
+                .iter()
+                .map(|s| format!("{:?}", s))
+                .collect::<Vec<_>>()
+                .join(", ")
+        );
 
-    if !inspection.metadata.is_empty() {
-        println!("\nmetadata:\n");
-        for (meta_key, meta_value) in &inspection.metadata {
-            println!("  {}: {}", meta_key, meta_value);
+        if !inspection.metadata.is_empty() {
+            println!("\nmetadata:\n");
+            for (meta_key, meta_value) in &inspection.metadata {
+                println!("  {}: {}", meta_key, meta_value);
+            }
         }
-    }
 
-    if let Some(tensors) = &inspection.tensors {
-        println!("\ntensors:\n");
+        if let Some(tensors) = &inspection.tensors {
+            println!("\ntensors:\n");
 
-        for tensor_info in tensors {
-            println!(
-                "  {}",
-                tensor_info
-                    .id
-                    .as_ref()
-                    .unwrap_or(&"<no tensor id>".to_string())
-            );
-            println!("    dtype: {:?}", tensor_info.dtype);
-            println!("    shape: {:?}", tensor_info.shape);
-            println!(
-                "    size: {} ({})",
-                humansize::format_size(tensor_info.size, humansize::DECIMAL),
-                tensor_info.size
-            );
+            for tensor_info in tensors {
+                println!(
+                    "  {}",
+                    tensor_info
+                        .id
+                        .as_ref()
+                        .unwrap_or(&"<no tensor id>".to_string())
+                );
+                println!("    dtype: {:?}", tensor_info.dtype);
+                println!("    shape: {:?}", tensor_info.shape);
+                println!(
+                    "    size: {} ({})",
+                    humansize::format_size(tensor_info.size, humansize::DECIMAL),
+                    tensor_info.size
+                );
 
-            if !tensor_info.metadata.is_empty() {
-                println!("    metadata:");
-                for (meta_key, meta_value) in &tensor_info.metadata {
-                    println!("      {}: {}", meta_key, meta_value);
+                if !tensor_info.metadata.is_empty() {
+                    println!("    metadata:");
+                    for (meta_key, meta_value) in &tensor_info.metadata {
+                        println!("      {}: {}", meta_key, meta_value);
+                    }
                 }
-            }
 
-            println!();
+                println!();
+            }
         }
     }
 
     if let Some(json_file_path) = &args.to_json {
         let json_str = serde_json::to_string_pretty(&inspection)?;
         std::fs::write(json_file_path, json_str)?;
 
-        println!("\nsaved to {:?}", json_file_path);
+        if !quiet {
+            println!("\nsaved to {:?}", json_file_path);
+        }
     }
 
     Ok(())

diff --git a/src/cli/mod.rs b/src/cli/mod.rs
@@ -54,6 +54,9 @@ pub(crate) struct InspectArgs {
     /// If the detail level is set to full, filter the tensors by this substring.
     #[clap(long, short = 'F')]
     filter: Option<String>,
+    /// Suppress inspection output.
+    #[clap(long, short = 'Q')]
+    quiet: Option<bool>,
     /// Save as JSON to the specified file.
     #[clap(long, short = 'J')]
     to_json: Option<PathBuf>,