zsfm_nn/
gguf.rs

1use std::io::{Read, Seek};
2
3use anyhow::{Context, Result};
4use candle_core::quantized::gguf_file;
5use candle_core::{DType, Device, Tensor};
6
7/// Dequantize a named GGUF tensor and cast it to `dtype`.
8pub fn load_tensor(
9    content: &gguf_file::Content,
10    reader: &mut (impl Read + Seek),
11    name: &str,
12    device: &Device,
13    dtype: DType,
14) -> Result<Tensor> {
15    let qt = content
16        .tensor(reader, name, device)
17        .with_context(|| format!("tensor '{name}' not found in GGUF"))?;
18    Ok(qt.dequantize(device)?.to_dtype(dtype)?)
19}
20
21/// Same as [`load_tensor`], but returns `Ok(None)` instead of erroring when the tensor is absent.
22pub fn try_load_tensor(
23    content: &gguf_file::Content,
24    reader: &mut (impl Read + Seek),
25    name: &str,
26    device: &Device,
27    dtype: DType,
28) -> Result<Option<Tensor>> {
29    match content.tensor(reader, name, device) {
30        Ok(qt) => Ok(Some(qt.dequantize(device)?.to_dtype(dtype)?)),
31        Err(_) => Ok(None),
32    }
33}
34
35/// Load an F32 weight PyTorch stores as `(d_out, d_in)`. Candle reverses the GGUF shape back to
36/// `(d_out, d_in)`; the `dim(0)` check guards against a stray transposed store (the Q8_0
37/// transpose trick some converters apply) by flipping the tensor back.
38pub fn load_weight(
39    content: &gguf_file::Content,
40    reader: &mut (impl Read + Seek),
41    name: &str,
42    expected_d_out: usize,
43    device: &Device,
44) -> Result<Tensor> {
45    let w = load_tensor(content, reader, name, device, DType::F32)?;
46    if w.dim(0)? != expected_d_out {
47        Ok(w.t()?.contiguous()?)
48    } else {
49        Ok(w)
50    }
51}
52
53/// Dequantize a named GGUF tensor to F32 and flatten it to a plain `Vec<f32>`.
54pub fn load_vec(
55    content: &gguf_file::Content,
56    reader: &mut (impl Read + Seek),
57    name: &str,
58    device: &Device,
59) -> Result<Vec<f32>> {
60    Ok(load_tensor(content, reader, name, device, DType::F32)?
61        .flatten_all()?
62        .to_vec1()?)
63}
zsfm_nn/gguf.rs

zsfm_nn/
gguf.rs
