forge/serialization/
mod.rs

1//! SafeTensors saving: named tensors -> file.
2//!
3//! There is deliberately no generic loader here. Model loading goes through
4//! [`crate::models::gpt2::Gpt2::from_safetensors_bytes`], which deserializes
5//! straight into the model's own parameter layout rather than materializing an
6//! intermediate `HashMap<String, Tensor>` of every tensor in the file.
7
8use std::path::Path;
9
10use safetensors::tensor::{Dtype, TensorView};
11
12use crate::error::{ForgeError, Result};
13
14/// Save named f32 tensors to a .safetensors file (checkpoints).
15/// Entries: (name, shape, host data).
16pub fn save_safetensors(
17    path: impl AsRef<Path>,
18    entries: &[(String, Vec<usize>, Vec<f32>)],
19) -> Result<()> {
20    let views: Vec<(String, TensorView<'_>)> = entries
21        .iter()
22        .map(|(name, shape, data)| {
23            TensorView::new(Dtype::F32, shape.clone(), bytemuck::cast_slice(data))
24                .map(|v| (name.clone(), v))
25                .map_err(|e| ForgeError::SafeTensors(format!("{name}: {e}")))
26        })
27        .collect::<Result<_>>()?;
28    safetensors::serialize_to_file(views, &None, path.as_ref())
29        .map_err(|e| ForgeError::SafeTensors(format!("{}: {e}", path.as_ref().display())))
30}
forge/serialization/mod.rs

forge/serialization/
mod.rs
