kernel/discovery/
modality_hints.rs

1//! Modality/capability hints derived from a Hugging Face model's `config.json`:
2//! a first guess (architecture list, a vision block, or a few speech config keys)
3//! at what a downloaded model can do, before full identification. A diffusers
4//! `model_index.json` is hinted through the pipeline-family registry
5//! ([`from_model_index`]).
6
7use std::collections::BTreeSet;
8use std::path::Path;
9
10use crate::records::{Capability, ExecutionMode, JsonValue, Modality};
11use crate::resolution::pipelines::{PipelineFamilyRegistry, diffusers_pipeline_class};
12
13/// A best-effort guess at a model's shape from its config metadata.
14#[derive(Debug, Clone, PartialEq)]
15pub struct Hint {
16    /// The guessed modality, if any.
17    pub modality: Option<Modality>,
18    /// The guessed capabilities.
19    pub capabilities: Vec<Capability>,
20    /// How the model executes.
21    pub execution: ExecutionMode,
22    /// A context-window hint pulled from the config.
23    pub context_length: Option<i64>,
24    /// The quantization the config names, as `4bit`, when it does.
25    pub quantization: Option<String>,
26}
27
28impl Hint {
29    fn new(modality: Modality, capabilities: Vec<Capability>, execution: ExecutionMode) -> Self {
30        Self {
31            modality: Some(modality),
32            capabilities,
33            execution,
34            context_length: None,
35            quantization: None,
36        }
37    }
38
39    /// An empty hint (no modality or capabilities) with the given execution shape
40    /// — the starting point before any rule matches.
41    pub fn unknown(execution: ExecutionMode) -> Self {
42        Self {
43            modality: None,
44            capabilities: Vec::new(),
45            execution,
46            context_length: None,
47            quantization: None,
48        }
49    }
50}
51
52/// A text-to-speech model.
53pub fn speech_hint() -> Hint {
54    Hint::new(
55        Modality::speech(),
56        vec![Capability::speak()],
57        ExecutionMode::Stream,
58    )
59}
60
61/// A speech-to-text (transcription) model.
62pub fn audio_hint() -> Hint {
63    Hint::new(
64        Modality::audio(),
65        vec![Capability::transcribe()],
66        ExecutionMode::Stream,
67    )
68}
69
70/// A text chat/completion model.
71pub fn text_hint() -> Hint {
72    Hint::new(
73        Modality::text(),
74        vec![Capability::chat(), Capability::complete()],
75        ExecutionMode::Stream,
76    )
77}
78
79/// An embedding model.
80pub fn embedding_hint() -> Hint {
81    Hint::new(
82        Modality::embedding(),
83        vec![Capability::embed()],
84        ExecutionMode::Stream,
85    )
86}
87
88/// A vision-capable chat model.
89pub fn vision_chat_hint() -> Hint {
90    Hint::new(
91        Modality::text(),
92        vec![
93            Capability::chat(),
94            Capability::complete(),
95            Capability::see(),
96        ],
97        ExecutionMode::Stream,
98    )
99}
100
101/// The default hint for a bare GGUF weight (text chat/completion).
102pub fn gguf_hint() -> Hint {
103    text_hint()
104}
105
106/// The default hint for a whisper `.bin` (transcription).
107pub fn whisper_bin_hint() -> Hint {
108    audio_hint()
109}
110
111/// The hint for a diffusers `model_index.json`: its pipeline family's modality and
112/// capabilities as a job. An unknown or absent `_class_name` yields an empty job
113/// hint (still a job — a diffusers bundle is never a streaming model).
114pub fn from_model_index(path: &Path) -> Hint {
115    if let Some(class) = diffusers_pipeline_class(path)
116        && let Some(family) = PipelineFamilyRegistry::shared().family(&class)
117    {
118        return Hint {
119            modality: Some(family.modality.clone()),
120            capabilities: family.capabilities.clone(),
121            execution: ExecutionMode::Job,
122            context_length: None,
123            quantization: None,
124        };
125    }
126    Hint::unknown(ExecutionMode::Job)
127}
128
129/// Architecture-name substrings that mark a vision-language model when a
130/// `vision_config` block is also present.
131const VISION_LANGUAGE_ARCHITECTURES: [&str; 5] =
132    ["Llava", "Qwen2VL", "Idefics", "PaliGemma", "Mllama"];
133
134/// The hint for a Hugging Face `config.json` file, or `None` if it is unreadable,
135/// unparseable, or matches no rule.
136pub fn from_config_json(path: &Path) -> Option<Hint> {
137    let bytes = std::fs::read(path).ok()?;
138    let json = serde_json::from_slice::<JsonValue>(&bytes).ok()?;
139    from_config(&json)
140}
141
142/// The hint for a parsed `config.json` value: a vision-language model (a
143/// `vision_config` block plus a matching architecture), else the first
144/// architecture that matches an [`architecture_hint`] rule, else a speech model
145/// recognized by its config keys, else `None`.
146pub fn from_config(json: &JsonValue) -> Option<Hint> {
147    let object = json.as_object()?;
148    let architectures: Vec<&str> = object
149        .get("architectures")
150        .and_then(JsonValue::as_array)
151        .map(|items| items.iter().filter_map(JsonValue::as_str).collect())
152        .unwrap_or_default();
153
154    // First present integer among the window keys, then require it positive
155    // (a present-but-zero value voids the hint rather than falling through).
156    let context_length = ["max_position_embeddings", "n_positions", "max_seq_len"]
157        .into_iter()
158        .find_map(|key| object.get(key).and_then(JsonValue::as_i64))
159        .filter(|value| *value > 0);
160    // MLX writes the bits it quantized to under `quantization`; a config without
161    // the block is unquantized, or not MLX's.
162    let quantization = object
163        .get("quantization")
164        .and_then(JsonValue::as_object)
165        .and_then(|block| block.get("bits"))
166        .and_then(JsonValue::as_i64)
167        .filter(|bits| *bits > 0)
168        .map(|bits| format!("{bits}bit"));
169
170    if object.contains_key("vision_config")
171        && architectures.iter().any(|architecture| {
172            architecture.ends_with("ForConditionalGeneration")
173                || VISION_LANGUAGE_ARCHITECTURES
174                    .iter()
175                    .any(|marker| architecture.contains(marker))
176        })
177    {
178        return Some(with_facts(
179            vision_chat_hint(),
180            context_length,
181            quantization.clone(),
182        ));
183    }
184
185    for architecture in &architectures {
186        if let Some(hint) = architecture_hint(architecture) {
187            return Some(with_facts(hint, context_length, quantization.clone()));
188        }
189    }
190
191    let keys: BTreeSet<&str> = object.keys().map(String::as_str).collect();
192    config_key_hint(&keys).map(|hint| with_facts(hint, context_length, quantization))
193}
194
195fn with_facts(mut hint: Hint, context_length: Option<i64>, quantization: Option<String>) -> Hint {
196    hint.context_length = context_length;
197    hint.quantization = quantization;
198    hint
199}
200
201/// The hint for a single architecture name, by substring/suffix rules (checked in
202/// priority order: speech, audio, embedding, then causal-LM text).
203fn architecture_hint(architecture: &str) -> Option<Hint> {
204    const SPEECH: [&str; 6] = ["Kokoro", "StyleTTS", "Bark", "ParlerTTS", "Vits", "Xtts"];
205    const EMBEDDING: [&str; 5] = [
206        "BertModel",
207        "NomicBertModel",
208        "ModernBertModel",
209        "XLMRobertaModel",
210        "MPNetModel",
211    ];
212    if SPEECH.iter().any(|marker| architecture.contains(marker)) {
213        return Some(speech_hint());
214    }
215    if architecture.contains("Whisper") {
216        return Some(audio_hint());
217    }
218    if EMBEDDING.iter().any(|marker| architecture.contains(marker)) {
219        return Some(embedding_hint());
220    }
221    if architecture.contains("LMHead") || architecture.ends_with("ForCausalLM") {
222        return Some(text_hint());
223    }
224    None
225}
226
227/// The hint for a set of config keys: a few speech models are recognized only by
228/// their config shape (all required keys present).
229fn config_key_hint(keys: &BTreeSet<&str>) -> Option<Hint> {
230    if keys.contains("istftnet") || keys.contains("plbert") {
231        return Some(speech_hint());
232    }
233    if keys.contains("style_dim") && keys.contains("n_mels") {
234        return Some(speech_hint());
235    }
236    None
237}
kernel/discovery/modality_hints.rs

kernel/discovery/
modality_hints.rs
