spark_runtime/weights/gguf/
config.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Build an atlas-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 atlas-core (which cannot see this crate's GGUF
7//! parser); this module bridges the two — it impls the atlas-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 atlas_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}
33
34/// Build a [`ModelConfig`] from the `.gguf` in `model_dir`, with no `config.json`.
35///
36/// vocab is taken from the `token_embd.weight` shape (ggml dims are
37/// `[embedding_length, vocab]`, so the trailing dim is the vocab), and tied
38/// embeddings are inferred from the absence of an `output.weight` tensor.
39pub fn config_from_gguf_dir(model_dir: &Path) -> Result<ModelConfig> {
40    let path = find_gguf(model_dir)
41        .with_context(|| format!("no .gguf file in {}", model_dir.display()))?;
42    let file =
43        std::fs::File::open(&path).with_context(|| format!("failed to open {}", path.display()))?;
44    // SAFETY: same mmap contract as the GGUF weight loader; the map outlives the
45    // borrow below and the file is not mutated concurrently.
46    let mmap = unsafe { memmap2::MmapOptions::new().map(&file)? };
47    let gguf = GgufFile::parse(&mmap)
48        .with_context(|| format!("failed to parse GGUF metadata: {}", path.display()))?;
49
50    let token_embd_vocab = gguf
51        .tensor("token_embd.weight")
52        .and_then(|t| t.dims.last().copied());
53    let has_output_weight = gguf.tensor("output.weight").is_some();
54
55    let inputs = GgufConfigInputs {
56        meta: &gguf,
57        token_embd_vocab,
58        has_output_weight,
59    };
60    config_from_gguf(&inputs).context("failed to build ModelConfig from GGUF metadata")
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    /// The GgufMeta bridge must forward to the inherent getters (no recursion)
68    /// and remap `get_arr_len` → `arr_len`.
69    #[test]
70    fn gguf_meta_bridge_forwards() {
71        // Hand-built minimal GGUF: one UINT32 KV + one STRING-array KV.
72        let mut b: Vec<u8> = Vec::new();
73        let push_u32 = |b: &mut Vec<u8>, v: u32| b.extend_from_slice(&v.to_le_bytes());
74        let push_u64 = |b: &mut Vec<u8>, v: u64| b.extend_from_slice(&v.to_le_bytes());
75        let push_str = |b: &mut Vec<u8>, s: &str| {
76            b.extend_from_slice(&(s.len() as u64).to_le_bytes());
77            b.extend_from_slice(s.as_bytes());
78        };
79        push_u32(&mut b, 0x4655_4747); // "GGUF"
80        push_u32(&mut b, 3); // version
81        push_u64(&mut b, 0); // tensor_count
82        push_u64(&mut b, 2); // kv_count
83        // key = "qwen3.block_count" : UINT32 = 28
84        push_str(&mut b, "qwen3.block_count");
85        push_u32(&mut b, 4); // UINT32
86        push_u32(&mut b, 28);
87        // key = "tokenizer.ggml.tokens" : ARRAY<STRING> len 3
88        push_str(&mut b, "tokenizer.ggml.tokens");
89        push_u32(&mut b, 9); // ARRAY
90        push_u32(&mut b, 8); // elem type STRING
91        push_u64(&mut b, 3); // len
92        for s in ["a", "bb", "ccc"] {
93            push_str(&mut b, s);
94        }
95        // Pad to the 32-byte alignment boundary so the (empty) tensor-data
96        // section start is within the buffer.
97        while !b.len().is_multiple_of(32) {
98            b.push(0);
99        }
100        let gguf = GgufFile::parse(&b).unwrap();
101        let m: &dyn GgufMeta = &gguf;
102        assert_eq!(m.get_u64("qwen3.block_count"), Some(28));
103        assert_eq!(m.get_arr_len("tokenizer.ggml.tokens"), Some(3));
104        assert_eq!(m.get_str("nonexistent"), None);
105    }
106}
spark_runtime/weights/gguf/config.rs

spark_runtime/weights/gguf/
config.rs
