zsfm_checkpoint/
hdf5_keras.rs

1//! HDF5 (`.h5` / `.hdf5`) and Keras v3 (`.keras`) checkpoint loading.
2//!
3//! HDF5 files are walked recursively: every numeric dataset becomes a tensor
4//! named by its group path with `/` separators replaced by `.` (e.g.
5//! `/model_weights/dense/kernel:0` → `model_weights.dense.kernel:0`).
6//! F32 and F16 datasets are preserved exactly; F64, integer, and bool datasets
7//! are cast to F32 with a note; non-numeric datasets (strings, compounds) are
8//! skipped with a note.
9//!
10//! `.keras` archives are zip files wrapping `model.weights.h5` — the embedded
11//! weights file is extracted to a temp location and parsed the same way.
12
13use std::path::Path;
14
15use anyhow::{Context, Result};
16use hdf5::types::{FloatSize, TypeDescriptor};
17
18use crate::cast::{f32_to_bytes, SrcDtype};
19use crate::read::{Checkpoint, RawTensor};
20
21pub fn load_hdf5(path: &Path) -> Result<Checkpoint> {
22    const MAX_HDF5_BYTES: u64 = 8 << 30;
23    const MAX_DATASETS: usize = 100_000;
24    let meta = std::fs::metadata(path).with_context(|| format!("stat {}", path.display()))?;
25    anyhow::ensure!(
26        meta.len() <= MAX_HDF5_BYTES,
27        "HDF5 file too large: {} bytes (max {MAX_HDF5_BYTES})",
28        meta.len()
29    );
30    let file =
31        hdf5::File::open(path).with_context(|| format!("open HDF5 file {}", path.display()))?;
32    let mut tensors = Vec::new();
33    walk_group(&file, &mut tensors)?;
34    anyhow::ensure!(
35        tensors.len() <= MAX_DATASETS,
36        "HDF5 has too many datasets: {} (max {MAX_DATASETS})",
37        tensors.len()
38    );
39    anyhow::ensure!(
40        !tensors.is_empty(),
41        "no numeric datasets found in {}",
42        path.display()
43    );
44    Ok(Checkpoint {
45        tensors,
46        metadata: Vec::new(),
47    })
48}
49
50fn walk_group(group: &hdf5::Group, out: &mut Vec<RawTensor>) -> Result<()> {
51    let mut datasets = group.datasets().unwrap_or_default();
52    datasets.sort_by_key(|d| d.name());
53    for ds in datasets {
54        match dataset_to_raw(&ds) {
55            Ok(Some(t)) => out.push(t),
56            Ok(None) => {}
57            Err(e) => eprintln!("note: skipping dataset {}: {e:#}", ds.name()),
58        }
59    }
60    let mut subgroups = group.groups().unwrap_or_default();
61    subgroups.sort_by_key(|g| g.name());
62    for sub in subgroups {
63        walk_group(&sub, out)?;
64    }
65    Ok(())
66}
67
68fn dataset_to_raw(ds: &hdf5::Dataset) -> Result<Option<RawTensor>> {
69    let name = ds.name().trim_start_matches('/').replace('/', ".");
70    let shape: Vec<u64> = ds.shape().iter().map(|&d| d as u64).collect();
71    let descriptor = ds.dtype()?.to_descriptor()?;
72
73    let (dtype, data): (SrcDtype, Vec<u8>) = match descriptor {
74        TypeDescriptor::Float(FloatSize::U4) => {
75            (SrcDtype::F32, f32_to_bytes(&ds.read_raw::<f32>()?))
76        }
77        #[allow(unreachable_patterns)] // U2 only exists with the f16 feature
78        TypeDescriptor::Float(FloatSize::U2) => {
79            let vals = ds.read_raw::<half::f16>()?;
80            (
81                SrcDtype::F16,
82                vals.iter()
83                    .flat_map(|v| v.to_bits().to_le_bytes())
84                    .collect(),
85            )
86        }
87        TypeDescriptor::Float(FloatSize::U8) => {
88            eprintln!("note: dataset {name}: casting F64 to F32");
89            let vals: Vec<f32> = ds.read_raw::<f64>()?.iter().map(|&v| v as f32).collect();
90            (SrcDtype::F32, f32_to_bytes(&vals))
91        }
92        TypeDescriptor::Integer(size) => {
93            eprintln!("note: dataset {name}: casting signed integer to F32");
94            let vals: Vec<f32> = match size {
95                hdf5::types::IntSize::U1 => {
96                    ds.read_raw::<i8>()?.iter().map(|&v| v as f32).collect()
97                }
98                hdf5::types::IntSize::U2 => {
99                    ds.read_raw::<i16>()?.iter().map(|&v| v as f32).collect()
100                }
101                hdf5::types::IntSize::U4 => {
102                    ds.read_raw::<i32>()?.iter().map(|&v| v as f32).collect()
103                }
104                hdf5::types::IntSize::U8 => {
105                    ds.read_raw::<i64>()?.iter().map(|&v| v as f32).collect()
106                }
107            };
108            (SrcDtype::F32, f32_to_bytes(&vals))
109        }
110        TypeDescriptor::Unsigned(size) => {
111            eprintln!("note: dataset {name}: casting unsigned integer to F32");
112            let vals: Vec<f32> = match size {
113                hdf5::types::IntSize::U1 => {
114                    ds.read_raw::<u8>()?.iter().map(|&v| v as f32).collect()
115                }
116                hdf5::types::IntSize::U2 => {
117                    ds.read_raw::<u16>()?.iter().map(|&v| v as f32).collect()
118                }
119                hdf5::types::IntSize::U4 => {
120                    ds.read_raw::<u32>()?.iter().map(|&v| v as f32).collect()
121                }
122                hdf5::types::IntSize::U8 => {
123                    ds.read_raw::<u64>()?.iter().map(|&v| v as f32).collect()
124                }
125            };
126            (SrcDtype::F32, f32_to_bytes(&vals))
127        }
128        TypeDescriptor::Boolean => {
129            eprintln!("note: dataset {name}: casting bool to F32");
130            let vals: Vec<f32> = ds
131                .read_raw::<bool>()?
132                .iter()
133                .map(|&v| v as u8 as f32)
134                .collect();
135            (SrcDtype::F32, f32_to_bytes(&vals))
136        }
137        other => {
138            eprintln!("note: skipping non-numeric dataset {name} ({other})");
139            return Ok(None);
140        }
141    };
142
143    Ok(Some(RawTensor {
144        name,
145        shape,
146        dtype,
147        data,
148    }))
149}
150
151// ---------------------------------------------------------------------------
152// Keras v3 (.keras = zip wrapping model.weights.h5)
153// ---------------------------------------------------------------------------
154
155pub fn load_keras(path: &Path) -> Result<Checkpoint> {
156    let file = std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?;
157    let mut archive = zip::ZipArchive::new(file).context("read .keras zip archive")?;
158
159    let weights_entry = archive
160        .file_names()
161        .find(|n| n.ends_with(".weights.h5"))
162        .or_else(|| archive.file_names().find(|n| n.ends_with(".h5")))
163        .map(String::from)
164        .context(".keras archive contains no .weights.h5 entry")?;
165
166    // libhdf5 needs a real file path, so extract the weights to a temp file.
167    let tmp = std::env::temp_dir().join(format!(
168        "zsfm-keras-{}-{}.h5",
169        std::process::id(),
170        std::time::SystemTime::now()
171            .duration_since(std::time::UNIX_EPOCH)
172            .unwrap_or_default()
173            .as_nanos()
174    ));
175    {
176        let mut entry = archive.by_name(&weights_entry)?;
177        let mut out = std::fs::File::create(&tmp)
178            .with_context(|| format!("create temp file {}", tmp.display()))?;
179        std::io::copy(&mut entry, &mut out).context("extract weights from .keras archive")?;
180    }
181    let result = load_hdf5(&tmp);
182    let _ = std::fs::remove_file(&tmp);
183    result.with_context(|| format!("parse {weights_entry} from .keras archive"))
184}
zsfm_checkpoint/hdf5_keras.rs

zsfm_checkpoint/
hdf5_keras.rs
