kernel/install/
file_selection.rs

1//! Choosing which files of a Hugging Face repo to download. A repo lists many
2//! siblings; this picks the model weights (preferring one GGUF quantization, or the
3//! safetensors set, dropping duplicate/converted formats) plus the small companion
4//! config/tokenizer files, and drops docs, images, and non-model directories.
5
6use std::collections::{BTreeMap, BTreeSet};
7
8use crate::discovery::gguf_shards;
9use crate::install::bytes::saturating_sum;
10
11const WEIGHT_EXTENSIONS: [&str; 6] = ["safetensors", "gguf", "bin", "ckpt", "pt", "pth"];
12const EXCLUDED_EXTENSIONS: [&str; 8] = ["md", "png", "jpg", "jpeg", "gif", "webp", "msgpack", "h5"];
13const EXCLUDED_DIRECTORIES: [&str; 3] = ["onnx", "openvino", "coreml"];
14/// Quantizations in descending preference; the first one present is chosen.
15const QUANT_PREFERENCE: [&str; 6] = ["q4_k_m", "q4_0", "q5_k_m", "q6_k", "q8_0", "f16"];
16/// Companion files kept alongside GGUF weights must be at most this big (10 MiB).
17const COMPANION_CAP: i64 = 10 << 20;
18/// Support files kept for a transformers model must be at most this big (100 MiB).
19const CONFIG_CAP: i64 = 100 << 20;
20
21/// One file listed in a Hugging Face repo, as returned by the hub API.
22#[derive(Debug, Clone, PartialEq, Eq, Hash)]
23pub struct HFSibling {
24    /// The file's path within the repo.
25    pub rfilename: String,
26    /// The file's size in bytes, if the listing reported it.
27    pub bytes: Option<i64>,
28    /// The LFS object's SHA-256 (`lfs.sha256`), when the file is stored in LFS. The
29    /// download path uses it as the content-addressed blob name and to verify the
30    /// bytes; plain (non-LFS) files don't report one.
31    pub sha256: Option<String>,
32}
33
34impl HFSibling {
35    /// A sibling for `rfilename` with no known LFS hash.
36    pub fn new(rfilename: impl Into<String>, bytes: Option<i64>) -> Self {
37        Self {
38            rfilename: rfilename.into(),
39            bytes,
40            sha256: None,
41        }
42    }
43
44    /// This sibling with its LFS SHA-256 set.
45    pub fn with_sha256(mut self, sha256: Option<String>) -> Self {
46        self.sha256 = sha256;
47        self
48    }
49
50    /// Whether this file is a model weight (by extension).
51    pub fn is_weight(&self) -> bool {
52        is_weight_path(&self.rfilename)
53    }
54}
55
56/// Whether `path`'s extension marks it a model weight file.
57pub fn is_weight_path(path: &str) -> bool {
58    WEIGHT_EXTENSIONS.contains(&file_extension(path).as_str())
59}
60
61/// The lowercased extension of `path` (the part after the last `.`), or empty when
62/// there is none or the only dot is a leading one (a dotfile has no extension).
63pub fn file_extension(path: &str) -> String {
64    match path.rfind('.') {
65        Some(dot) if dot > 0 => path[dot + 1..].to_lowercase(),
66        _ => String::new(),
67    }
68}
69
70/// Select the files to download from a repo's `siblings`: the eligible weights
71/// (one GGUF quant group, or the safetensors/pytorch set) plus small companions.
72pub fn select(siblings: &[HFSibling]) -> Vec<HFSibling> {
73    let kept: Vec<HFSibling> = siblings
74        .iter()
75        .filter(|s| is_eligible(s))
76        .cloned()
77        .collect();
78    let ggufs: Vec<HFSibling> = kept
79        .iter()
80        .filter(|s| s.rfilename.to_lowercase().ends_with(".gguf"))
81        .cloned()
82        .collect();
83    if !ggufs.is_empty() {
84        let others: Vec<HFSibling> = kept
85            .iter()
86            .filter(|s| !ggufs.contains(s))
87            .cloned()
88            .collect();
89        return gguf_selection(&ggufs, &others);
90    }
91    if kept.iter().any(|s| s.rfilename == "model_index.json") {
92        return diffusers_selection(&kept);
93    }
94    transformers_selection(&kept)
95}
96
97/// The path split into non-empty `/`-segments (empty segments are dropped).
98fn segments(path: &str) -> Vec<&str> {
99    path.split('/').filter(|part| !part.is_empty()).collect()
100}
101
102/// Whether a sibling is a plausible model file: not hidden, not a readme, not an
103/// excluded extension, not a flax/tf checkpoint, and not under an excluded dir.
104fn is_eligible(sibling: &HFSibling) -> bool {
105    let path = &sibling.rfilename;
106    let segments = segments(path);
107    let Some(filename) = segments.last().copied() else {
108        return false;
109    };
110    if segments.iter().any(|segment| segment.starts_with('.')) {
111        return false;
112    }
113    if filename.to_lowercase().starts_with("readme") {
114        return false;
115    }
116    if EXCLUDED_EXTENSIONS.contains(&file_extension(filename).as_str()) {
117        return false;
118    }
119    let stem = filename.to_lowercase();
120    if stem.starts_with("flax_model") || stem.starts_with("tf_model") {
121        return false;
122    }
123    if segments.len() > 1
124        && let Some(first) = segments.first()
125        && EXCLUDED_DIRECTORIES.contains(&first.to_lowercase().as_str())
126    {
127        return false;
128    }
129    true
130}
131
132#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
133struct GroupKey {
134    directory: String,
135    base: String,
136    total: usize,
137}
138
139/// Choose one complete GGUF quant group (plus any mmproj projector and small
140/// companions). Shards are grouped by `(directory, base, total)`; a sharded group
141/// is only a candidate once every shard index is present.
142fn gguf_selection(ggufs: &[HFSibling], others: &[HFSibling]) -> Vec<HFSibling> {
143    let mut groups: BTreeMap<GroupKey, Vec<HFSibling>> = BTreeMap::new();
144    let mut seen_indices: BTreeMap<GroupKey, BTreeSet<usize>> = BTreeMap::new();
145    for sibling in ggufs {
146        let segments = segments(&sibling.rfilename);
147        let filename = segments.last().copied().unwrap_or("");
148        let directory = segments[..segments.len().saturating_sub(1)].join("/");
149        if let Some(shard) = gguf_shards::parse(filename) {
150            let key = GroupKey {
151                directory,
152                base: shard.base,
153                total: shard.total,
154            };
155            groups.entry(key.clone()).or_default().push(sibling.clone());
156            seen_indices.entry(key).or_default().insert(shard.index);
157        } else {
158            // A non-sharded GGUF: its base is the name without the `.gguf` suffix
159            // (5 trailing bytes, regardless of case).
160            let stem = &filename[..filename.len() - ".gguf".len()];
161            let key = GroupKey {
162                directory,
163                base: stem.to_owned(),
164                total: 0,
165            };
166            groups.entry(key).or_default().push(sibling.clone());
167        }
168    }
169
170    let mmproj: Vec<HFSibling> = ggufs
171        .iter()
172        .filter(|s| s.rfilename.to_lowercase().contains("mmproj"))
173        .cloned()
174        .collect();
175
176    // `BTreeMap` iteration is already ordered by `(directory, base, total)`, so the
177    // surviving candidate groups come out in that sorted order.
178    let ordered: Vec<Vec<HFSibling>> = groups
179        .iter()
180        .filter(|(key, _)| {
181            !key.base.to_lowercase().contains("mmproj")
182                && (key.total == 0 || seen_indices.get(*key).map(BTreeSet::len) == Some(key.total))
183        })
184        .map(|(_, group)| group.clone())
185        .collect();
186
187    let chosen = pick_quant_group(&ordered);
188    if chosen.is_empty() {
189        return Vec::new();
190    }
191    let companions: Vec<HFSibling> = others
192        .iter()
193        .filter(|s| s.bytes.unwrap_or(0) <= COMPANION_CAP)
194        .cloned()
195        .collect();
196
197    let mut result = chosen.clone();
198    result.extend(mmproj.into_iter().filter(|s| !chosen.contains(s)));
199    result.extend(companions);
200    result
201}
202
203/// Pick the group matching the highest quant preference, else the smallest by total
204/// bytes.
205fn pick_quant_group(groups: &[Vec<HFSibling>]) -> Vec<HFSibling> {
206    if groups.is_empty() {
207        return Vec::new();
208    }
209    for token in QUANT_PREFERENCE {
210        if let Some(matched) = groups
211            .iter()
212            .find(|group| group.iter().any(|s| matches_quant(&s.rfilename, token)))
213        {
214            return matched.clone();
215        }
216    }
217    groups
218        .iter()
219        .min_by_key(|group| saturating_sum(group.iter().filter_map(|s| s.bytes)))
220        .cloned()
221        .unwrap_or_default()
222}
223
224/// Whether `rfilename` contains `token` as a whole quant word (bounded by non-quant
225/// characters on both sides), so `q4_0` matches `model-q4_0.gguf` but not `xq4_0y`.
226fn matches_quant(rfilename: &str, token: &str) -> bool {
227    let name = rfilename.to_lowercase();
228    let mut start = 0;
229    while let Some(offset) = name[start..].find(token) {
230        let at = start + offset;
231        let end = at + token.len();
232        let before_ok = match name[..at].chars().next_back() {
233            None => true,
234            Some(character) => !is_quant_character(character),
235        };
236        let after_ok = match name[end..].chars().next() {
237            None => true,
238            Some(character) => !is_quant_character(character),
239        };
240        if before_ok && after_ok {
241            return true;
242        }
243        start = end;
244    }
245    false
246}
247
248fn is_quant_character(character: char) -> bool {
249    character.is_alphanumeric() || character == '_'
250}
251
252/// Diffusers selection: drop root-level weights when a subtree carries them, drop a
253/// `.bin`/`.ckpt`/… when a `.safetensors` twin exists, and drop `.fp16.`/`.non_ema.`
254/// variants when the plain form is present.
255fn diffusers_selection(kept: &[HFSibling]) -> Vec<HFSibling> {
256    let paths: BTreeSet<&str> = kept.iter().map(|s| s.rfilename.as_str()).collect();
257    let tree_has_weights = kept
258        .iter()
259        .any(|s| s.rfilename.contains('/') && s.is_weight());
260    kept.iter()
261        .filter(|sibling| {
262            let path = &sibling.rfilename;
263            if tree_has_weights && !path.contains('/') && sibling.is_weight() {
264                return false;
265            }
266            let ext = file_extension(path);
267            if ["bin", "ckpt", "pt", "pth"].contains(&ext.as_str())
268                && let Some(stem) = path.get(..path.len().saturating_sub(ext.len() + 1))
269                && paths.contains(format!("{stem}.safetensors").as_str())
270            {
271                return false;
272            }
273            for variant in [".fp16.", ".non_ema."] {
274                if path.contains(variant) && paths.contains(path.replace(variant, ".").as_str()) {
275                    return false;
276                }
277            }
278            true
279        })
280        .cloned()
281        .collect()
282}
283
284/// Transformers selection: the root safetensors set (or the pytorch `.bin` set when
285/// none), plus small support files (config/tokenizer), excluding index sidecars.
286fn transformers_selection(kept: &[HFSibling]) -> Vec<HFSibling> {
287    let root: Vec<&HFSibling> = kept.iter().filter(|s| !s.rfilename.contains('/')).collect();
288    let safetensors: Vec<&HFSibling> = root
289        .iter()
290        .copied()
291        .filter(|s| {
292            file_extension(&s.rfilename) == "safetensors"
293                || s.rfilename.ends_with(".safetensors.index.json")
294        })
295        .collect();
296    let has_safetensors_weight = safetensors
297        .iter()
298        .any(|s| file_extension(&s.rfilename) == "safetensors");
299    let weights: Vec<HFSibling> = if has_safetensors_weight {
300        safetensors.iter().copied().cloned().collect()
301    } else {
302        root.iter()
303            .copied()
304            .filter(|s| {
305                s.rfilename.starts_with("pytorch_model")
306                    && (file_extension(&s.rfilename) == "bin"
307                        || s.rfilename.ends_with(".bin.index.json"))
308            })
309            .cloned()
310            .collect()
311    };
312    let support: Vec<HFSibling> = root
313        .iter()
314        .copied()
315        .filter(|s| {
316            !s.is_weight()
317                && !s.rfilename.ends_with(".index.json")
318                && s.bytes.unwrap_or(0) <= CONFIG_CAP
319        })
320        .cloned()
321        .collect();
322
323    let mut result = weights;
324    result.extend(support);
325    result
326}
kernel/install/file_selection.rs

kernel/install/
file_selection.rs
