diff --git a/src/serve.rs b/src/serve.rs index 1a5d9497..0392f198 100644 --- a/src/serve.rs +++ b/src/serve.rs @@ -78,13 +78,18 @@ pub fn list_models(dir: &Path) -> Vec<(String, PathBuf, u64)> { out } -/// The models among files given as (name, bytes): .gguf files, a split model once by its first shard with the -/// bytes of all its shards, vision projectors (`mmproj*`) left out. +/// The models among files given as (name, bytes): .gguf files, complete split models once by their first shard +/// with the bytes of their shards, vision projectors (`mmproj*`) left out. pub fn models(files: &[(String, u64)]) -> Vec<(String, u64)> { - let parts = |stem: &str| files.iter().filter(|(m, _)| shard(m).is_some_and(|(s, _)| s == stem)).map(|x| x.1).sum(); + let sizes: std::collections::HashMap<&str, u64> = files.iter().map(|(name, bytes)| (name.as_str(), *bytes)).collect(); files.iter().filter(|(n, _)| n.ends_with(".gguf") && !n.rsplit('/').next().unwrap_or(n).starts_with("mmproj")) - .filter_map(|(n, b)| match shard(n) { - Some((stem, 1)) => Some((n.clone(), parts(stem))), + .filter_map(|(n, b)| match shard_parts(n) { + Some((stem, 1, total)) if total > 0 && total as usize <= files.len() => { + (1..=total).try_fold(0u64, |bytes, i| { + sizes.get(format!("{stem}-{i:05}-of-{total:05}.gguf").as_str()) + .and_then(|part| bytes.checked_add(*part)) + }).map(|bytes| (n.clone(), bytes)) + } Some(_) => None, None => Some((n.clone(), *b)), }).collect() @@ -92,11 +97,15 @@ pub fn models(files: &[(String, u64)]) -> Vec<(String, u64)> { /// "dir/name-00002-of-00003.gguf" gives ("dir/name", 2). pub fn shard(n: &str) -> Option<(&str, u32)> { + shard_parts(n).map(|(stem, i, _)| (stem, i)) +} + +fn shard_parts(n: &str) -> Option<(&str, u32, u32)> { let (head, total) = n.strip_suffix(".gguf")?.rsplit_once("-of-")?; let (stem, i) = head.rsplit_once('-')?; let digits = |s: &str| s.len() == 5 && s.bytes().all(|b| b.is_ascii_digit()); if !digits(total) || !digits(i) { return None; } - Some((stem, i.parse().ok()?)) + Some((stem, i.parse().ok()?, total.parse().ok()?)) } /// The parts of an OpenAI chat completion request that strata uses. Sampling fields left out fall back to the diff --git a/tests/serve.rs b/tests/serve.rs index 03c13212..9eb0a4e5 100644 --- a/tests/serve.rs +++ b/tests/serve.rs @@ -1,5 +1,5 @@ use rusting_strata::chat::{Hold, Json, Msg, Template}; -use rusting_strata::serve::{chat_request, host_ok, key_ok, list_models, messages_body, read_request}; +use rusting_strata::serve::{Keys, Who, chat_request, host_ok, key_ok, list_models, messages_body, models, read_request}; /// Every request in tests/fixtures/chat renders byte-identical to llama.cpp's Jinja rendering of the model's own /// template (tools/gen_chat_fixtures.py); a request llama.cpp rejects must be rejected too. @@ -234,6 +234,26 @@ fn lists_models_once_per_split() { assert_eq!(l, [("a.gguf".into(), "a.gguf".into(), 3), ("org/repo/big-00001-of-00002.gguf".into(), "org/repo/big-00001-of-00002.gguf".into(), 12)]); } +#[test] +fn lists_only_complete_split_models() { + let files = vec![ + ("model-00001-of-00003.gguf".into(), 10), + ("model-00002-of-00003.gguf".into(), 20), + ("model-00001-of-00002.gguf".into(), 100), + ("model-00002-of-00002.gguf".into(), 20), + ("plain.gguf".into(), 7), + ]; + assert_eq!(models(&files), vec![("model-00001-of-00002.gguf".into(), 120), ("plain.gguf".into(), 7)]); + + let mut complete = files; + complete.push(("model-00003-of-00003.gguf".into(), 30)); + assert_eq!(models(&complete), vec![ + ("model-00001-of-00003.gguf".into(), 60), + ("model-00001-of-00002.gguf".into(), 120), + ("plain.gguf".into(), 7), + ]); +} + #[test] fn download_names_stay_in_the_models_folder() { use rusting_strata::hf::local_name;