huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-14 09:57:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Tensor tools print all (#2543)

Browse Source 
* Support whisper large-v3 turbo in the whisper-microphone example.

* Print all tensors when no argument is provided.
This commit is contained in:
Laurent Mazare
2024-10-05 10:05:14 +02:00
committed by GitHub
parent 410c89f72a
commit d2e432914e
1 changed files with 29 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
tensor-tools/src/main.rs
View File

@ -197,6 +197,11 @@ fn run_print(
match format { match format {
Format::Npz => { Format::Npz => {
let tensors = candle::npy::NpzTensors::new(file)?; let tensors = candle::npy::NpzTensors::new(file)?;
let names = if names.is_empty() {
tensors.names().into_iter().map(|v| v.to_string()).collect()
} else {
names
};
for name in names.iter() { for name in names.iter() {
println!("==== {name} ===="); println!("==== {name} ====");
match tensors.get(name)? { match tensors.get(name)? {
@ -209,6 +214,11 @@ fn run_print(
use candle::safetensors::Load; use candle::safetensors::Load;
let tensors = unsafe { candle::safetensors::MmapedSafetensors::new(file)? }; let tensors = unsafe { candle::safetensors::MmapedSafetensors::new(file)? };
let tensors: std::collections::HashMap<_, _> = tensors.tensors().into_iter().collect(); let tensors: std::collections::HashMap<_, _> = tensors.tensors().into_iter().collect();
let names = if names.is_empty() {
tensors.keys().map(|v| v.to_string()).collect()
} else {
names
};
for name in names.iter() { for name in names.iter() {
println!("==== {name} ===="); println!("==== {name} ====");
match tensors.get(name) { match tensors.get(name) {
@ -222,6 +232,15 @@ fn run_print(
} }
Format::Pth => { Format::Pth => {
let pth_file = candle::pickle::PthTensors::new(file, None)?; let pth_file = candle::pickle::PthTensors::new(file, None)?;
let names = if names.is_empty() {
pth_file
.tensor_infos()
.keys()
.map(|v| v.to_string())
.collect()
} else {
names
};
for name in names.iter() { for name in names.iter() {
println!("==== {name} ===="); println!("==== {name} ====");
match pth_file.get(name)? { match pth_file.get(name)? {
@ -238,6 +257,11 @@ fn run_print(
Format::Ggml => { Format::Ggml => {
let mut file = std::fs::File::open(file)?; let mut file = std::fs::File::open(file)?;
let content = candle::quantized::ggml_file::Content::read(&mut file, device)?; let content = candle::quantized::ggml_file::Content::read(&mut file, device)?;
let names = if names.is_empty() {
content.tensors.keys().map(|v| v.to_string()).collect()
} else {
names
};
for name in names.iter() { for name in names.iter() {
println!("==== {name} ===="); println!("==== {name} ====");
match content.tensors.get(name) { match content.tensors.get(name) {
@ -252,6 +276,11 @@ fn run_print(
Format::Gguf => { Format::Gguf => {
let mut file = std::fs::File::open(file)?; let mut file = std::fs::File::open(file)?;
let content = gguf_file::Content::read(&mut file)?; let content = gguf_file::Content::read(&mut file)?;
let names = if names.is_empty() {
content.tensor_infos.keys().map(|v| v.to_string()).collect()
} else {
names
};
for name in names.iter() { for name in names.iter() {
println!("==== {name} ===="); println!("==== {name} ====");
match content.tensor(&mut file, name, device) { match content.tensor(&mut file, name, device) {