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 reranker: text in, a relevance score out. It chats with no one and embeds
89/// nothing, so it claims no capability of its own; the runtime that can serve it
90/// supplies what it does.
91pub fn reranker_hint() -> Hint {
92    Hint::new(Modality::text(), Vec::new(), ExecutionMode::Stream)
93}
94
95/// What a sentence-transformers layout in a snapshot is.
96#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub enum SentenceTransformersLayout {
98    /// A bi-encoder, pooling a text into an embedding.
99    Embedder,
100    /// A cross-encoder, scoring a pair of texts against each other.
101    CrossEncoder,
102}
103
104/// The sentence-transformers layout in `dir`, if there is one. Both kinds ship a
105/// `config_sentence_transformers.json`; only its `model_type` tells a
106/// cross-encoder from an embedder, and one read as the other is served nonsense.
107pub fn sentence_transformers_layout(dir: &Path) -> Option<SentenceTransformersLayout> {
108    let config = dir.join("config_sentence_transformers.json");
109    let declared = std::fs::read(&config)
110        .ok()
111        .and_then(|bytes| serde_json::from_slice::<JsonValue>(&bytes).ok())
112        .and_then(|json| {
113            json.as_object()?
114                .get("model_type")
115                .and_then(JsonValue::as_str)
116                .map(str::to_owned)
117        });
118    if declared.as_deref() == Some("CrossEncoder") {
119        return Some(SentenceTransformersLayout::CrossEncoder);
120    }
121    (config.exists() || dir.join("1_Pooling").exists())
122        .then_some(SentenceTransformersLayout::Embedder)
123}
124
125/// A vision-capable chat model.
126pub fn vision_chat_hint() -> Hint {
127    Hint::new(
128        Modality::text(),
129        vec![
130            Capability::chat(),
131            Capability::complete(),
132            Capability::see(),
133        ],
134        ExecutionMode::Stream,
135    )
136}
137
138/// The default hint for a bare GGUF weight (text chat/completion).
139pub fn gguf_hint() -> Hint {
140    text_hint()
141}
142
143/// The default hint for a whisper `.bin` (transcription).
144pub fn whisper_bin_hint() -> Hint {
145    audio_hint()
146}
147
148/// The hint for a diffusers `model_index.json`: its pipeline family's modality and
149/// capabilities as a job. An unknown or absent `_class_name` yields an empty job
150/// hint (still a job — a diffusers bundle is never a streaming model).
151pub fn from_model_index(path: &Path) -> Hint {
152    if let Some(class) = diffusers_pipeline_class(path)
153        && let Some(family) = PipelineFamilyRegistry::shared().family(&class)
154    {
155        return Hint {
156            modality: Some(family.modality.clone()),
157            capabilities: family.capabilities.clone(),
158            execution: ExecutionMode::Job,
159            context_length: None,
160            quantization: None,
161        };
162    }
163    Hint::unknown(ExecutionMode::Job)
164}
165
166/// Architecture-name substrings that mark a vision-language model when a
167/// `vision_config` block is also present.
168const VISION_LANGUAGE_ARCHITECTURES: [&str; 5] =
169    ["Llava", "Qwen2VL", "Idefics", "PaliGemma", "Mllama"];
170
171/// The hint for a Hugging Face `config.json` file, or `None` if it is unreadable,
172/// unparseable, or matches no rule.
173pub fn from_config_json(path: &Path) -> Option<Hint> {
174    let bytes = std::fs::read(path).ok()?;
175    let json = serde_json::from_slice::<JsonValue>(&bytes).ok()?;
176    from_config(&json)
177}
178
179/// The hint for a parsed `config.json` value: a vision-language model (a
180/// `vision_config` block plus a matching architecture), else the first
181/// architecture that matches an [`architecture_hint`] rule, else a speech model
182/// recognized by its config keys, else `None`.
183pub fn from_config(json: &JsonValue) -> Option<Hint> {
184    let object = json.as_object()?;
185    let architectures: Vec<&str> = object
186        .get("architectures")
187        .and_then(JsonValue::as_array)
188        .map(|items| items.iter().filter_map(JsonValue::as_str).collect())
189        .unwrap_or_default();
190
191    // First present integer among the window keys, then require it positive
192    // (a present-but-zero value voids the hint rather than falling through).
193    let context_length = ["max_position_embeddings", "n_positions", "max_seq_len"]
194        .into_iter()
195        .find_map(|key| object.get(key).and_then(JsonValue::as_i64))
196        .filter(|value| *value > 0);
197    // MLX writes the bits it quantized to under `quantization`; a config without
198    // the block is unquantized, or not MLX's.
199    let quantization = object
200        .get("quantization")
201        .and_then(JsonValue::as_object)
202        .and_then(|block| block.get("bits"))
203        .and_then(JsonValue::as_i64)
204        .filter(|bits| *bits > 0)
205        .map(|bits| format!("{bits}bit"));
206
207    if object.contains_key("vision_config")
208        && architectures.iter().any(|architecture| {
209            architecture.ends_with("ForConditionalGeneration")
210                || VISION_LANGUAGE_ARCHITECTURES
211                    .iter()
212                    .any(|marker| architecture.contains(marker))
213        })
214    {
215        return Some(with_facts(
216            vision_chat_hint(),
217            context_length,
218            quantization.clone(),
219        ));
220    }
221
222    for architecture in &architectures {
223        if let Some(hint) = architecture_hint(architecture) {
224            return Some(with_facts(hint, context_length, quantization.clone()));
225        }
226    }
227
228    let keys: BTreeSet<&str> = object.keys().map(String::as_str).collect();
229    config_key_hint(&keys).map(|hint| with_facts(hint, context_length, quantization))
230}
231
232fn with_facts(mut hint: Hint, context_length: Option<i64>, quantization: Option<String>) -> Hint {
233    hint.context_length = context_length;
234    hint.quantization = quantization;
235    hint
236}
237
238/// The hint for a single architecture name, by substring/suffix rules (checked in
239/// priority order: speech, audio, embedding, then causal-LM text).
240fn architecture_hint(architecture: &str) -> Option<Hint> {
241    const SPEECH: [&str; 6] = ["Kokoro", "StyleTTS", "Bark", "ParlerTTS", "Vits", "Xtts"];
242    const EMBEDDING: [&str; 5] = [
243        "BertModel",
244        "NomicBertModel",
245        "ModernBertModel",
246        "XLMRobertaModel",
247        "MPNetModel",
248    ];
249    if SPEECH.iter().any(|marker| architecture.contains(marker)) {
250        return Some(speech_hint());
251    }
252    if architecture.contains("Whisper") {
253        return Some(audio_hint());
254    }
255    if EMBEDDING.iter().any(|marker| architecture.contains(marker)) {
256        return Some(embedding_hint());
257    }
258    if architecture.contains("LMHead") || architecture.ends_with("ForCausalLM") {
259        return Some(text_hint());
260    }
261    None
262}
263
264/// The hint for a set of config keys: a few speech models are recognized only by
265/// their config shape (all required keys present).
266fn config_key_hint(keys: &BTreeSet<&str>) -> Option<Hint> {
267    if keys.contains("istftnet") || keys.contains("plbert") {
268        return Some(speech_hint());
269    }
270    if keys.contains("style_dim") && keys.contains("n_mels") {
271        return Some(speech_hint());
272    }
273    None
274}
kernel/discovery/modality_hints.rs

kernel/discovery/
modality_hints.rs
