/home/noah/src/realizar/src/safetensors.rs

Source
//! Safetensors parser
//!
//! Pure Rust implementation of Safetensors format reader.
//! Used by `HuggingFace` for safe, zero-copy tensor storage.
//!
//! Format specification: <https://github.com/huggingface/safetensors>
//!
//! ## Format Overview
//!
//! ```text
//! Safetensors := HEADER METADATA TENSOR_DATA
//!
//! HEADER := {
//!   metadata_len: u64 (little-endian)
//! }
//!
//! METADATA := JSON {
//!   "tensor_name": {
//!     "dtype": "F32" | "F16" | "I32" | ...,
//!     "shape": [dim1, dim2, ...],
//!     "data_offsets": [start, end]
//!   },
//!   ...
//! }
//! ```

use std::{
    collections::HashMap,
    io::{Cursor, Read},
};

use serde::{Deserialize, Serialize};

use crate::error::{RealizarError, Result};
use crate::inference::simd_bf16_to_f32;

/// Safetensors data type
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
pub enum SafetensorsDtype {
    /// 32-bit float
    F32,
    /// 16-bit float
    F16,
    /// Brain float 16
    BF16,
    /// 32-bit signed integer
    I32,
    /// 64-bit signed integer
    I64,
    /// 8-bit unsigned integer
    U8,
    /// Boolean
    Bool,
}

/// JSON tensor metadata (internal)
#[derive(Debug, Deserialize)]
struct TensorMetadata {
    dtype: SafetensorsDtype,
    shape: Vec<usize>,
    data_offsets: [usize; 2],
}

/// Tensor metadata
#[derive(Debug, Clone, PartialEq)]
pub struct SafetensorsTensorInfo {
    /// Tensor name
    pub name: String,
    /// Data type
    pub dtype: SafetensorsDtype,
    /// Shape (dimensions)
    pub shape: Vec<usize>,
    /// Data offsets in file [start, end)
    pub data_offsets: [usize; 2],
}

/// Safetensors model container
#[derive(Debug, Clone)]
pub struct SafetensorsModel {
    /// Tensor metadata
    pub tensors: HashMap<String, SafetensorsTensorInfo>,
    /// Raw tensor data (not parsed yet)
    pub data: Vec<u8>,
}

impl SafetensorsModel {
    /// Parse Safetensors file from bytes
    ///
    /// # Arguments
    ///
    /// * `data` - Raw Safetensors file bytes
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - Invalid header
    /// - Malformed JSON metadata
    /// - Invalid data offsets
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let data = std::fs::read("model.safetensors")?;
    /// let model = SafetensorsModel::from_bytes(&data)?;
    /// println!("Loaded {} tensors", model.tensors.len());
    /// ```
    pub fn from_bytes(data: &[u8]) -> Result<Self> {
        let mut cursor = Cursor::new(data);

        // Parse header (8-byte metadata length)
        let metadata_len37 = Self::parse_header(&mut cursor)?7;

        // Parse JSON metadata
        let tensors30 = Self::parse_metadata(&mut cursor, metadata_len)?7;

        // Store remaining data
        let data_start =
            usize::try_from(8 + metadata_len).map_err(|_| RealizarError::UnsupportedOperation {
                operation: "convert_data_offset".to_string(),
                reason: format!(
                    "Data offset {} exceeds platform usize limit",
                    8 + metadata_len
                ),
            })?;
        let data = data[data_start..].to_vec();

        Ok(Self { tensors, data })
    }

    /// Extract F32 tensor data by name
    ///
    /// # Arguments
    ///
    /// * `name` - Tensor name to extract
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - Tensor not found
    /// - Tensor dtype is not F32
    /// - Data offsets are invalid
    ///
    /// # Panics
    ///
    /// Never panics. The `unwrap()` in byte conversion is safe because
    /// `chunks_exact(4)` guarantees exactly 4 bytes per chunk.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let model = SafetensorsModel::from_bytes(&data)?;
    /// let weights = model.get_tensor_f32("layer1.weight")?;
    /// println!("Weights: {:?}", weights);
    /// ```
    pub fn get_tensor_f32(&self, name: &str) -> Result<Vec<f32>> {
        // Find tensor metadata
        let tensor8 = self
            .tensors
            .get(name)
            .ok_or_else(|| RealizarError::UnsupportedOperation {
                operation: "get_tensor_f32".to_string(),
                reason: format!("Tensor '{name}' not found"),
            })?;

        // Verify dtype is F32
        if tensor.dtype != SafetensorsDtype::F32 {
            let dtype = &tensor.dtype;
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_f32".to_string(),
                reason: format!("Tensor '{name}' has dtype {dtype:?}, expected F32"),
            });
        }

        // Extract data slice
        let [start, end] = tensor.data_offsets;
        if end > self.data.len() {
            let data_len = self.data.len();
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_f32".to_string(),
                reason: format!("Data offset {end} exceeds data size {data_len}"),
            });
        }

        let bytes = &self.data[start..end];

        // Convert bytes to f32 vector
        if !bytes.len().is_multiple_of(4) {
            let len = bytes.len();
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_f32".to_string(),
                reason: format!("Data size {len} is not a multiple of 4"),
            });
        }

        let values = bytes
            .chunks_exact(4)
            .map5(|chunk| {
                f32::from_le_bytes(
                    chunk
                        .try_into()
                        .expect("chunks_exact(4) guarantees 4-byte slices"),
                )
            })
            .collect();

        Ok(values)
    }

    /// Parse header (8-byte metadata length)
    fn parse_header(cursor: &mut Cursor<&[u8]>) -> Result<u64> {
        let mut buf = [0u8; 8];
        cursor
            .read_exact(&mut buf)
            .map_err(|e| RealizarError::UnsupportedOperation {
                operation: "read_metadata_len".to_string(),
                reason: e.to_string(),
            })?;

        Ok(u64::from_le_bytes(buf))
    }

    /// Parse JSON metadata
    fn parse_metadata(
        cursor: &mut Cursor<&[u8]>,
        len: u64,
    ) -> Result<HashMap<String, SafetensorsTensorInfo>> {
        // Read JSON bytes
        let len_usize = usize::try_from(len).map_err(|_| RealizarError::UnsupportedOperation {
            operation: "convert_metadata_len".to_string(),
            reason: format!("Metadata length {len} exceeds platform usize limit"),
        })?;

        let mut json_bytes = vec![0u8; len_usize];
        cursor
            .read_exact(&mut json_bytes)
            .map_err(|e| RealizarError::UnsupportedOperation {
                operation: "read_metadata_json".to_string(),
                reason: e.to_string(),
            })?;

        // Parse JSON as generic Value first to handle __metadata__ and other special keys
        let json_value32: serde_json::Value32 = serde_json::from_slice(&json_bytes).map_err(|e| {3
            RealizarError::UnsupportedOperation {
                operation: "parse_json".to_string(),
                reason: e.to_string(),
            }
        })?;

        let json_map =
            json_value
                .as_object()
                .ok_or_else(|| RealizarError::UnsupportedOperation {
                    operation: "parse_json".to_string(),
                    reason: "Expected JSON object".to_string(),
                })?;

        // Convert to SafetensorsTensorInfo, skipping special keys like __metadata__
        let mut tensors = HashMap::new();
        for (name47, value47) in json_map {
            // Skip metadata keys (start with __)
            if name.starts_with("__") {
                continue;
            }

            // Parse tensor metadata
            let meta45: TensorMetadata45 = serde_json::from_value(value.clone()).map_err(|e| {1
                RealizarError::UnsupportedOperation {
                    operation: "parse_tensor_metadata".to_string(),
                    reason: format!("Failed to parse tensor '{name}': {e}"),
                }
            })?;

            tensors.insert(
                name.clone(),
                SafetensorsTensorInfo {
                    name: name.clone(),
                    dtype: meta.dtype,
                    shape: meta.shape,
                    data_offsets: meta.data_offsets,
                },
            );
        }

        Ok(tensors)
    }

    /// Get tensor data as F16 values (converts to F32)
    ///
    /// # Arguments
    ///
    /// * `name` - Tensor name to extract
    ///
    /// # Errors
    ///
    /// Returns error if tensor not found or dtype is not F16
    pub fn get_tensor_f16_as_f32(&self, name: &str) -> Result<Vec<f32>> {
        let tensor4 = self
            .tensors
            .get(name)
            .ok_or_else(|| RealizarError::UnsupportedOperation {
                operation: "get_tensor_f16_as_f32".to_string(),
                reason: format!("Tensor '{name}' not found"),
            })?;

        if tensor.dtype != SafetensorsDtype::F16 {
            let dtype = &tensor.dtype;
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_f16_as_f32".to_string(),
                reason: format!("Tensor '{name}' has dtype {dtype:?}, expected F16"),
            });
        }

        let [start, end] = tensor.data_offsets;
        if end > self.data.len() {
            let data_len = self.data.len();
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_f16_as_f32".to_string(),
                reason: format!("Data offset {end} exceeds data size {data_len}"),
            });
        }

        let bytes = &self.data[start..end];

        // Convert F16 bytes to F32
        let values: Vec<f32> = bytes
            .chunks_exact(2)
            .map2(|chunk| {
                let bits = u16::from_le_bytes([chunk[0], chunk[1]]);
                half::f16::from_bits(bits).to_f32()
            })
            .collect();

        Ok(values)
    }

    /// Get tensor data as BF16 values (converts to F32)
    ///
    /// # Arguments
    ///
    /// * `name` - Tensor name to extract
    ///
    /// # Errors
    ///
    /// Returns error if tensor not found or dtype is not BF16
    pub fn get_tensor_bf16_as_f32(&self, name: &str) -> Result<Vec<f32>> {
        let tensor4 = self
            .tensors
            .get(name)
            .ok_or_else(|| RealizarError::UnsupportedOperation {
                operation: "get_tensor_bf16_as_f32".to_string(),
                reason: format!("Tensor '{name}' not found"),
            })?;

        if tensor.dtype != SafetensorsDtype::BF16 {
            let dtype = &tensor.dtype;
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_bf16_as_f32".to_string(),
                reason: format!("Tensor '{name}' has dtype {dtype:?}, expected BF16"),
            });
        }

        let [start, end] = tensor.data_offsets;
        if end > self.data.len() {
            let data_len = self.data.len();
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_bf16_as_f32".to_string(),
                reason: format!("Data offset {end} exceeds data size {data_len}"),
            });
        }

        let bytes = &self.data[start..end];

        // Convert BF16 bytes to F32 using SIMD-accelerated conversion
        // This provides 3-4x speedup over scalar conversion
        let values = simd_bf16_to_f32(bytes);

        Ok(values)
    }

    /// Get tensor as F32 with automatic dtype conversion
    ///
    /// Supports F32, F16, and BF16 dtypes with automatic conversion to F32.
    ///
    /// # Arguments
    ///
    /// * `name` - Tensor name to extract
    ///
    /// # Errors
    ///
    /// Returns error if tensor not found or dtype is not supported
    pub fn get_tensor_auto(&self, name: &str) -> Result<Vec<f32>> {
        let tensor4 = self
            .tensors
            .get(name)
            .ok_or_else(|| RealizarError::UnsupportedOperation {
                operation: "get_tensor_auto".to_string(),
                reason: format!("Tensor '{name}' not found"),
            })?;

        match tensor.dtype {
            SafetensorsDtype::F32 => self.get_tensor_f32(name),
            SafetensorsDtype::F16 => self.get_tensor_f16_as_f32(name),
            SafetensorsDtype::BF16 => self.get_tensor_bf16_as_f32(name),
            _ => Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_auto".to_string(),
                reason: format!("Unsupported dtype {:?} for tensor '{name}'", tensor.dtype),
            }),
        }
    }

    /// Get list of tensor names
    #[must_use]
    pub fn tensor_names(&self) -> Vec<&str> {
        self.tensors.keys().map(String::as_str).collect()
    }

    /// Get tensor info by name
    #[must_use]
    pub fn get_tensor_info(&self, name: &str) -> Option<&SafetensorsTensorInfo> {
        self.tensors.get(name)
    }

    /// Check if model has a tensor with given name
    #[must_use]
    pub fn has_tensor(&self, name: &str) -> bool {
        self.tensors.contains_key(name)
    }
}

// ============================================================================
// SafeTensors Config loader (for sibling config.json)
// ============================================================================

/// Model configuration from config.json
#[derive(Debug, Clone, Deserialize)]
pub struct SafetensorsConfig {
    /// Hidden dimension
    #[serde(alias = "n_embd", alias = "d_model")]
    pub hidden_size: Option<usize>,
    /// Number of transformer layers
    #[serde(alias = "n_layer", alias = "num_layers")]
    pub num_hidden_layers: Option<usize>,
    /// Number of attention heads
    #[serde(alias = "n_head")]
    pub num_attention_heads: Option<usize>,
    /// Number of key-value heads (for GQA)
    pub num_key_value_heads: Option<usize>,
    /// Vocabulary size
    pub vocab_size: Option<usize>,
    /// Intermediate/FFN dimension
    #[serde(alias = "n_inner")]
    pub intermediate_size: Option<usize>,
    /// Maximum sequence length
    #[serde(alias = "n_positions", alias = "n_ctx")]
    pub max_position_embeddings: Option<usize>,
    /// RMSNorm epsilon
    pub rms_norm_eps: Option<f32>,
    /// RoPE theta
    pub rope_theta: Option<f32>,
    /// Model architecture name
    pub architectures: Option<Vec<String>>,
    /// Model type
    pub model_type: Option<String>,
    /// BOS token ID
    pub bos_token_id: Option<u32>,
    /// EOS token ID
    pub eos_token_id: Option<u32>,
}

impl SafetensorsConfig {
    /// Load config from sibling config.json file
    ///
    /// # Arguments
    ///
    /// * `model_path` - Path to the model file (config.json will be loaded from same directory)
    ///
    /// # Returns
    ///
    /// Config if found and parsed, None otherwise
    pub fn load_from_sibling(model_path: &std::path::Path) -> Option<Self> {
        let config_path = model_path.with_file_name("config.json");
        if !config_path.exists() {
            return None;
        }

        let content = std::fs::read_to_string(&config_path).ok()?0;
        serde_json::from_str(&content).ok()
    }

    /// Get number of key-value heads (defaults to num_attention_heads for MHA)
    #[must_use]
    pub fn num_kv_heads(&self) -> usize {
        self.num_key_value_heads
            .or(self.num_attention_heads)
            .unwrap_or(1)
    }

    /// Get model architecture string
    #[must_use]
    pub fn architecture(&self) -> String {
        self.architectures
            .as_ref()
            .and_then(|a| a6.first6())
            .cloned()
            .or_else(|| self.model_type9.clone9())
            .unwrap_or_else(|| "unknown"8.to_string8())
    }
}

// ============================================================================
// Zero-Copy Memory-Mapped SafeTensors Model (T-QA-020)
// ============================================================================

/// Zero-copy memory-mapped SafeTensors model container
///
/// Unlike `SafetensorsModel` which copies all tensor data to the heap,
/// `MappedSafeTensorsModel` uses memory-mapping (mmap) for true zero-copy
/// access to tensor data. This is critical for fast model loading (TTFT).
///
/// # Performance Characteristics
///
/// - **Loading time**: O(1) - only parses header/metadata, no data copy
/// - **Memory**: Only RSS grows as pages are accessed (demand paging)
/// - **TTFT target**: < 500ms for 3GB model
///
/// # Example
///
/// ```rust,ignore
/// let model = MappedSafeTensorsModel::load("model.safetensors")?;
/// let weights = model.get_tensor_bytes("layer1.weight")?;
/// // weights is a zero-copy slice into the mmap'd file
/// ```
#[cfg(not(target_arch = "wasm32"))]
#[derive(Debug)]
pub struct MappedSafeTensorsModel {
    /// Memory-mapped file data
    mmap: memmap2::Mmap,
    /// File path (for diagnostics)
    path: std::path::PathBuf,
    /// Tensor metadata (parsed from header)
    tensors: HashMap<String, SafetensorsTensorInfo>,
    /// Offset where tensor data begins (after header + JSON metadata)
    data_offset: usize,
}

#[cfg(not(target_arch = "wasm32"))]
impl MappedSafeTensorsModel {
    /// Load a SafeTensors file with zero-copy memory mapping
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the SafeTensors file
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - File cannot be opened
    /// - Memory mapping fails
    /// - Header/metadata parsing fails
    ///
    /// # Performance
    ///
    /// This method is O(1) with respect to file size - only the header
    /// and JSON metadata are parsed. Tensor data is not touched until
    /// `get_tensor_bytes()` is called.
    pub fn load<P: AsRef<std::path::Path>>(path: P) -> Result<Self> {
        let path = path.as_ref().to_path_buf();

        // Open file
        let file22 = std::fs::File::open(&path).map_err(|e| RealizarError::UnsupportedOperation {
            operation: "open_safetensors".to_string(),
            reason: format!("Failed to open file '{}': {}", path.display(), e),
        })?;

        // Memory-map the file (zero-copy)
        // SAFETY: File is opened read-only and we don't modify it
        let mmap = unsafe {
            memmap2::MmapOptions::new().map(&file).map_err(|e| {0
                RealizarError::UnsupportedOperation {
                    operation: "mmap_safetensors".to_string(),
                    reason: format!("Failed to mmap file '{}': {}", path.display(), e),
                }
            })?
        };

        // Parse header (8-byte metadata length)
        if mmap.len() < 8 {
            return Err(RealizarError::UnsupportedOperation {
                operation: "parse_safetensors_header".to_string(),
                reason: format!(
                    "File too small: {} bytes (minimum 8 for header)",
                    mmap.len()
                ),
            });
        }

        let metadata_len =
            u64::from_le_bytes(mmap[0..8].try_into().expect("slice is exactly 8 bytes"));

        let metadata_len_usize =
            usize::try_from(metadata_len).map_err(|_| RealizarError::UnsupportedOperation {
                operation: "parse_safetensors_header".to_string(),
                reason: format!("Metadata length {} exceeds platform limit", metadata_len),
            })?;

        // Verify we have enough data for metadata
        let data_offset = 8 + metadata_len_usize;
        if mmap.len() < data_offset {
            return Err(RealizarError::UnsupportedOperation {
                operation: "parse_safetensors_header".to_string(),
                reason: format!(
                    "File truncated: need {} bytes for header+metadata, have {}",
                    data_offset,
                    mmap.len()
                ),
            });
        }

        // Parse JSON metadata (from mmap'd memory, no copy)
        let json_bytes = &mmap[8..data_offset];
        let tensors = Self::parse_metadata(json_bytes)?0;

        Ok(Self {
            mmap,
            path,
            tensors,
            data_offset,
        })
    }

    /// Parse JSON metadata from bytes
    fn parse_metadata(json_bytes: &[u8]) -> Result<HashMap<String, SafetensorsTensorInfo>> {
        // Parse JSON as generic Value first to handle __metadata__ and other special keys
        let json_value: serde_json::Value = serde_json::from_slice(json_bytes).map_err(|e| {0
            RealizarError::UnsupportedOperation {
                operation: "parse_json".to_string(),
                reason: e.to_string(),
            }
        })?;

        let json_map =
            json_value
                .as_object()
                .ok_or_else(|| RealizarError::UnsupportedOperation {
                    operation: "parse_json".to_string(),
                    reason: "Expected JSON object".to_string(),
                })?;

        // Convert to SafetensorsTensorInfo, skipping special keys like __metadata__
        let mut tensors = HashMap::new();
        for (name54, value54) in json_map {
            // Skip metadata keys (start with __)
            if name.starts_with("__") {
                continue;
            }

            // Parse tensor metadata
            let meta: TensorMetadata = serde_json::from_value(value.clone()).map_err(|e| {0
                RealizarError::UnsupportedOperation {
                    operation: "parse_tensor_metadata".to_string(),
                    reason: format!("Failed to parse tensor '{name}': {e}"),
                }
            })?;

            tensors.insert(
                name.clone(),
                SafetensorsTensorInfo {
                    name: name.clone(),
                    dtype: meta.dtype,
                    shape: meta.shape,
                    data_offsets: meta.data_offsets,
                },
            );
        }

        Ok(tensors)
    }

    /// Get raw tensor bytes (zero-copy slice into mmap'd file)
    ///
    /// # Arguments
    ///
    /// * `name` - Tensor name
    ///
    /// # Returns
    ///
    /// Zero-copy slice into the memory-mapped file. The slice is valid
    /// as long as `self` is alive.
    ///
    /// # Errors
    ///
    /// Returns error if tensor not found or offsets are invalid.
    pub fn get_tensor_bytes(&self, name: &str) -> Result<&[u8]> {
        let tensor = self
            .tensors
            .get(name)
            .ok_or_else(|| RealizarError::UnsupportedOperation {
                operation: "get_tensor_bytes".to_string(),
                reason: format!("Tensor '{name}' not found"),
            })?;

        let [start, end] = tensor.data_offsets;
        let abs_start = self.data_offset + start;
        let abs_end = self.data_offset + end;

        if abs_end > self.mmap.len() {
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_bytes".to_string(),
                reason: format!(
                    "Tensor '{}' data offsets [{}, {}] exceed file size {}",
                    name,
                    abs_start,
                    abs_end,
                    self.mmap.len()
                ),
            });
        }

        Ok(&self.mmap[abs_start..abs_end])
    }

    /// Get tensor as F32 values (zero-copy bytes, then convert)
    ///
    /// # Arguments
    ///
    /// * `name` - Tensor name
    ///
    /// # Errors
    ///
    /// Returns error if tensor not found or dtype is not F32.
    pub fn get_tensor_f32(&self, name: &str) -> Result<Vec<f32>> {
        let tensor = self
            .tensors
            .get(name)
            .ok_or_else(|| RealizarError::UnsupportedOperation {
                operation: "get_tensor_f32".to_string(),
                reason: format!("Tensor '{name}' not found"),
            })?;

        if tensor.dtype != SafetensorsDtype::F32 {
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_f32".to_string(),
                reason: format!(
                    "Tensor '{}' has dtype {:?}, expected F32",
                    name, tensor.dtype
                ),
            });
        }

        let bytes = self.get_tensor_bytes(name)?0;

        if !bytes.len().is_multiple_of(4) {
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_f32".to_string(),
                reason: format!("Data size {} is not a multiple of 4", bytes.len()),
            });
        }

        let values = bytes
            .chunks_exact(4)
            .map54(|chunk| {
                f32::from_le_bytes(
                    chunk
                        .try_into()
                        .expect("chunks_exact(4) guarantees 4-byte slices"),
                )
            })
            .collect();

        Ok(values)
    }

    /// Get tensor as BF16 bytes (zero-copy, native format)
    ///
    /// Returns raw BF16 bytes for native SIMD processing without
    /// F32 conversion at boot time.
    ///
    /// # Arguments
    ///
    /// * `name` - Tensor name
    ///
    /// # Errors
    ///
    /// Returns error if tensor not found or dtype is not BF16.
    pub fn get_tensor_bf16_bytes(&self, name: &str) -> Result<&[u8]> {
        let tensor = self
            .tensors
            .get(name)
            .ok_or_else(|| RealizarError::UnsupportedOperation {
                operation: "get_tensor_bf16_bytes".to_string(),
                reason: format!("Tensor '{name}' not found"),
            })?;

        if tensor.dtype != SafetensorsDtype::BF16 {
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_bf16_bytes".to_string(),
                reason: format!(
                    "Tensor '{}' has dtype {:?}, expected BF16",
                    name, tensor.dtype
                ),
            });
        }

        self.get_tensor_bytes(name)
    }

    /// Get tensor as BF16 values converted to F32
    ///
    /// # Arguments
    ///
    /// * `name` - Tensor name
    ///
    /// # Errors
    ///
    /// Returns error if tensor not found or dtype is not BF16.
    pub fn get_tensor_bf16_as_f32(&self, name: &str) -> Result<Vec<f32>> {
        let bytes = self.get_tensor_bf16_bytes(name)?;

        // Convert BF16 bytes to F32 using SIMD-accelerated conversion
        // This provides 3-4x speedup over scalar conversion
        let values = simd_bf16_to_f32(bytes);

        Ok(values)
    }

    /// Get tensor as F16 bytes (zero-copy, native format)
    ///
    /// Returns raw F16 bytes for native SIMD processing.
    ///
    /// # Arguments
    ///
    /// * `name` - Tensor name
    ///
    /// # Errors
    ///
    /// Returns error if tensor not found or dtype is not F16.
    pub fn get_tensor_f16_bytes(&self, name: &str) -> Result<&[u8]> {
        let tensor = self
            .tensors
            .get(name)
            .ok_or_else(|| RealizarError::UnsupportedOperation {
                operation: "get_tensor_f16_bytes".to_string(),
                reason: format!("Tensor '{name}' not found"),
            })?;

        if tensor.dtype != SafetensorsDtype::F16 {
            return Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_f16_bytes".to_string(),
                reason: format!(
                    "Tensor '{}' has dtype {:?}, expected F16",
                    name, tensor.dtype
                ),
            });
        }

        self.get_tensor_bytes(name)
    }

    /// Get tensor as F16 values converted to F32
    pub fn get_tensor_f16_as_f32(&self, name: &str) -> Result<Vec<f32>> {
        let bytes = self.get_tensor_f16_bytes(name)?;

        let values: Vec<f32> = bytes
            .chunks_exact(2)
            .map(|chunk| {
                let bits = u16::from_le_bytes([chunk[0], chunk[1]]);
                half::f16::from_bits(bits).to_f32()
            })
            .collect();

        Ok(values)
    }

    /// Get tensor as F32 with automatic dtype conversion
    ///
    /// Supports F32, F16, and BF16 dtypes with automatic conversion to F32.
    pub fn get_tensor_auto(&self, name: &str) -> Result<Vec<f32>> {
        let tensor54 = self
            .tensors
            .get(name)
            .ok_or_else(|| RealizarError::UnsupportedOperation {
                operation: "get_tensor_auto".to_string(),
                reason: format!("Tensor '{name}' not found"),
            })?;

        match tensor.dtype {
            SafetensorsDtype::F32 => self.get_tensor_f32(name),
            SafetensorsDtype::F16 => self.get_tensor_f16_as_f32(name),
            SafetensorsDtype::BF16 => self.get_tensor_bf16_as_f32(name),
            _ => Err(RealizarError::UnsupportedOperation {
                operation: "get_tensor_auto".to_string(),
                reason: format!("Unsupported dtype {:?} for tensor '{}'", tensor.dtype, name),
            }),
        }
    }

    /// Get list of tensor names
    #[must_use]
    pub fn tensor_names(&self) -> Vec<&str> {
        self.tensors.keys().map(String::as_str).collect()
    }

    /// Get tensor info by name
    #[must_use]
    pub fn get_tensor_info(&self, name: &str) -> Option<&SafetensorsTensorInfo> {
        self.tensors.get(name)
    }

    /// Check if model has a tensor with given name
    #[must_use]
    pub fn has_tensor(&self, name: &str) -> bool {
        self.tensors.contains_key(name)
    }

    /// Get the file path
    #[must_use]
    pub fn path(&self) -> &std::path::Path {
        &self.path
    }

    /// Get the total file size in bytes
    #[must_use]
    pub fn file_size(&self) -> usize {
        self.mmap.len()
    }

    /// Get the number of tensors
    #[must_use]
    pub fn tensor_count(&self) -> usize {
        self.tensors.len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_empty_safetensors() {
        // Minimal valid Safetensors: 8-byte header + empty JSON "{}"
        let mut data = Vec::new();
        data.extend_from_slice(&2u64.to_le_bytes()); // metadata_len = 2
        data.extend_from_slice(b"{}"); // Empty JSON

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        assert_eq!(model.tensors.len(), 0);
        assert_eq!(model.data.len(), 0);
    }

    #[test]
    fn test_invalid_header_truncated() {
        // Only 4 bytes (should be 8)
        let data = [0u8; 4];
        let result = SafetensorsModel::from_bytes(&data);
        assert!(result.is_err());
    }

    #[test]
    fn test_empty_file() {
        let data = &[];
        let result = SafetensorsModel::from_bytes(data);
        assert!(result.is_err());
    }

    #[test]
    fn test_parse_single_tensor() {
        // Safetensors with one F32 tensor
        let json = r#"{"weight":{"dtype":"F32","shape":[2,3],"data_offsets":[0,24]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        // Add 24 bytes of dummy tensor data (2*3*4 = 24 bytes for F32)
        data.extend_from_slice(&[0u8; 24]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        assert_eq!(model.tensors.len(), 1);

        let tensor = model.tensors.get("weight").expect("test");
        assert_eq!(tensor.name, "weight");
        assert_eq!(tensor.dtype, SafetensorsDtype::F32);
        assert_eq!(tensor.shape, vec![2, 3]);
        assert_eq!(tensor.data_offsets, [0, 24]);
    }

    #[test]
    fn test_parse_multiple_tensors() {
        // Safetensors with multiple tensors of different types
        let json = r#"{
            "layer1.weight":{"dtype":"F32","shape":[128,256],"data_offsets":[0,131072]},
            "layer1.bias":{"dtype":"F32","shape":[128],"data_offsets":[131072,131584]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        // Add dummy tensor data
        data.extend_from_slice(&vec![0u8; 131_584]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        assert_eq!(model.tensors.len(), 2);

        let weight = model.tensors.get("layer1.weight").expect("test");
        assert_eq!(weight.dtype, SafetensorsDtype::F32);
        assert_eq!(weight.shape, vec![128, 256]);
        assert_eq!(weight.data_offsets, [0, 131_072]);

        let bias = model.tensors.get("layer1.bias").expect("test");
        assert_eq!(bias.dtype, SafetensorsDtype::F32);
        assert_eq!(bias.shape, vec![128]);
        assert_eq!(bias.data_offsets, [131_072, 131_584]);
    }

    #[test]
    fn test_parse_various_dtypes() {
        // Test different data types
        let json = r#"{
            "f32_tensor":{"dtype":"F32","shape":[2],"data_offsets":[0,8]},
            "i32_tensor":{"dtype":"I32","shape":[2],"data_offsets":[8,16]},
            "u8_tensor":{"dtype":"U8","shape":[4],"data_offsets":[16,20]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 20]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        assert_eq!(model.tensors.len(), 3);

        assert_eq!(
            model.tensors.get("f32_tensor").expect("test").dtype,
            SafetensorsDtype::F32
        );
        assert_eq!(
            model.tensors.get("i32_tensor").expect("test").dtype,
            SafetensorsDtype::I32
        );
        assert_eq!(
            model.tensors.get("u8_tensor").expect("test").dtype,
            SafetensorsDtype::U8
        );
    }

    #[test]
    fn test_invalid_json_error() {
        // Invalid JSON in metadata
        let mut data = Vec::new();
        data.extend_from_slice(&10u64.to_le_bytes()); // metadata_len = 10
        data.extend_from_slice(b"not json!!"); // Invalid JSON

        let result = SafetensorsModel::from_bytes(&data);
        assert!(result.is_err());
        assert!(matches!0(
            result.unwrap_err(),
            RealizarError::UnsupportedOperation { .. }
        ));
    }

    #[test]
    fn test_truncated_json_error() {
        // Header says JSON is longer than actual data
        let mut data = Vec::new();
        data.extend_from_slice(&100u64.to_le_bytes()); // metadata_len = 100
        data.extend_from_slice(b"{}"); // Only 2 bytes, not 100

        let result = SafetensorsModel::from_bytes(&data);
        assert!(result.is_err());
        assert!(matches!0(
            result.unwrap_err(),
            RealizarError::UnsupportedOperation { .. }
        ));
    }

    #[test]
    fn test_parse_all_dtypes() {
        // Test all supported data types
        let json = r#"{
            "f32":{"dtype":"F32","shape":[1],"data_offsets":[0,4]},
            "f16":{"dtype":"F16","shape":[1],"data_offsets":[4,6]},
            "bf16":{"dtype":"BF16","shape":[1],"data_offsets":[6,8]},
            "i32":{"dtype":"I32","shape":[1],"data_offsets":[8,12]},
            "i64":{"dtype":"I64","shape":[1],"data_offsets":[12,20]},
            "u8":{"dtype":"U8","shape":[1],"data_offsets":[20,21]},
            "bool":{"dtype":"Bool","shape":[1],"data_offsets":[21,22]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 22]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        assert_eq!(model.tensors.len(), 7);

        assert_eq!(
            model.tensors.get("f32").expect("test").dtype,
            SafetensorsDtype::F32
        );
        assert_eq!(
            model.tensors.get("f16").expect("test").dtype,
            SafetensorsDtype::F16
        );
        assert_eq!(
            model.tensors.get("bf16").expect("test").dtype,
            SafetensorsDtype::BF16
        );
        assert_eq!(
            model.tensors.get("i32").expect("test").dtype,
            SafetensorsDtype::I32
        );
        assert_eq!(
            model.tensors.get("i64").expect("test").dtype,
            SafetensorsDtype::I64
        );
        assert_eq!(
            model.tensors.get("u8").expect("test").dtype,
            SafetensorsDtype::U8
        );
        assert_eq!(
            model.tensors.get("bool").expect("test").dtype,
            SafetensorsDtype::Bool
        );
    }

    #[test]
    fn test_tensor_data_preserved() {
        // Verify tensor data is correctly preserved
        let json = r#"{"weight":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        // Add specific tensor data (two f32s: 1.0 and 2.0)
        data.extend_from_slice(&1.0f32.to_le_bytes());
        data.extend_from_slice(&2.0f32.to_le_bytes());

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        assert_eq!(model.data.len(), 8);

        // Verify we can read back the f32 values
        let val1 = f32::from_le_bytes(model.data[0..4].try_into().expect("test"));
        let val2 = f32::from_le_bytes(model.data[4..8].try_into().expect("test"));
        assert!((val1 - 1.0).abs() < 1e-6);
        assert!((val2 - 2.0).abs() < 1e-6);
    }

    #[test]
    fn test_multidimensional_shapes() {
        // Test tensors with various shapes
        let json = r#"{
            "scalar":{"dtype":"F32","shape":[],"data_offsets":[0,4]},
            "vector":{"dtype":"F32","shape":[10],"data_offsets":[4,44]},
            "matrix":{"dtype":"F32","shape":[3,4],"data_offsets":[44,92]},
            "tensor3d":{"dtype":"F32","shape":[2,3,4],"data_offsets":[92,188]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 188]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        assert_eq!(model.tensors.len(), 4);

        assert_eq!(
            model.tensors.get("scalar").expect("test").shape,
            Vec::<usize>::new()
        );
        assert_eq!(model.tensors.get("vector").expect("test").shape, vec![10]);
        assert_eq!(model.tensors.get("matrix").expect("test").shape, vec![3, 4]);
        assert_eq!(
            model.tensors.get("tensor3d").expect("test").shape,
            vec![2, 3, 4]
        );
    }

    #[test]
    fn test_aprender_linear_regression_format_compatibility() {
        // Test aprender LinearRegression SafeTensors format compatibility
        // Format: {"coefficients": [n_features], "intercept": [1]}
        // Example model: y = 2.0*x1 + 3.0*x2 + 1.5*x3 + 0.5

        let json = r#"{
            "coefficients":{"dtype":"F32","shape":[3],"data_offsets":[0,12]},
            "intercept":{"dtype":"F32","shape":[1],"data_offsets":[12,16]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();

        // Header
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());

        // Metadata
        data.extend_from_slice(json_bytes);

        // Tensor data: coefficients [2.0, 3.0, 1.5]
        data.extend_from_slice(&2.0f32.to_le_bytes());
        data.extend_from_slice(&3.0f32.to_le_bytes());
        data.extend_from_slice(&1.5f32.to_le_bytes());

        // intercept [0.5]
        data.extend_from_slice(&0.5f32.to_le_bytes());

        // Parse with realizar
        let model = SafetensorsModel::from_bytes(&data).expect("test");

        // Verify structure
        assert_eq!(model.tensors.len(), 2);

        // Check coefficients tensor
        let coef = model.tensors.get("coefficients").expect("test");
        assert_eq!(coef.dtype, SafetensorsDtype::F32);
        assert_eq!(coef.shape, vec![3]);
        assert_eq!(coef.data_offsets, [0, 12]);

        // Check intercept tensor
        let intercept = model.tensors.get("intercept").expect("test");
        assert_eq!(intercept.dtype, SafetensorsDtype::F32);
        assert_eq!(intercept.shape, vec![1]);
        assert_eq!(intercept.data_offsets, [12, 16]);

        // Verify we can extract the actual values
        let coef_vals: Vec<f32> = (0..3)
            .map1(|i| {
                let offset = i * 4;
                f32::from_le_bytes(model.data[offset..offset + 4].try_into().expect("test"))
            })
            .collect();
        assert!((coef_vals[0] - 2.0).abs() < 1e-6);
        assert!((coef_vals[1] - 3.0).abs() < 1e-6);
        assert!((coef_vals[2] - 1.5).abs() < 1e-6);

        let intercept_val = f32::from_le_bytes(model.data[12..16].try_into().expect("test"));
        assert!((intercept_val - 0.5).abs() < 1e-6);
    }

    #[test]
    fn test_get_tensor_f32_helper() {
        // Test the get_tensor_f32 helper method
        let json = r#"{
            "weights":{"dtype":"F32","shape":[4],"data_offsets":[0,16]},
            "bias":{"dtype":"F32","shape":[2],"data_offsets":[16,24]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);

        // weights: [1.0, 2.0, 3.0, 4.0]
        data.extend_from_slice(&1.0f32.to_le_bytes());
        data.extend_from_slice(&2.0f32.to_le_bytes());
        data.extend_from_slice(&3.0f32.to_le_bytes());
        data.extend_from_slice(&4.0f32.to_le_bytes());

        // bias: [0.5, 0.25]
        data.extend_from_slice(&0.5f32.to_le_bytes());
        data.extend_from_slice(&0.25f32.to_le_bytes());

        let model = SafetensorsModel::from_bytes(&data).expect("test");

        // Test extracting weights
        let weights = model.get_tensor_f32("weights").expect("test");
        assert_eq!(weights.len(), 4);
        assert!((weights[0] - 1.0).abs() < 1e-6);
        assert!((weights[1] - 2.0).abs() < 1e-6);
        assert!((weights[2] - 3.0).abs() < 1e-6);
        assert!((weights[3] - 4.0).abs() < 1e-6);

        // Test extracting bias
        let bias = model.get_tensor_f32("bias").expect("test");
        assert_eq!(bias.len(), 2);
        assert!((bias[0] - 0.5).abs() < 1e-6);
        assert!((bias[1] - 0.25).abs() < 1e-6);

        // Test error: tensor not found
        let result = model.get_tensor_f32("nonexistent");
        assert!(result.is_err());
    }

    #[test]
    fn test_get_tensor_f32_wrong_dtype() {
        // Test error when tensor has wrong dtype
        let json = r#"{
            "int_tensor":{"dtype":"I32","shape":[2],"data_offsets":[0,8]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&1i32.to_le_bytes());
        data.extend_from_slice(&2i32.to_le_bytes());

        let model = SafetensorsModel::from_bytes(&data).expect("test");

        // Should error because dtype is I32, not F32
        let result = model.get_tensor_f32("int_tensor");
        assert!(result.is_err());
    }

    #[test]
    fn test_get_tensor_f32_with_aprender_model() {
        // Use get_tensor_f32 with aprender LinearRegression format
        let json = r#"{
            "coefficients":{"dtype":"F32","shape":[3],"data_offsets":[0,12]},
            "intercept":{"dtype":"F32","shape":[1],"data_offsets":[12,16]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&2.0f32.to_le_bytes());
        data.extend_from_slice(&3.0f32.to_le_bytes());
        data.extend_from_slice(&1.5f32.to_le_bytes());
        data.extend_from_slice(&0.5f32.to_le_bytes());

        let model = SafetensorsModel::from_bytes(&data).expect("test");

        // Extract using helper method - much cleaner!
        let coefficients = model.get_tensor_f32("coefficients").expect("test");
        assert_eq!(coefficients, vec![2.0, 3.0, 1.5]);

        let intercept = model.get_tensor_f32("intercept").expect("test");
        assert_eq!(intercept, vec![0.5]);
    }

    // ========== Coverage tests for untested functions ==========

    #[test]
    fn test_cov_get_tensor_f16_as_f32() {
        let json = r#"{"weights":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        // Two F16 values: 1.0 and 2.0
        data.extend_from_slice(&half::f16::from_f32(1.0).to_le_bytes());
        data.extend_from_slice(&half::f16::from_f32(2.0).to_le_bytes());

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let weights = model.get_tensor_f16_as_f32("weights").expect("test");

        assert_eq!(weights.len(), 2);
        assert!((weights[0] - 1.0).abs() < 0.01);
        assert!((weights[1] - 2.0).abs() < 0.01);
    }

    #[test]
    fn test_cov_get_tensor_f16_not_found() {
        let json = r#"{"weights":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 4]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_f16_as_f32("nonexistent");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_get_tensor_f16_wrong_dtype() {
        let json = r#"{"weights":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 8]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_f16_as_f32("weights");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_get_tensor_f16_data_offset_exceeds() {
        let json = r#"{"weights":{"dtype":"F16","shape":[2],"data_offsets":[0,100]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 4]); // Only 4 bytes, offset says 100

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_f16_as_f32("weights");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_get_tensor_bf16_as_f32() {
        let json = r#"{"weights":{"dtype":"BF16","shape":[2],"data_offsets":[0,4]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        // Two BF16 values: 1.0 and 2.0
        data.extend_from_slice(&half::bf16::from_f32(1.0).to_le_bytes());
        data.extend_from_slice(&half::bf16::from_f32(2.0).to_le_bytes());

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let weights = model.get_tensor_bf16_as_f32("weights").expect("test");

        assert_eq!(weights.len(), 2);
        assert!((weights[0] - 1.0).abs() < 0.01);
        assert!((weights[1] - 2.0).abs() < 0.01);
    }

    #[test]
    fn test_cov_get_tensor_bf16_not_found() {
        let json = r#"{"weights":{"dtype":"BF16","shape":[2],"data_offsets":[0,4]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 4]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_bf16_as_f32("nonexistent");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_get_tensor_bf16_wrong_dtype() {
        let json = r#"{"weights":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 8]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_bf16_as_f32("weights");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_get_tensor_bf16_data_offset_exceeds() {
        let json = r#"{"weights":{"dtype":"BF16","shape":[2],"data_offsets":[0,100]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 4]); // Only 4 bytes

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_bf16_as_f32("weights");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_get_tensor_auto_f32() {
        let json = r#"{"weights":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&1.0f32.to_le_bytes());
        data.extend_from_slice(&2.0f32.to_le_bytes());

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let weights = model.get_tensor_auto("weights").expect("test");
        assert_eq!(weights, vec![1.0, 2.0]);
    }

    #[test]
    fn test_cov_get_tensor_auto_f16() {
        let json = r#"{"weights":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&half::f16::from_f32(1.0).to_le_bytes());
        data.extend_from_slice(&half::f16::from_f32(2.0).to_le_bytes());

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let weights = model.get_tensor_auto("weights").expect("test");
        assert_eq!(weights.len(), 2);
    }

    #[test]
    fn test_cov_get_tensor_auto_bf16() {
        let json = r#"{"weights":{"dtype":"BF16","shape":[2],"data_offsets":[0,4]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&half::bf16::from_f32(1.0).to_le_bytes());
        data.extend_from_slice(&half::bf16::from_f32(2.0).to_le_bytes());

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let weights = model.get_tensor_auto("weights").expect("test");
        assert_eq!(weights.len(), 2);
    }

    #[test]
    fn test_cov_get_tensor_auto_unsupported_dtype() {
        let json = r#"{"weights":{"dtype":"I32","shape":[2],"data_offsets":[0,8]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 8]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_auto("weights");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_get_tensor_auto_not_found() {
        let json = r#"{"weights":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 8]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_auto("nonexistent");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_tensor_names() {
        let json = r#"{
            "weight1":{"dtype":"F32","shape":[2],"data_offsets":[0,8]},
            "weight2":{"dtype":"F32","shape":[2],"data_offsets":[8,16]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 16]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let names = model.tensor_names();
        assert_eq!(names.len(), 2);
        assert!(names.contains(&"weight1"));
        assert!(names.contains(&"weight2"));
    }

    #[test]
    fn test_cov_get_tensor_info() {
        let json = r#"{"weight":{"dtype":"F32","shape":[2,3],"data_offsets":[0,24]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 24]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");

        let info = model.get_tensor_info("weight");
        assert!(info.is_some());
        let info = info.expect("operation failed");
        assert_eq!(info.shape, vec![2, 3]);
        assert_eq!(info.dtype, SafetensorsDtype::F32);

        assert!(model.get_tensor_info("nonexistent").is_none());
    }

    #[test]
    fn test_cov_has_tensor() {
        let json = r#"{"weight":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 8]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        assert!(model.has_tensor("weight"));
        assert!(!model.has_tensor("nonexistent"));
    }

    #[test]
    fn test_cov_get_tensor_f32_data_offset_exceeds() {
        let json = r#"{"weight":{"dtype":"F32","shape":[2],"data_offsets":[0,100]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 8]); // Only 8 bytes, offset says 100

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_f32("weight");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_get_tensor_f32_not_multiple_of_4() {
        let json = r#"{"weight":{"dtype":"F32","shape":[2],"data_offsets":[0,7]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 7]); // 7 bytes, not a multiple of 4

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        let result = model.get_tensor_f32("weight");
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_safetensors_config_num_kv_heads() {
        let config = SafetensorsConfig {
            hidden_size: Some(768),
            num_hidden_layers: Some(12),
            num_attention_heads: Some(12),
            num_key_value_heads: Some(4),
            vocab_size: Some(32000),
            intermediate_size: None,
            max_position_embeddings: None,
            rms_norm_eps: None,
            rope_theta: None,
            architectures: None,
            model_type: None,
            bos_token_id: None,
            eos_token_id: None,
        };

        assert_eq!(config.num_kv_heads(), 4);
    }

    #[test]
    fn test_cov_safetensors_config_num_kv_heads_default() {
        let config = SafetensorsConfig {
            hidden_size: Some(768),
            num_hidden_layers: Some(12),
            num_attention_heads: Some(12),
            num_key_value_heads: None, // Not set, should fall back to attention heads
            vocab_size: Some(32000),
            intermediate_size: None,
            max_position_embeddings: None,
            rms_norm_eps: None,
            rope_theta: None,
            architectures: None,
            model_type: None,
            bos_token_id: None,
            eos_token_id: None,
        };

        assert_eq!(config.num_kv_heads(), 12);
    }

    #[test]
    fn test_cov_safetensors_config_num_kv_heads_fallback() {
        let config = SafetensorsConfig {
            hidden_size: Some(768),
            num_hidden_layers: Some(12),
            num_attention_heads: None,
            num_key_value_heads: None,
            vocab_size: Some(32000),
            intermediate_size: None,
            max_position_embeddings: None,
            rms_norm_eps: None,
            rope_theta: None,
            architectures: None,
            model_type: None,
            bos_token_id: None,
            eos_token_id: None,
        };

        assert_eq!(config.num_kv_heads(), 1); // Fallback to 1
    }

    #[test]
    fn test_cov_safetensors_config_architecture_from_architectures() {
        let config = SafetensorsConfig {
            hidden_size: None,
            num_hidden_layers: None,
            num_attention_heads: None,
            num_key_value_heads: None,
            vocab_size: None,
            intermediate_size: None,
            max_position_embeddings: None,
            rms_norm_eps: None,
            rope_theta: None,
            architectures: Some(vec!["LlamaForCausalLM".to_string()]),
            model_type: Some("llama".to_string()),
            bos_token_id: None,
            eos_token_id: None,
        };

        assert_eq!(config.architecture(), "LlamaForCausalLM");
    }

    #[test]
    fn test_cov_safetensors_config_architecture_from_model_type() {
        let config = SafetensorsConfig {
            hidden_size: None,
            num_hidden_layers: None,
            num_attention_heads: None,
            num_key_value_heads: None,
            vocab_size: None,
            intermediate_size: None,
            max_position_embeddings: None,
            rms_norm_eps: None,
            rope_theta: None,
            architectures: None,
            model_type: Some("llama".to_string()),
            bos_token_id: None,
            eos_token_id: None,
        };

        assert_eq!(config.architecture(), "llama");
    }

    #[test]
    fn test_cov_safetensors_config_architecture_unknown() {
        let config = SafetensorsConfig {
            hidden_size: None,
            num_hidden_layers: None,
            num_attention_heads: None,
            num_key_value_heads: None,
            vocab_size: None,
            intermediate_size: None,
            max_position_embeddings: None,
            rms_norm_eps: None,
            rope_theta: None,
            architectures: None,
            model_type: None,
            bos_token_id: None,
            eos_token_id: None,
        };

        assert_eq!(config.architecture(), "unknown");
    }

    #[test]
    fn test_cov_safetensors_config_load_from_sibling_not_found() {
        let path = std::path::Path::new("/nonexistent/model.safetensors");
        let config = SafetensorsConfig::load_from_sibling(path);
        assert!(config.is_none());
    }

    #[test]
    fn test_cov_metadata_key_skipped() {
        // Test that __metadata__ key is skipped
        let json = r#"{
            "__metadata__":{"format":"pt"},
            "weight":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}
        }"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 8]);

        let model = SafetensorsModel::from_bytes(&data).expect("test");
        assert_eq!(model.tensors.len(), 1);
        assert!(model.tensors.contains_key("weight"));
        assert!(!model.tensors.contains_key("__metadata__"));
    }

    #[test]
    fn test_cov_json_not_object() {
        // JSON is an array, not an object
        let json = r"[]";
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);

        let result = SafetensorsModel::from_bytes(&data);
        assert!(result.is_err());
    }

    #[test]
    fn test_cov_tensor_metadata_parse_error() {
        // Tensor has invalid metadata (missing dtype)
        let json = r#"{"weight":{"shape":[2],"data_offsets":[0,8]}}"#;
        let json_bytes = json.as_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
        data.extend_from_slice(json_bytes);
        data.extend_from_slice(&[0u8; 8]);

        let result = SafetensorsModel::from_bytes(&data);
        assert!(result.is_err());
    }
}

Coverage Report

Created: 2026-01-25 15:05