teeny_quant/
read.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! Reading source `.safetensors` checkpoints. Opening/mmapping is delegated to
18//! [`teeny_data::safetensors::SafeTensors`]; this module only adds the "upcast whatever float
19//! dtype is on disk to `f32` for quantization math" step.
20
21use std::collections::HashMap;
22use std::io::Read as _;
23use std::path::Path;
24
25use half::{bf16, f16};
26use safetensors::Dtype;
27use safetensors::tensor::TensorView;
28
29use crate::error::{Error, Result};
30
31/// Reads `view`'s raw bytes as `f32`, upcasting from `F32`/`F16`/`BF16` as needed. Integer or
32/// boolean source tensors (already-quantized inputs, masks, etc.) aren't supported -- callers
33/// should pass those through unquantized rather than routing them through this function.
34pub fn read_f32(view: &TensorView<'_>, name: &str) -> Result<Vec<f32>> {
35    match view.dtype() {
36        Dtype::F32 => Ok(view
37            .data()
38            .chunks_exact(4)
39            .map(|b| f32::from_le_bytes(b.try_into().expect("chunks_exact(4)")))
40            .collect()),
41        Dtype::F16 => Ok(view
42            .data()
43            .chunks_exact(2)
44            .map(|b| f16::from_le_bytes(b.try_into().expect("chunks_exact(2)")).to_f32())
45            .collect()),
46        Dtype::BF16 => Ok(view
47            .data()
48            .chunks_exact(2)
49            .map(|b| bf16::from_le_bytes(b.try_into().expect("chunks_exact(2)")).to_f32())
50            .collect()),
51        other => Err(Error::UnsupportedDtype {
52            tensor: name.to_string(),
53            dtype: other,
54        }),
55    }
56}
57
58/// Whether `dtype` is one [`read_f32`] can upcast from.
59pub fn is_quantizable_float(dtype: Dtype) -> bool {
60    matches!(dtype, Dtype::F32 | Dtype::F16 | Dtype::BF16)
61}
62
63/// Reads a `.safetensors` file's `__metadata__` string map (e.g. the embedded
64/// `quantization_config` -- see [`crate::format`]) directly from `path`.
65///
66/// `safetensors::SafeTensors` (the type `teeny_data::safetensors::SafeTensors::tensors` hands
67/// back) doesn't expose this map publicly -- only the crate-internal `Metadata` header type
68/// does, via `SafeTensors::read_metadata`. That function insists the buffer's total length
69/// exactly match the header-declared data length, but doesn't care about the data *bytes*
70/// themselves for metadata extraction -- so this reads the real header off disk but only
71/// zero-pads out to the real file length for the (unread) tensor payload, avoiding an I/O read
72/// of the potentially multi-GB data section just to reach a handful of metadata strings.
73pub fn read_metadata(path: &Path) -> Result<HashMap<String, String>> {
74    let file_len = std::fs::metadata(path)?.len() as usize;
75    let mut file = std::fs::File::open(path)?;
76
77    let mut len_bytes = [0u8; 8];
78    file.read_exact(&mut len_bytes)?;
79    let header_len = u64::from_le_bytes(len_bytes) as usize;
80
81    let mut buffer = vec![0u8; file_len];
82    buffer[..8].copy_from_slice(&len_bytes);
83    file.read_exact(&mut buffer[8..8 + header_len])?;
84
85    let (_, metadata) = safetensors::SafeTensors::read_metadata(&buffer)?;
86    Ok(metadata.metadata().clone().unwrap_or_default())
87}
teeny_quant/read.rs

teeny_quant/
read.rs
