1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
//! Architecture registry: maps `metadata.architecture` to a model loader.
//! New architectures are additive — one module + one `register` call.
use std::collections::HashMap;
use burn::tensor::{Device, backend::Backend};
use combs_formats::ModelSource;
use crate::traits::GenerativeModel;
use crate::{ModelError, Result};
/// A constructor for a boxed model of some architecture.
pub type Loader<B> =
fn(&dyn ModelSource, &Device<B>) -> Result<Box<dyn GenerativeModel<B>>>;
/// Maps architecture identifiers (`config.json::model_type`, plus known
/// aliases) to loaders. Mirrors MLC's `model.py::MODELS` table.
pub struct ModelRegistry<B: Backend> {
loaders: HashMap<String, Loader<B>>,
}
impl<B: Backend> Default for ModelRegistry<B> {
fn default() -> Self {
Self::new()
}
}
impl<B: Backend> ModelRegistry<B> {
/// Creates a registry with the built-in architectures registered.
pub fn new() -> Self {
let mut r = ModelRegistry {
loaders: HashMap::new(),
};
// SmolLM2 reports model_type "llama" (older releases: "smollm2");
// both are Llama-structured.
r.register("llama", |source, device| {
Ok(Box::new(crate::llama::LlamaModel::<B>::load(source, device)?))
});
r.register("smollm2", |source, device| {
Ok(Box::new(crate::llama::LlamaModel::<B>::load(source, device)?))
});
// Qwen2/2.5 is llama-structured plus q/k/v projection biases (loaded
// by presence); an active upper-layer sliding partition
// (`use_sliding_window: true`) maps onto the per-layer attention
// layout (ArchSpec: first `max_window_layers` global, rest sliding).
r.register("qwen2", |source, device| {
Ok(Box::new(crate::llama::LlamaModel::<B>::load(source, device)?))
});
// Qwen3 is qwen2 plus per-head q/k RMSNorm (ArchSpec `qk_norm`) and
// an explicit `head_dim` decoupled from hidden/heads.
r.register("qwen3", |source, device| {
Ok(Box::new(crate::llama::LlamaModel::<B>::load(source, device)?))
});
// Mistral v0.3+/Nemo report `sliding_window: null` and are plain
// llama; v0.1's all-layer sliding window rides the same per-layer
// layout phi3 uses.
r.register("mistral", |source, device| {
Ok(Box::new(crate::llama::LlamaModel::<B>::load(source, device)?))
});
// Phi-3/3.5/4-mini: llama-structured with fused qkv/gate_up
// projections (split at load) and an all-layer sliding window that
// the per-layer attention layout expresses directly — every shipped
// mini config activates it, so no sliding guard here.
r.register("phi3", |source, device| {
Ok(Box::new(crate::llama::LlamaModel::<B>::load(source, device)?))
});
// SmolVLM reports model_type "idefics3" (SigLIP + pixel-shuffle + SmolLM2).
r.register("idefics3", |source, device| {
Ok(Box::new(crate::smolvlm::SmolVlmModel::<B>::load(
source, device,
)?))
});
// Gemma 3 text ("gemma3_text") and the text trunk of multimodal
// Gemma 3 ("gemma3") ride the universal decoder: ArchSpec supplies
// the (1+w) norm flavor, qk/sandwich norms, sqrt(hidden) embed
// scale, query_pre_attn_scalar, dual-RoPE local theta, and the
// every-Nth-global sliding layout.
r.register("gemma3_text", |source, device| {
Ok(Box::new(crate::llama::LlamaModel::<B>::load(source, device)?))
});
r.register("gemma3", |source, device| {
Ok(Box::new(crate::llama::LlamaModel::<B>::load(source, device)?))
});
r
}
/// Registers (or replaces) the loader for an architecture id.
pub fn register(&mut self, architecture: &str, loader: Loader<B>) {
self.loaders.insert(architecture.to_string(), loader);
}
/// Whether an architecture id has a loader.
pub fn supports(&self, architecture: &str) -> bool {
self.loaders.contains_key(architecture)
}
/// Registered architecture ids.
pub fn architectures(&self) -> Vec<&str> {
let mut v: Vec<&str> = self.loaders.keys().map(String::as_str).collect();
v.sort();
v
}
/// Loads the model described by `source`'s metadata.
pub fn load(
&self,
source: &dyn ModelSource,
device: &Device<B>,
) -> Result<Box<dyn GenerativeModel<B>>> {
let arch = &source.metadata().architecture;
let loader = self
.loaders
.get(arch)
.ok_or_else(|| ModelError::UnsupportedArchitecture(arch.clone()))?;
loader(source, device)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn registry_has_llama_aliases() {
let r = ModelRegistry::<burn::backend::NdArray<f32>>::new();
assert!(r.supports("llama"));
assert!(r.supports("smollm2"));
assert!(r.supports("qwen2"));
assert!(r.supports("mistral"));
assert!(r.supports("qwen3"));
}
}