1use std::collections::{BTreeMap, BTreeSet};
8use std::io::{Read, Seek};
9
10use candle_core::quantized::gguf_file::{Content, Value};
11use candle_core::{Error, Result};
12use serde::Serialize;
13
14use super::{block_quantization_format, GgmlDType, GgufHadamard};
15
16mod header;
17mod metadata;
18use header::{Header, TensorHeader};
19pub use metadata::GgufModelMetadata;
20
21#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
22pub struct GgufTensorInventory {
23 pub name: String,
24 pub dtype: String,
25 pub ggml_type: u32,
26 pub candle_dtype_available: bool,
28 pub quantization_format: Option<String>,
29 pub dimensions: Vec<u64>,
31 pub block_axis: usize,
32 pub logical_values_per_block: u64,
33 pub bytes_per_block: u64,
34 pub elements: u64,
35 pub absolute_offset: u64,
36 pub bytes: u64,
37}
38
39#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
40pub struct GgufSplitInventory {
41 pub index: u64,
42 pub count: u64,
43 pub total_tensors: u64,
44}
45
46#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
47pub struct GgufInventory {
48 pub schema_version: u32,
49 pub architecture: String,
50 pub quantization_version: Option<u64>,
51 #[serde(skip_serializing_if = "Option::is_none")]
53 pub hadamard: Option<GgufHadamard>,
54 pub declared_file_bytes: u64,
55 pub tensor_data_offset: u64,
56 pub tensor_payload_bytes: u64,
57 pub split: Option<GgufSplitInventory>,
58 pub tensor_counts_by_dtype: BTreeMap<String, usize>,
59 pub tensors: Vec<GgufTensorInventory>,
61}
62
63impl GgufInventory {
64 pub fn read<R: Read + Seek>(reader: &mut R, declared_file_bytes: u64) -> Result<Self> {
67 Self::from_header(Header::read(reader)?, declared_file_bytes)
68 }
69
70 pub fn from_content(content: &Content, declared_file_bytes: u64) -> Result<Self> {
71 let header = Header {
72 metadata: content
73 .metadata
74 .iter()
75 .filter(|(key, _)| {
76 matches!(
77 key.as_str(),
78 "general.architecture"
79 | "general.alignment"
80 | "general.quantization_version"
81 | "split.no"
82 | "split.count"
83 | "split.tensors.count"
84 ) || super::hadamard::declares_transform(key)
85 || super::hadamard::is_geometry_key(key)
86 })
87 .map(|(key, value)| (key.clone(), value.clone()))
88 .collect(),
89 tensors: content
90 .tensor_infos
91 .iter()
92 .map(|(name, info)| {
93 Ok(TensorHeader {
94 name: name.clone(),
95 dimensions: info.shape.dims().iter().map(|&n| n as u64).collect(),
96 dtype: CANDLE_DTYPES
97 .iter()
98 .find(|(_, dtype)| *dtype == info.ggml_dtype)
99 .ok_or_else(|| {
100 Error::Msg(format!(
101 "missing GGML file ID for {:?}",
102 info.ggml_dtype
103 ))
104 })?
105 .0,
106 offset: info.offset,
107 })
108 })
109 .collect::<Result<_>>()?,
110 data_offset: content.tensor_data_offset,
111 };
112 Self::from_header(header, declared_file_bytes)
113 }
114
115 fn from_header(content: Header, declared_file_bytes: u64) -> Result<Self> {
116 let invalid = |reason: String| Error::Msg(format!("invalid GGUF inventory: {reason}"));
117 let unsupported: BTreeSet<_> = content
118 .tensors
119 .iter()
120 .filter(|tensor| block_abi(tensor.dtype).is_err())
121 .map(|tensor| tensor.dtype)
122 .collect();
123 if !unsupported.is_empty() {
124 return Err(invalid(format!(
125 "missing block ABIs for GGML types {unsupported:?}"
126 )));
127 }
128 let architecture = content
129 .metadata
130 .get("general.architecture")
131 .ok_or_else(|| invalid("missing general.architecture".into()))?
132 .to_string()?
133 .to_owned();
134 if architecture.is_empty() || content.tensors.is_empty() {
135 return Err(invalid(
136 "architecture and tensor table must be nonempty".into(),
137 ));
138 }
139 let hadamard = GgufHadamard::parse(
140 &content.metadata,
141 &architecture,
142 content
143 .tensors
144 .iter()
145 .map(|tensor| (tensor.name.as_str(), tensor.dimensions.as_slice())),
146 )?;
147 if content.data_offset > declared_file_bytes {
148 return Err(invalid(
149 "tensor data starts beyond the declared file length".into(),
150 ));
151 }
152 let alignment = match content.metadata.get("general.alignment") {
153 Some(value) => metadata_integer(value)?,
154 None => 32,
155 };
156 if alignment == 0 || !alignment.is_multiple_of(8) {
157 return Err(invalid(
158 "alignment must be a nonzero multiple of eight".into(),
159 ));
160 }
161 if !content.data_offset.is_multiple_of(alignment) {
162 return Err(invalid("tensor data offset is misaligned".into()));
163 }
164 let mut tensors = Vec::with_capacity(content.tensors.len());
165 let mut tensor_counts_by_dtype = BTreeMap::new();
166 let mut tensor_payload_bytes = 0_u64;
167 for info in &content.tensors {
168 let name = &info.name;
169 let dimensions = info.dimensions.clone();
170 if name.is_empty() || !(1..=4).contains(&dimensions.len()) || dimensions.contains(&0) {
171 return Err(invalid(format!(
172 "{name:?} needs a nonempty name and one to four nonzero dimensions"
173 )));
174 }
175 let block_axis = dimensions.len() - 1;
176 let abi = block_abi(info.dtype)?;
177 let logical_values_per_block = abi.values;
178 let bytes_per_block = abi.bytes;
179 if !dimensions[block_axis].is_multiple_of(logical_values_per_block) {
180 return Err(invalid(format!(
181 "{name:?} has an incomplete row quantization block"
182 )));
183 }
184 let elements = dimensions
185 .iter()
186 .try_fold(1_u64, |count, &dimension| count.checked_mul(dimension))
187 .ok_or_else(|| invalid(format!("{name:?} element count overflows u64")))?;
188 let bytes = (elements / logical_values_per_block)
189 .checked_mul(bytes_per_block)
190 .ok_or_else(|| invalid(format!("{name:?} payload size overflows u64")))?;
191 let absolute_offset = content
192 .data_offset
193 .checked_add(info.offset)
194 .ok_or_else(|| invalid(format!("{name:?} absolute offset overflows u64")))?;
195 let end = absolute_offset
196 .checked_add(bytes)
197 .ok_or_else(|| invalid(format!("{name:?} payload end overflows u64")))?;
198 if !info.offset.is_multiple_of(alignment) || end > declared_file_bytes {
199 return Err(invalid(format!(
200 "{name:?} is misaligned or extends beyond the declared file length"
201 )));
202 }
203 tensor_payload_bytes = tensor_payload_bytes
204 .checked_add(bytes)
205 .ok_or_else(|| invalid("total tensor bytes overflow u64".into()))?;
206 let dtype = abi.name;
207 *tensor_counts_by_dtype.entry(dtype.clone()).or_insert(0) += 1;
208 tensors.push(GgufTensorInventory {
209 name: name.clone(),
210 dtype,
211 ggml_type: info.dtype,
212 candle_dtype_available: abi.candle_dtype_available,
213 quantization_format: abi.format.map(str::to_owned),
214 dimensions,
215 block_axis,
216 logical_values_per_block,
217 bytes_per_block,
218 elements,
219 absolute_offset,
220 bytes,
221 });
222 }
223 tensors.sort_by_key(|tensor| tensor.absolute_offset);
224 for pair in tensors.windows(2) {
225 if pair[0].absolute_offset + pair[0].bytes > pair[1].absolute_offset {
226 return Err(invalid(format!(
227 "tensor payloads {:?} and {:?} overlap",
228 pair[0].name, pair[1].name
229 )));
230 }
231 }
232 tensors.sort_by(|a, b| a.name.cmp(&b.name));
233 let split = match (
234 content.metadata.get("split.no"),
235 content.metadata.get("split.count"),
236 content.metadata.get("split.tensors.count"),
237 ) {
238 (None, None, None) => None,
239 (Some(index), Some(count), Some(total)) => {
240 let split = GgufSplitInventory {
241 index: metadata_integer(index)?,
242 count: metadata_integer(count)?,
243 total_tensors: metadata_integer(total)?,
244 };
245 if split.count == 0
246 || split.index >= split.count
247 || split.total_tensors < tensors.len() as u64
248 {
249 return Err(invalid(
250 "inconsistent split index, count, or tensor total".into(),
251 ));
252 }
253 Some(split)
254 }
255 _ => return Err(invalid("incomplete split metadata".into())),
256 };
257 Ok(Self {
258 schema_version: 1,
259 architecture,
260 quantization_version: content
261 .metadata
262 .get("general.quantization_version")
263 .map(metadata_integer)
264 .transpose()?,
265 hadamard,
266 declared_file_bytes,
267 tensor_data_offset: content.data_offset,
268 tensor_payload_bytes,
269 split,
270 tensor_counts_by_dtype,
271 tensors,
272 })
273 }
274}
275
276const CANDLE_DTYPES: &[(u32, GgmlDType)] = &[
277 (0, GgmlDType::F32),
278 (1, GgmlDType::F16),
279 (2, GgmlDType::Q4_0),
280 (3, GgmlDType::Q4_1),
281 (6, GgmlDType::Q5_0),
282 (7, GgmlDType::Q5_1),
283 (8, GgmlDType::Q8_0),
284 (9, GgmlDType::Q8_1),
285 (10, GgmlDType::Q2K),
286 (11, GgmlDType::Q3K),
287 (12, GgmlDType::Q4K),
288 (13, GgmlDType::Q5K),
289 (14, GgmlDType::Q6K),
290 (15, GgmlDType::Q8K),
291 (30, GgmlDType::BF16),
292];
293
294pub(super) struct BlockAbi {
295 name: String,
296 pub(super) format: Option<&'static str>,
297 pub(super) values: u64,
298 pub(super) bytes: u64,
299 candle_dtype_available: bool,
300}
301
302pub(super) fn block_abi(code: u32) -> Result<BlockAbi> {
303 if let Some((_, dtype)) = CANDLE_DTYPES.iter().find(|(id, _)| *id == code) {
304 return Ok(BlockAbi {
305 name: format!("{dtype:?}"),
306 format: block_quantization_format(*dtype),
307 values: dtype.block_size() as u64,
308 bytes: dtype.type_size() as u64,
309 candle_dtype_available: true,
310 });
311 }
312 match code {
316 142 => {
320 return Ok(BlockAbi {
321 name: "PQ2_0".into(),
322 format: Some("quantization.gguf.pq2-0"),
323 values: 128,
324 bytes: 34,
325 candle_dtype_available: false,
326 })
327 }
328 23 => {
329 return Ok(BlockAbi {
330 name: "IQ4_XS".into(),
331 format: Some("quantization.gguf.iq4-xs"),
332 values: 256,
333 bytes: 136,
334 candle_dtype_available: false,
335 })
336 }
337 20 => {
339 return Ok(BlockAbi {
340 name: "IQ4_NL".into(),
341 format: Some("quantization.gguf.iq4-nl"),
342 values: 32,
343 bytes: 18,
344 candle_dtype_available: false,
345 })
346 }
347 21 => {
350 return Ok(BlockAbi {
351 name: "IQ3_S".into(),
352 format: Some("quantization.gguf.iq3-s"),
353 values: 256,
354 bytes: 110,
355 candle_dtype_available: false,
356 })
357 }
358 _ => {}
359 }
360 Err(Error::Msg(format!(
361 "GGUF inventory has no block ABI for GGML type {code}"
362 )))
363}
364
365fn metadata_integer(value: &Value) -> Result<u64> {
366 match value {
367 Value::U8(value) => Ok((*value).into()),
368 Value::U16(value) => Ok((*value).into()),
369 Value::U32(value) => Ok((*value).into()),
370 Value::U64(value) => Ok(*value),
371 Value::I8(value) => u64::try_from(*value).map_err(Error::wrap),
372 Value::I16(value) => u64::try_from(*value).map_err(Error::wrap),
373 Value::I32(value) => u64::try_from(*value).map_err(Error::wrap),
374 Value::I64(value) => u64::try_from(*value).map_err(Error::wrap),
375 _ => Err(Error::Msg(
376 "GGUF inventory metadata must be an integer".into(),
377 )),
378 }
379}