1use super::*;
5const BYTES_PER_MIB: usize = 1024 * 1024;
6use anyhow::{Context, Result, ensure};
7use serde_json::{Value, json};
8
9pub(crate) fn inspect(map: &[u8]) -> Result<Inventory> {
10 let mut reader = Header { bytes: map, pos: 0 };
11
12 ensure!(reader.take(4)? == b"GGUF", "expected a GGUF file");
13
14 let version = reader.u32()?;
15
16 ensure!(
17 matches!(version, 2 | 3),
18 "unsupported GGUF version {version}"
19 );
20
21 let tensor_count = reader.count()?;
22 let metadata_count = reader.count()?;
23 let mut metadata = serde_json::Map::new();
24 let mut alignment = 32_u64;
25
26 for _ in 0..metadata_count {
27 let key = reader.string()?.to_owned();
28 let kind = reader.u32()?;
29 let value = reader.value(kind, 0)?;
30
31 if key == "general.alignment" {
32 ensure!(kind == 4, "GGUF alignment must be uint32");
33
34 alignment = value.as_u64().context("invalid GGUF alignment")?;
35 }
36
37 ensure!(
38 !metadata.contains_key(&key),
39 "duplicate GGUF metadata {key}"
40 );
41 metadata.insert(key, json!({"type": kind, "value": value}));
42 }
43
44 let mut tensors = Vec::new();
45
46 for _ in 0..tensor_count {
47 let name = reader.string()?.to_owned();
48 let rank = reader.u32()?;
49
50 ensure!((1..=4).contains(&rank), "invalid GGUF tensor rank");
51
52 let mut shape = (0..rank)
53 .map(|_| reader.u64())
54 .collect::<Result<Vec<_>>>()?;
55
56 shape.reverse(); let encoding = codec(reader.u32()?);
59 let offset = reader.u64()?;
60
61 tensors.push((name, shape, encoding, offset));
62 }
63
64 ensure!(
65 alignment > 0 && alignment.is_power_of_two(),
66 "invalid GGUF alignment"
67 );
68
69 let data_start = (reader.pos as u64)
70 .checked_add(alignment - 1)
71 .context("GGUF alignment overflow")?
72 / alignment
73 * alignment;
74
75 ensure!(data_start <= map.len() as u64, "truncated GGUF padding");
76 tensors.sort_by_key(|t| t.3);
77
78 let mut described = Vec::new();
79
80 for (i, (name, shape, encoding, relative)) in tensors.iter().enumerate() {
81 let offset = data_start
82 .checked_add(*relative)
83 .context("GGUF tensor offset overflow")?;
84 let end = match tensors.get(i + 1) {
85 Some(next) => data_start
86 .checked_add(next.3)
87 .context("GGUF tensor offset overflow")?,
88 None => map.len() as u64,
89 };
90
91 ensure!(
92 relative.is_multiple_of(alignment) && offset <= end && end <= map.len() as u64,
93 "GGUF tensor range out of bounds"
94 );
95
96 let length = encoded_len(*encoding, shape)?.unwrap_or(end - offset);
97
98 ensure!(length <= end - offset, "truncated GGUF tensor {name}");
99 described.push(Tensor::new(
100 name.clone(),
101 TensorRole::Opaque,
102 Some(shape.clone()),
103 TensorEncoding::Ggml {
104 encoding: *encoding,
105 data: DataSpan {
106 object: ObjectId(0),
107 offset,
108 length,
109 },
110 },
111 ));
112 }
113
114 Ok(Inventory {
115 format: ContainerFormat::Gguf { version },
116 metadata: Value::Object(metadata),
117 tensors: described,
118 })
119}
120
121fn encoded_len(encoding: GgmlEncoding, shape: &[u64]) -> Result<Option<u64>> {
122 let (block, bytes) = match encoding {
123 GgmlEncoding::F32 => (1, 4),
124 GgmlEncoding::F16 | GgmlEncoding::Bf16 => (1, 2),
125 GgmlEncoding::Q4_0 => (32, 18),
126 GgmlEncoding::Q4_1 => (32, 20),
127 GgmlEncoding::Q8_0 => (32, 34),
128 GgmlEncoding::Q4K => (256, 144),
129 GgmlEncoding::Q5K => (256, 176),
130 GgmlEncoding::Q6K => (256, 210),
131 GgmlEncoding::Opaque { .. } => return Ok(None),
132 };
133 let width = *shape.last().context("GGUF tensor has no dimensions")?;
134
135 ensure!(
136 width > 0 && width.is_multiple_of(block),
137 "GGUF row does not fit its quantization block"
138 );
139
140 let elements = shape.iter().try_fold(1_u64, |size, &dim| {
141 ensure!(dim > 0, "empty GGUF tensor dimension");
142
143 size.checked_mul(dim).context("GGUF tensor size overflow")
144 })?;
145
146 Ok(Some(
147 (elements / block)
148 .checked_mul(bytes)
149 .context("GGUF tensor size overflow")?,
150 ))
151}
152
153fn codec(code: u32) -> GgmlEncoding {
154 match code {
155 0 => GgmlEncoding::F32,
156 1 => GgmlEncoding::F16,
157 2 => GgmlEncoding::Q4_0,
158 3 => GgmlEncoding::Q4_1,
159 8 => GgmlEncoding::Q8_0,
160 12 => GgmlEncoding::Q4K,
161 13 => GgmlEncoding::Q5K,
162 14 => GgmlEncoding::Q6K,
163 30 => GgmlEncoding::Bf16,
164 code => GgmlEncoding::Opaque { code },
165 }
166}
167
168struct Header<'a> {
169 bytes: &'a [u8],
170 pos: usize,
171}
172
173impl<'a> Header<'a> {
174 fn take(&mut self, count: usize) -> Result<&'a [u8]> {
175 let end = self
176 .pos
177 .checked_add(count)
178 .context("GGUF header overflow")?;
179
180 ensure!(end <= 64 * BYTES_PER_MIB, "GGUF header exceeds 64 MiB");
181
182 let bytes = self
183 .bytes
184 .get(self.pos..end)
185 .context("truncated GGUF header")?;
186 self.pos = end;
187
188 Ok(bytes)
189 }
190
191 fn u32(&mut self) -> Result<u32> {
192 Ok(u32::from_le_bytes(self.take(4)?.try_into()?))
193 }
194
195 fn u64(&mut self) -> Result<u64> {
196 Ok(u64::from_le_bytes(self.take(8)?.try_into()?))
197 }
198
199 fn count(&mut self) -> Result<u64> {
200 let count = self.u64()?;
201
202 ensure!(
203 count <= 1_000_000,
204 "GGUF entry count exceeds inspection limit"
205 );
206
207 Ok(count)
208 }
209
210 fn string(&mut self) -> Result<&'a str> {
211 let len = usize::try_from(self.u64()?).context("GGUF string size overflow")?;
212
213 std::str::from_utf8(self.take(len)?).context("invalid UTF-8 in GGUF header")
214 }
215
216 fn value(&mut self, kind: u32, depth: usize) -> Result<Value> {
217 ensure!(depth <= 8, "GGUF metadata nesting limit exceeded");
218
219 Ok(match kind {
220 0 => json!(self.take(1)?[0]),
221 1 => json!(self.take(1)?[0] as i8),
222 2 => json!(u16::from_le_bytes(self.take(2)?.try_into()?)),
223 3 => json!(i16::from_le_bytes(self.take(2)?.try_into()?)),
224 4 => json!(self.u32()?),
225 5 => json!(self.u32()? as i32),
226 6 => {
227 let bits = self.u32()?;
228
229 float_value(f64::from(f32::from_bits(bits)), u64::from(bits))
230 }
231 7 => {
232 let value = self.take(1)?[0];
233
234 ensure!(value <= 1, "invalid GGUF boolean");
235
236 json!(value == 1)
237 }
238 8 => json!(self.string()?),
239 9 => {
240 let element = self.u32()?;
241 let count = self.count()?;
242 let values = (0..count)
243 .map(|_| self.value(element, depth + 1))
244 .collect::<Result<Vec<_>>>()?;
245
246 json!({"element_type": element, "values": values})
247 }
248 10 => json!(self.u64()?),
249 11 => json!(self.u64()? as i64),
250 12 => {
251 let bits = self.u64()?;
252
253 float_value(f64::from_bits(bits), bits)
254 }
255 _ => anyhow::bail!("unknown GGUF metadata type {kind}"),
256 })
257 }
258}
259
260fn float_value(value: f64, bits: u64) -> Value {
261 serde_json::Number::from_f64(value)
262 .map(Value::Number)
263 .unwrap_or_else(|| json!({"nonfinite_bits": bits}))
264}