kernel/resolution/
gguf.rs

1//! A minimal GGUF header reader: enough to pull the architecture, context
2//! length, and chat-template presence without loading the weights.
3//!
4//! Values are little-endian. The reader streams over a buffered file handle and
5//! seeks past values it does not need, so it never reads the tensor data.
6
7use std::collections::BTreeMap;
8use std::fs::File;
9use std::io::{BufReader, Read, Seek, SeekFrom};
10use std::path::Path;
11
12use crate::resolution::format::GgufFacts;
13
14const MAX_KV_PAIRS: u64 = 512;
15const MAX_STRING_LEN: u64 = 1 << 16;
16const MAX_STRING_ARRAY_LEN: u64 = 1 << 24;
17
18const TYPE_UINT32: u32 = 4;
19const TYPE_INT32: u32 = 5;
20const TYPE_STRING: u32 = 8;
21const TYPE_ARRAY: u32 = 9;
22const TYPE_UINT64: u32 = 10;
23const TYPE_INT64: u32 = 11;
24
25/// Whether the file begins with the GGUF magic bytes.
26pub fn has_gguf_magic(path: &Path) -> bool {
27    first_four(path) == Some(*b"GGUF")
28}
29
30/// Whether the file begins with the legacy GGML magic bytes (`lmgg` is the
31/// little-endian byte order of the `ggml` magic).
32pub fn has_ggml_magic(path: &Path) -> bool {
33    first_four(path) == Some(*b"lmgg")
34}
35
36/// The `general.architecture` value from a GGUF header, if any.
37pub fn gguf_general_architecture(path: &Path) -> Option<String> {
38    gguf_facts(path)?.architecture
39}
40
41/// Read the architecture, context length, and chat-template presence from a GGUF
42/// header. Returns `None` if the file is not a valid GGUF (v2+) header.
43pub fn gguf_facts(path: &Path) -> Option<GgufFacts> {
44    let mut reader = Reader::open(path)?;
45    if reader.read_array::<4>()? != *b"GGUF" {
46        return None;
47    }
48    let version = reader.read_u32()?;
49    if version < 2 {
50        return None;
51    }
52    let _tensor_count = reader.read_u64()?;
53    let kv_count = reader.read_u64()?;
54
55    let mut architecture: Option<String> = None;
56    let mut context_lengths: BTreeMap<String, i64> = BTreeMap::new();
57    let mut has_chat_template = false;
58
59    for _ in 0..kv_count.min(MAX_KV_PAIRS) {
60        let Some(key) = reader.read_string() else {
61            break;
62        };
63        let Some(value_type) = reader.read_u32() else {
64            break;
65        };
66
67        if key == "general.architecture" {
68            if value_type == TYPE_STRING {
69                let Some(value) = reader.read_string() else {
70                    break;
71                };
72                architecture = Some(value);
73            } else if !reader.skip_value(value_type) {
74                break;
75            }
76        } else if key == "tokenizer.chat_template" {
77            has_chat_template = true;
78            if !reader.skip_value(value_type) {
79                break;
80            }
81        } else if key.ends_with(".context_length") {
82            match read_integer(&mut reader, value_type) {
83                Some(value) => {
84                    if value > 0 {
85                        context_lengths.insert(key, value);
86                    }
87                }
88                None => {
89                    if !reader.skip_value(value_type) {
90                        break;
91                    }
92                }
93            }
94        } else if !reader.skip_value(value_type) {
95            break;
96        }
97    }
98
99    let context_length = architecture
100        .as_ref()
101        .and_then(|arch| {
102            context_lengths
103                .get(&format!("{arch}.context_length"))
104                .copied()
105        })
106        .or_else(|| {
107            if context_lengths.len() == 1 {
108                context_lengths.values().copied().next()
109            } else {
110                None
111            }
112        });
113
114    Some(GgufFacts {
115        architecture,
116        context_length,
117        has_chat_template,
118    })
119}
120
121fn first_four(path: &Path) -> Option<[u8; 4]> {
122    let mut file = File::open(path).ok()?;
123    let mut buffer = [0u8; 4];
124    file.read_exact(&mut buffer).ok()?;
125    Some(buffer)
126}
127
128fn read_integer(reader: &mut Reader, value_type: u32) -> Option<i64> {
129    match value_type {
130        TYPE_UINT32 => reader.read_u32().map(i64::from),
131        TYPE_INT32 => reader.read_i32().map(i64::from),
132        TYPE_UINT64 => reader
133            .read_u64()
134            .map(|value| value.min(i64::MAX as u64) as i64),
135        TYPE_INT64 => reader.read_i64(),
136        _ => None,
137    }
138}
139
140fn scalar_width(value_type: u32) -> Option<u64> {
141    match value_type {
142        0 | 1 | 7 => Some(1), // uint8 / int8 / bool
143        2..=3 => Some(2),     // uint16 / int16
144        4..=6 => Some(4),     // uint32 / int32 / float32
145        10..=12 => Some(8),   // uint64 / int64 / float64
146        _ => None,
147    }
148}
149
150struct Reader {
151    inner: BufReader<File>,
152}
153
154impl Reader {
155    fn open(path: &Path) -> Option<Self> {
156        Some(Self {
157            inner: BufReader::new(File::open(path).ok()?),
158        })
159    }
160
161    fn read_bytes(&mut self, count: usize) -> Option<Vec<u8>> {
162        let mut buffer = vec![0u8; count];
163        self.inner.read_exact(&mut buffer).ok()?;
164        Some(buffer)
165    }
166
167    fn read_array<const N: usize>(&mut self) -> Option<[u8; N]> {
168        let mut buffer = [0u8; N];
169        self.inner.read_exact(&mut buffer).ok()?;
170        Some(buffer)
171    }
172
173    fn read_u32(&mut self) -> Option<u32> {
174        self.read_array::<4>().map(u32::from_le_bytes)
175    }
176
177    fn read_i32(&mut self) -> Option<i32> {
178        self.read_array::<4>().map(i32::from_le_bytes)
179    }
180
181    fn read_u64(&mut self) -> Option<u64> {
182        self.read_array::<8>().map(u64::from_le_bytes)
183    }
184
185    fn read_i64(&mut self) -> Option<i64> {
186        self.read_array::<8>().map(i64::from_le_bytes)
187    }
188
189    fn read_string(&mut self) -> Option<String> {
190        let length = self.read_u64()?;
191        if length > MAX_STRING_LEN {
192            return None;
193        }
194        let bytes = self.read_bytes(length as usize)?;
195        Some(String::from_utf8_lossy(&bytes).into_owned())
196    }
197
198    fn skip(&mut self, count: u64) -> bool {
199        if count == 0 {
200            return true;
201        }
202        if count > i64::MAX as u64 {
203            return false;
204        }
205        self.inner.seek(SeekFrom::Current(count as i64)).is_ok()
206    }
207
208    fn skip_value(&mut self, value_type: u32) -> bool {
209        if let Some(width) = scalar_width(value_type) {
210            return self.skip(width);
211        }
212        match value_type {
213            TYPE_STRING => match self.read_u64() {
214                Some(length) => self.skip(length),
215                None => false,
216            },
217            TYPE_ARRAY => self.skip_array(),
218            _ => false,
219        }
220    }
221
222    fn skip_array(&mut self) -> bool {
223        let Some(element_type) = self.read_u32() else {
224            return false;
225        };
226        let Some(count) = self.read_u64() else {
227            return false;
228        };
229        if let Some(width) = scalar_width(element_type) {
230            return match count.checked_mul(width) {
231                Some(total) => self.skip(total),
232                None => false,
233            };
234        }
235        if element_type != TYPE_STRING || count > MAX_STRING_ARRAY_LEN {
236            return false;
237        }
238        for _ in 0..count {
239            let Some(length) = self.read_u64() else {
240                return false;
241            };
242            if !self.skip(length) {
243                return false;
244            }
245        }
246        true
247    }
248}
kernel/resolution/gguf.rs

kernel/resolution/
gguf.rs
