cherenkov/model/
inspect.rs

1use super::qwen::tensor_role as qwen_role;
2use super::*;
3use crate::qwen4_exp::Manifest;
4use anyhow::{Context, Result, ensure};
5use cherenkov_model_data as data;
6use std::path::Path;
7
8mod packed;
9
10/// A checkpoint description and the immutable objects backing its tensor views.
11/// Pass the packed directory explicitly to inspect prepared data.
12pub struct Checkpoint {
13    /// Architecture, tensor roles, and storage descriptors inferred by the adapter.
14    pub description: ModelDescription,
15    /// Backing files addressed by the description's object IDs.
16    pub objects: MappedObjects,
17}
18
19impl Checkpoint {
20    /// Inspect a raw checkpoint or a prepared directory containing `manifest.json`.
21    /// Backing files must remain unchanged while this checkpoint or its views live.
22    pub fn open(path: &Path) -> Result<Self> {
23        if path.join("manifest.json").is_file() {
24            return packed::open(path);
25        }
26
27        let source = data::Checkpoint::open(path)?;
28        let description = describe_raw(&source, &source.inventory.metadata["config"])?;
29
30        Ok(Self {
31            description,
32            objects: source.objects,
33        })
34    }
35
36    /// Return a retained byte view after checking the span's object ID and bounds.
37    /// A view can outlive this checkpoint, but must not outlive its artifact lease.
38    pub fn map(&self, span: &DataSpan) -> Result<MappedBytes> {
39        self.objects.view(span)
40    }
41}
42
43pub(crate) fn describe_raw(
44    source: &data::Checkpoint,
45    config: &serde_json::Value,
46) -> Result<ModelDescription> {
47    describe_inventory(&source.inventory, &source.objects, config)
48}
49
50pub(crate) fn describe_inventory(
51    inventory: &data::Inventory,
52    source: &dyn ByteSource,
53    config: &serde_json::Value,
54) -> Result<ModelDescription> {
55    let (container, conventions, architecture) = formats(inventory, config);
56    let mlx = matches!(
57        conventions,
58        CheckpointConventions::Qwen4ExpMlx | CheckpointConventions::MlxAffine
59    );
60    let interpret_mlx = mlx
61        || (matches!(architecture, Architecture::Qwen4Exp)
62            && matches!(container, ContainerFormat::Safetensors));
63    let mut tensors = if interpret_mlx {
64        data::mlx::tensors(inventory, config)?
65    } else {
66        inventory.tensors.clone()
67    };
68
69    if matches!(architecture, Architecture::Qwen4Exp)
70        && matches!(container, ContainerFormat::Safetensors)
71    {
72        for tensor in &mut tensors {
73            tensor.role = qwen_role(&tensor.name, tensor.shape().map_or(0, <[u64]>::len));
74        }
75    }
76
77    let ngram = ngram_table(inventory, source, config, &tensors)?;
78
79    Ok(ModelDescription {
80        schema_version: 1,
81        format: CheckpointFormat {
82            container,
83            conventions,
84        },
85        architecture,
86        metadata: inventory.metadata.clone(),
87        tensors,
88        ngram,
89    })
90}
91
92fn formats(
93    inventory: &data::Inventory,
94    config: &serde_json::Value,
95) -> (ContainerFormat, CheckpointConventions, Architecture) {
96    if let data::ContainerFormat::Gguf { version } = inventory.format {
97        let name = inventory.metadata["general.architecture"]["value"]
98            .as_str()
99            .unwrap_or("unknown");
100        let architecture = match name {
101            "qwen4exp" | "qwen4_exp" => Architecture::Qwen4Exp,
102            _ => Architecture::Opaque { name: name.into() },
103        };
104
105        return (
106            ContainerFormat::Gguf { version },
107            CheckpointConventions::Gguf,
108            architecture,
109        );
110    }
111
112    let architecture = architecture(config);
113    let native = inventory
114        .tensors
115        .iter()
116        .any(|t| t.name.ends_with(".mlp.experts.gate_up_proj"));
117    let mlx = inventory
118        .tensors
119        .iter()
120        .any(|t| t.name.contains(".mlp.switch_mlp."));
121    let conventions = match (&architecture, native, mlx) {
122        (Architecture::Qwen4Exp, true, _) => CheckpointConventions::Qwen4ExpHuggingFace,
123        (Architecture::Qwen4Exp, _, true) => CheckpointConventions::Qwen4ExpMlx,
124        _ if declares_mlx(config) => CheckpointConventions::MlxAffine,
125        _ => CheckpointConventions::Opaque {
126            name: "unrecognized checkpoint conventions".into(),
127        },
128    };
129
130    (ContainerFormat::Safetensors, conventions, architecture)
131}
132
133fn declares_mlx(config: &serde_json::Value) -> bool {
134    let Some(quant) = config
135        .get("quantization")
136        .or_else(|| config.get("quantization_config"))
137    else {
138        return false;
139    };
140
141    quant.get("quant_method").is_none()
142        && quant["bits"].is_u64()
143        && quant["mode"].as_str().unwrap_or("affine") == "affine"
144}
145
146fn architecture(config: &serde_json::Value) -> Architecture {
147    match config["model_type"].as_str() {
148        Some("qwen4_exp" | "qwen4_exp_text") => Architecture::Qwen4Exp,
149        name => Architecture::Opaque {
150            name: name.unwrap_or("unknown").into(),
151        },
152    }
153}
154
155fn ngram_table(
156    inventory: &data::Inventory,
157    source: &dyn ByteSource,
158    config: &serde_json::Value,
159    tensors: &[Tensor],
160) -> Result<Option<NgramTable>> {
161    let mut entries: Vec<_> = tensors
162        .iter()
163        .enumerate()
164        .filter(|(_, t)| t.role == TensorRole::NgramEmbedding)
165        .collect();
166
167    if entries.is_empty() {
168        return Ok(None);
169    }
170
171    // Numeric shard order matters: lexicographic order places shard_10 before 2.
172    entries.sort_by_key(|(_, t)| shard_number(&t.name));
173
174    let name = &entries[0].1.name;
175    let prefix = name
176        .split(".ngram_embedding.")
177        .next()
178        .context("n-gram prefix missing")?;
179    let read = |suffix| -> Result<Vec<i64>> {
180        let entry = inventory
181            .tensors
182            .iter()
183            .find(|t| t.name == format!("{prefix}.{suffix}"))
184            .with_context(|| format!("missing n-gram {suffix}"))?;
185        let TensorEncoding::Dense { tensor } = &entry.encoding else {
186            anyhow::bail!("n-gram metadata must be dense I64");
187        };
188
189        ensure!(tensor.dtype == Dtype::I64, "n-gram metadata must be I64");
190
191        ensure!(
192            tensor.data.length <= 1024 * 1024,
193            "n-gram metadata exceeds 1 MiB"
194        );
195
196        let mut bytes = Vec::new();
197
198        source.read(&tensor.data, &mut bytes)?;
199
200        Ok(bytes
201            .as_chunks::<8>()
202            .0
203            .iter()
204            .map(|&b| i64::from_le_bytes(b))
205            .collect())
206    };
207    let text = config.get("text_config").unwrap_or(config);
208    let hashing = NgramHash::Qwen4Exp {
209        ngram_size: text["ngram_size"].as_u64().unwrap_or(3),
210        heads_per_ngram: text["heads_per_ngram"].as_u64().unwrap_or(8),
211        head_offsets: unsigned(read("ngram_heads_offsets")?)?,
212        head_vocab_sizes: unsigned(read("ngram_heads_vocab_sizes")?)?,
213        layer_multipliers: read("layer_multipliers")?,
214    };
215    let mut shards = Vec::new();
216    let mut first_row = 0_u64;
217
218    for (expected, (id, tensor)) in entries.into_iter().enumerate() {
219        ensure!(
220            shard_number(&tensor.name) == Some(expected),
221            "n-gram shards must be contiguous"
222        );
223
224        let rows = *tensor
225            .shape()
226            .and_then(|s| s.first())
227            .context("n-gram rows missing")?;
228
229        shards.push(TableShard {
230            first_row,
231            rows,
232            tensor: TensorId(id),
233        });
234
235        first_row = first_row
236            .checked_add(rows)
237            .context("n-gram row count overflow")?;
238    }
239
240    Ok(Some(NgramTable { hashing, shards }))
241}
242
243fn unsigned(values: Vec<i64>) -> Result<Vec<u64>> {
244    values
245        .into_iter()
246        .map(|v| u64::try_from(v).context("negative n-gram index metadata"))
247        .collect()
248}
249
250fn shard_number(name: &str) -> Option<usize> {
251    let tail = name.split(".ngram_embedding.").nth(1)?;
252    let number = tail
253        .strip_prefix("shards.")
254        .or_else(|| tail.strip_prefix("shard_"))?;
255
256    number.split('.').next()?.parse().ok()
257}
cherenkov/model/inspect.rs

cherenkov/model/
inspect.rs
