skippy_model/
imatrix.rs

1use std::collections::BTreeMap;
2use std::fs;
3use std::io::{Cursor, Read};
4use std::path::{Path, PathBuf};
5
6use anyhow::{Context, Result, anyhow, ensure};
7
8type ByteReader<'a> = Cursor<&'a [u8]>;
9
10const GGUF_MAGIC: &[u8; 4] = b"GGUF";
11const GGUF_ALIGNMENT_DEFAULT: u64 = 32;
12const GGML_TYPE_F32: u32 = 0;
13const TENSOR_SUMS_SUFFIX: &str = ".in_sum2";
14const TENSOR_COUNTS_SUFFIX: &str = ".counts";
15
16const GGUF_TYPE_UINT32: u32 = 4;
17const GGUF_TYPE_FLOAT32: u32 = 6;
18const GGUF_TYPE_BOOL: u32 = 7;
19const GGUF_TYPE_STRING: u32 = 8;
20const GGUF_TYPE_ARRAY: u32 = 9;
21const GGUF_TYPE_UINT64: u32 = 10;
22const GGUF_TYPE_INT64: u32 = 11;
23const GGUF_TYPE_FLOAT64: u32 = 12;
24
25const KV_GENERAL_ALIGNMENT: &str = "general.alignment";
26const KV_IMATRIX_DATASETS: &str = "imatrix.datasets";
27const KV_IMATRIX_CHUNK_COUNT: &str = "imatrix.chunk_count";
28
29#[derive(Debug, Clone, PartialEq)]
30pub struct ImatrixEntry {
31    pub name: String,
32    pub values: Vec<f32>,
33}
34
35#[derive(Debug, Clone, PartialEq)]
36pub struct Imatrix {
37    entries: Vec<ImatrixEntry>,
38    source_path: PathBuf,
39    dataset: Option<String>,
40    chunk_count: i32,
41}
42
43impl Imatrix {
44    pub fn load(
45        path: &Path,
46        include_weights: &[String],
47        exclude_weights: &[String],
48    ) -> Result<Self> {
49        let bytes =
50            fs::read(path).with_context(|| format!("read imatrix file {}", path.display()))?;
51        Self::from_bytes(path, &bytes, include_weights, exclude_weights)
52    }
53
54    /// Parse an importance matrix from caller-owned bytes.
55    ///
56    /// This lets callers bind verification and parsing to the same immutable
57    /// snapshot instead of reopening a path after it has been hashed.
58    pub fn from_bytes(
59        path: &Path,
60        bytes: &[u8],
61        include_weights: &[String],
62        exclude_weights: &[String],
63    ) -> Result<Self> {
64        let loaded = if bytes.starts_with(GGUF_MAGIC) {
65            load_gguf_imatrix(bytes, path)?
66        } else {
67            load_legacy_imatrix(bytes, path)?
68        };
69        Self::from_loaded(path, loaded, include_weights, exclude_weights)
70    }
71
72    pub fn entries(&self) -> &[ImatrixEntry] {
73        &self.entries
74    }
75
76    pub fn source_path(&self) -> &Path {
77        &self.source_path
78    }
79
80    pub fn dataset(&self) -> Option<&str> {
81        self.dataset.as_deref()
82    }
83
84    pub fn chunk_count(&self) -> i32 {
85        self.chunk_count
86    }
87
88    pub fn entry_count(&self) -> usize {
89        self.entries.len()
90    }
91
92    fn from_loaded(
93        path: &Path,
94        loaded: LoadedImatrix,
95        include_weights: &[String],
96        exclude_weights: &[String],
97    ) -> Result<Self> {
98        let mut selected = loaded
99            .entries
100            .into_iter()
101            .filter(|entry| include_exclude_match(&entry.name, include_weights, exclude_weights))
102            .collect::<Vec<_>>();
103        selected.sort_by(|a, b| a.name.cmp(&b.name));
104        ensure!(
105            !selected.is_empty(),
106            "imatrix filters removed all entries from {}",
107            path.display()
108        );
109
110        Ok(Self {
111            entries: selected,
112            source_path: path.to_path_buf(),
113            dataset: loaded.dataset,
114            chunk_count: loaded.chunk_count,
115        })
116    }
117}
118
119struct LoadedImatrix {
120    entries: Vec<ImatrixEntry>,
121    dataset: Option<String>,
122    chunk_count: i32,
123}
124
125fn load_legacy_imatrix(bytes: &[u8], path: &Path) -> Result<LoadedImatrix> {
126    let mut reader = Cursor::new(bytes);
127    let entry_count = read_i32(&mut reader)
128        .with_context(|| format!("read imatrix entry count from {}", path.display()))?;
129    ensure!(
130        entry_count > 0,
131        "imatrix file has no entries: {}",
132        path.display()
133    );
134    let entry_count = usize::try_from(entry_count).context("imatrix entry count overflow")?;
135    ensure!(
136        entry_count <= remaining_bytes(&reader) / 12,
137        "imatrix entry count exceeds the available file data"
138    );
139    let mut entries = Vec::with_capacity(entry_count);
140    for index in 0..entry_count {
141        entries.push(read_legacy_imatrix_entry(&mut reader, index)?);
142    }
143    let (chunk_count, dataset) = read_legacy_imatrix_trailer(&mut reader)?;
144    Ok(LoadedImatrix {
145        entries,
146        dataset,
147        chunk_count,
148    })
149}
150
151fn read_legacy_imatrix_entry(reader: &mut ByteReader<'_>, index: usize) -> Result<ImatrixEntry> {
152    let name_len = read_i32(reader).with_context(|| format!("read imatrix name length {index}"))?;
153    ensure!(name_len > 0, "imatrix entry {index} has empty name");
154    let name_len = usize::try_from(name_len).context("imatrix name length overflow")?;
155    ensure!(
156        name_len <= remaining_bytes(reader),
157        "imatrix entry {index} name extends past end of file"
158    );
159    let mut name_bytes = vec![0_u8; name_len];
160    reader
161        .read_exact(&mut name_bytes)
162        .with_context(|| format!("read imatrix name {index}"))?;
163    let name = String::from_utf8(name_bytes).with_context(|| format!("imatrix name {index}"))?;
164    let ncall = read_i32(reader).with_context(|| format!("read imatrix ncall for {name}"))?;
165    let value_count =
166        read_i32(reader).with_context(|| format!("read imatrix value count for {name}"))?;
167    ensure!(value_count > 0, "imatrix entry {name} has no values");
168    let value_count = usize::try_from(value_count).context("imatrix value count overflow")?;
169    let value_bytes = value_count
170        .checked_mul(std::mem::size_of::<f32>())
171        .context("imatrix value byte count overflow")?;
172    ensure!(
173        value_bytes <= remaining_bytes(reader),
174        "imatrix entry {name} values extend past end of file"
175    );
176    let mut values = (0..value_count)
177        .map(|_| read_f32(reader))
178        .collect::<Result<Vec<_>>>()?;
179    if ncall > 0 {
180        let denom = ncall as f32;
181        for value in &mut values {
182            *value /= denom;
183        }
184    }
185    Ok(ImatrixEntry { name, values })
186}
187
188fn read_legacy_imatrix_trailer(reader: &mut ByteReader<'_>) -> Result<(i32, Option<String>)> {
189    if reader.position() as usize >= reader.get_ref().len() {
190        return Ok((0, None));
191    }
192    let chunk_count = read_i32(reader)?;
193    if reader.position() as usize >= reader.get_ref().len() {
194        return Ok((chunk_count, None));
195    }
196    let len = read_i32(reader)?;
197    if len <= 0 {
198        return Ok((chunk_count, None));
199    }
200    let len = usize::try_from(len).context("imatrix dataset length overflow")?;
201    ensure!(
202        len <= remaining_bytes(reader),
203        "imatrix dataset extends past end of file"
204    );
205    let mut dataset = vec![0_u8; len];
206    reader.read_exact(&mut dataset)?;
207    Ok((chunk_count, Some(String::from_utf8(dataset)?)))
208}
209
210fn load_gguf_imatrix(bytes: &[u8], path: &Path) -> Result<LoadedImatrix> {
211    let mut reader = Cursor::new(bytes);
212    let header = GgufHeader::read(&mut reader)?;
213    ensure!(
214        header.version >= 2,
215        "unsupported GGUF imatrix version {} in {}",
216        header.version,
217        path.display()
218    );
219    ensure!(
220        header.metadata_count <= remaining_bytes(&reader) as u64 / 12,
221        "GGUF metadata count exceeds the available file data"
222    );
223    let mut metadata = GgufMetadata::default();
224    for _ in 0..header.metadata_count {
225        let key = read_gguf_string(&mut reader)?;
226        let value_type = read_u32(&mut reader)?;
227        metadata.read_value(&mut reader, &key, value_type)?;
228    }
229
230    ensure!(
231        header.tensor_count <= remaining_bytes(&reader) as u64 / 32,
232        "GGUF tensor count exceeds the available file data"
233    );
234    let tensor_count =
235        usize::try_from(header.tensor_count).context("GGUF tensor count overflow")?;
236    let mut tensors = Vec::with_capacity(tensor_count);
237    for _ in 0..header.tensor_count {
238        tensors.push(GgufTensorInfo::read(&mut reader)?);
239    }
240    let data_start = align_to(reader.position(), metadata.alignment)
241        .context("GGUF imatrix data offset overflow")?;
242    let entries = read_gguf_entries(bytes, &tensors, data_start)?;
243    ensure!(
244        !entries.is_empty(),
245        "GGUF imatrix has no paired tensors in {}",
246        path.display()
247    );
248    Ok(LoadedImatrix {
249        entries,
250        dataset: metadata.datasets.into_iter().next(),
251        chunk_count: metadata.chunk_count.unwrap_or(0) as i32,
252    })
253}
254
255struct GgufHeader {
256    version: u32,
257    tensor_count: u64,
258    metadata_count: u64,
259}
260
261impl GgufHeader {
262    fn read(reader: &mut ByteReader<'_>) -> Result<Self> {
263        let mut magic = [0_u8; 4];
264        reader.read_exact(&mut magic)?;
265        ensure!(&magic == GGUF_MAGIC, "not a GGUF file");
266        Ok(Self {
267            version: read_u32(reader)?,
268            tensor_count: read_u64(reader)?,
269            metadata_count: read_u64(reader)?,
270        })
271    }
272}
273
274#[derive(Default)]
275struct GgufMetadata {
276    alignment: u64,
277    datasets: Vec<String>,
278    chunk_count: Option<u32>,
279}
280
281impl GgufMetadata {
282    fn read_value(
283        &mut self,
284        reader: &mut ByteReader<'_>,
285        key: &str,
286        value_type: u32,
287    ) -> Result<()> {
288        match (key, value_type) {
289            (KV_GENERAL_ALIGNMENT, GGUF_TYPE_UINT32) => {
290                self.alignment = read_u32(reader)? as u64;
291            }
292            (KV_GENERAL_ALIGNMENT, GGUF_TYPE_UINT64) => {
293                self.alignment = read_u64(reader)?;
294            }
295            (KV_IMATRIX_CHUNK_COUNT, GGUF_TYPE_UINT32) => {
296                self.chunk_count = Some(read_u32(reader)?);
297            }
298            (KV_IMATRIX_DATASETS, GGUF_TYPE_ARRAY) => {
299                self.datasets = read_string_array(reader)?;
300            }
301            _ => skip_gguf_value(reader, value_type)?,
302        }
303        if self.alignment == 0 {
304            self.alignment = GGUF_ALIGNMENT_DEFAULT;
305        }
306        Ok(())
307    }
308}
309
310struct GgufTensorInfo {
311    name: String,
312    dims: Vec<u64>,
313    tensor_type: u32,
314    offset: u64,
315}
316
317impl GgufTensorInfo {
318    fn read(reader: &mut ByteReader<'_>) -> Result<Self> {
319        let name = read_gguf_string(reader)?;
320        let n_dims = read_u32(reader)?;
321        ensure!(n_dims > 0, "GGUF tensor {name} has no dimensions");
322        ensure!(
323            n_dims as usize <= remaining_bytes(reader) / 8,
324            "GGUF tensor {name} dimensions extend past end of file"
325        );
326        let dims = (0..n_dims)
327            .map(|_| read_u64(reader))
328            .collect::<Result<Vec<_>>>()?;
329        let tensor_type = read_u32(reader)?;
330        let offset = read_u64(reader)?;
331        Ok(Self {
332            name,
333            dims,
334            tensor_type,
335            offset,
336        })
337    }
338
339    fn element_count(&self) -> Result<usize> {
340        self.dims.iter().try_fold(1_usize, |acc, dim| {
341            let dim = usize::try_from(*dim)
342                .with_context(|| format!("GGUF tensor {} dimension overflow", self.name))?;
343            acc.checked_mul(dim)
344                .with_context(|| format!("GGUF tensor {} is too large", self.name))
345        })
346    }
347}
348
349fn read_gguf_entries(
350    bytes: &[u8],
351    tensors: &[GgufTensorInfo],
352    data_start: u64,
353) -> Result<Vec<ImatrixEntry>> {
354    let mut sums = BTreeMap::<String, Vec<f32>>::new();
355    let mut counts = BTreeMap::<String, Vec<f32>>::new();
356    for tensor in tensors {
357        if tensor.tensor_type != GGML_TYPE_F32 {
358            continue;
359        }
360        let Some((base_name, kind)) = imatrix_tensor_name(&tensor.name) else {
361            continue;
362        };
363        let values = read_gguf_f32_tensor(bytes, tensor, data_start)?;
364        match kind {
365            ImatrixTensorKind::Sums => {
366                sums.insert(base_name.to_string(), values);
367            }
368            ImatrixTensorKind::Counts => {
369                counts.insert(base_name.to_string(), values);
370            }
371        }
372    }
373    let mut entries = Vec::new();
374    for (name, sum_values) in sums {
375        let count_values = counts
376            .remove(&name)
377            .with_context(|| format!("GGUF imatrix tensor {name} is missing counts"))?;
378        entries.push(ImatrixEntry {
379            name,
380            values: normalize_gguf_imatrix_values(&sum_values, &count_values)?,
381        });
382    }
383    Ok(entries)
384}
385
386enum ImatrixTensorKind {
387    Sums,
388    Counts,
389}
390
391fn imatrix_tensor_name(name: &str) -> Option<(&str, ImatrixTensorKind)> {
392    if let Some(base) = name.strip_suffix(TENSOR_SUMS_SUFFIX) {
393        return Some((base, ImatrixTensorKind::Sums));
394    }
395    name.strip_suffix(TENSOR_COUNTS_SUFFIX)
396        .map(|base| (base, ImatrixTensorKind::Counts))
397}
398
399fn normalize_gguf_imatrix_values(sums: &[f32], counts: &[f32]) -> Result<Vec<f32>> {
400    ensure!(!counts.is_empty(), "GGUF imatrix entry has no counts");
401    ensure!(
402        sums.len().is_multiple_of(counts.len()),
403        "GGUF imatrix sums/counts shape mismatch"
404    );
405    let values_per_count = sums.len() / counts.len();
406    let mut output = vec![1.0_f32; sums.len()];
407    for (count_index, count) in counts.iter().enumerate() {
408        if *count <= 0.0 {
409            continue;
410        }
411        let offset = count_index * values_per_count;
412        for i in 0..values_per_count {
413            output[offset + i] = sums[offset + i] / count;
414        }
415    }
416    Ok(output)
417}
418
419fn read_gguf_f32_tensor(
420    bytes: &[u8],
421    tensor: &GgufTensorInfo,
422    data_start: u64,
423) -> Result<Vec<f32>> {
424    let element_count = tensor.element_count()?;
425    let byte_len = element_count
426        .checked_mul(std::mem::size_of::<f32>())
427        .with_context(|| format!("GGUF tensor {} is too large", tensor.name))?;
428    let offset = data_start
429        .checked_add(tensor.offset)
430        .with_context(|| format!("GGUF tensor {} offset overflow", tensor.name))?;
431    let offset = usize::try_from(offset)
432        .with_context(|| format!("GGUF tensor {} offset exceeds address space", tensor.name))?;
433    let end = offset
434        .checked_add(byte_len)
435        .with_context(|| format!("GGUF tensor {} byte range overflow", tensor.name))?;
436    ensure!(
437        end <= bytes.len(),
438        "GGUF tensor {} extends past end of file",
439        tensor.name
440    );
441    bytes[offset..end]
442        .as_chunks::<4>()
443        .0
444        .iter()
445        .map(|chunk| Ok(f32::from_le_bytes(*chunk)))
446        .collect::<Result<Vec<_>>>()
447}
448
449fn include_exclude_match(
450    name: &str,
451    include_weights: &[String],
452    exclude_weights: &[String],
453) -> bool {
454    if !exclude_weights.is_empty() {
455        return !exclude_weights.iter().any(|filter| name.contains(filter));
456    }
457    if !include_weights.is_empty() {
458        return include_weights.iter().any(|filter| name.contains(filter));
459    }
460    true
461}
462
463fn read_string_array(reader: &mut ByteReader<'_>) -> Result<Vec<String>> {
464    let element_type = read_u32(reader)?;
465    ensure!(
466        element_type == GGUF_TYPE_STRING,
467        "expected GGUF string array, found type {element_type}"
468    );
469    let len = read_u64(reader)?;
470    ensure!(
471        len <= remaining_bytes(reader) as u64 / 8,
472        "GGUF string array count exceeds the available file data"
473    );
474    (0..len).map(|_| read_gguf_string(reader)).collect()
475}
476
477fn skip_gguf_value(reader: &mut ByteReader<'_>, value_type: u32) -> Result<()> {
478    match value_type {
479        0 | 1 => skip_bytes(reader, 1),
480        2 | 3 => skip_bytes(reader, 2),
481        GGUF_TYPE_UINT32 | 5 | GGUF_TYPE_FLOAT32 => skip_bytes(reader, 4),
482        GGUF_TYPE_BOOL => skip_bytes(reader, 1),
483        GGUF_TYPE_STRING => {
484            let _ = read_gguf_string(reader)?;
485            Ok(())
486        }
487        GGUF_TYPE_ARRAY => skip_gguf_array(reader),
488        GGUF_TYPE_UINT64 | GGUF_TYPE_INT64 | GGUF_TYPE_FLOAT64 => skip_bytes(reader, 8),
489        _ => Err(anyhow!("unsupported GGUF metadata type {value_type}")),
490    }
491}
492
493fn skip_gguf_array(reader: &mut ByteReader<'_>) -> Result<()> {
494    let element_type = read_u32(reader)?;
495    let len = read_u64(reader)?;
496    ensure!(
497        len <= remaining_bytes(reader) as u64,
498        "GGUF array count exceeds the available file data"
499    );
500    for _ in 0..len {
501        skip_gguf_value(reader, element_type)?;
502    }
503    Ok(())
504}
505
506fn skip_bytes(reader: &mut ByteReader<'_>, len: u64) -> Result<()> {
507    let position = reader.position();
508    let next = position
509        .checked_add(len)
510        .context("GGUF metadata offset overflow")?;
511    ensure!(
512        next <= reader.get_ref().len() as u64,
513        "GGUF metadata extends past end of file"
514    );
515    reader.set_position(next);
516    Ok(())
517}
518
519fn read_gguf_string(reader: &mut ByteReader<'_>) -> Result<String> {
520    let len = read_u64(reader)?;
521    let len = usize::try_from(len).context("GGUF string length overflow")?;
522    ensure!(
523        len <= remaining_bytes(reader),
524        "GGUF string extends past end of file"
525    );
526    let mut bytes = vec![0_u8; len];
527    reader.read_exact(&mut bytes)?;
528    String::from_utf8(bytes).context("GGUF string is not UTF-8")
529}
530
531fn read_i32(reader: &mut ByteReader<'_>) -> Result<i32> {
532    let mut bytes = [0_u8; 4];
533    reader.read_exact(&mut bytes)?;
534    Ok(i32::from_le_bytes(bytes))
535}
536
537fn read_u32(reader: &mut ByteReader<'_>) -> Result<u32> {
538    let mut bytes = [0_u8; 4];
539    reader.read_exact(&mut bytes)?;
540    Ok(u32::from_le_bytes(bytes))
541}
542
543fn read_u64(reader: &mut ByteReader<'_>) -> Result<u64> {
544    let mut bytes = [0_u8; 8];
545    reader.read_exact(&mut bytes)?;
546    Ok(u64::from_le_bytes(bytes))
547}
548
549fn read_f32(reader: &mut ByteReader<'_>) -> Result<f32> {
550    let mut bytes = [0_u8; 4];
551    reader.read_exact(&mut bytes)?;
552    Ok(f32::from_le_bytes(bytes))
553}
554
555fn remaining_bytes(reader: &ByteReader<'_>) -> usize {
556    reader
557        .get_ref()
558        .len()
559        .saturating_sub(reader.position() as usize)
560}
561
562fn align_to(value: u64, alignment: u64) -> Option<u64> {
563    if alignment <= 1 {
564        return Some(value);
565    }
566    value.div_ceil(alignment).checked_mul(alignment)
567}
568
569#[cfg(test)]
570mod tests {
571    use super::*;
572
573    #[test]
574    fn loads_legacy_imatrix_with_include_filter() {
575        let root = unique_temp_dir();
576        fs::create_dir_all(&root).unwrap();
577        let imatrix_path = root.join("imatrix.dat");
578        write_legacy_imatrix(
579            &imatrix_path,
580            &[
581                ("blk.0.attn_q.weight", 2, &[2.0, 4.0]),
582                ("blk.0.ffn_down.weight", 1, &[9.0, 12.0]),
583            ],
584        );
585
586        let imatrix = Imatrix::load(&imatrix_path, &["attn_q".to_string()], &Vec::new()).unwrap();
587
588        assert_eq!(imatrix.entry_count(), 1);
589        assert_eq!(imatrix.entries[0].values, vec![1.0, 2.0]);
590        assert_eq!(imatrix.dataset(), Some("dataset.txt"));
591        assert_eq!(imatrix.chunk_count(), 3);
592        fs::remove_dir_all(root).unwrap();
593    }
594
595    #[test]
596    fn loads_gguf_imatrix_with_normalized_counts() {
597        let root = unique_temp_dir();
598        fs::create_dir_all(&root).unwrap();
599        let imatrix_path = root.join("imatrix.gguf");
600        write_gguf_imatrix(&imatrix_path);
601
602        let imatrix = Imatrix::load(&imatrix_path, &["attn_q".to_string()], &Vec::new()).unwrap();
603
604        assert_eq!(imatrix.entry_count(), 1);
605        assert_eq!(imatrix.entries[0].values, vec![1.0, 2.0, 1.0, 1.0]);
606        assert_eq!(imatrix.dataset(), Some("calibration.txt"));
607        assert_eq!(imatrix.chunk_count(), 7);
608        fs::remove_dir_all(root).unwrap();
609    }
610
611    #[test]
612    fn rejects_alignment_overflow() {
613        assert_eq!(align_to(u64::MAX, 32), None);
614    }
615
616    fn unique_temp_dir() -> PathBuf {
617        static NEXT_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
618        let nanos = std::time::SystemTime::now()
619            .duration_since(std::time::UNIX_EPOCH)
620            .unwrap()
621            .as_nanos();
622        let id = NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
623        std::env::temp_dir().join(format!("skippy-imatrix-{nanos}-{id}"))
624    }
625
626    fn write_legacy_imatrix(path: &Path, entries: &[(&str, i32, &[f32])]) {
627        let mut bytes = Vec::new();
628        bytes.extend_from_slice(&(entries.len() as i32).to_le_bytes());
629        for (name, ncall, values) in entries {
630            bytes.extend_from_slice(&(name.len() as i32).to_le_bytes());
631            bytes.extend_from_slice(name.as_bytes());
632            bytes.extend_from_slice(&ncall.to_le_bytes());
633            bytes.extend_from_slice(&(values.len() as i32).to_le_bytes());
634            for value in *values {
635                bytes.extend_from_slice(&value.to_le_bytes());
636            }
637        }
638        bytes.extend_from_slice(&3_i32.to_le_bytes());
639        let dataset = "dataset.txt";
640        bytes.extend_from_slice(&(dataset.len() as i32).to_le_bytes());
641        bytes.extend_from_slice(dataset.as_bytes());
642        fs::write(path, bytes).unwrap();
643    }
644
645    fn write_gguf_imatrix(path: &Path) {
646        let sums_name = "blk.0.attn_q.weight.in_sum2";
647        let counts_name = "blk.0.attn_q.weight.counts";
648        let mut bytes = Vec::new();
649        bytes.extend_from_slice(GGUF_MAGIC);
650        bytes.extend_from_slice(&3_u32.to_le_bytes());
651        bytes.extend_from_slice(&2_u64.to_le_bytes());
652        bytes.extend_from_slice(&4_u64.to_le_bytes());
653        write_gguf_kv_string(&mut bytes, "general.type", "imatrix");
654        write_gguf_kv_u32(&mut bytes, KV_GENERAL_ALIGNMENT, 32);
655        write_gguf_kv_u32(&mut bytes, KV_IMATRIX_CHUNK_COUNT, 7);
656        write_gguf_kv_string_array(&mut bytes, KV_IMATRIX_DATASETS, &["calibration.txt"]);
657        write_gguf_tensor_info(&mut bytes, sums_name, &[2, 2], GGML_TYPE_F32, 0);
658        write_gguf_tensor_info(&mut bytes, counts_name, &[1, 2], GGML_TYPE_F32, 16);
659        while bytes.len() % 32 != 0 {
660            bytes.push(0);
661        }
662        for value in [2.0_f32, 4.0, 9.0, 11.0] {
663            bytes.extend_from_slice(&value.to_le_bytes());
664        }
665        for value in [2.0_f32, 0.0] {
666            bytes.extend_from_slice(&value.to_le_bytes());
667        }
668        fs::write(path, bytes).unwrap();
669    }
670
671    fn write_gguf_kv_string(bytes: &mut Vec<u8>, key: &str, value: &str) {
672        write_gguf_string(bytes, key);
673        bytes.extend_from_slice(&GGUF_TYPE_STRING.to_le_bytes());
674        write_gguf_string(bytes, value);
675    }
676
677    fn write_gguf_kv_u32(bytes: &mut Vec<u8>, key: &str, value: u32) {
678        write_gguf_string(bytes, key);
679        bytes.extend_from_slice(&GGUF_TYPE_UINT32.to_le_bytes());
680        bytes.extend_from_slice(&value.to_le_bytes());
681    }
682
683    fn write_gguf_kv_string_array(bytes: &mut Vec<u8>, key: &str, values: &[&str]) {
684        write_gguf_string(bytes, key);
685        bytes.extend_from_slice(&GGUF_TYPE_ARRAY.to_le_bytes());
686        bytes.extend_from_slice(&GGUF_TYPE_STRING.to_le_bytes());
687        bytes.extend_from_slice(&(values.len() as u64).to_le_bytes());
688        for value in values {
689            write_gguf_string(bytes, value);
690        }
691    }
692
693    fn write_gguf_tensor_info(
694        bytes: &mut Vec<u8>,
695        name: &str,
696        dims: &[u64],
697        tensor_type: u32,
698        offset: u64,
699    ) {
700        write_gguf_string(bytes, name);
701        bytes.extend_from_slice(&(dims.len() as u32).to_le_bytes());
702        for dim in dims {
703            bytes.extend_from_slice(&dim.to_le_bytes());
704        }
705        bytes.extend_from_slice(&tensor_type.to_le_bytes());
706        bytes.extend_from_slice(&offset.to_le_bytes());
707    }
708
709    fn write_gguf_string(bytes: &mut Vec<u8>, value: &str) {
710        bytes.extend_from_slice(&(value.len() as u64).to_le_bytes());
711        bytes.extend_from_slice(value.as_bytes());
712    }
713}
skippy_model/imatrix.rs

skippy_model/
imatrix.rs
