Skip to main content

kernel/discovery/
gguf_models.rs

1//! Turning a flat list of GGUF files into [`DiscoveredModel`]s: loose files
2//! become one model each, shard sets become a single model keyed by their shared
3//! base (flagged as still-downloading when incomplete). Shared by the file-tree
4//! scanners (LM Studio, loose files).
5
6use std::collections::HashMap;
7use std::path::{Path, PathBuf};
8
9use crate::discovery::gguf_shards::group;
10use crate::discovery::modality_hints::gguf_hint;
11use crate::discovery::scanner::DiscoveredModel;
12use crate::records::{ModelSource, SourceKind};
13
14/// Whether `name` is a GGUF file's, whatever the case of its extension. A
15/// name that is only the extension has no stem and is not one.
16pub(crate) fn is_gguf_name(name: &str) -> bool {
17    name.rsplit_once('.')
18        .is_some_and(|(stem, extension)| !stem.is_empty() && extension.eq_ignore_ascii_case("gguf"))
19}
20
21/// Whether a filename looks like a multimodal projector (never the primary
22/// weight). Case-insensitive `mmproj` substring.
23pub fn is_mmproj_name(name: &str) -> bool {
24    name.to_ascii_lowercase().contains("mmproj")
25}
26
27/// Build discovered models from `files` (each a GGUF path and its byte size).
28/// `repo` derives the repository label from a file's path (e.g. a relative
29/// `<org>/<model>` prefix), or `None`. Returns the models and any per-shard-set
30/// issues (a shard set missing its first part is skipped, not emitted).
31pub fn discovered_models(
32    files: &[(PathBuf, i64)],
33    kind: &SourceKind,
34    repo: impl Fn(&Path) -> Option<String>,
35) -> (Vec<DiscoveredModel>, Vec<String>) {
36    let (groups, loose) = group(files);
37    let mut bytes_by_path: HashMap<&Path, i64> = HashMap::new();
38    for (path, bytes) in files {
39        bytes_by_path.entry(path.as_path()).or_insert(*bytes);
40    }
41    let hint = gguf_hint();
42    let mut discovered = Vec::new();
43    let mut issues = Vec::new();
44
45    for path in &loose {
46        let name = path
47            .file_stem()
48            .and_then(|stem| stem.to_str())
49            .unwrap_or_default()
50            .to_owned();
51        let mut model = DiscoveredModel::new(name, source(kind, path, &repo));
52        apply_hint(&mut model, &hint);
53        model.footprint_bytes = bytes_by_path.get(path.as_path()).copied().unwrap_or(0);
54        model.primary_weight_path = Some(display(path));
55        discovered.push(model);
56    }
57
58    for shard_group in groups {
59        let Some(first) = shard_group.first_shard() else {
60            issues.push(format!(
61                "sharded model {} is missing its first part — skipped",
62                shard_group.base
63            ));
64            continue;
65        };
66        let mut model = DiscoveredModel::new(shard_group.base.clone(), source(kind, first, &repo));
67        apply_hint(&mut model, &hint);
68        model.footprint_bytes = shard_group.footprint_bytes();
69        model.primary_weight_path = Some(display(first));
70        model.downloading = !shard_group.complete();
71        discovered.push(model);
72    }
73
74    (discovered, issues)
75}
76
77fn source(kind: &SourceKind, path: &Path, repo: &impl Fn(&Path) -> Option<String>) -> ModelSource {
78    let mut source = ModelSource::new(kind.clone(), &display(path));
79    source.repo = repo(path);
80    source
81}
82
83fn apply_hint(model: &mut DiscoveredModel, hint: &crate::discovery::modality_hints::Hint) {
84    model.modality_hint = hint.modality.clone();
85    model.capabilities_hint = hint.capabilities.clone();
86    model.execution_hint = hint.execution;
87}
88
89fn display(path: &Path) -> String {
90    path.to_string_lossy().into_owned()
91}