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

kernel/discovery/
modality_hints.rs
