use crate::quant::{bf16_to_f32, f16_to_f32};
use crate::session::CeraError;
use serde::Deserialize;
use std::collections::BTreeMap;
#[derive(Debug, Clone, Deserialize)]
pub struct SafeTensorInfo {
pub dtype: String,
pub shape: Vec<usize>,
pub data_offsets: (usize, usize),
}
#[derive(Debug, Clone)]
pub struct SafeTensorsHeader {
pub header_size_bytes: usize,
pub tensors: BTreeMap<String, SafeTensorInfo>,
}
impl SafeTensorsHeader {
pub fn parse_from_bytes(bytes: &[u8]) -> Result<Self, CeraError> {
if bytes.len() < 8 {
return Err(CeraError::Backend(
"invalid safetensors: buffer smaller than 8 bytes".into(),
));
}
let len_bytes: [u8; 8] = bytes[0..8].try_into().map_err(|_| {
CeraError::Backend("failed to read 8-byte safetensors header length".into())
})?;
let header_len_u64 = u64::from_le_bytes(len_bytes);
if header_len_u64 == 0 || header_len_u64 > 100_000_000 {
return Err(CeraError::Backend(format!(
"invalid safetensors header length: {header_len_u64}"
)));
}
let header_len = header_len_u64 as usize;
let total_header = match 8usize.checked_add(header_len) {
Some(t) => t,
None => {
return Err(CeraError::Backend(
"safetensors header length overflows usize".into(),
));
}
};
if bytes.len() < total_header {
return Err(CeraError::Backend(format!(
"invalid safetensors: buffer size ({}) smaller than 8 + header size ({header_len})",
bytes.len()
)));
}
let json_str = std::str::from_utf8(&bytes[8..total_header])
.map_err(|e| CeraError::Backend(format!("invalid utf8 in safetensors header: {e}")))?;
Self::parse_from_json_str(json_str, total_header)
}
pub fn parse_from_reader<R: std::io::Read>(reader: &mut R) -> Result<Self, CeraError> {
let mut len_bytes = [0u8; 8];
reader.read_exact(&mut len_bytes).map_err(|e| {
CeraError::Backend(format!(
"failed to read 8-byte safetensors header length: {e}"
))
})?;
let header_len_u64 = u64::from_le_bytes(len_bytes);
if header_len_u64 == 0 || header_len_u64 > 100_000_000 {
return Err(CeraError::Backend(format!(
"invalid safetensors header length: {header_len_u64}"
)));
}
let header_len = header_len_u64 as usize;
let mut json_bytes = vec![0u8; header_len];
reader.read_exact(&mut json_bytes).map_err(|e| {
CeraError::Backend(format!(
"failed to read {header_len} bytes of safetensors header: {e}"
))
})?;
let json_str = std::str::from_utf8(&json_bytes)
.map_err(|e| CeraError::Backend(format!("invalid utf8 in safetensors header: {e}")))?;
let total_header = match 8usize.checked_add(header_len) {
Some(t) => t,
None => {
return Err(CeraError::Backend(
"safetensors header length overflows usize".into(),
));
}
};
Self::parse_from_json_str(json_str, total_header)
}
pub fn parse_from_json_str(json_str: &str, total_header: usize) -> Result<Self, CeraError> {
let raw_map: BTreeMap<String, serde_json::Value> =
serde_json::from_str(json_str).map_err(|e| {
CeraError::Backend(format!("failed to parse safetensors json header: {e}"))
})?;
let mut tensors = BTreeMap::new();
for (name, val) in raw_map {
if name == "__metadata__" {
continue;
}
let info: SafeTensorInfo = serde_json::from_value(val).map_err(|e| {
CeraError::Backend(format!("failed to parse tensor `{name}` metadata: {e}"))
})?;
tensors.insert(name, info);
}
Ok(Self {
header_size_bytes: total_header,
tensors,
})
}
}
fn map_whisper_block_sub(sub: &str) -> &str {
match sub {
"self_attn.q_proj.weight" => "attn.query.weight",
"self_attn.q_proj.bias" => "attn.query.bias",
"self_attn.k_proj.weight" => "attn.key.weight",
"self_attn.k_proj.bias" => "attn.key.bias",
"self_attn.v_proj.weight" => "attn.value.weight",
"self_attn.v_proj.bias" => "attn.value.bias",
"self_attn.out_proj.weight" => "attn.out.weight",
"self_attn.out_proj.bias" => "attn.out.bias",
"self_attn_layer_norm.weight" => "attn_ln.weight",
"self_attn_layer_norm.bias" => "attn_ln.bias",
"encoder_attn.q_proj.weight" => "cross_attn.query.weight",
"encoder_attn.q_proj.bias" => "cross_attn.query.bias",
"encoder_attn.k_proj.weight" => "cross_attn.key.weight",
"encoder_attn.k_proj.bias" => "cross_attn.key.bias",
"encoder_attn.v_proj.weight" => "cross_attn.value.weight",
"encoder_attn.v_proj.bias" => "cross_attn.value.bias",
"encoder_attn.out_proj.weight" => "cross_attn.out.weight",
"encoder_attn.out_proj.bias" => "cross_attn.out.bias",
"encoder_attn_layer_norm.weight" => "cross_attn_ln.weight",
"encoder_attn_layer_norm.bias" => "cross_attn_ln.bias",
"fc1.weight" => "mlp.0.weight",
"fc1.bias" => "mlp.0.bias",
"fc2.weight" => "mlp.2.weight",
"fc2.bias" => "mlp.2.bias",
"final_layer_norm.weight" => "mlp_ln.weight",
"final_layer_norm.bias" => "mlp_ln.bias",
_ => sub,
}
}
pub fn translate_hf_to_gguf_tensor_name(hf_name: &str) -> String {
if hf_name == "model.embed_tokens.weight"
|| hf_name == "lfm2.embed_tokens.weight"
|| hf_name == "transformer.wte.weight"
|| hf_name == "embeddings.word_embeddings.weight"
{
return "token_embd.weight".to_string();
}
if hf_name == "lfm2.embedding_norm.weight" {
return "token_embd_norm.weight".to_string();
}
if hf_name == "model.norm.weight"
|| hf_name == "lfm2.norm.weight"
|| hf_name == "transformer.ln_f.weight"
|| hf_name == "ln_f.weight"
{
return "output_norm.weight".to_string();
}
if hf_name == "lm_head.weight" || hf_name == "proj_out.weight" {
return "output.weight".to_string();
}
if let Some(rest) = hf_name
.strip_prefix("model.encoder.")
.or_else(|| hf_name.strip_prefix("encoder."))
{
match rest {
"conv1.weight" => return "encoder.conv1.weight".to_string(),
"conv1.bias" => return "encoder.conv1.bias".to_string(),
"conv2.weight" => return "encoder.conv2.weight".to_string(),
"conv2.bias" => return "encoder.conv2.bias".to_string(),
"embed_positions.weight" | "positional_embedding" => {
return "encoder.positional_embedding".to_string();
}
"layer_norm.weight" | "ln_post.weight" => {
return "encoder.ln_post.weight".to_string();
}
"layer_norm.bias" | "ln_post.bias" => {
return "encoder.ln_post.bias".to_string();
}
_ => {}
}
if let Some(stripped) = rest.strip_prefix("layers.")
&& let Some((idx, sub)) = stripped.split_once('.')
{
let suffix = map_whisper_block_sub(sub);
return format!("encoder.blocks.{idx}.{suffix}");
}
}
if let Some(rest) = hf_name
.strip_prefix("model.decoder.")
.or_else(|| hf_name.strip_prefix("decoder."))
{
match rest {
"embed_tokens.weight" | "token_embeddings.weight" => {
return "decoder.token_embeddings.weight".to_string();
}
"embed_positions.weight" | "positional_embedding" => {
return "decoder.positional_embedding".to_string();
}
"layer_norm.weight" | "ln.weight" | "ln_post.weight" => {
return "decoder.ln_post.weight".to_string();
}
"layer_norm.bias" | "ln.bias" | "ln_post.bias" => {
return "decoder.ln_post.bias".to_string();
}
_ => {}
}
if let Some(stripped) = rest.strip_prefix("layers.")
&& let Some((idx, sub)) = stripped.split_once('.')
{
let suffix = map_whisper_block_sub(sub);
return format!("decoder.blocks.{idx}.{suffix}");
}
}
let layer_rest = hf_name
.strip_prefix("model.layers.")
.or_else(|| hf_name.strip_prefix("lfm2.layers."))
.or_else(|| hf_name.strip_prefix("transformer.h."))
.or_else(|| hf_name.strip_prefix("layers."));
if let Some((layer_idx, sub_name)) = layer_rest.and_then(|r| r.split_once('.')) {
let gguf_suffix = match sub_name {
"self_attn.q_proj.weight" => "attn_q.weight",
"self_attn.q_proj.bias" => "attn_q.bias",
"self_attn.k_proj.weight" => "attn_k.weight",
"self_attn.k_proj.bias" => "attn_k.bias",
"self_attn.v_proj.weight" => "attn_v.weight",
"self_attn.v_proj.bias" => "attn_v.bias",
"self_attn.o_proj.weight" | "self_attn.out_proj.weight" => "attn_output.weight",
"self_attn.o_proj.bias" | "self_attn.out_proj.bias" => "attn_output.bias",
"self_attn.qkv_proj.weight" => "attn_qkv.weight",
"self_attn.qkv_proj.bias" => "attn_qkv.bias",
"self_attn.q_norm.weight" | "self_attn.q_layernorm.weight" => "attn_q_norm.weight",
"self_attn.q_norm.bias" | "self_attn.q_layernorm.bias" => "attn_q_norm.bias",
"self_attn.k_norm.weight" | "self_attn.k_layernorm.weight" => "attn_k_norm.weight",
"self_attn.k_norm.bias" | "self_attn.k_layernorm.bias" => "attn_k_norm.bias",
"mlp.gate_proj.weight" | "feed_forward.w1.weight" => "ffn_gate.weight",
"mlp.gate_proj.bias" | "feed_forward.w1.bias" => "ffn_gate.bias",
"mlp.up_proj.weight" | "feed_forward.w3.weight" => "ffn_up.weight",
"mlp.up_proj.bias" | "feed_forward.w3.bias" => "ffn_up.bias",
"mlp.down_proj.weight" | "feed_forward.w2.weight" => "ffn_down.weight",
"mlp.down_proj.bias" | "feed_forward.w2.bias" => "ffn_down.bias",
"input_layernorm.weight" | "operator_norm.weight" => "attn_norm.weight",
"input_layernorm.bias" | "operator_norm.bias" => "attn_norm.bias",
"post_attention_layernorm.weight" => "ffn_norm.weight",
"post_attention_layernorm.bias" => "ffn_norm.bias",
"operator.conv.weight" | "conv.conv.weight" => "shortconv.conv.weight",
"operator.conv.bias" | "conv.conv.bias" => "shortconv.conv.bias",
"operator.in_proj.weight" | "conv.in_proj.weight" => "shortconv.in_proj.weight",
"operator.in_proj.bias" | "conv.in_proj.bias" => "shortconv.in_proj.bias",
"operator.out_proj.weight" | "conv.out_proj.weight" => "shortconv.out_proj.weight",
"operator.out_proj.bias" | "conv.out_proj.bias" => "shortconv.out_proj.bias",
_ => sub_name,
};
return format!("blk.{layer_idx}.{gguf_suffix}");
}
hf_name.to_string()
}
pub fn decode_safetensor_to_f32_into(
raw_bytes: &[u8],
dtype: &str,
out: &mut Vec<f32>,
) -> Result<(), CeraError> {
out.clear();
match dtype.to_ascii_uppercase().as_str() {
"F32" => {
if !raw_bytes.len().is_multiple_of(4) {
return Err(CeraError::Backend(
"F32 tensor bytes not multiple of 4".into(),
));
}
out.reserve(raw_bytes.len() / 4);
let (chunks, _) = raw_bytes.as_chunks::<4>();
out.extend(chunks.iter().map(|c| f32::from_le_bytes(*c)));
Ok(())
}
"F64" => {
if !raw_bytes.len().is_multiple_of(8) {
return Err(CeraError::Backend(
"F64 tensor bytes not multiple of 8".into(),
));
}
out.reserve(raw_bytes.len() / 8);
let (chunks, _) = raw_bytes.as_chunks::<8>();
out.extend(chunks.iter().map(|c| f64::from_le_bytes(*c) as f32));
Ok(())
}
"BF16" => {
if !raw_bytes.len().is_multiple_of(2) {
return Err(CeraError::Backend(
"BF16 tensor bytes not multiple of 2".into(),
));
}
out.reserve(raw_bytes.len() / 2);
let (chunks, _) = raw_bytes.as_chunks::<2>();
out.extend(chunks.iter().map(|c| bf16_to_f32(u16::from_le_bytes(*c))));
Ok(())
}
"F16" => {
if !raw_bytes.len().is_multiple_of(2) {
return Err(CeraError::Backend(
"F16 tensor bytes not multiple of 2".into(),
));
}
out.reserve(raw_bytes.len() / 2);
let (chunks, _) = raw_bytes.as_chunks::<2>();
out.extend(chunks.iter().map(|c| f16_to_f32(u16::from_le_bytes(*c))));
Ok(())
}
"I32" => {
if !raw_bytes.len().is_multiple_of(4) {
return Err(CeraError::Backend(
"I32 tensor bytes not multiple of 4".into(),
));
}
out.reserve(raw_bytes.len() / 4);
let (chunks, _) = raw_bytes.as_chunks::<4>();
out.extend(chunks.iter().map(|c| i32::from_le_bytes(*c) as f32));
Ok(())
}
"I64" => {
if !raw_bytes.len().is_multiple_of(8) {
return Err(CeraError::Backend(
"I64 tensor bytes not multiple of 8".into(),
));
}
out.reserve(raw_bytes.len() / 8);
let (chunks, _) = raw_bytes.as_chunks::<8>();
out.extend(chunks.iter().map(|c| i64::from_le_bytes(*c) as f32));
Ok(())
}
"U32" => {
if !raw_bytes.len().is_multiple_of(4) {
return Err(CeraError::Backend(
"U32 tensor bytes not multiple of 4".into(),
));
}
out.reserve(raw_bytes.len() / 4);
let (chunks, _) = raw_bytes.as_chunks::<4>();
out.extend(chunks.iter().map(|c| u32::from_le_bytes(*c) as f32));
Ok(())
}
"U8" => {
out.reserve(raw_bytes.len());
out.extend(raw_bytes.iter().map(|&b| b as f32));
Ok(())
}
"I8" => {
out.reserve(raw_bytes.len());
out.extend(raw_bytes.iter().map(|&b| b as i8 as f32));
Ok(())
}
"BOOL" => {
out.reserve(raw_bytes.len());
out.extend(raw_bytes.iter().map(|&b| if b != 0 { 1.0 } else { 0.0 }));
Ok(())
}
other => Err(CeraError::Backend(format!(
"unsupported safetensors dtype: `{other}`"
))),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tensor_name_translation() {
assert_eq!(
translate_hf_to_gguf_tensor_name("model.embed_tokens.weight"),
"token_embd.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("lm_head.weight"),
"output.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("model.norm.weight"),
"output_norm.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("model.layers.0.self_attn.q_proj.weight"),
"blk.0.attn_q.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("model.layers.15.mlp.down_proj.weight"),
"blk.15.ffn_down.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("model.encoder.conv1.weight"),
"encoder.conv1.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("model.encoder.layers.2.self_attn.k_proj.weight"),
"encoder.blocks.2.attn.key.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("model.decoder.layers.3.encoder_attn.q_proj.weight"),
"decoder.blocks.3.cross_attn.query.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("model.decoder.embed_tokens.weight"),
"decoder.token_embeddings.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("proj_out.weight"),
"output.weight"
);
assert_eq!(
translate_hf_to_gguf_tensor_name("model.decoder.layer_norm.weight"),
"decoder.ln_post.weight"
);
}
}