use std::io::Write;
use std::path::PathBuf;
fn f32_to_bytes(values: &[f32]) -> Vec<u8> {
values.iter().flat_map(|f| f.to_le_bytes()).collect()
}
fn create_safetensors(path: &std::path::Path, tensors: Vec<(&str, &str, Vec<usize>, Vec<u8>)>) {
let mut sorted: Vec<_> = tensors.into_iter().collect();
sorted.sort_by(|a, b| a.0.cmp(b.0));
let mut header = serde_json::Map::new();
header.insert(
"__metadata__".to_string(),
serde_json::Value::Object(serde_json::Map::new()),
);
let mut offset = 0usize;
let mut entries = Vec::new();
for (name, dtype, shape, data) in sorted {
let end = offset + data.len();
entries.push((
name.to_string(),
dtype.to_string(),
shape,
data,
offset,
end,
));
offset = end;
}
for (name, dtype, shape, _, start, end) in &entries {
let mut obj = serde_json::Map::new();
obj.insert(
"dtype".to_string(),
serde_json::Value::String(dtype.clone()),
);
obj.insert("shape".to_string(), serde_json::json!(shape));
obj.insert(
"data_offsets".to_string(),
serde_json::json!([*start, *end]),
);
header.insert(name.clone(), serde_json::Value::Object(obj));
}
let header_json = serde_json::to_string(&header).unwrap();
let header_bytes = header_json.as_bytes();
let mut file = std::fs::File::create(path).unwrap();
file.write_all(&(header_bytes.len() as u64).to_le_bytes())
.unwrap();
file.write_all(header_bytes).unwrap();
for (_, _, _, data, _, _) in &entries {
file.write_all(data).unwrap();
}
}
fn main() {
let args: Vec<String> = std::env::args().collect();
let output_path = if args.len() > 1 {
PathBuf::from(&args[1])
} else {
PathBuf::from("simple_model.safetensors")
};
let hidden = 8;
let layer1_weight = f32_to_bytes(
&(0..hidden * hidden)
.map(|i| (i as f32) / (hidden * hidden) as f32)
.collect::<Vec<_>>(),
);
let layer1_bias = f32_to_bytes(&vec![0.1; hidden]);
let layer2_weight = f32_to_bytes(
&(0..hidden)
.map(|i| (i as f32) / hidden as f32)
.collect::<Vec<_>>(),
);
let layer2_bias = f32_to_bytes(&[0.0]);
create_safetensors(
&output_path,
vec![
("layer1.bias", "F32", vec![hidden], layer1_bias),
(
"layer1.weight",
"F32",
vec![hidden, hidden],
layer1_weight.clone(),
),
("layer2.bias", "F32", vec![1], layer2_bias),
("layer2.weight", "F32", vec![1, hidden], layer2_weight),
],
);
let param_count = hidden * hidden + hidden + hidden + 1;
let bytes_approx = (hidden * hidden + hidden + hidden + 1) * std::mem::size_of::<f32>();
eprintln!(
"Wrote {} (~{} float params, ~{:.1} KiB of F32 weights). Try:\n modelc inspect {:?}\n modelc compile {:?} -o ./simple_mlp_serve",
output_path.display(),
param_count,
bytes_approx as f64 / 1024.0,
output_path,
output_path,
);
}