Add the tensor-tools binary. (#510)

candle-core/examples/tensor-tools.rs

@@ -0,0 +1,72 @@
+use candle_core::Result;
+use clap::{Parser, Subcommand};
+
+#[derive(Subcommand, Debug, Clone)]
+enum Command {
+    Ls { files: Vec<std::path::PathBuf> },
+}
+
+#[derive(Parser, Debug, Clone)]
+struct Args {
+    /// Enable verbose mode.
+    #[arg(short, long)]
+    verbose: bool,
+
+    #[command(subcommand)]
+    command: Command,
+}
+
+fn run_ls(file: &std::path::PathBuf) -> Result<()> {
+    match file.extension().and_then(|e| e.to_str()) {
+        Some("npz") => {
+            let tensors = candle_core::npy::NpzTensors::new(file)?;
+            let mut names = tensors.names();
+            names.sort();
+            for name in names {
+                let shape_dtype = match tensors.get_shape_and_dtype(name) {
+                    Ok((shape, dtype)) => format!("[{shape:?}; {dtype:?}]"),
+                    Err(err) => err.to_string(),
+                };
+                println!("{name}: {shape_dtype}")
+            }
+        }
+        Some("safetensor") | Some("safetensors") => {
+            let tensors = unsafe { candle_core::safetensors::MmapedFile::new(file)? };
+            let tensors = tensors.deserialize()?;
+            let mut tensors = tensors.tensors();
+            tensors.sort_by(|a, b| a.0.cmp(&b.0));
+            for (name, view) in tensors.iter() {
+                let dtype = view.dtype();
+                let dtype = match candle_core::DType::try_from(dtype) {
+                    Ok(dtype) => format!("{dtype:?}"),
+                    Err(_) => format!("{dtype:?}"),
+                };
+                let shape = view.shape();
+                println!("{name}: [{shape:?}; {dtype}]")
+            }
+        }
+        Some(_) => {
+            println!("{file:?}: unsupported file extension")
+        }
+        None => {
+            println!("{file:?}: no file extension")
+        }
+    }
+    Ok(())
+}
+
+fn main() -> anyhow::Result<()> {
+    let args = Args::parse();
+    match args.command {
+        Command::Ls { files } => {
+            let multiple_files = files.len() > 1;
+            for file in files.iter() {
+                if multiple_files {
+                    println!("--- {file:?} ---");
+                }
+                run_ls(file)?
+            }
+        }
+    }
+    Ok(())
+}