pub use crate::std::layers_std::LayerSpec;
const RMD1_HEADER_SIZE: usize = 20;
const LAYER_META_SIZE: usize = 20;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelFormatStdError {
Truncated,
BadMagic,
BadVersion,
DtypeMismatch,
BadHeader,
CapacityTooSmall,
InvalidActivation,
InvalidLayer,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DecodedCountsStd {
pub layers: usize,
pub weights: usize,
pub biases: usize,
}
pub type DecodedCounts = DecodedCountsStd;
pub fn encoded_size(layer_count: usize, weights_len: usize, biases_len: usize) -> Option<usize> {
encoded_size_v1(layer_count, weights_len, biases_len)
}
pub fn encoded_size_v1(layer_count: usize, weights_len: usize, biases_len: usize) -> Option<usize> {
let layers_bytes = layer_count.checked_mul(LAYER_META_SIZE)?;
let weights_bytes = weights_len.checked_mul(core::mem::size_of::<f32>())?;
let biases_bytes = biases_len.checked_mul(core::mem::size_of::<f32>())?;
RMD1_HEADER_SIZE
.checked_add(layers_bytes)?
.checked_add(weights_bytes)?
.checked_add(biases_bytes)
}
pub fn encode_model(
layers: &[LayerSpec],
weights: &[f32],
biases: &[f32],
out: &mut [u8],
) -> Result<usize, ModelFormatStdError> {
encode_model_v1(layers, weights, biases, out)
}
pub fn encode_model_v1(
layers: &[LayerSpec],
weights: &[f32],
biases: &[f32],
out: &mut [u8],
) -> Result<usize, ModelFormatStdError> {
let needed = encoded_size_v1(layers.len(), weights.len(), biases.len())
.ok_or(ModelFormatStdError::BadHeader)?;
if out.len() < needed {
return Err(ModelFormatStdError::CapacityTooSmall);
}
write_header(out, layers.len(), weights.len(), biases.len())?;
let mut cursor = RMD1_HEADER_SIZE;
for layer in layers {
let LayerSpec::Dense(dense) = *layer;
validate_layer(dense.into(), weights.len(), biases.len())?;
write_u32(out, &mut cursor, dense.input_size)?;
write_u32(out, &mut cursor, dense.output_size)?;
write_u32(out, &mut cursor, dense.weight_offset)?;
write_u32(out, &mut cursor, dense.bias_offset)?;
out[cursor] = dense.activation.to_u8();
cursor += 1;
out[cursor..cursor + 3].copy_from_slice(&[0u8; 3]);
cursor += 3;
}
for &w in weights {
out[cursor..cursor + 4].copy_from_slice(&w.to_le_bytes());
cursor += 4;
}
for &b in biases {
out[cursor..cursor + 4].copy_from_slice(&b.to_le_bytes());
cursor += 4;
}
Ok(cursor)
}
pub fn decode_model(
bytes: &[u8],
layers_out: &mut [LayerSpec],
weights_out: &mut [f32],
biases_out: &mut [f32],
) -> Result<usize, ModelFormatStdError> {
let decoded = decode_model_v1(bytes, layers_out, weights_out, biases_out)?;
Ok(decoded.layers)
}
pub fn decode_model_v1(
bytes: &[u8],
layers_out: &mut [LayerSpec],
weights_out: &mut [f32],
biases_out: &mut [f32],
) -> Result<DecodedCountsStd, ModelFormatStdError> {
let (layer_count, weights_len, biases_len) = parse_header(bytes)?;
if layers_out.len() < layer_count
|| weights_out.len() < weights_len
|| biases_out.len() < biases_len
{
return Err(ModelFormatStdError::CapacityTooSmall);
}
let expected = encoded_size_v1(layer_count, weights_len, biases_len)
.ok_or(ModelFormatStdError::BadHeader)?;
if bytes.len() < expected {
return Err(ModelFormatStdError::Truncated);
}
let mut cursor = RMD1_HEADER_SIZE;
for slot in layers_out.iter_mut().take(layer_count) {
let input_size = read_u32(bytes, &mut cursor)? as usize;
let output_size = read_u32(bytes, &mut cursor)? as usize;
let weight_offset = read_u32(bytes, &mut cursor)? as usize;
let bias_offset = read_u32(bytes, &mut cursor)? as usize;
let activation_u8 = *bytes.get(cursor).ok_or(ModelFormatStdError::Truncated)?;
cursor += 1;
cursor = cursor
.checked_add(3)
.ok_or(ModelFormatStdError::BadHeader)?;
let activation = native_neural_network::activations::ActivationKind::from_u8(activation_u8)
.ok_or(ModelFormatStdError::InvalidActivation)?;
let dense = native_neural_network::layers::DenseLayerDesc {
input_size,
output_size,
weight_offset,
bias_offset,
activation,
};
validate_layer(dense, weights_len, biases_len)?;
*slot = LayerSpec::Dense(dense.into());
}
for w in &mut weights_out[..weights_len] {
let chunk = bytes
.get(cursor..cursor + 4)
.ok_or(ModelFormatStdError::Truncated)?;
*w = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
cursor += 4;
}
for b in &mut biases_out[..biases_len] {
let chunk = bytes
.get(cursor..cursor + 4)
.ok_or(ModelFormatStdError::Truncated)?;
*b = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
cursor += 4;
}
Ok(DecodedCountsStd {
layers: layer_count,
weights: weights_len,
biases: biases_len,
})
}
fn write_header(
out: &mut [u8],
layer_count: usize,
weights_len: usize,
biases_len: usize,
) -> Result<(), ModelFormatStdError> {
if out.len() < RMD1_HEADER_SIZE {
return Err(ModelFormatStdError::CapacityTooSmall);
}
out[0..4].copy_from_slice(b"RMD1");
out[4..6].copy_from_slice(&1u16.to_le_bytes());
out[6] = 0u8;
out[7] = 0u8;
out[8..12].copy_from_slice(
&(u32::try_from(layer_count).map_err(|_| ModelFormatStdError::BadHeader)?).to_le_bytes(),
);
out[12..16].copy_from_slice(
&(u32::try_from(weights_len).map_err(|_| ModelFormatStdError::BadHeader)?).to_le_bytes(),
);
out[16..20].copy_from_slice(
&(u32::try_from(biases_len).map_err(|_| ModelFormatStdError::BadHeader)?).to_le_bytes(),
);
Ok(())
}
fn parse_header(bytes: &[u8]) -> Result<(usize, usize, usize), ModelFormatStdError> {
if bytes.len() < RMD1_HEADER_SIZE {
return Err(ModelFormatStdError::Truncated);
}
if &bytes[0..4] != b"RMD1" {
return Err(ModelFormatStdError::BadMagic);
}
let version = u16::from_le_bytes([bytes[4], bytes[5]]);
if version != 1 {
return Err(ModelFormatStdError::BadVersion);
}
if bytes[6] != 0 {
return Err(ModelFormatStdError::DtypeMismatch);
}
let layer_count = u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize;
let weights_len = u32::from_le_bytes([bytes[12], bytes[13], bytes[14], bytes[15]]) as usize;
let biases_len = u32::from_le_bytes([bytes[16], bytes[17], bytes[18], bytes[19]]) as usize;
if layer_count == 0 {
return Err(ModelFormatStdError::BadHeader);
}
Ok((layer_count, weights_len, biases_len))
}
fn write_u32(out: &mut [u8], cursor: &mut usize, value: usize) -> Result<(), ModelFormatStdError> {
let value = u32::try_from(value).map_err(|_| ModelFormatStdError::BadHeader)?;
out[*cursor..*cursor + 4].copy_from_slice(&value.to_le_bytes());
*cursor += 4;
Ok(())
}
fn read_u32(bytes: &[u8], cursor: &mut usize) -> Result<u32, ModelFormatStdError> {
let chunk = bytes
.get(*cursor..*cursor + 4)
.ok_or(ModelFormatStdError::Truncated)?;
*cursor += 4;
Ok(u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
}
fn validate_layer(
layer: native_neural_network::layers::DenseLayerDesc,
weights_len: usize,
biases_len: usize,
) -> Result<(), ModelFormatStdError> {
if layer.input_size == 0 || layer.output_size == 0 {
return Err(ModelFormatStdError::InvalidLayer);
}
let w_len = layer
.input_size
.checked_mul(layer.output_size)
.ok_or(ModelFormatStdError::InvalidLayer)?;
let w_end = layer
.weight_offset
.checked_add(w_len)
.ok_or(ModelFormatStdError::InvalidLayer)?;
let b_end = layer
.bias_offset
.checked_add(layer.output_size)
.ok_or(ModelFormatStdError::InvalidLayer)?;
if w_end > weights_len || b_end > biases_len {
return Err(ModelFormatStdError::InvalidLayer);
}
Ok(())
}
impl core::fmt::Display for ModelFormatStdError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "ModelFormatStdError::{:?}", self)
}
}
impl std::error::Error for ModelFormatStdError {}