## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

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

    Line
    Count
    Source
    1
    //! Safetensors parser
    2
    //!
    3
    //! Pure Rust implementation of Safetensors format reader.
    4
    //! Used by `HuggingFace` for safe, zero-copy tensor storage.
    5
    //!
    6
    //! Format specification: <https://github.com/huggingface/safetensors>
    7
    //!
    8
    //! ## Format Overview
    9
    //!
    10
    //! ```text
    11
    //! Safetensors := HEADER METADATA TENSOR_DATA
    12
    //!
    13
    //! HEADER := {
    14
    //!   metadata_len: u64 (little-endian)
    15
    //! }
    16
    //!
    17
    //! METADATA := JSON {
    18
    //!   "tensor_name": {
    19
    //!     "dtype": "F32" | "F16" | "I32" | ...,
    20
    //!     "shape": [dim1, dim2, ...],
    21
    //!     "data_offsets": [start, end]
    22
    //!   },
    23
    //!   ...
    24
    //! }
    25
    //! ```
    26
    27
    use std::{
    28
        collections::HashMap,
    29
        io::{Cursor, Read},
    30
    };
    31
    32
    use serde::{Deserialize, Serialize};
    33
    34
    use crate::error::{RealizarError, Result};
    35
    use crate::inference::simd_bf16_to_f32;
    36
    37
    /// Safetensors data type
    38
    #[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
    39
    pub enum SafetensorsDtype {
    40
        /// 32-bit float
    41
        F32,
    42
        /// 16-bit float
    43
        F16,
    44
        /// Brain float 16
    45
        BF16,
    46
        /// 32-bit signed integer
    47
        I32,
    48
        /// 64-bit signed integer
    49
        I64,
    50
        /// 8-bit unsigned integer
    51
        U8,
    52
        /// Boolean
    53
        Bool,
    54
    }
    55
    56
    /// JSON tensor metadata (internal)
    57
    #[derive(Debug, Deserialize)]
    58
    struct TensorMetadata {
    59
        dtype: SafetensorsDtype,
    60
        shape: Vec<usize>,
    61
        data_offsets: [usize; 2],
    62
    }
    63
    64
    /// Tensor metadata
    65
    #[derive(Debug, Clone, PartialEq)]
    66
    pub struct SafetensorsTensorInfo {
    67
        /// Tensor name
    68
        pub name: String,
    69
        /// Data type
    70
        pub dtype: SafetensorsDtype,
    71
        /// Shape (dimensions)
    72
        pub shape: Vec<usize>,
    73
        /// Data offsets in file [start, end)
    74
        pub data_offsets: [usize; 2],
    75
    }
    76
    77
    /// Safetensors model container
    78
    #[derive(Debug, Clone)]
    79
    pub struct SafetensorsModel {
    80
        /// Tensor metadata
    81
        pub tensors: HashMap<String, SafetensorsTensorInfo>,
    82
        /// Raw tensor data (not parsed yet)
    83
        pub data: Vec<u8>,
    84
    }
    85
    86
    impl SafetensorsModel {
    87
        /// Parse Safetensors file from bytes
    88
        ///
    89
        /// # Arguments
    90
        ///
    91
        /// * `data` - Raw Safetensors file bytes
    92
        ///
    93
        /// # Errors
    94
        ///
    95
        /// Returns error if:
    96
        /// - Invalid header
    97
        /// - Malformed JSON metadata
    98
        /// - Invalid data offsets
    99
        ///
    100
        /// # Examples
    101
        ///
    102
        /// ```rust,ignore
    103
        /// let data = std::fs::read("model.safetensors")?;
    104
        /// let model = SafetensorsModel::from_bytes(&data)?;
    105
        /// println!("Loaded {} tensors", model.tensors.len());
    106
        /// ```
    107
    44
        pub fn from_bytes(data: &[u8]) -> Result<Self> {
    108
    44
            let mut cursor = Cursor::new(data);
    109
    110
            // Parse header (8-byte metadata length)
    111
    44
            let metadata_len37 = Self::parse_header(&mut cursor)?7;
    112
    113
            // Parse JSON metadata
    114
    37
            let tensors30 = Self::parse_metadata(&mut cursor, metadata_len)?7;
    115
    116
            // Store remaining data
    117
    30
            let data_start =
    118
    30
                usize::try_from(8 + metadata_len).map_err(|_| RealizarError::UnsupportedOperation {
    119
    0
                    operation: "convert_data_offset".to_string(),
    120
    0
                    reason: format!(
    121
    0
                        "Data offset {} exceeds platform usize limit",
    122
    0
                        8 + metadata_len
    123
                    ),
    124
    0
                })?;
    125
    30
            let data = data[data_start..].to_vec();
    126
    127
    30
            Ok(Self { tensors, data })
    128
    44
        }
    129
    130
        /// Extract F32 tensor data by name
    131
        ///
    132
        /// # Arguments
    133
        ///
    134
        /// * `name` - Tensor name to extract
    135
        ///
    136
        /// # Errors
    137
        ///
    138
        /// Returns error if:
    139
        /// - Tensor not found
    140
        /// - Tensor dtype is not F32
    141
        /// - Data offsets are invalid
    142
        ///
    143
        /// # Panics
    144
        ///
    145
        /// Never panics. The `unwrap()` in byte conversion is safe because
    146
        /// `chunks_exact(4)` guarantees exactly 4 bytes per chunk.
    147
        ///
    148
        /// # Examples
    149
        ///
    150
        /// ```rust,ignore
    151
        /// let model = SafetensorsModel::from_bytes(&data)?;
    152
        /// let weights = model.get_tensor_f32("layer1.weight")?;
    153
        /// println!("Weights: {:?}", weights);
    154
        /// ```
    155
    9
        pub fn get_tensor_f32(&self, name: &str) -> Result<Vec<f32>> {
    156
            // Find tensor metadata
    157
    9
            let tensor8 = self
    158
    9
                .tensors
    159
    9
                .get(name)
    160
    9
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    161
    1
                    operation: "get_tensor_f32".to_string(),
    162
    1
                    reason: format!("Tensor '{name}' not found"),
    163
    1
                })?;
    164
    165
            // Verify dtype is F32
    166
    8
            if tensor.dtype != SafetensorsDtype::F32 {
    167
    1
                let dtype = &tensor.dtype;
    168
    1
                return Err(RealizarError::UnsupportedOperation {
    169
    1
                    operation: "get_tensor_f32".to_string(),
    170
    1
                    reason: format!("Tensor '{name}' has dtype {dtype:?}, expected F32"),
    171
    1
                });
    172
    7
            }
    173
    174
            // Extract data slice
    175
    7
            let [start, end] = tensor.data_offsets;
    176
    7
            if end > self.data.len() {
    177
    1
                let data_len = self.data.len();
    178
    1
                return Err(RealizarError::UnsupportedOperation {
    179
    1
                    operation: "get_tensor_f32".to_string(),
    180
    1
                    reason: format!("Data offset {end} exceeds data size {data_len}"),
    181
    1
                });
    182
    6
            }
    183
    184
    6
            let bytes = &self.data[start..end];
    185
    186
            // Convert bytes to f32 vector
    187
    6
            if !bytes.len().is_multiple_of(4) {
    188
    1
                let len = bytes.len();
    189
    1
                return Err(RealizarError::UnsupportedOperation {
    190
    1
                    operation: "get_tensor_f32".to_string(),
    191
    1
                    reason: format!("Data size {len} is not a multiple of 4"),
    192
    1
                });
    193
    5
            }
    194
    195
    5
            let values = bytes
    196
    5
                .chunks_exact(4)
    197
    12
                .map5(|chunk| {
    198
    12
                    f32::from_le_bytes(
    199
    12
                        chunk
    200
    12
                            .try_into()
    201
    12
                            .expect("chunks_exact(4) guarantees 4-byte slices"),
    202
                    )
    203
    12
                })
    204
    5
                .collect();
    205
    206
    5
            Ok(values)
    207
    9
        }
    208
    209
        /// Parse header (8-byte metadata length)
    210
    44
        fn parse_header(cursor: &mut Cursor<&[u8]>) -> Result<u64> {
    211
    44
            let mut buf = [0u8; 8];
    212
    44
            cursor
    213
    44
                .read_exact(&mut buf)
    214
    44
                .map_err(|e| RealizarError::UnsupportedOperation {
    215
    7
                    operation: "read_metadata_len".to_string(),
    216
    7
                    reason: e.to_string(),
    217
    7
                })?;
    218
    219
    37
            Ok(u64::from_le_bytes(buf))
    220
    44
        }
    221
    222
        /// Parse JSON metadata
    223
    37
        fn parse_metadata(
    224
    37
            cursor: &mut Cursor<&[u8]>,
    225
    37
            len: u64,
    226
    37
        ) -> Result<HashMap<String, SafetensorsTensorInfo>> {
    227
            // Read JSON bytes
    228
    37
            let len_usize = usize::try_from(len).map_err(|_| RealizarError::UnsupportedOperation {
    229
    0
                operation: "convert_metadata_len".to_string(),
    230
    0
                reason: format!("Metadata length {len} exceeds platform usize limit"),
    231
    0
            })?;
    232
    233
    37
            let mut json_bytes = vec![0u8; len_usize];
    234
    37
            cursor
    235
    37
                .read_exact(&mut json_bytes)
    236
    37
                .map_err(|e| RealizarError::UnsupportedOperation {
    237
    2
                    operation: "read_metadata_json".to_string(),
    238
    2
                    reason: e.to_string(),
    239
    2
                })?;
    240
    241
            // Parse JSON as generic Value first to handle __metadata__ and other special keys
    242
    35
            let json_value32: serde_json::Value32 = serde_json::from_slice(&json_bytes).map_err(|e| {3
    243
    3
                RealizarError::UnsupportedOperation {
    244
    3
                    operation: "parse_json".to_string(),
    245
    3
                    reason: e.to_string(),
    246
    3
                }
    247
    3
            })?;
    248
    249
    31
            let json_map =
    250
    32
                json_value
    251
    32
                    .as_object()
    252
    32
                    .ok_or_else(|| RealizarError::UnsupportedOperation {
    253
    1
                        operation: "parse_json".to_string(),
    254
    1
                        reason: "Expected JSON object".to_string(),
    255
    1
                    })?;
    256
    257
            // Convert to SafetensorsTensorInfo, skipping special keys like __metadata__
    258
    31
            let mut tensors = HashMap::new();
    259
    77
            for (name47, value47) in json_map {
    260
                // Skip metadata keys (start with __)
    261
    47
                if name.starts_with("__") {
    262
    1
                    continue;
    263
    46
                }
    264
    265
                // Parse tensor metadata
    266
    46
                let meta45: TensorMetadata45 = serde_json::from_value(value.clone()).map_err(|e| {1
    267
    1
                    RealizarError::UnsupportedOperation {
    268
    1
                        operation: "parse_tensor_metadata".to_string(),
    269
    1
                        reason: format!("Failed to parse tensor '{name}': {e}"),
    270
    1
                    }
    271
    1
                })?;
    272
    273
    45
                tensors.insert(
    274
    45
                    name.clone(),
    275
    45
                    SafetensorsTensorInfo {
    276
    45
                        name: name.clone(),
    277
    45
                        dtype: meta.dtype,
    278
    45
                        shape: meta.shape,
    279
    45
                        data_offsets: meta.data_offsets,
    280
    45
                    },
    281
                );
    282
            }
    283
    284
    30
            Ok(tensors)
    285
    37
        }
    286
    287
        /// Get tensor data as F16 values (converts to F32)
    288
        ///
    289
        /// # Arguments
    290
        ///
    291
        /// * `name` - Tensor name to extract
    292
        ///
    293
        /// # Errors
    294
        ///
    295
        /// Returns error if tensor not found or dtype is not F16
    296
    5
        pub fn get_tensor_f16_as_f32(&self, name: &str) -> Result<Vec<f32>> {
    297
    5
            let tensor4 = self
    298
    5
                .tensors
    299
    5
                .get(name)
    300
    5
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    301
    1
                    operation: "get_tensor_f16_as_f32".to_string(),
    302
    1
                    reason: format!("Tensor '{name}' not found"),
    303
    1
                })?;
    304
    305
    4
            if tensor.dtype != SafetensorsDtype::F16 {
    306
    1
                let dtype = &tensor.dtype;
    307
    1
                return Err(RealizarError::UnsupportedOperation {
    308
    1
                    operation: "get_tensor_f16_as_f32".to_string(),
    309
    1
                    reason: format!("Tensor '{name}' has dtype {dtype:?}, expected F16"),
    310
    1
                });
    311
    3
            }
    312
    313
    3
            let [start, end] = tensor.data_offsets;
    314
    3
            if end > self.data.len() {
    315
    1
                let data_len = self.data.len();
    316
    1
                return Err(RealizarError::UnsupportedOperation {
    317
    1
                    operation: "get_tensor_f16_as_f32".to_string(),
    318
    1
                    reason: format!("Data offset {end} exceeds data size {data_len}"),
    319
    1
                });
    320
    2
            }
    321
    322
    2
            let bytes = &self.data[start..end];
    323
    324
            // Convert F16 bytes to F32
    325
    2
            let values: Vec<f32> = bytes
    326
    2
                .chunks_exact(2)
    327
    4
                .map2(|chunk| {
    328
    4
                    let bits = u16::from_le_bytes([chunk[0], chunk[1]]);
    329
    4
                    half::f16::from_bits(bits).to_f32()
    330
    4
                })
    331
    2
                .collect();
    332
    333
    2
            Ok(values)
    334
    5
        }
    335
    336
        /// Get tensor data as BF16 values (converts to F32)
    337
        ///
    338
        /// # Arguments
    339
        ///
    340
        /// * `name` - Tensor name to extract
    341
        ///
    342
        /// # Errors
    343
        ///
    344
        /// Returns error if tensor not found or dtype is not BF16
    345
    5
        pub fn get_tensor_bf16_as_f32(&self, name: &str) -> Result<Vec<f32>> {
    346
    5
            let tensor4 = self
    347
    5
                .tensors
    348
    5
                .get(name)
    349
    5
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    350
    1
                    operation: "get_tensor_bf16_as_f32".to_string(),
    351
    1
                    reason: format!("Tensor '{name}' not found"),
    352
    1
                })?;
    353
    354
    4
            if tensor.dtype != SafetensorsDtype::BF16 {
    355
    1
                let dtype = &tensor.dtype;
    356
    1
                return Err(RealizarError::UnsupportedOperation {
    357
    1
                    operation: "get_tensor_bf16_as_f32".to_string(),
    358
    1
                    reason: format!("Tensor '{name}' has dtype {dtype:?}, expected BF16"),
    359
    1
                });
    360
    3
            }
    361
    362
    3
            let [start, end] = tensor.data_offsets;
    363
    3
            if end > self.data.len() {
    364
    1
                let data_len = self.data.len();
    365
    1
                return Err(RealizarError::UnsupportedOperation {
    366
    1
                    operation: "get_tensor_bf16_as_f32".to_string(),
    367
    1
                    reason: format!("Data offset {end} exceeds data size {data_len}"),
    368
    1
                });
    369
    2
            }
    370
    371
    2
            let bytes = &self.data[start..end];
    372
    373
            // Convert BF16 bytes to F32 using SIMD-accelerated conversion
    374
            // This provides 3-4x speedup over scalar conversion
    375
    2
            let values = simd_bf16_to_f32(bytes);
    376
    377
    2
            Ok(values)
    378
    5
        }
    379
    380
        /// Get tensor as F32 with automatic dtype conversion
    381
        ///
    382
        /// Supports F32, F16, and BF16 dtypes with automatic conversion to F32.
    383
        ///
    384
        /// # Arguments
    385
        ///
    386
        /// * `name` - Tensor name to extract
    387
        ///
    388
        /// # Errors
    389
        ///
    390
        /// Returns error if tensor not found or dtype is not supported
    391
    5
        pub fn get_tensor_auto(&self, name: &str) -> Result<Vec<f32>> {
    392
    5
            let tensor4 = self
    393
    5
                .tensors
    394
    5
                .get(name)
    395
    5
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    396
    1
                    operation: "get_tensor_auto".to_string(),
    397
    1
                    reason: format!("Tensor '{name}' not found"),
    398
    1
                })?;
    399
    400
    4
            match tensor.dtype {
    401
    1
                SafetensorsDtype::F32 => self.get_tensor_f32(name),
    402
    1
                SafetensorsDtype::F16 => self.get_tensor_f16_as_f32(name),
    403
    1
                SafetensorsDtype::BF16 => self.get_tensor_bf16_as_f32(name),
    404
    1
                _ => Err(RealizarError::UnsupportedOperation {
    405
    1
                    operation: "get_tensor_auto".to_string(),
    406
    1
                    reason: format!("Unsupported dtype {:?} for tensor '{name}'", tensor.dtype),
    407
    1
                }),
    408
            }
    409
    5
        }
    410
    411
        /// Get list of tensor names
    412
        #[must_use]
    413
    1
        pub fn tensor_names(&self) -> Vec<&str> {
    414
    1
            self.tensors.keys().map(String::as_str).collect()
    415
    1
        }
    416
    417
        /// Get tensor info by name
    418
        #[must_use]
    419
    2
        pub fn get_tensor_info(&self, name: &str) -> Option<&SafetensorsTensorInfo> {
    420
    2
            self.tensors.get(name)
    421
    2
        }
    422
    423
        /// Check if model has a tensor with given name
    424
        #[must_use]
    425
    2
        pub fn has_tensor(&self, name: &str) -> bool {
    426
    2
            self.tensors.contains_key(name)
    427
    2
        }
    428
    }
    429
    430
    // ============================================================================
    431
    // SafeTensors Config loader (for sibling config.json)
    432
    // ============================================================================
    433
    434
    /// Model configuration from config.json
    435
    #[derive(Debug, Clone, Deserialize)]
    436
    pub struct SafetensorsConfig {
    437
        /// Hidden dimension
    438
        #[serde(alias = "n_embd", alias = "d_model")]
    439
        pub hidden_size: Option<usize>,
    440
        /// Number of transformer layers
    441
        #[serde(alias = "n_layer", alias = "num_layers")]
    442
        pub num_hidden_layers: Option<usize>,
    443
        /// Number of attention heads
    444
        #[serde(alias = "n_head")]
    445
        pub num_attention_heads: Option<usize>,
    446
        /// Number of key-value heads (for GQA)
    447
        pub num_key_value_heads: Option<usize>,
    448
        /// Vocabulary size
    449
        pub vocab_size: Option<usize>,
    450
        /// Intermediate/FFN dimension
    451
        #[serde(alias = "n_inner")]
    452
        pub intermediate_size: Option<usize>,
    453
        /// Maximum sequence length
    454
        #[serde(alias = "n_positions", alias = "n_ctx")]
    455
        pub max_position_embeddings: Option<usize>,
    456
        /// RMSNorm epsilon
    457
        pub rms_norm_eps: Option<f32>,
    458
        /// RoPE theta
    459
        pub rope_theta: Option<f32>,
    460
        /// Model architecture name
    461
        pub architectures: Option<Vec<String>>,
    462
        /// Model type
    463
        pub model_type: Option<String>,
    464
        /// BOS token ID
    465
        pub bos_token_id: Option<u32>,
    466
        /// EOS token ID
    467
        pub eos_token_id: Option<u32>,
    468
    }
    469
    470
    impl SafetensorsConfig {
    471
        /// Load config from sibling config.json file
    472
        ///
    473
        /// # Arguments
    474
        ///
    475
        /// * `model_path` - Path to the model file (config.json will be loaded from same directory)
    476
        ///
    477
        /// # Returns
    478
        ///
    479
        /// Config if found and parsed, None otherwise
    480
    19
        pub fn load_from_sibling(model_path: &std::path::Path) -> Option<Self> {
    481
    19
            let config_path = model_path.with_file_name("config.json");
    482
    19
            if !config_path.exists() {
    483
    3
                return None;
    484
    16
            }
    485
    486
    16
            let content = std::fs::read_to_string(&config_path).ok()?0;
    487
    16
            serde_json::from_str(&content).ok()
    488
    19
        }
    489
    490
        /// Get number of key-value heads (defaults to num_attention_heads for MHA)
    491
        #[must_use]
    492
    16
        pub fn num_kv_heads(&self) -> usize {
    493
    16
            self.num_key_value_heads
    494
    16
                .or(self.num_attention_heads)
    495
    16
                .unwrap_or(1)
    496
    16
        }
    497
    498
        /// Get model architecture string
    499
        #[must_use]
    500
    15
        pub fn architecture(&self) -> String {
    501
    15
            self.architectures
    502
    15
                .as_ref()
    503
    15
                .and_then(|a| a6.first6())
    504
    15
                .cloned()
    505
    15
                .or_else(|| self.model_type9.clone9())
    506
    15
                .unwrap_or_else(|| "unknown"8.to_string8())
    507
    15
        }
    508
    }
    509
    510
    // ============================================================================
    511
    // Zero-Copy Memory-Mapped SafeTensors Model (T-QA-020)
    512
    // ============================================================================
    513
    514
    /// Zero-copy memory-mapped SafeTensors model container
    515
    ///
    516
    /// Unlike `SafetensorsModel` which copies all tensor data to the heap,
    517
    /// `MappedSafeTensorsModel` uses memory-mapping (mmap) for true zero-copy
    518
    /// access to tensor data. This is critical for fast model loading (TTFT).
    519
    ///
    520
    /// # Performance Characteristics
    521
    ///
    522
    /// - **Loading time**: O(1) - only parses header/metadata, no data copy
    523
    /// - **Memory**: Only RSS grows as pages are accessed (demand paging)
    524
    /// - **TTFT target**: < 500ms for 3GB model
    525
    ///
    526
    /// # Example
    527
    ///
    528
    /// ```rust,ignore
    529
    /// let model = MappedSafeTensorsModel::load("model.safetensors")?;
    530
    /// let weights = model.get_tensor_bytes("layer1.weight")?;
    531
    /// // weights is a zero-copy slice into the mmap'd file
    532
    /// ```
    533
    #[cfg(not(target_arch = "wasm32"))]
    534
    #[derive(Debug)]
    535
    pub struct MappedSafeTensorsModel {
    536
        /// Memory-mapped file data
    537
        mmap: memmap2::Mmap,
    538
        /// File path (for diagnostics)
    539
        path: std::path::PathBuf,
    540
        /// Tensor metadata (parsed from header)
    541
        tensors: HashMap<String, SafetensorsTensorInfo>,
    542
        /// Offset where tensor data begins (after header + JSON metadata)
    543
        data_offset: usize,
    544
    }
    545
    546
    #[cfg(not(target_arch = "wasm32"))]
    547
    impl MappedSafeTensorsModel {
    548
        /// Load a SafeTensors file with zero-copy memory mapping
    549
        ///
    550
        /// # Arguments
    551
        ///
    552
        /// * `path` - Path to the SafeTensors file
    553
        ///
    554
        /// # Errors
    555
        ///
    556
        /// Returns error if:
    557
        /// - File cannot be opened
    558
        /// - Memory mapping fails
    559
        /// - Header/metadata parsing fails
    560
        ///
    561
        /// # Performance
    562
        ///
    563
        /// This method is O(1) with respect to file size - only the header
    564
        /// and JSON metadata are parsed. Tensor data is not touched until
    565
        /// `get_tensor_bytes()` is called.
    566
    23
        pub fn load<P: AsRef<std::path::Path>>(path: P) -> Result<Self> {
    567
    23
            let path = path.as_ref().to_path_buf();
    568
    569
            // Open file
    570
    23
            let file22 = std::fs::File::open(&path).map_err(|e| RealizarError::UnsupportedOperation {
    571
    1
                operation: "open_safetensors".to_string(),
    572
    1
                reason: format!("Failed to open file '{}': {}", path.display(), e),
    573
    1
            })?;
    574
    575
            // Memory-map the file (zero-copy)
    576
            // SAFETY: File is opened read-only and we don't modify it
    577
    22
            let mmap = unsafe {
    578
    22
                memmap2::MmapOptions::new().map(&file).map_err(|e| {0
    579
    0
                    RealizarError::UnsupportedOperation {
    580
    0
                        operation: "mmap_safetensors".to_string(),
    581
    0
                        reason: format!("Failed to mmap file '{}': {}", path.display(), e),
    582
    0
                    }
    583
    0
                })?
    584
            };
    585
    586
            // Parse header (8-byte metadata length)
    587
    22
            if mmap.len() < 8 {
    588
    0
                return Err(RealizarError::UnsupportedOperation {
    589
    0
                    operation: "parse_safetensors_header".to_string(),
    590
    0
                    reason: format!(
    591
    0
                        "File too small: {} bytes (minimum 8 for header)",
    592
    0
                        mmap.len()
    593
    0
                    ),
    594
    0
                });
    595
    22
            }
    596
    597
    22
            let metadata_len =
    598
    22
                u64::from_le_bytes(mmap[0..8].try_into().expect("slice is exactly 8 bytes"));
    599
    600
    22
            let metadata_len_usize =
    601
    22
                usize::try_from(metadata_len).map_err(|_| RealizarError::UnsupportedOperation {
    602
    0
                    operation: "parse_safetensors_header".to_string(),
    603
    0
                    reason: format!("Metadata length {} exceeds platform limit", metadata_len),
    604
    0
                })?;
    605
    606
            // Verify we have enough data for metadata
    607
    22
            let data_offset = 8 + metadata_len_usize;
    608
    22
            if mmap.len() < data_offset {
    609
    1
                return Err(RealizarError::UnsupportedOperation {
    610
    1
                    operation: "parse_safetensors_header".to_string(),
    611
    1
                    reason: format!(
    612
    1
                        "File truncated: need {} bytes for header+metadata, have {}",
    613
    1
                        data_offset,
    614
    1
                        mmap.len()
    615
    1
                    ),
    616
    1
                });
    617
    21
            }
    618
    619
            // Parse JSON metadata (from mmap'd memory, no copy)
    620
    21
            let json_bytes = &mmap[8..data_offset];
    621
    21
            let tensors = Self::parse_metadata(json_bytes)?0;
    622
    623
    21
            Ok(Self {
    624
    21
                mmap,
    625
    21
                path,
    626
    21
                tensors,
    627
    21
                data_offset,
    628
    21
            })
    629
    23
        }
    630
    631
        /// Parse JSON metadata from bytes
    632
    21
        fn parse_metadata(json_bytes: &[u8]) -> Result<HashMap<String, SafetensorsTensorInfo>> {
    633
            // Parse JSON as generic Value first to handle __metadata__ and other special keys
    634
    21
            let json_value: serde_json::Value = serde_json::from_slice(json_bytes).map_err(|e| {0
    635
    0
                RealizarError::UnsupportedOperation {
    636
    0
                    operation: "parse_json".to_string(),
    637
    0
                    reason: e.to_string(),
    638
    0
                }
    639
    0
            })?;
    640
    641
    21
            let json_map =
    642
    21
                json_value
    643
    21
                    .as_object()
    644
    21
                    .ok_or_else(|| RealizarError::UnsupportedOperation {
    645
    0
                        operation: "parse_json".to_string(),
    646
    0
                        reason: "Expected JSON object".to_string(),
    647
    0
                    })?;
    648
    649
            // Convert to SafetensorsTensorInfo, skipping special keys like __metadata__
    650
    21
            let mut tensors = HashMap::new();
    651
    75
            for (name54, value54) in json_map {
    652
                // Skip metadata keys (start with __)
    653
    54
                if name.starts_with("__") {
    654
    0
                    continue;
    655
    54
                }
    656
    657
                // Parse tensor metadata
    658
    54
                let meta: TensorMetadata = serde_json::from_value(value.clone()).map_err(|e| {0
    659
    0
                    RealizarError::UnsupportedOperation {
    660
    0
                        operation: "parse_tensor_metadata".to_string(),
    661
    0
                        reason: format!("Failed to parse tensor '{name}': {e}"),
    662
    0
                    }
    663
    0
                })?;
    664
    665
    54
                tensors.insert(
    666
    54
                    name.clone(),
    667
    54
                    SafetensorsTensorInfo {
    668
    54
                        name: name.clone(),
    669
    54
                        dtype: meta.dtype,
    670
    54
                        shape: meta.shape,
    671
    54
                        data_offsets: meta.data_offsets,
    672
    54
                    },
    673
                );
    674
            }
    675
    676
    21
            Ok(tensors)
    677
    21
        }
    678
    679
        /// Get raw tensor bytes (zero-copy slice into mmap'd file)
    680
        ///
    681
        /// # Arguments
    682
        ///
    683
        /// * `name` - Tensor name
    684
        ///
    685
        /// # Returns
    686
        ///
    687
        /// Zero-copy slice into the memory-mapped file. The slice is valid
    688
        /// as long as `self` is alive.
    689
        ///
    690
        /// # Errors
    691
        ///
    692
        /// Returns error if tensor not found or offsets are invalid.
    693
    54
        pub fn get_tensor_bytes(&self, name: &str) -> Result<&[u8]> {
    694
    54
            let tensor = self
    695
    54
                .tensors
    696
    54
                .get(name)
    697
    54
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    698
    0
                    operation: "get_tensor_bytes".to_string(),
    699
    0
                    reason: format!("Tensor '{name}' not found"),
    700
    0
                })?;
    701
    702
    54
            let [start, end] = tensor.data_offsets;
    703
    54
            let abs_start = self.data_offset + start;
    704
    54
            let abs_end = self.data_offset + end;
    705
    706
    54
            if abs_end > self.mmap.len() {
    707
    0
                return Err(RealizarError::UnsupportedOperation {
    708
    0
                    operation: "get_tensor_bytes".to_string(),
    709
    0
                    reason: format!(
    710
    0
                        "Tensor '{}' data offsets [{}, {}] exceed file size {}",
    711
    0
                        name,
    712
    0
                        abs_start,
    713
    0
                        abs_end,
    714
    0
                        self.mmap.len()
    715
    0
                    ),
    716
    0
                });
    717
    54
            }
    718
    719
    54
            Ok(&self.mmap[abs_start..abs_end])
    720
    54
        }
    721
    722
        /// Get tensor as F32 values (zero-copy bytes, then convert)
    723
        ///
    724
        /// # Arguments
    725
        ///
    726
        /// * `name` - Tensor name
    727
        ///
    728
        /// # Errors
    729
        ///
    730
        /// Returns error if tensor not found or dtype is not F32.
    731
    54
        pub fn get_tensor_f32(&self, name: &str) -> Result<Vec<f32>> {
    732
    54
            let tensor = self
    733
    54
                .tensors
    734
    54
                .get(name)
    735
    54
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    736
    0
                    operation: "get_tensor_f32".to_string(),
    737
    0
                    reason: format!("Tensor '{name}' not found"),
    738
    0
                })?;
    739
    740
    54
            if tensor.dtype != SafetensorsDtype::F32 {
    741
    0
                return Err(RealizarError::UnsupportedOperation {
    742
    0
                    operation: "get_tensor_f32".to_string(),
    743
    0
                    reason: format!(
    744
    0
                        "Tensor '{}' has dtype {:?}, expected F32",
    745
    0
                        name, tensor.dtype
    746
    0
                    ),
    747
    0
                });
    748
    54
            }
    749
    750
    54
            let bytes = self.get_tensor_bytes(name)?0;
    751
    752
    54
            if !bytes.len().is_multiple_of(4) {
    753
    0
                return Err(RealizarError::UnsupportedOperation {
    754
    0
                    operation: "get_tensor_f32".to_string(),
    755
    0
                    reason: format!("Data size {} is not a multiple of 4", bytes.len()),
    756
    0
                });
    757
    54
            }
    758
    759
    54
            let values = bytes
    760
    54
                .chunks_exact(4)
    761
    66.1k
                .map54(|chunk| {
    762
    66.1k
                    f32::from_le_bytes(
    763
    66.1k
                        chunk
    764
    66.1k
                            .try_into()
    765
    66.1k
                            .expect("chunks_exact(4) guarantees 4-byte slices"),
    766
                    )
    767
    66.1k
                })
    768
    54
                .collect();
    769
    770
    54
            Ok(values)
    771
    54
        }
    772
    773
        /// Get tensor as BF16 bytes (zero-copy, native format)
    774
        ///
    775
        /// Returns raw BF16 bytes for native SIMD processing without
    776
        /// F32 conversion at boot time.
    777
        ///
    778
        /// # Arguments
    779
        ///
    780
        /// * `name` - Tensor name
    781
        ///
    782
        /// # Errors
    783
        ///
    784
        /// Returns error if tensor not found or dtype is not BF16.
    785
    0
        pub fn get_tensor_bf16_bytes(&self, name: &str) -> Result<&[u8]> {
    786
    0
            let tensor = self
    787
    0
                .tensors
    788
    0
                .get(name)
    789
    0
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    790
    0
                    operation: "get_tensor_bf16_bytes".to_string(),
    791
    0
                    reason: format!("Tensor '{name}' not found"),
    792
    0
                })?;
    793
    794
    0
            if tensor.dtype != SafetensorsDtype::BF16 {
    795
    0
                return Err(RealizarError::UnsupportedOperation {
    796
    0
                    operation: "get_tensor_bf16_bytes".to_string(),
    797
    0
                    reason: format!(
    798
    0
                        "Tensor '{}' has dtype {:?}, expected BF16",
    799
    0
                        name, tensor.dtype
    800
    0
                    ),
    801
    0
                });
    802
    0
            }
    803
    804
    0
            self.get_tensor_bytes(name)
    805
    0
        }
    806
    807
        /// Get tensor as BF16 values converted to F32
    808
        ///
    809
        /// # Arguments
    810
        ///
    811
        /// * `name` - Tensor name
    812
        ///
    813
        /// # Errors
    814
        ///
    815
        /// Returns error if tensor not found or dtype is not BF16.
    816
    0
        pub fn get_tensor_bf16_as_f32(&self, name: &str) -> Result<Vec<f32>> {
    817
    0
            let bytes = self.get_tensor_bf16_bytes(name)?;
    818
    819
            // Convert BF16 bytes to F32 using SIMD-accelerated conversion
    820
            // This provides 3-4x speedup over scalar conversion
    821
    0
            let values = simd_bf16_to_f32(bytes);
    822
    823
    0
            Ok(values)
    824
    0
        }
    825
    826
        /// Get tensor as F16 bytes (zero-copy, native format)
    827
        ///
    828
        /// Returns raw F16 bytes for native SIMD processing.
    829
        ///
    830
        /// # Arguments
    831
        ///
    832
        /// * `name` - Tensor name
    833
        ///
    834
        /// # Errors
    835
        ///
    836
        /// Returns error if tensor not found or dtype is not F16.
    837
    0
        pub fn get_tensor_f16_bytes(&self, name: &str) -> Result<&[u8]> {
    838
    0
            let tensor = self
    839
    0
                .tensors
    840
    0
                .get(name)
    841
    0
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    842
    0
                    operation: "get_tensor_f16_bytes".to_string(),
    843
    0
                    reason: format!("Tensor '{name}' not found"),
    844
    0
                })?;
    845
    846
    0
            if tensor.dtype != SafetensorsDtype::F16 {
    847
    0
                return Err(RealizarError::UnsupportedOperation {
    848
    0
                    operation: "get_tensor_f16_bytes".to_string(),
    849
    0
                    reason: format!(
    850
    0
                        "Tensor '{}' has dtype {:?}, expected F16",
    851
    0
                        name, tensor.dtype
    852
    0
                    ),
    853
    0
                });
    854
    0
            }
    855
    856
    0
            self.get_tensor_bytes(name)
    857
    0
        }
    858
    859
        /// Get tensor as F16 values converted to F32
    860
    0
        pub fn get_tensor_f16_as_f32(&self, name: &str) -> Result<Vec<f32>> {
    861
    0
            let bytes = self.get_tensor_f16_bytes(name)?;
    862
    863
    0
            let values: Vec<f32> = bytes
    864
    0
                .chunks_exact(2)
    865
    0
                .map(|chunk| {
    866
    0
                    let bits = u16::from_le_bytes([chunk[0], chunk[1]]);
    867
    0
                    half::f16::from_bits(bits).to_f32()
    868
    0
                })
    869
    0
                .collect();
    870
    871
    0
            Ok(values)
    872
    0
        }
    873
    874
        /// Get tensor as F32 with automatic dtype conversion
    875
        ///
    876
        /// Supports F32, F16, and BF16 dtypes with automatic conversion to F32.
    877
    61
        pub fn get_tensor_auto(&self, name: &str) -> Result<Vec<f32>> {
    878
    61
            let tensor54 = self
    879
    61
                .tensors
    880
    61
                .get(name)
    881
    61
                .ok_or_else(|| RealizarError::UnsupportedOperation {
    882
    7
                    operation: "get_tensor_auto".to_string(),
    883
    7
                    reason: format!("Tensor '{name}' not found"),
    884
    7
                })?;
    885
    886
    54
            match tensor.dtype {
    887
    54
                SafetensorsDtype::F32 => self.get_tensor_f32(name),
    888
    0
                SafetensorsDtype::F16 => self.get_tensor_f16_as_f32(name),
    889
    0
                SafetensorsDtype::BF16 => self.get_tensor_bf16_as_f32(name),
    890
    0
                _ => Err(RealizarError::UnsupportedOperation {
    891
    0
                    operation: "get_tensor_auto".to_string(),
    892
    0
                    reason: format!("Unsupported dtype {:?} for tensor '{}'", tensor.dtype, name),
    893
    0
                }),
    894
            }
    895
    61
        }
    896
    897
        /// Get list of tensor names
    898
        #[must_use]
    899
    0
        pub fn tensor_names(&self) -> Vec<&str> {
    900
    0
            self.tensors.keys().map(String::as_str).collect()
    901
    0
        }
    902
    903
        /// Get tensor info by name
    904
        #[must_use]
    905
    0
        pub fn get_tensor_info(&self, name: &str) -> Option<&SafetensorsTensorInfo> {
    906
    0
            self.tensors.get(name)
    907
    0
        }
    908
    909
        /// Check if model has a tensor with given name
    910
        #[must_use]
    911
    11
        pub fn has_tensor(&self, name: &str) -> bool {
    912
    11
            self.tensors.contains_key(name)
    913
    11
        }
    914
    915
        /// Get the file path
    916
        #[must_use]
    917
    0
        pub fn path(&self) -> &std::path::Path {
    918
    0
            &self.path
    919
    0
        }
    920
    921
        /// Get the total file size in bytes
    922
        #[must_use]
    923
    0
        pub fn file_size(&self) -> usize {
    924
    0
            self.mmap.len()
    925
    0
        }
    926
    927
        /// Get the number of tensors
    928
        #[must_use]
    929
    0
        pub fn tensor_count(&self) -> usize {
    930
    0
            self.tensors.len()
    931
    0
        }
    932
    }
    933
    934
    #[cfg(test)]
    935
    mod tests {
    936
        use super::*;
    937
    938
        #[test]
    939
    1
        fn test_parse_empty_safetensors() {
    940
            // Minimal valid Safetensors: 8-byte header + empty JSON "{}"
    941
    1
            let mut data = Vec::new();
    942
    1
            data.extend_from_slice(&2u64.to_le_bytes()); // metadata_len = 2
    943
    1
            data.extend_from_slice(b"{}"); // Empty JSON
    944
    945
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    946
    1
            assert_eq!(model.tensors.len(), 0);
    947
    1
            assert_eq!(model.data.len(), 0);
    948
    1
        }
    949
    950
        #[test]
    951
    1
        fn test_invalid_header_truncated() {
    952
            // Only 4 bytes (should be 8)
    953
    1
            let data = [0u8; 4];
    954
    1
            let result = SafetensorsModel::from_bytes(&data);
    955
    1
            assert!(result.is_err());
    956
    1
        }
    957
    958
        #[test]
    959
    1
        fn test_empty_file() {
    960
    1
            let data = &[];
    961
    1
            let result = SafetensorsModel::from_bytes(data);
    962
    1
            assert!(result.is_err());
    963
    1
        }
    964
    965
        #[test]
    966
    1
        fn test_parse_single_tensor() {
    967
            // Safetensors with one F32 tensor
    968
    1
            let json = r#"{"weight":{"dtype":"F32","shape":[2,3],"data_offsets":[0,24]}}"#;
    969
    1
            let json_bytes = json.as_bytes();
    970
    971
    1
            let mut data = Vec::new();
    972
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    973
    1
            data.extend_from_slice(json_bytes);
    974
            // Add 24 bytes of dummy tensor data (2*3*4 = 24 bytes for F32)
    975
    1
            data.extend_from_slice(&[0u8; 24]);
    976
    977
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    978
    1
            assert_eq!(model.tensors.len(), 1);
    979
    980
    1
            let tensor = model.tensors.get("weight").expect("test");
    981
    1
            assert_eq!(tensor.name, "weight");
    982
    1
            assert_eq!(tensor.dtype, SafetensorsDtype::F32);
    983
    1
            assert_eq!(tensor.shape, vec![2, 3]);
    984
    1
            assert_eq!(tensor.data_offsets, [0, 24]);
    985
    1
        }
    986
    987
        #[test]
    988
    1
        fn test_parse_multiple_tensors() {
    989
            // Safetensors with multiple tensors of different types
    990
    1
            let json = r#"{
    991
    1
                "layer1.weight":{"dtype":"F32","shape":[128,256],"data_offsets":[0,131072]},
    992
    1
                "layer1.bias":{"dtype":"F32","shape":[128],"data_offsets":[131072,131584]}
    993
    1
            }"#;
    994
    1
            let json_bytes = json.as_bytes();
    995
    996
    1
            let mut data = Vec::new();
    997
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    998
    1
            data.extend_from_slice(json_bytes);
    999
            // Add dummy tensor data
    1000
    1
            data.extend_from_slice(&vec![0u8; 131_584]);
    1001
    1002
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1003
    1
            assert_eq!(model.tensors.len(), 2);
    1004
    1005
    1
            let weight = model.tensors.get("layer1.weight").expect("test");
    1006
    1
            assert_eq!(weight.dtype, SafetensorsDtype::F32);
    1007
    1
            assert_eq!(weight.shape, vec![128, 256]);
    1008
    1
            assert_eq!(weight.data_offsets, [0, 131_072]);
    1009
    1010
    1
            let bias = model.tensors.get("layer1.bias").expect("test");
    1011
    1
            assert_eq!(bias.dtype, SafetensorsDtype::F32);
    1012
    1
            assert_eq!(bias.shape, vec![128]);
    1013
    1
            assert_eq!(bias.data_offsets, [131_072, 131_584]);
    1014
    1
        }
    1015
    1016
        #[test]
    1017
    1
        fn test_parse_various_dtypes() {
    1018
            // Test different data types
    1019
    1
            let json = r#"{
    1020
    1
                "f32_tensor":{"dtype":"F32","shape":[2],"data_offsets":[0,8]},
    1021
    1
                "i32_tensor":{"dtype":"I32","shape":[2],"data_offsets":[8,16]},
    1022
    1
                "u8_tensor":{"dtype":"U8","shape":[4],"data_offsets":[16,20]}
    1023
    1
            }"#;
    1024
    1
            let json_bytes = json.as_bytes();
    1025
    1026
    1
            let mut data = Vec::new();
    1027
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1028
    1
            data.extend_from_slice(json_bytes);
    1029
    1
            data.extend_from_slice(&[0u8; 20]);
    1030
    1031
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1032
    1
            assert_eq!(model.tensors.len(), 3);
    1033
    1034
    1
            assert_eq!(
    1035
    1
                model.tensors.get("f32_tensor").expect("test").dtype,
    1036
                SafetensorsDtype::F32
    1037
            );
    1038
    1
            assert_eq!(
    1039
    1
                model.tensors.get("i32_tensor").expect("test").dtype,
    1040
                SafetensorsDtype::I32
    1041
            );
    1042
    1
            assert_eq!(
    1043
    1
                model.tensors.get("u8_tensor").expect("test").dtype,
    1044
                SafetensorsDtype::U8
    1045
            );
    1046
    1
        }
    1047
    1048
        #[test]
    1049
    1
        fn test_invalid_json_error() {
    1050
            // Invalid JSON in metadata
    1051
    1
            let mut data = Vec::new();
    1052
    1
            data.extend_from_slice(&10u64.to_le_bytes()); // metadata_len = 10
    1053
    1
            data.extend_from_slice(b"not json!!"); // Invalid JSON
    1054
    1055
    1
            let result = SafetensorsModel::from_bytes(&data);
    1056
    1
            assert!(result.is_err());
    1057
    1
            assert!(matches!0(
    1058
    1
                result.unwrap_err(),
    1059
                RealizarError::UnsupportedOperation { .. }
    1060
            ));
    1061
    1
        }
    1062
    1063
        #[test]
    1064
    1
        fn test_truncated_json_error() {
    1065
            // Header says JSON is longer than actual data
    1066
    1
            let mut data = Vec::new();
    1067
    1
            data.extend_from_slice(&100u64.to_le_bytes()); // metadata_len = 100
    1068
    1
            data.extend_from_slice(b"{}"); // Only 2 bytes, not 100
    1069
    1070
    1
            let result = SafetensorsModel::from_bytes(&data);
    1071
    1
            assert!(result.is_err());
    1072
    1
            assert!(matches!0(
    1073
    1
                result.unwrap_err(),
    1074
                RealizarError::UnsupportedOperation { .. }
    1075
            ));
    1076
    1
        }
    1077
    1078
        #[test]
    1079
    1
        fn test_parse_all_dtypes() {
    1080
            // Test all supported data types
    1081
    1
            let json = r#"{
    1082
    1
                "f32":{"dtype":"F32","shape":[1],"data_offsets":[0,4]},
    1083
    1
                "f16":{"dtype":"F16","shape":[1],"data_offsets":[4,6]},
    1084
    1
                "bf16":{"dtype":"BF16","shape":[1],"data_offsets":[6,8]},
    1085
    1
                "i32":{"dtype":"I32","shape":[1],"data_offsets":[8,12]},
    1086
    1
                "i64":{"dtype":"I64","shape":[1],"data_offsets":[12,20]},
    1087
    1
                "u8":{"dtype":"U8","shape":[1],"data_offsets":[20,21]},
    1088
    1
                "bool":{"dtype":"Bool","shape":[1],"data_offsets":[21,22]}
    1089
    1
            }"#;
    1090
    1
            let json_bytes = json.as_bytes();
    1091
    1092
    1
            let mut data = Vec::new();
    1093
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1094
    1
            data.extend_from_slice(json_bytes);
    1095
    1
            data.extend_from_slice(&[0u8; 22]);
    1096
    1097
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1098
    1
            assert_eq!(model.tensors.len(), 7);
    1099
    1100
    1
            assert_eq!(
    1101
    1
                model.tensors.get("f32").expect("test").dtype,
    1102
                SafetensorsDtype::F32
    1103
            );
    1104
    1
            assert_eq!(
    1105
    1
                model.tensors.get("f16").expect("test").dtype,
    1106
                SafetensorsDtype::F16
    1107
            );
    1108
    1
            assert_eq!(
    1109
    1
                model.tensors.get("bf16").expect("test").dtype,
    1110
                SafetensorsDtype::BF16
    1111
            );
    1112
    1
            assert_eq!(
    1113
    1
                model.tensors.get("i32").expect("test").dtype,
    1114
                SafetensorsDtype::I32
    1115
            );
    1116
    1
            assert_eq!(
    1117
    1
                model.tensors.get("i64").expect("test").dtype,
    1118
                SafetensorsDtype::I64
    1119
            );
    1120
    1
            assert_eq!(
    1121
    1
                model.tensors.get("u8").expect("test").dtype,
    1122
                SafetensorsDtype::U8
    1123
            );
    1124
    1
            assert_eq!(
    1125
    1
                model.tensors.get("bool").expect("test").dtype,
    1126
                SafetensorsDtype::Bool
    1127
            );
    1128
    1
        }
    1129
    1130
        #[test]
    1131
    1
        fn test_tensor_data_preserved() {
    1132
            // Verify tensor data is correctly preserved
    1133
    1
            let json = r#"{"weight":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
    1134
    1
            let json_bytes = json.as_bytes();
    1135
    1136
    1
            let mut data = Vec::new();
    1137
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1138
    1
            data.extend_from_slice(json_bytes);
    1139
            // Add specific tensor data (two f32s: 1.0 and 2.0)
    1140
    1
            data.extend_from_slice(&1.0f32.to_le_bytes());
    1141
    1
            data.extend_from_slice(&2.0f32.to_le_bytes());
    1142
    1143
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1144
    1
            assert_eq!(model.data.len(), 8);
    1145
    1146
            // Verify we can read back the f32 values
    1147
    1
            let val1 = f32::from_le_bytes(model.data[0..4].try_into().expect("test"));
    1148
    1
            let val2 = f32::from_le_bytes(model.data[4..8].try_into().expect("test"));
    1149
    1
            assert!((val1 - 1.0).abs() < 1e-6);
    1150
    1
            assert!((val2 - 2.0).abs() < 1e-6);
    1151
    1
        }
    1152
    1153
        #[test]
    1154
    1
        fn test_multidimensional_shapes() {
    1155
            // Test tensors with various shapes
    1156
    1
            let json = r#"{
    1157
    1
                "scalar":{"dtype":"F32","shape":[],"data_offsets":[0,4]},
    1158
    1
                "vector":{"dtype":"F32","shape":[10],"data_offsets":[4,44]},
    1159
    1
                "matrix":{"dtype":"F32","shape":[3,4],"data_offsets":[44,92]},
    1160
    1
                "tensor3d":{"dtype":"F32","shape":[2,3,4],"data_offsets":[92,188]}
    1161
    1
            }"#;
    1162
    1
            let json_bytes = json.as_bytes();
    1163
    1164
    1
            let mut data = Vec::new();
    1165
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1166
    1
            data.extend_from_slice(json_bytes);
    1167
    1
            data.extend_from_slice(&[0u8; 188]);
    1168
    1169
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1170
    1
            assert_eq!(model.tensors.len(), 4);
    1171
    1172
    1
            assert_eq!(
    1173
    1
                model.tensors.get("scalar").expect("test").shape,
    1174
    1
                Vec::<usize>::new()
    1175
            );
    1176
    1
            assert_eq!(model.tensors.get("vector").expect("test").shape, vec![10]);
    1177
    1
            assert_eq!(model.tensors.get("matrix").expect("test").shape, vec![3, 4]);
    1178
    1
            assert_eq!(
    1179
    1
                model.tensors.get("tensor3d").expect("test").shape,
    1180
    1
                vec![2, 3, 4]
    1181
            );
    1182
    1
        }
    1183
    1184
        #[test]
    1185
    1
        fn test_aprender_linear_regression_format_compatibility() {
    1186
            // Test aprender LinearRegression SafeTensors format compatibility
    1187
            // Format: {"coefficients": [n_features], "intercept": [1]}
    1188
            // Example model: y = 2.0*x1 + 3.0*x2 + 1.5*x3 + 0.5
    1189
    1190
    1
            let json = r#"{
    1191
    1
                "coefficients":{"dtype":"F32","shape":[3],"data_offsets":[0,12]},
    1192
    1
                "intercept":{"dtype":"F32","shape":[1],"data_offsets":[12,16]}
    1193
    1
            }"#;
    1194
    1
            let json_bytes = json.as_bytes();
    1195
    1196
    1
            let mut data = Vec::new();
    1197
    1198
            // Header
    1199
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1200
    1201
            // Metadata
    1202
    1
            data.extend_from_slice(json_bytes);
    1203
    1204
            // Tensor data: coefficients [2.0, 3.0, 1.5]
    1205
    1
            data.extend_from_slice(&2.0f32.to_le_bytes());
    1206
    1
            data.extend_from_slice(&3.0f32.to_le_bytes());
    1207
    1
            data.extend_from_slice(&1.5f32.to_le_bytes());
    1208
    1209
            // intercept [0.5]
    1210
    1
            data.extend_from_slice(&0.5f32.to_le_bytes());
    1211
    1212
            // Parse with realizar
    1213
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1214
    1215
            // Verify structure
    1216
    1
            assert_eq!(model.tensors.len(), 2);
    1217
    1218
            // Check coefficients tensor
    1219
    1
            let coef = model.tensors.get("coefficients").expect("test");
    1220
    1
            assert_eq!(coef.dtype, SafetensorsDtype::F32);
    1221
    1
            assert_eq!(coef.shape, vec![3]);
    1222
    1
            assert_eq!(coef.data_offsets, [0, 12]);
    1223
    1224
            // Check intercept tensor
    1225
    1
            let intercept = model.tensors.get("intercept").expect("test");
    1226
    1
            assert_eq!(intercept.dtype, SafetensorsDtype::F32);
    1227
    1
            assert_eq!(intercept.shape, vec![1]);
    1228
    1
            assert_eq!(intercept.data_offsets, [12, 16]);
    1229
    1230
            // Verify we can extract the actual values
    1231
    1
            let coef_vals: Vec<f32> = (0..3)
    1232
    3
                .map1(|i| {
    1233
    3
                    let offset = i * 4;
    1234
    3
                    f32::from_le_bytes(model.data[offset..offset + 4].try_into().expect("test"))
    1235
    3
                })
    1236
    1
                .collect();
    1237
    1
            assert!((coef_vals[0] - 2.0).abs() < 1e-6);
    1238
    1
            assert!((coef_vals[1] - 3.0).abs() < 1e-6);
    1239
    1
            assert!((coef_vals[2] - 1.5).abs() < 1e-6);
    1240
    1241
    1
            let intercept_val = f32::from_le_bytes(model.data[12..16].try_into().expect("test"));
    1242
    1
            assert!((intercept_val - 0.5).abs() < 1e-6);
    1243
    1
        }
    1244
    1245
        #[test]
    1246
    1
        fn test_get_tensor_f32_helper() {
    1247
            // Test the get_tensor_f32 helper method
    1248
    1
            let json = r#"{
    1249
    1
                "weights":{"dtype":"F32","shape":[4],"data_offsets":[0,16]},
    1250
    1
                "bias":{"dtype":"F32","shape":[2],"data_offsets":[16,24]}
    1251
    1
            }"#;
    1252
    1
            let json_bytes = json.as_bytes();
    1253
    1254
    1
            let mut data = Vec::new();
    1255
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1256
    1
            data.extend_from_slice(json_bytes);
    1257
    1258
            // weights: [1.0, 2.0, 3.0, 4.0]
    1259
    1
            data.extend_from_slice(&1.0f32.to_le_bytes());
    1260
    1
            data.extend_from_slice(&2.0f32.to_le_bytes());
    1261
    1
            data.extend_from_slice(&3.0f32.to_le_bytes());
    1262
    1
            data.extend_from_slice(&4.0f32.to_le_bytes());
    1263
    1264
            // bias: [0.5, 0.25]
    1265
    1
            data.extend_from_slice(&0.5f32.to_le_bytes());
    1266
    1
            data.extend_from_slice(&0.25f32.to_le_bytes());
    1267
    1268
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1269
    1270
            // Test extracting weights
    1271
    1
            let weights = model.get_tensor_f32("weights").expect("test");
    1272
    1
            assert_eq!(weights.len(), 4);
    1273
    1
            assert!((weights[0] - 1.0).abs() < 1e-6);
    1274
    1
            assert!((weights[1] - 2.0).abs() < 1e-6);
    1275
    1
            assert!((weights[2] - 3.0).abs() < 1e-6);
    1276
    1
            assert!((weights[3] - 4.0).abs() < 1e-6);
    1277
    1278
            // Test extracting bias
    1279
    1
            let bias = model.get_tensor_f32("bias").expect("test");
    1280
    1
            assert_eq!(bias.len(), 2);
    1281
    1
            assert!((bias[0] - 0.5).abs() < 1e-6);
    1282
    1
            assert!((bias[1] - 0.25).abs() < 1e-6);
    1283
    1284
            // Test error: tensor not found
    1285
    1
            let result = model.get_tensor_f32("nonexistent");
    1286
    1
            assert!(result.is_err());
    1287
    1
        }
    1288
    1289
        #[test]
    1290
    1
        fn test_get_tensor_f32_wrong_dtype() {
    1291
            // Test error when tensor has wrong dtype
    1292
    1
            let json = r#"{
    1293
    1
                "int_tensor":{"dtype":"I32","shape":[2],"data_offsets":[0,8]}
    1294
    1
            }"#;
    1295
    1
            let json_bytes = json.as_bytes();
    1296
    1297
    1
            let mut data = Vec::new();
    1298
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1299
    1
            data.extend_from_slice(json_bytes);
    1300
    1
            data.extend_from_slice(&1i32.to_le_bytes());
    1301
    1
            data.extend_from_slice(&2i32.to_le_bytes());
    1302
    1303
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1304
    1305
            // Should error because dtype is I32, not F32
    1306
    1
            let result = model.get_tensor_f32("int_tensor");
    1307
    1
            assert!(result.is_err());
    1308
    1
        }
    1309
    1310
        #[test]
    1311
    1
        fn test_get_tensor_f32_with_aprender_model() {
    1312
            // Use get_tensor_f32 with aprender LinearRegression format
    1313
    1
            let json = r#"{
    1314
    1
                "coefficients":{"dtype":"F32","shape":[3],"data_offsets":[0,12]},
    1315
    1
                "intercept":{"dtype":"F32","shape":[1],"data_offsets":[12,16]}
    1316
    1
            }"#;
    1317
    1
            let json_bytes = json.as_bytes();
    1318
    1319
    1
            let mut data = Vec::new();
    1320
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1321
    1
            data.extend_from_slice(json_bytes);
    1322
    1
            data.extend_from_slice(&2.0f32.to_le_bytes());
    1323
    1
            data.extend_from_slice(&3.0f32.to_le_bytes());
    1324
    1
            data.extend_from_slice(&1.5f32.to_le_bytes());
    1325
    1
            data.extend_from_slice(&0.5f32.to_le_bytes());
    1326
    1327
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1328
    1329
            // Extract using helper method - much cleaner!
    1330
    1
            let coefficients = model.get_tensor_f32("coefficients").expect("test");
    1331
    1
            assert_eq!(coefficients, vec![2.0, 3.0, 1.5]);
    1332
    1333
    1
            let intercept = model.get_tensor_f32("intercept").expect("test");
    1334
    1
            assert_eq!(intercept, vec![0.5]);
    1335
    1
        }
    1336
    1337
        // ========== Coverage tests for untested functions ==========
    1338
    1339
        #[test]
    1340
    1
        fn test_cov_get_tensor_f16_as_f32() {
    1341
    1
            let json = r#"{"weights":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}}"#;
    1342
    1
            let json_bytes = json.as_bytes();
    1343
    1344
    1
            let mut data = Vec::new();
    1345
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1346
    1
            data.extend_from_slice(json_bytes);
    1347
            // Two F16 values: 1.0 and 2.0
    1348
    1
            data.extend_from_slice(&half::f16::from_f32(1.0).to_le_bytes());
    1349
    1
            data.extend_from_slice(&half::f16::from_f32(2.0).to_le_bytes());
    1350
    1351
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1352
    1
            let weights = model.get_tensor_f16_as_f32("weights").expect("test");
    1353
    1354
    1
            assert_eq!(weights.len(), 2);
    1355
    1
            assert!((weights[0] - 1.0).abs() < 0.01);
    1356
    1
            assert!((weights[1] - 2.0).abs() < 0.01);
    1357
    1
        }
    1358
    1359
        #[test]
    1360
    1
        fn test_cov_get_tensor_f16_not_found() {
    1361
    1
            let json = r#"{"weights":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}}"#;
    1362
    1
            let json_bytes = json.as_bytes();
    1363
    1364
    1
            let mut data = Vec::new();
    1365
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1366
    1
            data.extend_from_slice(json_bytes);
    1367
    1
            data.extend_from_slice(&[0u8; 4]);
    1368
    1369
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1370
    1
            let result = model.get_tensor_f16_as_f32("nonexistent");
    1371
    1
            assert!(result.is_err());
    1372
    1
        }
    1373
    1374
        #[test]
    1375
    1
        fn test_cov_get_tensor_f16_wrong_dtype() {
    1376
    1
            let json = r#"{"weights":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
    1377
    1
            let json_bytes = json.as_bytes();
    1378
    1379
    1
            let mut data = Vec::new();
    1380
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1381
    1
            data.extend_from_slice(json_bytes);
    1382
    1
            data.extend_from_slice(&[0u8; 8]);
    1383
    1384
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1385
    1
            let result = model.get_tensor_f16_as_f32("weights");
    1386
    1
            assert!(result.is_err());
    1387
    1
        }
    1388
    1389
        #[test]
    1390
    1
        fn test_cov_get_tensor_f16_data_offset_exceeds() {
    1391
    1
            let json = r#"{"weights":{"dtype":"F16","shape":[2],"data_offsets":[0,100]}}"#;
    1392
    1
            let json_bytes = json.as_bytes();
    1393
    1394
    1
            let mut data = Vec::new();
    1395
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1396
    1
            data.extend_from_slice(json_bytes);
    1397
    1
            data.extend_from_slice(&[0u8; 4]); // Only 4 bytes, offset says 100
    1398
    1399
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1400
    1
            let result = model.get_tensor_f16_as_f32("weights");
    1401
    1
            assert!(result.is_err());
    1402
    1
        }
    1403
    1404
        #[test]
    1405
    1
        fn test_cov_get_tensor_bf16_as_f32() {
    1406
    1
            let json = r#"{"weights":{"dtype":"BF16","shape":[2],"data_offsets":[0,4]}}"#;
    1407
    1
            let json_bytes = json.as_bytes();
    1408
    1409
    1
            let mut data = Vec::new();
    1410
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1411
    1
            data.extend_from_slice(json_bytes);
    1412
            // Two BF16 values: 1.0 and 2.0
    1413
    1
            data.extend_from_slice(&half::bf16::from_f32(1.0).to_le_bytes());
    1414
    1
            data.extend_from_slice(&half::bf16::from_f32(2.0).to_le_bytes());
    1415
    1416
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1417
    1
            let weights = model.get_tensor_bf16_as_f32("weights").expect("test");
    1418
    1419
    1
            assert_eq!(weights.len(), 2);
    1420
    1
            assert!((weights[0] - 1.0).abs() < 0.01);
    1421
    1
            assert!((weights[1] - 2.0).abs() < 0.01);
    1422
    1
        }
    1423
    1424
        #[test]
    1425
    1
        fn test_cov_get_tensor_bf16_not_found() {
    1426
    1
            let json = r#"{"weights":{"dtype":"BF16","shape":[2],"data_offsets":[0,4]}}"#;
    1427
    1
            let json_bytes = json.as_bytes();
    1428
    1429
    1
            let mut data = Vec::new();
    1430
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1431
    1
            data.extend_from_slice(json_bytes);
    1432
    1
            data.extend_from_slice(&[0u8; 4]);
    1433
    1434
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1435
    1
            let result = model.get_tensor_bf16_as_f32("nonexistent");
    1436
    1
            assert!(result.is_err());
    1437
    1
        }
    1438
    1439
        #[test]
    1440
    1
        fn test_cov_get_tensor_bf16_wrong_dtype() {
    1441
    1
            let json = r#"{"weights":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
    1442
    1
            let json_bytes = json.as_bytes();
    1443
    1444
    1
            let mut data = Vec::new();
    1445
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1446
    1
            data.extend_from_slice(json_bytes);
    1447
    1
            data.extend_from_slice(&[0u8; 8]);
    1448
    1449
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1450
    1
            let result = model.get_tensor_bf16_as_f32("weights");
    1451
    1
            assert!(result.is_err());
    1452
    1
        }
    1453
    1454
        #[test]
    1455
    1
        fn test_cov_get_tensor_bf16_data_offset_exceeds() {
    1456
    1
            let json = r#"{"weights":{"dtype":"BF16","shape":[2],"data_offsets":[0,100]}}"#;
    1457
    1
            let json_bytes = json.as_bytes();
    1458
    1459
    1
            let mut data = Vec::new();
    1460
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1461
    1
            data.extend_from_slice(json_bytes);
    1462
    1
            data.extend_from_slice(&[0u8; 4]); // Only 4 bytes
    1463
    1464
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1465
    1
            let result = model.get_tensor_bf16_as_f32("weights");
    1466
    1
            assert!(result.is_err());
    1467
    1
        }
    1468
    1469
        #[test]
    1470
    1
        fn test_cov_get_tensor_auto_f32() {
    1471
    1
            let json = r#"{"weights":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
    1472
    1
            let json_bytes = json.as_bytes();
    1473
    1474
    1
            let mut data = Vec::new();
    1475
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1476
    1
            data.extend_from_slice(json_bytes);
    1477
    1
            data.extend_from_slice(&1.0f32.to_le_bytes());
    1478
    1
            data.extend_from_slice(&2.0f32.to_le_bytes());
    1479
    1480
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1481
    1
            let weights = model.get_tensor_auto("weights").expect("test");
    1482
    1
            assert_eq!(weights, vec![1.0, 2.0]);
    1483
    1
        }
    1484
    1485
        #[test]
    1486
    1
        fn test_cov_get_tensor_auto_f16() {
    1487
    1
            let json = r#"{"weights":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}}"#;
    1488
    1
            let json_bytes = json.as_bytes();
    1489
    1490
    1
            let mut data = Vec::new();
    1491
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1492
    1
            data.extend_from_slice(json_bytes);
    1493
    1
            data.extend_from_slice(&half::f16::from_f32(1.0).to_le_bytes());
    1494
    1
            data.extend_from_slice(&half::f16::from_f32(2.0).to_le_bytes());
    1495
    1496
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1497
    1
            let weights = model.get_tensor_auto("weights").expect("test");
    1498
    1
            assert_eq!(weights.len(), 2);
    1499
    1
        }
    1500
    1501
        #[test]
    1502
    1
        fn test_cov_get_tensor_auto_bf16() {
    1503
    1
            let json = r#"{"weights":{"dtype":"BF16","shape":[2],"data_offsets":[0,4]}}"#;
    1504
    1
            let json_bytes = json.as_bytes();
    1505
    1506
    1
            let mut data = Vec::new();
    1507
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1508
    1
            data.extend_from_slice(json_bytes);
    1509
    1
            data.extend_from_slice(&half::bf16::from_f32(1.0).to_le_bytes());
    1510
    1
            data.extend_from_slice(&half::bf16::from_f32(2.0).to_le_bytes());
    1511
    1512
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1513
    1
            let weights = model.get_tensor_auto("weights").expect("test");
    1514
    1
            assert_eq!(weights.len(), 2);
    1515
    1
        }
    1516
    1517
        #[test]
    1518
    1
        fn test_cov_get_tensor_auto_unsupported_dtype() {
    1519
    1
            let json = r#"{"weights":{"dtype":"I32","shape":[2],"data_offsets":[0,8]}}"#;
    1520
    1
            let json_bytes = json.as_bytes();
    1521
    1522
    1
            let mut data = Vec::new();
    1523
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1524
    1
            data.extend_from_slice(json_bytes);
    1525
    1
            data.extend_from_slice(&[0u8; 8]);
    1526
    1527
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1528
    1
            let result = model.get_tensor_auto("weights");
    1529
    1
            assert!(result.is_err());
    1530
    1
        }
    1531
    1532
        #[test]
    1533
    1
        fn test_cov_get_tensor_auto_not_found() {
    1534
    1
            let json = r#"{"weights":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
    1535
    1
            let json_bytes = json.as_bytes();
    1536
    1537
    1
            let mut data = Vec::new();
    1538
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1539
    1
            data.extend_from_slice(json_bytes);
    1540
    1
            data.extend_from_slice(&[0u8; 8]);
    1541
    1542
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1543
    1
            let result = model.get_tensor_auto("nonexistent");
    1544
    1
            assert!(result.is_err());
    1545
    1
        }
    1546
    1547
        #[test]
    1548
    1
        fn test_cov_tensor_names() {
    1549
    1
            let json = r#"{
    1550
    1
                "weight1":{"dtype":"F32","shape":[2],"data_offsets":[0,8]},
    1551
    1
                "weight2":{"dtype":"F32","shape":[2],"data_offsets":[8,16]}
    1552
    1
            }"#;
    1553
    1
            let json_bytes = json.as_bytes();
    1554
    1555
    1
            let mut data = Vec::new();
    1556
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1557
    1
            data.extend_from_slice(json_bytes);
    1558
    1
            data.extend_from_slice(&[0u8; 16]);
    1559
    1560
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1561
    1
            let names = model.tensor_names();
    1562
    1
            assert_eq!(names.len(), 2);
    1563
    1
            assert!(names.contains(&"weight1"));
    1564
    1
            assert!(names.contains(&"weight2"));
    1565
    1
        }
    1566
    1567
        #[test]
    1568
    1
        fn test_cov_get_tensor_info() {
    1569
    1
            let json = r#"{"weight":{"dtype":"F32","shape":[2,3],"data_offsets":[0,24]}}"#;
    1570
    1
            let json_bytes = json.as_bytes();
    1571
    1572
    1
            let mut data = Vec::new();
    1573
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1574
    1
            data.extend_from_slice(json_bytes);
    1575
    1
            data.extend_from_slice(&[0u8; 24]);
    1576
    1577
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1578
    1579
    1
            let info = model.get_tensor_info("weight");
    1580
    1
            assert!(info.is_some());
    1581
    1
            let info = info.expect("operation failed");
    1582
    1
            assert_eq!(info.shape, vec![2, 3]);
    1583
    1
            assert_eq!(info.dtype, SafetensorsDtype::F32);
    1584
    1585
    1
            assert!(model.get_tensor_info("nonexistent").is_none());
    1586
    1
        }
    1587
    1588
        #[test]
    1589
    1
        fn test_cov_has_tensor() {
    1590
    1
            let json = r#"{"weight":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}}"#;
    1591
    1
            let json_bytes = json.as_bytes();
    1592
    1593
    1
            let mut data = Vec::new();
    1594
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1595
    1
            data.extend_from_slice(json_bytes);
    1596
    1
            data.extend_from_slice(&[0u8; 8]);
    1597
    1598
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1599
    1
            assert!(model.has_tensor("weight"));
    1600
    1
            assert!(!model.has_tensor("nonexistent"));
    1601
    1
        }
    1602
    1603
        #[test]
    1604
    1
        fn test_cov_get_tensor_f32_data_offset_exceeds() {
    1605
    1
            let json = r#"{"weight":{"dtype":"F32","shape":[2],"data_offsets":[0,100]}}"#;
    1606
    1
            let json_bytes = json.as_bytes();
    1607
    1608
    1
            let mut data = Vec::new();
    1609
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1610
    1
            data.extend_from_slice(json_bytes);
    1611
    1
            data.extend_from_slice(&[0u8; 8]); // Only 8 bytes, offset says 100
    1612
    1613
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1614
    1
            let result = model.get_tensor_f32("weight");
    1615
    1
            assert!(result.is_err());
    1616
    1
        }
    1617
    1618
        #[test]
    1619
    1
        fn test_cov_get_tensor_f32_not_multiple_of_4() {
    1620
    1
            let json = r#"{"weight":{"dtype":"F32","shape":[2],"data_offsets":[0,7]}}"#;
    1621
    1
            let json_bytes = json.as_bytes();
    1622
    1623
    1
            let mut data = Vec::new();
    1624
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1625
    1
            data.extend_from_slice(json_bytes);
    1626
    1
            data.extend_from_slice(&[0u8; 7]); // 7 bytes, not a multiple of 4
    1627
    1628
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1629
    1
            let result = model.get_tensor_f32("weight");
    1630
    1
            assert!(result.is_err());
    1631
    1
        }
    1632
    1633
        #[test]
    1634
    1
        fn test_cov_safetensors_config_num_kv_heads() {
    1635
    1
            let config = SafetensorsConfig {
    1636
    1
                hidden_size: Some(768),
    1637
    1
                num_hidden_layers: Some(12),
    1638
    1
                num_attention_heads: Some(12),
    1639
    1
                num_key_value_heads: Some(4),
    1640
    1
                vocab_size: Some(32000),
    1641
    1
                intermediate_size: None,
    1642
    1
                max_position_embeddings: None,
    1643
    1
                rms_norm_eps: None,
    1644
    1
                rope_theta: None,
    1645
    1
                architectures: None,
    1646
    1
                model_type: None,
    1647
    1
                bos_token_id: None,
    1648
    1
                eos_token_id: None,
    1649
    1
            };
    1650
    1651
    1
            assert_eq!(config.num_kv_heads(), 4);
    1652
    1
        }
    1653
    1654
        #[test]
    1655
    1
        fn test_cov_safetensors_config_num_kv_heads_default() {
    1656
    1
            let config = SafetensorsConfig {
    1657
    1
                hidden_size: Some(768),
    1658
    1
                num_hidden_layers: Some(12),
    1659
    1
                num_attention_heads: Some(12),
    1660
    1
                num_key_value_heads: None, // Not set, should fall back to attention heads
    1661
    1
                vocab_size: Some(32000),
    1662
    1
                intermediate_size: None,
    1663
    1
                max_position_embeddings: None,
    1664
    1
                rms_norm_eps: None,
    1665
    1
                rope_theta: None,
    1666
    1
                architectures: None,
    1667
    1
                model_type: None,
    1668
    1
                bos_token_id: None,
    1669
    1
                eos_token_id: None,
    1670
    1
            };
    1671
    1672
    1
            assert_eq!(config.num_kv_heads(), 12);
    1673
    1
        }
    1674
    1675
        #[test]
    1676
    1
        fn test_cov_safetensors_config_num_kv_heads_fallback() {
    1677
    1
            let config = SafetensorsConfig {
    1678
    1
                hidden_size: Some(768),
    1679
    1
                num_hidden_layers: Some(12),
    1680
    1
                num_attention_heads: None,
    1681
    1
                num_key_value_heads: None,
    1682
    1
                vocab_size: Some(32000),
    1683
    1
                intermediate_size: None,
    1684
    1
                max_position_embeddings: None,
    1685
    1
                rms_norm_eps: None,
    1686
    1
                rope_theta: None,
    1687
    1
                architectures: None,
    1688
    1
                model_type: None,
    1689
    1
                bos_token_id: None,
    1690
    1
                eos_token_id: None,
    1691
    1
            };
    1692
    1693
    1
            assert_eq!(config.num_kv_heads(), 1); // Fallback to 1
    1694
    1
        }
    1695
    1696
        #[test]
    1697
    1
        fn test_cov_safetensors_config_architecture_from_architectures() {
    1698
    1
            let config = SafetensorsConfig {
    1699
    1
                hidden_size: None,
    1700
    1
                num_hidden_layers: None,
    1701
    1
                num_attention_heads: None,
    1702
    1
                num_key_value_heads: None,
    1703
    1
                vocab_size: None,
    1704
    1
                intermediate_size: None,
    1705
    1
                max_position_embeddings: None,
    1706
    1
                rms_norm_eps: None,
    1707
    1
                rope_theta: None,
    1708
    1
                architectures: Some(vec!["LlamaForCausalLM".to_string()]),
    1709
    1
                model_type: Some("llama".to_string()),
    1710
    1
                bos_token_id: None,
    1711
    1
                eos_token_id: None,
    1712
    1
            };
    1713
    1714
    1
            assert_eq!(config.architecture(), "LlamaForCausalLM");
    1715
    1
        }
    1716
    1717
        #[test]
    1718
    1
        fn test_cov_safetensors_config_architecture_from_model_type() {
    1719
    1
            let config = SafetensorsConfig {
    1720
    1
                hidden_size: None,
    1721
    1
                num_hidden_layers: None,
    1722
    1
                num_attention_heads: None,
    1723
    1
                num_key_value_heads: None,
    1724
    1
                vocab_size: None,
    1725
    1
                intermediate_size: None,
    1726
    1
                max_position_embeddings: None,
    1727
    1
                rms_norm_eps: None,
    1728
    1
                rope_theta: None,
    1729
    1
                architectures: None,
    1730
    1
                model_type: Some("llama".to_string()),
    1731
    1
                bos_token_id: None,
    1732
    1
                eos_token_id: None,
    1733
    1
            };
    1734
    1735
    1
            assert_eq!(config.architecture(), "llama");
    1736
    1
        }
    1737
    1738
        #[test]
    1739
    1
        fn test_cov_safetensors_config_architecture_unknown() {
    1740
    1
            let config = SafetensorsConfig {
    1741
    1
                hidden_size: None,
    1742
    1
                num_hidden_layers: None,
    1743
    1
                num_attention_heads: None,
    1744
    1
                num_key_value_heads: None,
    1745
    1
                vocab_size: None,
    1746
    1
                intermediate_size: None,
    1747
    1
                max_position_embeddings: None,
    1748
    1
                rms_norm_eps: None,
    1749
    1
                rope_theta: None,
    1750
    1
                architectures: None,
    1751
    1
                model_type: None,
    1752
    1
                bos_token_id: None,
    1753
    1
                eos_token_id: None,
    1754
    1
            };
    1755
    1756
    1
            assert_eq!(config.architecture(), "unknown");
    1757
    1
        }
    1758
    1759
        #[test]
    1760
    1
        fn test_cov_safetensors_config_load_from_sibling_not_found() {
    1761
    1
            let path = std::path::Path::new("/nonexistent/model.safetensors");
    1762
    1
            let config = SafetensorsConfig::load_from_sibling(path);
    1763
    1
            assert!(config.is_none());
    1764
    1
        }
    1765
    1766
        #[test]
    1767
    1
        fn test_cov_metadata_key_skipped() {
    1768
            // Test that __metadata__ key is skipped
    1769
    1
            let json = r#"{
    1770
    1
                "__metadata__":{"format":"pt"},
    1771
    1
                "weight":{"dtype":"F32","shape":[2],"data_offsets":[0,8]}
    1772
    1
            }"#;
    1773
    1
            let json_bytes = json.as_bytes();
    1774
    1775
    1
            let mut data = Vec::new();
    1776
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1777
    1
            data.extend_from_slice(json_bytes);
    1778
    1
            data.extend_from_slice(&[0u8; 8]);
    1779
    1780
    1
            let model = SafetensorsModel::from_bytes(&data).expect("test");
    1781
    1
            assert_eq!(model.tensors.len(), 1);
    1782
    1
            assert!(model.tensors.contains_key("weight"));
    1783
    1
            assert!(!model.tensors.contains_key("__metadata__"));
    1784
    1
        }
    1785
    1786
        #[test]
    1787
    1
        fn test_cov_json_not_object() {
    1788
            // JSON is an array, not an object
    1789
    1
            let json = r"[]";
    1790
    1
            let json_bytes = json.as_bytes();
    1791
    1792
    1
            let mut data = Vec::new();
    1793
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1794
    1
            data.extend_from_slice(json_bytes);
    1795
    1796
    1
            let result = SafetensorsModel::from_bytes(&data);
    1797
    1
            assert!(result.is_err());
    1798
    1
        }
    1799
    1800
        #[test]
    1801
    1
        fn test_cov_tensor_metadata_parse_error() {
    1802
            // Tensor has invalid metadata (missing dtype)
    1803
    1
            let json = r#"{"weight":{"shape":[2],"data_offsets":[0,8]}}"#;
    1804
    1
            let json_bytes = json.as_bytes();
    1805
    1806
    1
            let mut data = Vec::new();
    1807
    1
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1808
    1
            data.extend_from_slice(json_bytes);
    1809
    1
            data.extend_from_slice(&[0u8; 8]);
    1810
    1811
    1
            let result = SafetensorsModel::from_bytes(&data);
    1812
    1
            assert!(result.is_err());
    1813
    1
        }
    1814
    }