qualia_core_db/inference/
safetensor.rs

1//! Phase 6 / task #12 — **safetensor (+ MLX) source parsing + dtype gate** for the streaming
2//! transcoder (`p64_weight::transcode_safetensor_to_p64`).
3//!
4//! This module only **parses + validates** a source; the streaming **emit** to the P64 container
5//! lives in `p64_weight` (it owns the container's private serializers). The split keeps the format
6//! writer encapsulated.
7//!
8//! ## Scope (honest)
9//! * **safetensor** — the on-disk layout is parsed here: an 8-byte little-endian header length,
10//!   then a JSON header `{ name: { dtype, shape, data_offsets:[begin,end] }, … }`, then the raw
11//!   tensor bytes. The JSON header is small (KBs); the tensor bytes are **never** read here — only
12//!   their offsets — so a multi-GB file is *not* loaded to plan the transcode.
13//! * **MLX** — Apple MLX exports are safetensor-format (often with `__metadata__.format = "mlx"`);
14//!   they parse through this path. MLX `.npz` archives are **deferred** (a different container).
15//! * **high-fidelity only** — [`is_high_fidelity_ggml`] accepts `F32 / F16 / BF16 / Q8_0` and
16//!   **rejects** `Q4_*` and other low-precision quant types (the "Q4 rejected" rail).
17
18/// A detected model source container.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum SourceFormat {
21    /// safetensor (incl. MLX-safetensor).
22    Safetensor,
23    /// GGUF (the established path — `p64_weight::compile_gguf_to_p64`).
24    Gguf,
25    /// Unrecognised.
26    Unknown,
27}
28
29/// Sniff the container format from the first bytes. GGUF starts with the ASCII magic `GGUF`;
30/// safetensor starts with an 8-byte LE header length immediately followed by a `{` (the JSON).
31pub fn detect_format(head: &[u8]) -> SourceFormat {
32    if head.len() >= 4 && &head[0..4] == b"GGUF" {
33        return SourceFormat::Gguf;
34    }
35    if head.len() >= 9 {
36        let hlen = u64::from_le_bytes(head[0..8].try_into().unwrap()) as usize;
37        // a sane JSON header length, and the JSON object opens right after the 8-byte length.
38        if hlen > 0 && hlen < (1 << 30) && head[8] == b'{' {
39            return SourceFormat::Safetensor;
40        }
41    }
42    SourceFormat::Unknown
43}
44
45/// One tensor declared in a safetensor header.
46#[derive(Debug, Clone, PartialEq, Eq)]
47pub struct SafeTensorEntry {
48    pub name: String,
49    pub dtype: String,
50    pub shape: Vec<usize>,
51    /// Byte range of the tensor within the data region (relative to `data_start`).
52    pub begin: usize,
53    pub end: usize,
54}
55
56impl SafeTensorEntry {
57    #[inline]
58    pub fn byte_len(&self) -> usize {
59        self.end - self.begin
60    }
61}
62
63/// The parsed plan: every tensor's metadata + the absolute offset where tensor data begins.
64#[derive(Debug, Clone)]
65pub struct SafeTensorPlan {
66    pub tensors: Vec<SafeTensorEntry>,
67    /// Absolute byte offset (in the source) of the start of the tensor-data region.
68    pub data_start: usize,
69    /// `true` if `__metadata__.format == "mlx"`.
70    pub is_mlx: bool,
71}
72
73/// Parse a safetensor header (the small JSON prefix only — never the tensor bytes). Validates that
74/// every declared byte range lies within `src`.
75pub fn parse_safetensor_header(src: &[u8]) -> Result<SafeTensorPlan, String> {
76    if src.len() < 8 {
77        return Err("safetensor: too small for 8-byte header length".to_string());
78    }
79    let hlen = u64::from_le_bytes(src[0..8].try_into().unwrap()) as usize;
80    let data_start = 8usize
81        .checked_add(hlen)
82        .ok_or("safetensor: header length overflow")?;
83    if data_start > src.len() {
84        return Err("safetensor: header length exceeds file".to_string());
85    }
86    let json: serde_json::Value = serde_json::from_slice(&src[8..data_start])
87        .map_err(|e| format!("safetensor: header JSON parse error: {e}"))?;
88    let obj = json
89        .as_object()
90        .ok_or("safetensor: header is not a JSON object")?;
91
92    let is_mlx = obj
93        .get("__metadata__")
94        .and_then(|m| m.get("format"))
95        .and_then(|f| f.as_str())
96        .map(|s| s.eq_ignore_ascii_case("mlx"))
97        .unwrap_or(false);
98
99    let data_len = src.len() - data_start;
100    let mut tensors = Vec::new();
101    for (name, spec) in obj {
102        if name == "__metadata__" {
103            continue;
104        }
105        let dtype = spec
106            .get("dtype")
107            .and_then(|d| d.as_str())
108            .ok_or_else(|| format!("safetensor: tensor '{name}' missing dtype"))?
109            .to_string();
110        let shape = spec
111            .get("shape")
112            .and_then(|s| s.as_array())
113            .ok_or_else(|| format!("safetensor: tensor '{name}' missing shape"))?
114            .iter()
115            .map(|v| v.as_u64().unwrap_or(0) as usize)
116            .collect::<Vec<_>>();
117        let offs = spec
118            .get("data_offsets")
119            .and_then(|o| o.as_array())
120            .ok_or_else(|| format!("safetensor: tensor '{name}' missing data_offsets"))?;
121        if offs.len() != 2 {
122            return Err(format!(
123                "safetensor: tensor '{name}' data_offsets must be [begin,end]"
124            ));
125        }
126        let begin = offs[0].as_u64().unwrap_or(0) as usize;
127        let end = offs[1].as_u64().unwrap_or(0) as usize;
128        if end < begin || end > data_len {
129            return Err(format!("safetensor: tensor '{name}' offsets out of bounds"));
130        }
131        tensors.push(SafeTensorEntry {
132            name: name.clone(),
133            dtype,
134            shape,
135            begin,
136            end,
137        });
138    }
139    // Deterministic order (header JSON object order is not guaranteed).
140    tensors.sort_by(|a, b| a.begin.cmp(&b.begin));
141    Ok(SafeTensorPlan {
142        tensors,
143        data_start,
144        is_mlx,
145    })
146}
147
148// ── dtype gate (high-fidelity only) ──────────────────────────────────────────────────────────────
149
150/// GGML element type codes used here (mirrors `gguf_sharder`): `0=F32, 1=F16, 8=Q8_0, 30=BF16`;
151/// low-precision quants are `Q4_0=2, Q4_1=3, Q4_K=12, …`.
152pub const GGML_F32: u32 = 0;
153pub const GGML_F16: u32 = 1;
154pub const GGML_Q8_0: u32 = 8;
155pub const GGML_BF16: u32 = 30;
156
157/// Map a safetensor dtype string to a GGML element type. `None` for anything not a supported
158/// high-fidelity weight dtype (so unknown / low-precision safetensor dtypes are rejected upstream).
159pub fn safetensor_dtype_to_ggml(dtype: &str) -> Option<u32> {
160    match dtype {
161        "F32" => Some(GGML_F32),
162        "F16" => Some(GGML_F16),
163        "BF16" => Some(GGML_BF16),
164        _ => None, // F8/I8/U8/BOOL/F64/… are not high-fidelity weight inputs for this path
165    }
166}
167
168/// Whether a GGML element type is a **high-fidelity** source this versioned path accepts.
169/// Accepts `F32 / F16 / BF16 / Q8_0`; **rejects** `Q4_*` and every other low-precision quant — the
170/// "ingest high-fidelity sources only; Q4 rejected/warned" rail.
171pub fn is_high_fidelity_ggml(ggml_type: u32) -> bool {
172    matches!(ggml_type, GGML_F32 | GGML_F16 | GGML_Q8_0 | GGML_BF16)
173}
174
175/// Bytes per element for the dtypes this path accepts (used to validate declared tensor sizes).
176pub fn ggml_elem_bytes(ggml_type: u32) -> Option<usize> {
177    match ggml_type {
178        GGML_F32 => Some(4),
179        GGML_F16 | GGML_BF16 => Some(2),
180        GGML_Q8_0 => Some(1), // block-quantised; treated per-byte for verbatim repackage
181        _ => None,
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    /// Build a minimal safetensor: two F16 tensors of given element counts.
190    fn synth_safetensor(t: &[(&str, &str, Vec<usize>, usize)]) -> Vec<u8> {
191        // assign contiguous byte ranges
192        let mut entries = serde_json::Map::new();
193        let mut cursor = 0usize;
194        for (name, dtype, shape, nbytes) in t {
195            let begin = cursor;
196            let end = cursor + nbytes;
197            cursor = end;
198            entries.insert(
199                (*name).to_string(),
200                serde_json::json!({ "dtype": dtype, "shape": shape, "data_offsets": [begin, end] }),
201            );
202        }
203        let header = serde_json::Value::Object(entries);
204        let header_bytes = serde_json::to_vec(&header).unwrap();
205        let mut out = Vec::new();
206        out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
207        out.extend_from_slice(&header_bytes);
208        out.resize(out.len() + cursor, 0u8); // zeroed tensor data
209        out
210    }
211
212    #[test]
213    fn detects_formats() {
214        let st = synth_safetensor(&[("w", "F16", vec![2, 2], 8)]);
215        assert_eq!(detect_format(&st), SourceFormat::Safetensor);
216        assert_eq!(detect_format(b"GGUF\0\0\0\0"), SourceFormat::Gguf);
217        assert_eq!(detect_format(b"not a model"), SourceFormat::Unknown);
218    }
219
220    #[test]
221    fn parses_header_and_offsets() {
222        let st = synth_safetensor(&[("a", "F16", vec![4], 8), ("b", "F32", vec![2, 2], 16)]);
223        let plan = parse_safetensor_header(&st).unwrap();
224        assert_eq!(plan.tensors.len(), 2);
225        assert_eq!(plan.tensors[0].name, "a");
226        assert_eq!(plan.tensors[0].byte_len(), 8);
227        assert_eq!(plan.tensors[1].dtype, "F32");
228        assert_eq!(plan.tensors[1].byte_len(), 16);
229        assert!(!plan.is_mlx);
230    }
231
232    #[test]
233    fn dtype_gate_accepts_high_fidelity_rejects_q4() {
234        assert_eq!(safetensor_dtype_to_ggml("F16"), Some(GGML_F16));
235        assert_eq!(safetensor_dtype_to_ggml("BF16"), Some(GGML_BF16));
236        assert_eq!(safetensor_dtype_to_ggml("U8"), None);
237
238        assert!(is_high_fidelity_ggml(GGML_F16));
239        assert!(is_high_fidelity_ggml(GGML_Q8_0));
240        assert!(!is_high_fidelity_ggml(12)); // Q4_K — rejected
241        assert!(!is_high_fidelity_ggml(2)); // Q4_0 — rejected
242    }
243}
qualia_core_db/inference/safetensor.rs

qualia_core_db/inference/
safetensor.rs
