diff --git a/Cargo.lock b/Cargo.lock index b7494ac..606e820 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,9 +4,9 @@ version = 3 [[package]] name = "anstream" -version = "0.6.14" +version = "0.6.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "418c75fa768af9c03be99d17643f93f79bbba589895012a80e3452a19ddda15b" +checksum = "64e15c1ab1f89faffbf04a634d5e1962e9074f2741eef6d97f3c4e322426d526" dependencies = [ "anstyle", "anstyle-parse", @@ -19,33 +19,33 @@ dependencies = [ [[package]] name = "anstyle" -version = "1.0.7" +version = "1.0.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "038dfcf04a5feb68e9c60b21c9625a54c2c0616e79b72b0fd87075a056ae1d1b" +checksum = "1bec1de6f59aedf83baf9ff929c98f2ad654b97c9510f4e70cf6f661d49fd5b1" [[package]] name = "anstyle-parse" -version = "0.2.4" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c03a11a9034d92058ceb6ee011ce58af4a9bf61491aa7e1e59ecd24bd40d22d4" +checksum = "eb47de1e80c2b463c735db5b217a0ddc39d612e7ac9e2e96a5aed1f57616c1cb" dependencies = [ "utf8parse", ] [[package]] name = "anstyle-query" -version = "1.1.0" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ad186efb764318d35165f1758e7dcef3b10628e26d41a44bc5550652e6804391" +checksum = "6d36fc52c7f6c869915e99412912f22093507da8d9e942ceaf66fe4b7c14422a" dependencies = [ "windows-sys", ] [[package]] name = "anstyle-wincon" -version = "3.0.3" +version = "3.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61a38449feb7068f52bb06c12759005cf459ee52bb4adc1d5a7c4322d716fb19" +checksum = "5bf74e1b6e971609db8ca7a9ce79fd5768ab6ae46441c572e46cf596f59e57f8" dependencies = [ "anstyle", "windows-sys", @@ -53,9 +53,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.10" +version = "4.5.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f6b81fb3c84f5563d509c59b5a48d935f689e993afa90fe39047f05adef9142" +checksum = "0fbb260a053428790f3de475e304ff84cdbc4face759ea7a3e64c1edd938a7fc" dependencies = [ "clap_builder", "clap_derive", @@ -63,9 +63,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.10" +version = "4.5.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca6706fd5224857d9ac5eb9355f6683563cc0541c7cd9d014043b57cbec78ac" +checksum = "64b17d7ea74e9f833c7dbf2cbe4fb12ff26783eda4782a8975b72f895c9b4d99" dependencies = [ "anstream", "anstyle", @@ -75,9 +75,9 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.5.8" +version = "4.5.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2bac35c6dafb060fd4d275d9a4ffae97917c13a6327903a8be2153cd964f7085" +checksum = "501d359d5f3dcaf6ecdeee48833ae73ec6e42723a1e52419c79abf9507eec0a0" dependencies = [ "heck", "proc-macro2", @@ -87,15 +87,15 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.7.1" +version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b82cf0babdbd58558212896d1a4272303a57bdb245c2bf1147185fb45640e70" +checksum = "1462739cb27611015575c0c11df5df7601141071f07518d56fcc1be504cbec97" [[package]] name = "colorchoice" -version = "1.0.1" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b6a852b24ab71dffc585bcb46eaf7959d175cb865a7152e35b348d1b2960422" +checksum = "d3fd119d74b830634cea2a0f58bbd0d54540518a14397557951e79340abc28c0" [[package]] name = "equivalent" @@ -124,9 +124,9 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "indexmap" -version = "2.2.6" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "168fb715dda47215e360912c096649d23d58bf392ac62f73919e831745e40f26" +checksum = "de3fc2e30ba82dd1b3911c8de1ffc143c74a914a14e99514d7637e3099df5ea0" dependencies = [ "equivalent", "hashbrown", @@ -134,9 +134,9 @@ dependencies = [ [[package]] name = "is_terminal_polyfill" -version = "1.70.0" +version = "1.70.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8478577c03552c21db0e2724ffb8986a5ce7af88107e6be5d2ee6e158c12800" +checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" [[package]] name = "libc" diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 2e6ab70..8be9ff7 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -7,7 +7,9 @@ mod split; #[macro_use] extern crate clap; + use clap::Parser; +use std::collections::HashSet; fn main() { use Commands::*; @@ -35,3 +37,12 @@ enum Commands { const YES: &str = "✔️ "; const ERR: &str = "❌ "; + +fn split_keys(arg: &Option) -> Option> { + arg.as_ref().map(|s| { + s.split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .collect::>() + }) +} diff --git a/xtask/src/show.rs b/xtask/src/show.rs index 6ff9890..377e44e 100644 --- a/xtask/src/show.rs +++ b/xtask/src/show.rs @@ -1,9 +1,9 @@ -use crate::{loose_shards::LooseShards, ERR, YES}; +use crate::{loose_shards::LooseShards, split_keys, ERR, YES}; use ggus::{ GGufFileHeader, GGufMetaDataValueType, GGufMetaKV, GGufMetaKVPairs, GGufReadError, GGufReader, GGufTensors, }; -use std::{fmt, fs::File, path::PathBuf}; +use std::{collections::HashSet, fmt, fs::File, path::PathBuf}; #[derive(Args, Default)] pub struct ShowArgs { @@ -12,19 +12,45 @@ pub struct ShowArgs { /// If set, show all shards in the directory #[clap(long)] shards: bool, + /// How many elements to show in arrays, `all` for all elements + #[clap(long, short = 'n', default_value = "8")] + array_detail: String, + /// Meta to show (split with `,`) + #[clap(long, short = 'm')] + filter_meta: Option, + /// Tensors to show (split with `,`) + #[clap(long, short = 't')] + filter_tensor: Option, } struct Failed; impl ShowArgs { pub fn show(self) { - let files = if self.shards { - LooseShards::from(&*self.file) + let Self { + file, + shards, + array_detail, + filter_meta, + filter_tensor, + } = self; + + let detail = match array_detail.trim().to_lowercase().as_str() { + "all" => usize::MAX, + s => s + .parse() + .expect("Invalid array detail, should be an integer or `all`"), + }; + let filter_meta = split_keys(&filter_meta); + let filter_tensor = split_keys(&filter_tensor); + + let files = if shards { + LooseShards::from(&*file) .into_iter() .filter(|p| p.is_file()) .collect::>() - } else if self.file.is_file() { - vec![self.file] + } else if file.is_file() { + vec![file] } else { vec![] }; @@ -64,7 +90,7 @@ impl ShowArgs { continue; } }; - if let Err(Failed) = show_meta_kvs(&kvs) { + if show_meta_kvs(&kvs, &filter_meta, detail).is_err() { println!(); continue; } @@ -78,7 +104,7 @@ impl ShowArgs { continue; } }; - let _ = show_tensors(&tensors); + let _ = show_tensors(&tensors, &filter_tensor); println!(); } } @@ -129,11 +155,24 @@ fn show_header(header: &GGufFileHeader) -> Result<(), Failed> { Ok(()) } -fn show_meta_kvs(kvs: &GGufMetaKVPairs) -> Result<(), Failed> { - if let Some(width) = kvs.keys().map(|k| k.len()).max() { +fn show_meta_kvs<'a>( + kvs: &GGufMetaKVPairs, + filter: &Option>, + detail: usize, +) -> Result<(), Failed> { + let kvs = filter.as_ref().map_or_else( + || kvs.kvs().collect::>(), + |to_keep| { + kvs.kvs() + .filter(move |m| to_keep.contains(m.key())) + .collect::>() + }, + ); + + if let Some(width) = kvs.iter().map(|kv| kv.key().len()).max() { show_title("Meta KV"); - for kv in kvs.kvs() { - show_meta_kv(kv, width)?; + for kv in kvs { + show_meta_kv(kv, width, detail)?; } println!(); } @@ -141,12 +180,12 @@ fn show_meta_kvs(kvs: &GGufMetaKVPairs) -> Result<(), Failed> { Ok(()) } -fn show_meta_kv(kv: GGufMetaKV, width: usize) -> Result<(), Failed> { +fn show_meta_kv(kv: GGufMetaKV, width: usize, detail: usize) -> Result<(), Failed> { let key = kv.key(); let ty = kv.ty(); let mut reader = kv.value_reader(); let mut buf = String::new(); - match fmt_meta_val(&mut reader, ty, 1, &mut buf) { + match fmt_meta_val(&mut reader, ty, 1, detail, &mut buf) { Ok(()) => { println!("{YES}{key:·( reader: &mut GGufReader<'a>, ty: GGufMetaDataValueType, len: usize, + detail: usize, buf: &mut String, ) -> Result<(), GGufReadError<'a>> { struct MultiLines<'a>(&'a str); @@ -212,34 +252,47 @@ fn fmt_meta_val<'a>( } T::Array => { let (ty, len) = reader.read_arr_header()?; - fmt_meta_val(reader, ty, len, buf)?; + fmt_meta_val(reader, ty, len, detail, buf)?; } } } - _ if len <= 8 => { + _ if len <= detail => { buf.push('['); for i in 0..len { if i > 0 { buf.push_str(", "); } - fmt_meta_val(reader, ty, 1, buf)?; + fmt_meta_val(reader, ty, 1, detail, buf)?; } buf.push(']'); } _ => { buf.push('['); - for _ in 0..8 { - fmt_meta_val(reader, ty, 1, buf)?; + for _ in 0..detail { + fmt_meta_val(reader, ty, 1, detail, buf)?; buf.push_str(", "); } - buf.push_str(&format!("...({} more)]", len - 8)); + buf.push_str(&format!("...({} more)]", len - detail)); } } Ok(()) } -fn show_tensors(tensors: &GGufTensors) -> Result<(), Failed> { - if let Some(name_width) = tensors.names().map(|k| k.len()).max() { +fn show_tensors<'a>( + tensors: &GGufTensors, + filter: &Option>, +) -> Result<(), Failed> { + let tensors = filter.as_ref().map_or_else( + || tensors.iter().collect::>(), + |to_keep| { + tensors + .iter() + .filter(move |t| to_keep.contains(t.name())) + .collect::>() + }, + ); + + if let Some(name_width) = tensors.iter().map(|t| t.name().len()).max() { show_title("Tensors"); let tensors = tensors.iter().collect::>(); let off_width = tensors.last().unwrap().offset().to_string().len() + 1;