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    let file = hdf5::File::open(path)
23        .with_context(|| format!("open HDF5 file {}", path.display()))?;
24    let mut tensors = Vec::new();
25    walk_group(&file, &mut tensors)?;
26    anyhow::ensure!(
27        !tensors.is_empty(),
28        "no numeric datasets found in {}",
29        path.display()
30    );
31    Ok(Checkpoint { tensors, metadata: Vec::new() })
32}
33
34fn walk_group(group: &hdf5::Group, out: &mut Vec<RawTensor>) -> Result<()> {
35    let mut datasets = group.datasets().unwrap_or_default();
36    datasets.sort_by_key(|d| d.name());
37    for ds in datasets {
38        match dataset_to_raw(&ds) {
39            Ok(Some(t)) => out.push(t),
40            Ok(None) => {}
41            Err(e) => eprintln!("note: skipping dataset {}: {e:#}", ds.name()),
42        }
43    }
44    let mut subgroups = group.groups().unwrap_or_default();
45    subgroups.sort_by_key(|g| g.name());
46    for sub in subgroups {
47        walk_group(&sub, out)?;
48    }
49    Ok(())
50}
51
52fn dataset_to_raw(ds: &hdf5::Dataset) -> Result<Option<RawTensor>> {
53    let name = ds.name().trim_start_matches('/').replace('/', ".");
54    let shape: Vec<u64> = ds.shape().iter().map(|&d| d as u64).collect();
55    let descriptor = ds.dtype()?.to_descriptor()?;
56
57    let (dtype, data): (SrcDtype, Vec<u8>) = match descriptor {
58        TypeDescriptor::Float(FloatSize::U4) => {
59            (SrcDtype::F32, f32_to_bytes(&ds.read_raw::<f32>()?))
60        }
61        #[allow(unreachable_patterns)] // U2 only exists with the f16 feature
62        TypeDescriptor::Float(FloatSize::U2) => {
63            let vals = ds.read_raw::<half::f16>()?;
64            (
65                SrcDtype::F16,
66                vals.iter().flat_map(|v| v.to_bits().to_le_bytes()).collect(),
67            )
68        }
69        TypeDescriptor::Float(FloatSize::U8) => {
70            eprintln!("note: dataset {name}: casting F64 to F32");
71            let vals: Vec<f32> = ds.read_raw::<f64>()?.iter().map(|&v| v as f32).collect();
72            (SrcDtype::F32, f32_to_bytes(&vals))
73        }
74        TypeDescriptor::Integer(size) => {
75            eprintln!("note: dataset {name}: casting signed integer to F32");
76            let vals: Vec<f32> = match size {
77                hdf5::types::IntSize::U1 => {
78                    ds.read_raw::<i8>()?.iter().map(|&v| v as f32).collect()
79                }
80                hdf5::types::IntSize::U2 => {
81                    ds.read_raw::<i16>()?.iter().map(|&v| v as f32).collect()
82                }
83                hdf5::types::IntSize::U4 => {
84                    ds.read_raw::<i32>()?.iter().map(|&v| v as f32).collect()
85                }
86                hdf5::types::IntSize::U8 => {
87                    ds.read_raw::<i64>()?.iter().map(|&v| v as f32).collect()
88                }
89            };
90            (SrcDtype::F32, f32_to_bytes(&vals))
91        }
92        TypeDescriptor::Unsigned(size) => {
93            eprintln!("note: dataset {name}: casting unsigned integer to F32");
94            let vals: Vec<f32> = match size {
95                hdf5::types::IntSize::U1 => {
96                    ds.read_raw::<u8>()?.iter().map(|&v| v as f32).collect()
97                }
98                hdf5::types::IntSize::U2 => {
99                    ds.read_raw::<u16>()?.iter().map(|&v| v as f32).collect()
100                }
101                hdf5::types::IntSize::U4 => {
102                    ds.read_raw::<u32>()?.iter().map(|&v| v as f32).collect()
103                }
104                hdf5::types::IntSize::U8 => {
105                    ds.read_raw::<u64>()?.iter().map(|&v| v as f32).collect()
106                }
107            };
108            (SrcDtype::F32, f32_to_bytes(&vals))
109        }
110        TypeDescriptor::Boolean => {
111            eprintln!("note: dataset {name}: casting bool to F32");
112            let vals: Vec<f32> = ds
113                .read_raw::<bool>()?
114                .iter()
115                .map(|&v| v as u8 as f32)
116                .collect();
117            (SrcDtype::F32, f32_to_bytes(&vals))
118        }
119        other => {
120            eprintln!("note: skipping non-numeric dataset {name} ({other})");
121            return Ok(None);
122        }
123    };
124
125    Ok(Some(RawTensor { name, shape, dtype, data }))
126}
127
128// ---------------------------------------------------------------------------
129// Keras v3 (.keras = zip wrapping model.weights.h5)
130// ---------------------------------------------------------------------------
131
132pub fn load_keras(path: &Path) -> Result<Checkpoint> {
133    let file = std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?;
134    let mut archive = zip::ZipArchive::new(file).context("read .keras zip archive")?;
135
136    let weights_entry = archive
137        .file_names()
138        .find(|n| n.ends_with(".weights.h5"))
139        .or_else(|| archive.file_names().find(|n| n.ends_with(".h5")))
140        .map(String::from)
141        .context(".keras archive contains no .weights.h5 entry")?;
142
143    // libhdf5 needs a real file path, so extract the weights to a temp file.
144    let tmp = std::env::temp_dir().join(format!(
145        "zsfm-keras-{}-{}.h5",
146        std::process::id(),
147        std::time::SystemTime::now()
148            .duration_since(std::time::UNIX_EPOCH)
149            .unwrap_or_default()
150            .as_nanos()
151    ));
152    {
153        let mut entry = archive.by_name(&weights_entry)?;
154        let mut out = std::fs::File::create(&tmp)
155            .with_context(|| format!("create temp file {}", tmp.display()))?;
156        std::io::copy(&mut entry, &mut out).context("extract weights from .keras archive")?;
157    }
158    let result = load_hdf5(&tmp);
159    let _ = std::fs::remove_file(&tmp);
160    result.with_context(|| format!("parse {weights_entry} from .keras archive"))
161}
zsfm_checkpoint/hdf5_keras.rs

zsfm_checkpoint/
hdf5_keras.rs
