combs_models/
registry.rs

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

combs_models/
registry.rs
