kernel/discovery/
gguf_shards.rs

1//! Recognizing and grouping multi-part GGUF weight files named like
2//! `model-00001-of-00005.gguf`.
3
4use std::collections::{BTreeMap, BTreeSet};
5use std::path::{Path, PathBuf};
6
7/// The parts of a sharded GGUF filename.
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct ShardName {
10    /// The base name shared by every shard in the set.
11    pub base: String,
12    /// This shard's 1-based index.
13    pub index: usize,
14    /// The total number of shards declared in the name.
15    pub total: usize,
16}
17
18/// One member file of a shard group.
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct Member {
21    /// The shard's 1-based index.
22    pub index: usize,
23    /// The file's path.
24    pub path: PathBuf,
25    /// The file's size in bytes.
26    pub bytes: i64,
27}
28
29/// A set of shards that together form one model.
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub struct ShardGroup {
32    /// The shared base name.
33    pub base: String,
34    /// The declared total shard count.
35    pub total: usize,
36    /// The member files present, sorted by index.
37    pub members: Vec<Member>,
38}
39
40impl ShardGroup {
41    /// The path of the first shard (index 1), if present.
42    pub fn first_shard(&self) -> Option<&Path> {
43        self.members
44            .iter()
45            .find(|member| member.index == 1)
46            .map(|member| member.path.as_path())
47    }
48
49    /// The total size of all present members.
50    pub fn footprint_bytes(&self) -> i64 {
51        self.members.iter().map(|member| member.bytes).sum()
52    }
53
54    /// Whether every declared shard is present.
55    pub fn complete(&self) -> bool {
56        self.members.len() == self.total
57    }
58}
59
60/// Parse a `<base>-<index>-of-<total>.gguf` filename. The index and total fields
61/// are exactly five digits, `1 <= index <= total`, and the base is non-empty.
62pub fn parse(filename: &str) -> Option<ShardName> {
63    let cut = filename.len().checked_sub(".gguf".len())?;
64    if !filename.get(cut..)?.eq_ignore_ascii_case(".gguf") {
65        return None;
66    }
67    let stem = &filename[..cut];
68    let of_position = stem.rfind("-of-")?;
69    let total_field = &stem[of_position + "-of-".len()..];
70    let total = five_digit_number(total_field)?;
71    if total == 0 {
72        return None;
73    }
74    let head = &stem[..of_position];
75    let dash_position = head.rfind('-')?;
76    let index_field = &head[dash_position + 1..];
77    let index = five_digit_number(index_field)?;
78    if index == 0 || index > total {
79        return None;
80    }
81    let base = &head[..dash_position];
82    if base.is_empty() {
83        return None;
84    }
85    Some(ShardName {
86        base: base.to_owned(),
87        index,
88        total,
89    })
90}
91
92/// Build the canonical filename for a shard.
93pub fn shard_filename(base: &str, index: usize, total: usize) -> String {
94    format!("{base}-{index:05}-of-{total:05}.gguf")
95}
96
97/// Split files into shard groups and loose (non-sharded) files. Shards are
98/// grouped by directory, base name, and declared total, and each group's members
99/// are sorted by index.
100pub fn group(files: &[(PathBuf, i64)]) -> (Vec<ShardGroup>, Vec<PathBuf>) {
101    let mut buckets: BTreeMap<(String, String, usize), Vec<Member>> = BTreeMap::new();
102    let mut loose: Vec<PathBuf> = Vec::new();
103
104    for (path, bytes) in files {
105        let shard = path
106            .file_name()
107            .and_then(|name| name.to_str())
108            .and_then(parse);
109        match shard {
110            Some(shard) => {
111                let directory = path
112                    .parent()
113                    .map(|dir| dir.to_string_lossy().into_owned())
114                    .unwrap_or_default();
115                buckets
116                    .entry((directory, shard.base, shard.total))
117                    .or_default()
118                    .push(Member {
119                        index: shard.index,
120                        path: path.clone(),
121                        bytes: *bytes,
122                    });
123            }
124            None => loose.push(path.clone()),
125        }
126    }
127
128    let groups = buckets
129        .into_iter()
130        .map(|((_, base, total), members)| {
131            let mut seen = BTreeSet::new();
132            let mut members: Vec<Member> = members
133                .into_iter()
134                .filter(|member| seen.insert(member.index))
135                .collect();
136            members.sort_by_key(|member| member.index);
137            ShardGroup {
138                base,
139                total,
140                members,
141            }
142        })
143        .collect();
144    (groups, loose)
145}
146
147fn five_digit_number(field: &str) -> Option<usize> {
148    if field.len() != 5 || !field.bytes().all(|byte| byte.is_ascii_digit()) {
149        return None;
150    }
151    field.parse().ok()
152}
kernel/discovery/gguf_shards.rs

kernel/discovery/
gguf_shards.rs
