cherenkov/qwen4_exp/
config.rs

1//! Model configuration for the qwen4_exp text engine.
2
3use anyhow::{Context, Result};
4use serde::Deserialize;
5use std::path::Path;
6
7/// Text-model configuration, parsed from config.json's `text_config`.
8#[derive(Debug, Clone, Deserialize)]
9pub struct Qwen4ExpConfig {
10    pub hidden_size: usize,
11    pub num_hidden_layers: usize,
12    pub layer_types: Vec<String>,
13    pub num_attention_heads: usize,
14    pub num_key_value_heads: usize,
15    pub head_dim: usize,
16    pub linear_num_key_heads: usize,
17    pub linear_num_value_heads: usize,
18    pub linear_key_head_dim: usize,
19    pub linear_value_head_dim: usize,
20    pub linear_conv_kernel_dim: usize,
21    pub rms_norm_eps: f64,
22    pub vocab_size: usize,
23    pub partial_rotary_factor: f64,
24    pub rope_parameters: RopeParams,
25    #[serde(default = "four")]
26    pub hc_count: usize,
27    #[serde(default = "d320")]
28    pub hc_lowrank: usize,
29    pub num_experts: usize,
30    pub num_experts_per_tok: usize,
31    pub moe_intermediate_size: usize,
32    pub shared_expert_intermediate_size: usize,
33    #[serde(default = "yes")]
34    pub norm_topk_prob: bool,
35    #[serde(default)]
36    pub ple_layer_ids: Vec<usize>,
37    #[serde(default = "four")]
38    pub ple_conv_kernel_size: usize,
39    pub ple_embed_dim: usize,
40    #[serde(default = "three")]
41    pub ngram_size: usize,
42    #[serde(default = "eight")]
43    pub heads_per_ngram: usize,
44    pub indexer_n_heads: usize,
45    pub indexer_kv_heads: usize,
46    pub indexer_head_dim: usize,
47    pub indexer_budget: usize,
48    pub indexer_compress_ratio: usize,
49    #[serde(default = "silu")]
50    pub output_gate_type: String,
51    pub eos_token_id: u32,
52    #[serde(default)]
53    pub mtp_num_hidden_layers: usize,
54}
55
56#[derive(Debug, Clone, Deserialize)]
57pub struct RopeParams {
58    pub rope_theta: f64,
59}
60
61fn four() -> usize {
62    4
63}
64
65fn three() -> usize {
66    3
67}
68
69fn eight() -> usize {
70    8
71}
72
73fn d320() -> usize {
74    320
75}
76
77fn yes() -> bool {
78    true
79}
80
81fn silu() -> String {
82    "silu".into()
83}
84
85impl Qwen4ExpConfig {
86    pub fn load(model_dir: &Path) -> Result<Self> {
87        let path = model_dir.join("config.json");
88        let bytes = std::fs::read(&path).with_context(|| format!("reading {}", path.display()))?;
89        let root: serde_json::Value = serde_json::from_slice(&bytes).context("config.json")?;
90        let model_type = root["model_type"].as_str().context("model_type missing")?;
91
92        anyhow::ensure!(
93            matches!(model_type, "qwen4_exp" | "qwen4_exp_text"),
94            "unsupported model_type {model_type}; expected qwen4_exp"
95        );
96
97        let text = root.get("text_config").cloned().unwrap_or(root);
98
99        anyhow::ensure!(
100            text["model_type"] == "qwen4_exp_text",
101            "expected qwen4_exp_text configuration"
102        );
103
104        let cfg: Qwen4ExpConfig = serde_json::from_value(text).context("text_config")?;
105
106        anyhow::ensure!(
107            cfg.layer_types.len() == cfg.num_hidden_layers,
108            "layer_types length mismatch"
109        );
110        anyhow::ensure!(cfg.hc_count > 1, "hc_count must exceed 1");
111
112        Ok(cfg)
113    }
114
115    pub fn is_linear(&self, layer: usize) -> bool {
116        self.layer_types[layer] == "linear_attention"
117    }
118
119    /// Zero-based layer index carrying the n-gram (PLE) block, if any.
120    pub fn ple_layer(&self) -> Option<usize> {
121        self.ple_layer_ids.first().map(|id| id - 1)
122    }
123
124    pub fn hc_hidden(&self) -> usize {
125        self.hc_count * self.hidden_size
126    }
127}
128
129#[cfg(test)]
130#[path = "../../tests/unit/qwen4_exp/config.rs"]
131mod tests;
cherenkov/qwen4_exp/config.rs

cherenkov/qwen4_exp/
config.rs
