spark_runtime/weights/gguf/
config.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Build an avarok-core [`ModelConfig`] from a bare GGUF file's metadata, so a
4//! directory containing only a `.gguf` (no `config.json`) can be served.
5//!
6//! `config_from_gguf` lives in avarok-core (which cannot see this crate's GGUF
7//! parser); this module bridges the two — it impls the avarok-core [`GgufMeta`]
8//! accessor over [`GgufFile`] and supplies the two tensor-section facts the
9//! builder needs (vocab rows + presence of an untied `output.weight`).
10
11use std::path::Path;
12
13use anyhow::{Context, Result};
14use avarok_core::config::{GgufConfigInputs, GgufMeta, ModelConfig, config_from_gguf};
15
16use super::container::GgufFile;
17use super::find_gguf;
18
19impl GgufMeta for GgufFile {
20    fn get_u64(&self, key: &str) -> Option<u64> {
21        GgufFile::get_u64(self, key)
22    }
23    fn get_f64(&self, key: &str) -> Option<f64> {
24        GgufFile::get_f64(self, key)
25    }
26    fn get_str(&self, key: &str) -> Option<&str> {
27        GgufFile::get_str(self, key)
28    }
29    fn get_arr_len(&self, key: &str) -> Option<usize> {
30        GgufFile::arr_len(self, key)
31    }
32    /// Integer array, every element widened to u64.
33    ///
34    /// Accepts any integer element type: DeepSeek-V4.1 ships
35    /// `attention.compress_ratios` and `engram.layer_ids` as int32 but
36    /// `engram.{multipliers,primes,offsets}` as uint64, and a caller asking for
37    /// "the numbers in this array" should not have to know which. Returns
38    /// `None` if the key is absent or ANY element is not an integer, rather
39    /// than silently dropping elements.
40    fn get_u64_arr(&self, key: &str) -> Option<Vec<u64>> {
41        let arr = GgufFile::get(self, key)?.as_array()?;
42        arr.iter()
43            .map(|v| {
44                v.as_u64()
45                    .or_else(|| v.as_i64().and_then(|i| u64::try_from(i).ok()))
46            })
47            .collect()
48    }
49    /// Float array, every element widened to f64. Same all-or-nothing rule.
50    fn get_f64_arr(&self, key: &str) -> Option<Vec<f64>> {
51        let arr = GgufFile::get(self, key)?.as_array()?;
52        arr.iter().map(|v| v.as_f64()).collect()
53    }
54}
55
56/// Build a [`ModelConfig`] from the `.gguf` in `model_dir`, with no `config.json`.
57///
58/// vocab is taken from the `token_embd.weight` shape (ggml dims are
59/// `[embedding_length, vocab]`, so the trailing dim is the vocab), and tied
60/// embeddings are inferred from the absence of an `output.weight` tensor.
61pub fn config_from_gguf_dir(model_dir: &Path) -> Result<ModelConfig> {
62    let path = find_gguf(model_dir)
63        .with_context(|| format!("no .gguf file in {}", model_dir.display()))?;
64    let file =
65        std::fs::File::open(&path).with_context(|| format!("failed to open {}", path.display()))?;
66    // SAFETY: same mmap contract as the GGUF weight loader; the map outlives the
67    // borrow below and the file is not mutated concurrently.
68    let mmap = unsafe { memmap2::MmapOptions::new().map(&file)? };
69    let gguf = GgufFile::parse(&mmap)
70        .with_context(|| format!("failed to parse GGUF metadata: {}", path.display()))?;
71
72    let token_embd_vocab = gguf
73        .tensor("token_embd.weight")
74        .and_then(|t| t.dims.last().copied());
75    let has_output_weight = gguf.tensor("output.weight").is_some();
76
77    let inputs = GgufConfigInputs {
78        meta: &gguf,
79        token_embd_vocab,
80        has_output_weight,
81    };
82    config_from_gguf(&inputs).context("failed to build ModelConfig from GGUF metadata")
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    /// The GgufMeta bridge must forward to the inherent getters (no recursion)
90    /// and remap `get_arr_len` → `arr_len`.
91    #[test]
92    fn gguf_meta_bridge_forwards() {
93        // Hand-built minimal GGUF: one UINT32 KV + one STRING-array KV.
94        let mut b: Vec<u8> = Vec::new();
95        let push_u32 = |b: &mut Vec<u8>, v: u32| b.extend_from_slice(&v.to_le_bytes());
96        let push_u64 = |b: &mut Vec<u8>, v: u64| b.extend_from_slice(&v.to_le_bytes());
97        let push_str = |b: &mut Vec<u8>, s: &str| {
98            b.extend_from_slice(&(s.len() as u64).to_le_bytes());
99            b.extend_from_slice(s.as_bytes());
100        };
101        push_u32(&mut b, 0x4655_4747); // "GGUF"
102        push_u32(&mut b, 3); // version
103        push_u64(&mut b, 0); // tensor_count
104        push_u64(&mut b, 2); // kv_count
105        // key = "qwen3.block_count" : UINT32 = 28
106        push_str(&mut b, "qwen3.block_count");
107        push_u32(&mut b, 4); // UINT32
108        push_u32(&mut b, 28);
109        // key = "tokenizer.ggml.tokens" : ARRAY<STRING> len 3
110        push_str(&mut b, "tokenizer.ggml.tokens");
111        push_u32(&mut b, 9); // ARRAY
112        push_u32(&mut b, 8); // elem type STRING
113        push_u64(&mut b, 3); // len
114        for s in ["a", "bb", "ccc"] {
115            push_str(&mut b, s);
116        }
117        // Pad to the 32-byte alignment boundary so the (empty) tensor-data
118        // section start is within the buffer.
119        while !b.len().is_multiple_of(32) {
120            b.push(0);
121        }
122        let gguf = GgufFile::parse(&b).unwrap();
123        let m: &dyn GgufMeta = &gguf;
124        assert_eq!(m.get_u64("qwen3.block_count"), Some(28));
125        assert_eq!(m.get_arr_len("tokenizer.ggml.tokens"), Some(3));
126        assert_eq!(m.get_str("nonexistent"), None);
127    }
128}
spark_runtime/weights/gguf/config.rs

spark_runtime/weights/gguf/
config.rs
