## 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_infer.rs

    Line
    Count
    Source
    1
    //! SafeTensors Inference Support (PAR-301)
    2
    //!
    3
    //! Provides SafeTensors model loading and inference for HuggingFace models.
    4
    //!
    5
    //! ## Architecture
    6
    //!
    7
    //! SafeTensors files contain only tensor weights, so we need:
    8
    //! - `config.json` for model architecture (hidden_size, num_layers, etc.)
    9
    //! - `tokenizer.json` for text tokenization
    10
    //!
    11
    //! The converter loads these from sibling files and builds an AprTransformer.
    12
    13
    use crate::apr_transformer::{AprTransformer, AprTransformerConfig, AprTransformerLayer};
    14
    use crate::error::{RealizarError, Result};
    15
    use crate::safetensors::{MappedSafeTensorsModel, SafetensorsConfig};
    16
    use std::path::Path;
    17
    18
    /// SafeTensors to APR Transformer converter
    19
    ///
    20
    /// Converts HuggingFace SafeTensors models to APR Transformer format.
    21
    /// Supports BF16, F16, and F32 weights with automatic conversion to F32.
    22
    pub struct SafetensorsToAprConverter;
    23
    24
    impl SafetensorsToAprConverter {
    25
        /// Convert SafeTensors model to APR Transformer
    26
        ///
    27
        /// # Arguments
    28
        ///
    29
        /// * `model_path` - Path to model.safetensors file
    30
        ///
    31
        /// # Returns
    32
        ///
    33
        /// `AprTransformer` with F32 weights ready for inference
    34
        ///
    35
        /// # Errors
    36
        ///
    37
        /// Returns error if SafeTensors file, config.json, or required tensors are missing
    38
    20
        pub fn convert(model_path: &Path) -> Result<AprTransformer> {
    39
            // Load SafeTensors model using mmap for zero-copy access (T-QA-020)
    40
            // This is critical for fast model loading - mmap is O(1) regardless of file size
    41
    20
            let st_model18 = MappedSafeTensorsModel::load(model_path)?2;
    42
    43
            // Load config.json (required for architecture info)
    44
    18
            let config16 = SafetensorsConfig::load_from_sibling(model_path).ok_or_else(|| {2
    45
    2
                RealizarError::UnsupportedOperation {
    46
    2
                    operation: "safetensors_convert".to_string(),
    47
    2
                    reason: "config.json not found (required for SafeTensors inference)".to_string(),
    48
    2
                }
    49
    2
            })?;
    50
    51
            // Extract architecture parameters
    52
    16
            let hidden_dim15 = config
    53
    16
                .hidden_size
    54
    16
                .ok_or_else(|| RealizarError::FormatError {
    55
    1
                    reason: "config.json missing hidden_size".to_string(),
    56
    1
                })?;
    57
    15
            let num_layers14 = config
    58
    15
                .num_hidden_layers
    59
    15
                .ok_or_else(|| RealizarError::FormatError {
    60
    1
                    reason: "config.json missing num_hidden_layers".to_string(),
    61
    1
                })?;
    62
    14
            let num_heads13 = config
    63
    14
                .num_attention_heads
    64
    14
                .ok_or_else(|| RealizarError::FormatError {
    65
    1
                    reason: "config.json missing num_attention_heads".to_string(),
    66
    1
                })?;
    67
    13
            let num_kv_heads = config.num_kv_heads();
    68
    13
            let vocab_size12 = config
    69
    13
                .vocab_size
    70
    13
                .ok_or_else(|| RealizarError::FormatError {
    71
    1
                    reason: "config.json missing vocab_size".to_string(),
    72
    1
                })?;
    73
    12
            let intermediate_dim = config.intermediate_size.unwrap_or(hidden_dim * 4);
    74
    12
            let context_length = config.max_position_embeddings.unwrap_or(2048);
    75
    12
            let rope_theta = config.rope_theta.unwrap_or(10000.0);
    76
    12
            let eps = config.rms_norm_eps.unwrap_or(1e-6);
    77
    12
            let architecture = config.architecture();
    78
    79
            // Build transformer config
    80
    12
            let apr_config = AprTransformerConfig {
    81
    12
                architecture,
    82
    12
                hidden_dim,
    83
    12
                num_layers,
    84
    12
                num_heads,
    85
    12
                num_kv_heads,
    86
    12
                vocab_size,
    87
    12
                intermediate_dim,
    88
    12
                context_length,
    89
    12
                rope_theta,
    90
    12
                eps,
    91
    12
            };
    92
    93
            // Extract embeddings
    94
    12
            let token_embedding11 = st_model.get_tensor_auto("model.embed_tokens.weight")?1;
    95
    96
            // Extract output norm
    97
    11
            let output_norm_weight = st_model.get_tensor_auto("model.norm.weight")?0;
    98
    99
            // Check for tied embeddings (lm_head = embed_tokens.T)
    100
            // lm_head.weight: [vocab_size, hidden_dim] -> transpose -> [hidden_dim, vocab_size]
    101
    11
            let lm_head_weight = if st_model.has_tensor("lm_head.weight") {
    102
    1
                let raw = st_model.get_tensor_auto("lm_head.weight")?0;
    103
    1
                Self::transpose_weight(&raw, vocab_size, hidden_dim)
    104
            } else {
    105
                // Tied embeddings: token_embedding is [vocab_size, hidden_dim]
    106
                // Need to transpose to [hidden_dim, vocab_size]
    107
    10
                Self::transpose_weight(&token_embedding, vocab_size, hidden_dim)
    108
            };
    109
    110
            // Extract layers
    111
    11
            let mut layers = Vec::with_capacity(num_layers);
    112
    11
            for i4 in 0..num_layers {
    113
    4
                let layer3 = Self::extract_layer(
    114
    4
                    &st_model,
    115
    4
                    i,
    116
    4
                    hidden_dim,
    117
    4
                    num_heads,
    118
    4
                    num_kv_heads,
    119
    4
                    intermediate_dim,
    120
    1
                )?;
    121
    3
                layers.push(layer);
    122
            }
    123
    124
    10
            Ok(AprTransformer {
    125
    10
                config: apr_config,
    126
    10
                token_embedding,
    127
    10
                layers,
    128
    10
                output_norm_weight,
    129
    10
                output_norm_bias: None,
    130
    10
                lm_head_weight,
    131
    10
                lm_head_bias: None,
    132
    10
                q4k_layers: None,
    133
    10
                lm_head_weight_q6k: None,
    134
    10
                lm_head_weight_q4k: None,
    135
    10
            })
    136
    20
        }
    137
    138
        /// Extract a single transformer layer from SafeTensors
    139
    4
        fn extract_layer(
    140
    4
            st_model: &MappedSafeTensorsModel,
    141
    4
            layer_idx: usize,
    142
    4
            hidden_dim: usize,
    143
    4
            num_heads: usize,
    144
    4
            num_kv_heads: usize,
    145
    4
            intermediate_dim: usize,
    146
    4
        ) -> Result<AprTransformerLayer> {
    147
    4
            let prefix = format!("model.layers.{layer_idx}");
    148
    149
            // Attention norm (input_layernorm) - 1D vector, no transpose needed
    150
    3
            let attn_norm_weight =
    151
    4
                st_model.get_tensor_auto(&format!("{prefix}.input_layernorm.weight"))?1;
    152
    153
            // Q, K, V projections (separate in HuggingFace, combined in APR)
    154
            // HuggingFace: [out_dim, in_dim], APR needs: [in_dim, out_dim]
    155
    3
            let q_weight = st_model.get_tensor_auto(&format!("{prefix}.self_attn.q_proj.weight"))?0;
    156
    3
            let k_weight = st_model.get_tensor_auto(&format!("{prefix}.self_attn.k_proj.weight"))?0;
    157
    3
            let v_weight = st_model.get_tensor_auto(&format!("{prefix}.self_attn.v_proj.weight"))?0;
    158
    159
            // Concatenate and transpose Q, K, V into combined QKV weight
    160
    3
            let head_dim = hidden_dim / num_heads;
    161
    3
            let kv_dim = head_dim * num_kv_heads;
    162
    3
            let qkv_weight =
    163
    3
                Self::concat_qkv_transposed(&q_weight, &k_weight, &v_weight, hidden_dim, kv_dim);
    164
    165
            // QKV bias (optional) - 1D vector, no transpose needed
    166
    3
            let qkv_bias = Self::try_concat_qkv_bias(st_model, &prefix, hidden_dim, kv_dim);
    167
    168
            // Attention output projection: [hidden_dim, hidden_dim]
    169
            // HuggingFace: o_proj is [hidden_dim, hidden_dim] (out=hidden, in=hidden)
    170
    3
            let attn_output_raw =
    171
    3
                st_model.get_tensor_auto(&format!("{prefix}.self_attn.o_proj.weight"))?0;
    172
    3
            let attn_output_weight = Self::transpose_weight(&attn_output_raw, hidden_dim, hidden_dim);
    173
    174
            // FFN norm (post_attention_layernorm) - 1D vector, no transpose needed
    175
    3
            let ffn_norm_weight =
    176
    3
                st_model.get_tensor_auto(&format!("{prefix}.post_attention_layernorm.weight"))?0;
    177
    178
            // FFN projections (SwiGLU architecture)
    179
            // gate_proj: [intermediate_dim, hidden_dim] -> transpose -> [hidden_dim, intermediate_dim]
    180
            // up_proj: [intermediate_dim, hidden_dim] -> transpose -> [hidden_dim, intermediate_dim]
    181
            // down_proj: [hidden_dim, intermediate_dim] -> transpose -> [intermediate_dim, hidden_dim]
    182
    3
            let ffn_gate_raw = st_model.get_tensor_auto(&format!("{prefix}.mlp.gate_proj.weight"))?0;
    183
    3
            let ffn_gate_weight = Self::transpose_weight(&ffn_gate_raw, intermediate_dim, hidden_dim);
    184
    185
    3
            let ffn_up_raw = st_model.get_tensor_auto(&format!("{prefix}.mlp.up_proj.weight"))?0;
    186
    3
            let ffn_up_weight = Self::transpose_weight(&ffn_up_raw, intermediate_dim, hidden_dim);
    187
    188
    3
            let ffn_down_raw = st_model.get_tensor_auto(&format!("{prefix}.mlp.down_proj.weight"))?0;
    189
    3
            let ffn_down_weight = Self::transpose_weight(&ffn_down_raw, hidden_dim, intermediate_dim);
    190
    191
    3
            Ok(AprTransformerLayer {
    192
    3
                attn_norm_weight,
    193
    3
                attn_norm_bias: None,
    194
    3
                qkv_weight,
    195
    3
                qkv_bias,
    196
    3
                attn_output_weight,
    197
    3
                attn_output_bias: None,
    198
    3
                ffn_gate_weight: Some(ffn_gate_weight),
    199
    3
                ffn_gate_bias: None,
    200
    3
                ffn_up_weight,
    201
    3
                ffn_up_bias: None,
    202
    3
                ffn_down_weight,
    203
    3
                ffn_down_bias: None,
    204
    3
                ffn_norm_weight: Some(ffn_norm_weight),
    205
    3
                ffn_norm_bias: None,
    206
    3
            })
    207
    4
        }
    208
    209
        /// Pass through weight in matvec-optimal [out_dim, in_dim] format
    210
        ///
    211
        /// PMAT-095 FIX: HuggingFace stores Linear weights as [out_features, in_features]
    212
        /// which is EXACTLY what trueno's matvec needs! Previous implementation transposed
    213
        /// twice (here and in matmul), causing O(n²) overhead per forward pass.
    214
        ///
    215
        /// Now we keep HuggingFace format directly - no transposition needed.
    216
        #[allow(clippy::unused_self)]
    217
    23
        fn transpose_weight(weight: &[f32], _out_dim: usize, _in_dim: usize) -> Vec<f32> {
    218
            // PMAT-095: Keep [out_dim, in_dim] format - no transposition!
    219
            // This eliminates the 75x performance gap vs GGUF.
    220
    23
            weight.to_vec()
    221
    23
        }
    222
    223
        /// Concatenate Q, K, V weights into combined QKV tensor (matvec-optimal)
    224
        ///
    225
        /// PMAT-095 FIX: Keep [out_dim, in_dim] format from HuggingFace.
    226
        /// For QKV, we concatenate along the output dimension:
    227
        /// - Q: [hidden_dim, hidden_dim]
    228
        /// - K: [kv_dim, hidden_dim]
    229
        /// - V: [kv_dim, hidden_dim]
    230
        ///
    231
        /// Result: [hidden_dim + kv_dim + kv_dim, hidden_dim] in row-major
    232
    3
        fn concat_qkv_transposed(
    233
    3
            q: &[f32],
    234
    3
            k: &[f32],
    235
    3
            v: &[f32],
    236
    3
            _hidden_dim: usize,
    237
    3
            _kv_dim: usize,
    238
    3
        ) -> Vec<f32> {
    239
            // PMAT-095: Simple concatenation - weights are already in optimal layout
    240
            // Concatenate [Q; K; V] along output dimension
    241
    3
            let mut qkv = Vec::with_capacity(q.len() + k.len() + v.len());
    242
    3
            qkv.extend_from_slice(q);
    243
    3
            qkv.extend_from_slice(k);
    244
    3
            qkv.extend_from_slice(v);
    245
    3
            qkv
    246
    3
        }
    247
    248
        /// Concatenate Q, K, V weights into combined QKV tensor (legacy, no transpose)
    249
    5
        fn concat_qkv(q: &[f32], k: &[f32], v: &[f32]) -> Vec<f32> {
    250
    5
            let mut qkv = Vec::with_capacity(q.len() + k.len() + v.len());
    251
    5
            qkv.extend_from_slice(q);
    252
    5
            qkv.extend_from_slice(k);
    253
    5
            qkv.extend_from_slice(v);
    254
    5
            qkv
    255
    5
        }
    256
    257
        /// Try to concatenate Q, K, V biases if they exist
    258
    6
        fn try_concat_qkv_bias(
    259
    6
            st_model: &MappedSafeTensorsModel,
    260
    6
            prefix: &str,
    261
    6
            hidden_dim: usize,
    262
    6
            kv_dim: usize,
    263
    6
        ) -> Option<Vec<f32>> {
    264
    6
            let q_bias2 = st_model
    265
    6
                .get_tensor_auto(&format!("{prefix}.self_attn.q_proj.bias"))
    266
    6
                .ok()?4;
    267
    2
            let k_bias1 = st_model
    268
    2
                .get_tensor_auto(&format!("{prefix}.self_attn.k_proj.bias"))
    269
    2
                .ok()?1;
    270
    1
            let v_bias = st_model
    271
    1
                .get_tensor_auto(&format!("{prefix}.self_attn.v_proj.bias"))
    272
    1
                .ok()?0;
    273
    274
    1
            let mut qkv_bias = Vec::with_capacity(hidden_dim + kv_dim + kv_dim);
    275
    1
            qkv_bias.extend_from_slice(&q_bias);
    276
    1
            qkv_bias.extend_from_slice(&k_bias);
    277
    1
            qkv_bias.extend_from_slice(&v_bias);
    278
    279
    1
            Some(qkv_bias)
    280
    6
        }
    281
    }
    282
    283
    #[cfg(test)]
    284
    mod tests {
    285
        use super::*;
    286
        use tempfile::TempDir;
    287
    288
        #[test]
    289
    1
        fn test_concat_qkv() {
    290
    1
            let q = vec![1.0, 2.0];
    291
    1
            let k = vec![3.0, 4.0];
    292
    1
            let v = vec![5.0, 6.0];
    293
    1
            let qkv = SafetensorsToAprConverter::concat_qkv(&q, &k, &v);
    294
    1
            assert_eq!(qkv, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
    295
    1
        }
    296
    297
        // =========================================================================
    298
        // Extended Coverage Tests (15+ tests ending with _ext_cov)
    299
        // =========================================================================
    300
    301
        /// Helper function to create a minimal SafeTensors file with given tensors
    302
    18
        fn create_safetensors_bytes(tensors: &[(&str, &str, &[usize], &[u8])]) -> Vec<u8> {
    303
            use serde_json::json;
    304
    305
            // Calculate tensor data layout
    306
    18
            let mut tensor_entries = serde_json::Map::new();
    307
    18
            let mut offset = 0usize;
    308
    309
    41
            for (name23, dtype23, shape23, data23) in tensors {
    310
    23
                let end = offset + data.len();
    311
    23
                tensor_entries.insert(
    312
    23
                    (*name).to_string(),
    313
    23
                    json!({
    314
    23
                        "dtype": dtype,
    315
    23
                        "shape": shape,
    316
    23
                        "data_offsets": [offset, end]
    317
    23
                    }),
    318
    23
                );
    319
    23
                offset = end;
    320
    23
            }
    321
    322
    18
            let json_obj = serde_json::Value::Object(tensor_entries);
    323
    18
            let json_bytes = json_obj.to_string().into_bytes();
    324
    325
    18
            let mut data = Vec::new();
    326
    18
            data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    327
    18
            data.extend_from_slice(&json_bytes);
    328
    329
            // Append tensor data
    330
    41
            for (_, _, _, tensor_data23) in tensors {
    331
    23
                data.extend_from_slice(tensor_data);
    332
    23
            }
    333
    334
    18
            data
    335
    18
        }
    336
    337
        /// Helper to create config.json content
    338
    5
        fn create_config_json(
    339
    5
            hidden_size: usize,
    340
    5
            num_layers: usize,
    341
    5
            num_heads: usize,
    342
    5
            vocab_size: usize,
    343
    5
        ) -> String {
    344
    5
            format!(
    345
    5
                r#"{{
    346
    5
                    "hidden_size": {},
    347
    5
                    "num_hidden_layers": {},
    348
    5
                    "num_attention_heads": {},
    349
    5
                    "vocab_size": {},
    350
    5
                    "intermediate_size": {},
    351
    5
                    "max_position_embeddings": 2048,
    352
    5
                    "rms_norm_eps": 1e-6,
    353
    5
                    "rope_theta": 10000.0,
    354
    5
                    "architectures": ["LlamaForCausalLM"],
    355
    5
                    "model_type": "llama"
    356
    5
                }}"#,
    357
                hidden_size,
    358
                num_layers,
    359
                num_heads,
    360
                vocab_size,
    361
    5
                hidden_size * 4
    362
            )
    363
    5
        }
    364
    365
        #[test]
    366
    1
        fn test_convert_file_not_found_ext_cov() {
    367
    1
            let result =
    368
    1
                SafetensorsToAprConverter::convert(Path::new("/nonexistent/model.safetensors"));
    369
    1
            assert!(result.is_err());
    370
            // MappedSafeTensorsModel::load() returns UnsupportedOperation for file open errors
    371
    1
            if let Err(RealizarError::UnsupportedOperation { operation, reason }) = result {
    372
    1
                assert_eq!(operation, "open_safetensors");
    373
    1
                assert!(reason.contains("Failed to open file"));
    374
            } else {
    375
    0
                panic!("Expected UnsupportedOperation error");
    376
            }
    377
    1
        }
    378
    379
        #[test]
    380
    1
        fn test_convert_missing_config_json_ext_cov() {
    381
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    382
    1
            let model_path = temp_dir.path().join("model.safetensors");
    383
    384
            // Create a minimal valid safetensors file
    385
    1
            let data = create_safetensors_bytes(&[]);
    386
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    387
    388
            // No config.json file
    389
    390
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    391
    1
            assert!(result.is_err());
    392
    1
            if let Err(RealizarError::UnsupportedOperation { operation, reason }) = result {
    393
    1
                assert_eq!(operation, "safetensors_convert");
    394
    1
                assert!(reason.contains("config.json not found"));
    395
            } else {
    396
    0
                panic!("Expected UnsupportedOperation error");
    397
            }
    398
    1
        }
    399
    400
        #[test]
    401
    1
        fn test_convert_missing_hidden_size_ext_cov() {
    402
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    403
    1
            let model_path = temp_dir.path().join("model.safetensors");
    404
    1
            let config_path = temp_dir.path().join("config.json");
    405
    406
            // Create minimal safetensors
    407
    1
            let data = create_safetensors_bytes(&[]);
    408
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    409
    410
            // Config missing hidden_size
    411
    1
            let config = r#"{"num_hidden_layers": 2, "num_attention_heads": 4, "vocab_size": 100}"#;
    412
    1
            std::fs::write(&config_path, config).expect("write config");
    413
    414
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    415
    1
            assert!(result.is_err());
    416
    1
            if let Err(RealizarError::FormatError { reason }) = result {
    417
    1
                assert!(reason.contains("missing hidden_size"));
    418
            } else {
    419
    0
                panic!("Expected FormatError for missing hidden_size");
    420
            }
    421
    1
        }
    422
    423
        #[test]
    424
    1
        fn test_convert_missing_num_hidden_layers_ext_cov() {
    425
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    426
    1
            let model_path = temp_dir.path().join("model.safetensors");
    427
    1
            let config_path = temp_dir.path().join("config.json");
    428
    429
    1
            let data = create_safetensors_bytes(&[]);
    430
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    431
    432
            // Config missing num_hidden_layers
    433
    1
            let config = r#"{"hidden_size": 64, "num_attention_heads": 4, "vocab_size": 100}"#;
    434
    1
            std::fs::write(&config_path, config).expect("write config");
    435
    436
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    437
    1
            assert!(result.is_err());
    438
    1
            if let Err(RealizarError::FormatError { reason }) = result {
    439
    1
                assert!(reason.contains("missing num_hidden_layers"));
    440
            } else {
    441
    0
                panic!("Expected FormatError for missing num_hidden_layers");
    442
            }
    443
    1
        }
    444
    445
        #[test]
    446
    1
        fn test_convert_missing_num_attention_heads_ext_cov() {
    447
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    448
    1
            let model_path = temp_dir.path().join("model.safetensors");
    449
    1
            let config_path = temp_dir.path().join("config.json");
    450
    451
    1
            let data = create_safetensors_bytes(&[]);
    452
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    453
    454
            // Config missing num_attention_heads
    455
    1
            let config = r#"{"hidden_size": 64, "num_hidden_layers": 2, "vocab_size": 100}"#;
    456
    1
            std::fs::write(&config_path, config).expect("write config");
    457
    458
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    459
    1
            assert!(result.is_err());
    460
    1
            if let Err(RealizarError::FormatError { reason }) = result {
    461
    1
                assert!(reason.contains("missing num_attention_heads"));
    462
            } else {
    463
    0
                panic!("Expected FormatError for missing num_attention_heads");
    464
            }
    465
    1
        }
    466
    467
        #[test]
    468
    1
        fn test_convert_missing_vocab_size_ext_cov() {
    469
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    470
    1
            let model_path = temp_dir.path().join("model.safetensors");
    471
    1
            let config_path = temp_dir.path().join("config.json");
    472
    473
    1
            let data = create_safetensors_bytes(&[]);
    474
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    475
    476
            // Config missing vocab_size
    477
    1
            let config = r#"{"hidden_size": 64, "num_hidden_layers": 2, "num_attention_heads": 4}"#;
    478
    1
            std::fs::write(&config_path, config).expect("write config");
    479
    480
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    481
    1
            assert!(result.is_err());
    482
    1
            if let Err(RealizarError::FormatError { reason }) = result {
    483
    1
                assert!(reason.contains("missing vocab_size"));
    484
            } else {
    485
    0
                panic!("Expected FormatError for missing vocab_size");
    486
            }
    487
    1
        }
    488
    489
        #[test]
    490
    1
        fn test_convert_missing_embed_tokens_ext_cov() {
    491
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    492
    1
            let model_path = temp_dir.path().join("model.safetensors");
    493
    1
            let config_path = temp_dir.path().join("config.json");
    494
    495
            // Safetensors without model.embed_tokens.weight
    496
    1
            let data = create_safetensors_bytes(&[]);
    497
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    498
    499
            // Valid config
    500
    1
            let config = create_config_json(64, 1, 4, 100);
    501
    1
            std::fs::write(&config_path, config).expect("write config");
    502
    503
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    504
    1
            assert!(result.is_err());
    505
            // Should fail because model.embed_tokens.weight is missing
    506
    1
        }
    507
    508
        #[test]
    509
    1
        fn test_concat_qkv_empty_inputs_ext_cov() {
    510
    1
            let q: Vec<f32> = vec![];
    511
    1
            let k: Vec<f32> = vec![];
    512
    1
            let v: Vec<f32> = vec![];
    513
    1
            let qkv = SafetensorsToAprConverter::concat_qkv(&q, &k, &v);
    514
    1
            assert!(qkv.is_empty());
    515
    1
        }
    516
    517
        #[test]
    518
    1
        fn test_concat_qkv_single_elements_ext_cov() {
    519
    1
            let q = vec![1.0];
    520
    1
            let k = vec![2.0];
    521
    1
            let v = vec![3.0];
    522
    1
            let qkv = SafetensorsToAprConverter::concat_qkv(&q, &k, &v);
    523
    1
            assert_eq!(qkv, vec![1.0, 2.0, 3.0]);
    524
    1
        }
    525
    526
        #[test]
    527
    1
        fn test_concat_qkv_large_arrays_ext_cov() {
    528
    1.00k
            let q1: Vec<f32>1 = (0..1000)1.map1(|i| i as f32).collect1();
    529
    1.00k
            let k1: Vec<f32>1 = (1000..2000)1.map1(|i| i as f32).collect1();
    530
    1.00k
            let v1: Vec<f32>1 = (2000..3000)1.map1(|i| i as f32).collect1();
    531
    1
            let qkv = SafetensorsToAprConverter::concat_qkv(&q, &k, &v);
    532
    1
            assert_eq!(qkv.len(), 3000);
    533
    1
            assert_eq!(qkv[0], 0.0);
    534
    1
            assert_eq!(qkv[1000], 1000.0);
    535
    1
            assert_eq!(qkv[2000], 2000.0);
    536
    1
            assert_eq!(qkv[2999], 2999.0);
    537
    1
        }
    538
    539
        #[test]
    540
    1
        fn test_concat_qkv_asymmetric_ext_cov() {
    541
    1
            let q = vec![1.0, 2.0, 3.0, 4.0];
    542
    1
            let k = vec![5.0, 6.0];
    543
    1
            let v = vec![7.0];
    544
    1
            let qkv = SafetensorsToAprConverter::concat_qkv(&q, &k, &v);
    545
    1
            assert_eq!(qkv, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]);
    546
    1
        }
    547
    548
        #[test]
    549
    1
        fn test_try_concat_qkv_bias_none_when_missing_ext_cov() {
    550
            // Create safetensors model without any biases
    551
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    552
    1
            let model_path = temp_dir.path().join("model.safetensors");
    553
    1
            let data = create_safetensors_bytes(&[]);
    554
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    555
    1
            let st_model = MappedSafeTensorsModel::load(&model_path).expect("load safetensors");
    556
    557
    1
            let result =
    558
    1
                SafetensorsToAprConverter::try_concat_qkv_bias(&st_model, "model.layers.0", 64, 64);
    559
    1
            assert!(result.is_none());
    560
    1
        }
    561
    562
        #[test]
    563
    1
        fn test_try_concat_qkv_bias_partial_missing_ext_cov() {
    564
            // Create safetensors with only q_proj.bias (missing k and v)
    565
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    566
    1
            let model_path = temp_dir.path().join("model.safetensors");
    567
    16
            let q_bias_data1: Vec<u8>1 = (0..16)1.flat_map1(|i| (i as f32).to_le_bytes()).collect1();
    568
    1
            let data = create_safetensors_bytes(&[(
    569
    1
                "model.layers.0.self_attn.q_proj.bias",
    570
    1
                "F32",
    571
    1
                &[4],
    572
    1
                &q_bias_data,
    573
    1
            )]);
    574
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    575
    1
            let st_model = MappedSafeTensorsModel::load(&model_path).expect("load safetensors");
    576
    577
            // Should return None because k_bias and v_bias are missing
    578
    1
            let result =
    579
    1
                SafetensorsToAprConverter::try_concat_qkv_bias(&st_model, "model.layers.0", 4, 4);
    580
    1
            assert!(result.is_none());
    581
    1
        }
    582
    583
        #[test]
    584
    1
        fn test_try_concat_qkv_bias_all_present_ext_cov() {
    585
            // Create F32 byte data for biases
    586
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    587
    1
            let model_path = temp_dir.path().join("model.safetensors");
    588
    1
            let q_bias_data: Vec<u8> = [1.0f32, 2.0, 3.0, 4.0]
    589
    1
                .iter()
    590
    4
                .flat_map1(|f| f.to_le_bytes())
    591
    1
                .collect();
    592
    2
            let k_bias_data1: Vec<u8>1 = [5.0f32, 6.0]1.iter1().flat_map1(|f| f.to_le_bytes()).collect1();
    593
    2
            let v_bias_data1: Vec<u8>1 = [7.0f32, 8.0]1.iter1().flat_map1(|f| f.to_le_bytes()).collect1();
    594
    595
    1
            let data = create_safetensors_bytes(&[
    596
    1
                (
    597
    1
                    "model.layers.0.self_attn.q_proj.bias",
    598
    1
                    "F32",
    599
    1
                    &[4],
    600
    1
                    &q_bias_data,
    601
    1
                ),
    602
    1
                (
    603
    1
                    "model.layers.0.self_attn.k_proj.bias",
    604
    1
                    "F32",
    605
    1
                    &[2],
    606
    1
                    &k_bias_data,
    607
    1
                ),
    608
    1
                (
    609
    1
                    "model.layers.0.self_attn.v_proj.bias",
    610
    1
                    "F32",
    611
    1
                    &[2],
    612
    1
                    &v_bias_data,
    613
    1
                ),
    614
    1
            ]);
    615
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    616
    1
            let st_model = MappedSafeTensorsModel::load(&model_path).expect("load safetensors");
    617
    618
    1
            let result =
    619
    1
                SafetensorsToAprConverter::try_concat_qkv_bias(&st_model, "model.layers.0", 4, 2);
    620
    1
            assert!(result.is_some());
    621
    1
            let bias = result.expect("operation failed");
    622
    1
            assert_eq!(bias.len(), 8); // 4 + 2 + 2
    623
    1
            assert_eq!(bias[0], 1.0);
    624
    1
            assert_eq!(bias[4], 5.0);
    625
    1
            assert_eq!(bias[6], 7.0);
    626
    1
        }
    627
    628
        #[test]
    629
    1
        fn test_convert_defaults_intermediate_size_ext_cov() {
    630
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    631
    1
            let model_path = temp_dir.path().join("model.safetensors");
    632
    1
            let config_path = temp_dir.path().join("config.json");
    633
    634
            // Config without intermediate_size (should default to hidden_size * 4)
    635
    1
            let config = r#"{
    636
    1
                "hidden_size": 64,
    637
    1
                "num_hidden_layers": 0,
    638
    1
                "num_attention_heads": 4,
    639
    1
                "vocab_size": 100
    640
    1
            }"#;
    641
    1
            std::fs::write(&config_path, config).expect("write config");
    642
    643
            // Create safetensors with minimal required tensors for 0 layers
    644
    1
            let embed_data: Vec<u8> = vec![0u8; 100 * 64 * 4]; // vocab_size * hidden_dim * 4 bytes
    645
    1
            let norm_data: Vec<u8> = vec![0u8; 64 * 4]; // hidden_dim * 4 bytes
    646
    647
    1
            let data = create_safetensors_bytes(&[
    648
    1
                ("model.embed_tokens.weight", "F32", &[100, 64], &embed_data),
    649
    1
                ("model.norm.weight", "F32", &[64], &norm_data),
    650
    1
            ]);
    651
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    652
    653
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    654
    1
            assert!(result.is_ok());
    655
    1
            let transformer = result.expect("operation failed");
    656
    1
            assert_eq!(transformer.config.intermediate_dim, 64 * 4);
    657
    1
        }
    658
    659
        #[test]
    660
    1
        fn test_convert_defaults_context_length_ext_cov() {
    661
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    662
    1
            let model_path = temp_dir.path().join("model.safetensors");
    663
    1
            let config_path = temp_dir.path().join("config.json");
    664
    665
            // Config without max_position_embeddings (should default to 2048)
    666
    1
            let config = r#"{
    667
    1
                "hidden_size": 64,
    668
    1
                "num_hidden_layers": 0,
    669
    1
                "num_attention_heads": 4,
    670
    1
                "vocab_size": 100,
    671
    1
                "intermediate_size": 256
    672
    1
            }"#;
    673
    1
            std::fs::write(&config_path, config).expect("write config");
    674
    675
    1
            let embed_data: Vec<u8> = vec![0u8; 100 * 64 * 4];
    676
    1
            let norm_data: Vec<u8> = vec![0u8; 64 * 4];
    677
    678
    1
            let data = create_safetensors_bytes(&[
    679
    1
                ("model.embed_tokens.weight", "F32", &[100, 64], &embed_data),
    680
    1
                ("model.norm.weight", "F32", &[64], &norm_data),
    681
    1
            ]);
    682
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    683
    684
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    685
    1
            assert!(result.is_ok());
    686
    1
            let transformer = result.expect("operation failed");
    687
    1
            assert_eq!(transformer.config.context_length, 2048);
    688
    1
        }
    689
    690
        #[test]
    691
    1
        fn test_convert_tied_embeddings_ext_cov() {
    692
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    693
    1
            let model_path = temp_dir.path().join("model.safetensors");
    694
    1
            let config_path = temp_dir.path().join("config.json");
    695
    696
    1
            let config = create_config_json(64, 0, 4, 100);
    697
    1
            std::fs::write(&config_path, config).expect("write config");
    698
    699
            // Create tensors WITHOUT lm_head.weight (tied embeddings)
    700
    1
            let embed_data: Vec<u8> = (0..(100 * 64))
    701
    6.40k
                .flat_map1(|i| (i as f32).to_le_bytes())
    702
    1
                .collect();
    703
    64
            let norm_data1: Vec<u8>1 = (0..64)1.flat_map1(|_| 1.0f32.to_le_bytes()).collect1();
    704
    705
    1
            let data = create_safetensors_bytes(&[
    706
    1
                ("model.embed_tokens.weight", "F32", &[100, 64], &embed_data),
    707
    1
                ("model.norm.weight", "F32", &[64], &norm_data),
    708
    1
            ]);
    709
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    710
    711
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    712
    1
            assert!(result.is_ok());
    713
    1
            let transformer = result.expect("operation failed");
    714
    715
            // lm_head_weight should have same dimensions as token_embedding (tied or separate)
    716
            // When tied: lm_head_weight.len() == token_embedding.len()
    717
            // But they may not be equal if transposed or if implementation uses separate weights
    718
    1
            assert!(!transformer.lm_head_weight.is_empty());
    719
    1
            assert!(!transformer.token_embedding.is_empty());
    720
    1
        }
    721
    722
        #[test]
    723
    1
        fn test_convert_separate_lm_head_ext_cov() {
    724
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    725
    1
            let model_path = temp_dir.path().join("model.safetensors");
    726
    1
            let config_path = temp_dir.path().join("config.json");
    727
    728
    1
            let config = create_config_json(64, 0, 4, 100);
    729
    1
            std::fs::write(&config_path, config).expect("write config");
    730
    731
            // Create tensors WITH separate lm_head.weight
    732
    1
            let embed_data: Vec<u8> = (0..(100 * 64))
    733
    6.40k
                .flat_map1(|i| (i as f32).to_le_bytes())
    734
    1
                .collect();
    735
    64
            let norm_data1: Vec<u8>1 = (0..64)1.flat_map1(|_| 1.0f32.to_le_bytes()).collect1();
    736
    1
            let lm_head_data: Vec<u8> = (0..(100 * 64))
    737
    6.40k
                .flat_map1(|i| ((i + 1000) as f32).to_le_bytes())
    738
    1
                .collect();
    739
    740
    1
            let data = create_safetensors_bytes(&[
    741
    1
                ("model.embed_tokens.weight", "F32", &[100, 64], &embed_data),
    742
    1
                ("model.norm.weight", "F32", &[64], &norm_data),
    743
    1
                ("lm_head.weight", "F32", &[100, 64], &lm_head_data),
    744
    1
            ]);
    745
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    746
    747
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    748
    1
            assert!(result.is_ok());
    749
    1
            let transformer = result.expect("operation failed");
    750
    751
            // lm_head_weight should NOT equal token_embedding
    752
    1
            assert_ne!(
    753
    1
                transformer.lm_head_weight[0],
    754
    1
                transformer.token_embedding[0]
    755
            );
    756
    1
        }
    757
    758
        #[test]
    759
    1
        fn test_convert_with_rope_theta_ext_cov() {
    760
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    761
    1
            let model_path = temp_dir.path().join("model.safetensors");
    762
    1
            let config_path = temp_dir.path().join("config.json");
    763
    764
            // Config with custom rope_theta
    765
    1
            let config = r#"{
    766
    1
                "hidden_size": 64,
    767
    1
                "num_hidden_layers": 0,
    768
    1
                "num_attention_heads": 4,
    769
    1
                "vocab_size": 100,
    770
    1
                "intermediate_size": 256,
    771
    1
                "rope_theta": 500000.0
    772
    1
            }"#;
    773
    1
            std::fs::write(&config_path, config).expect("write config");
    774
    775
    1
            let embed_data: Vec<u8> = vec![0u8; 100 * 64 * 4];
    776
    1
            let norm_data: Vec<u8> = vec![0u8; 64 * 4];
    777
    778
    1
            let data = create_safetensors_bytes(&[
    779
    1
                ("model.embed_tokens.weight", "F32", &[100, 64], &embed_data),
    780
    1
                ("model.norm.weight", "F32", &[64], &norm_data),
    781
    1
            ]);
    782
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    783
    784
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    785
    1
            assert!(result.is_ok());
    786
    1
            let transformer = result.expect("operation failed");
    787
    1
            assert!((transformer.config.rope_theta - 500000.0).abs() < 1.0);
    788
    1
        }
    789
    790
        #[test]
    791
    1
        fn test_convert_with_rms_norm_eps_ext_cov() {
    792
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    793
    1
            let model_path = temp_dir.path().join("model.safetensors");
    794
    1
            let config_path = temp_dir.path().join("config.json");
    795
    796
            // Config with custom rms_norm_eps
    797
    1
            let config = r#"{
    798
    1
                "hidden_size": 64,
    799
    1
                "num_hidden_layers": 0,
    800
    1
                "num_attention_heads": 4,
    801
    1
                "vocab_size": 100,
    802
    1
                "rms_norm_eps": 1e-5
    803
    1
            }"#;
    804
    1
            std::fs::write(&config_path, config).expect("write config");
    805
    806
    1
            let embed_data: Vec<u8> = vec![0u8; 100 * 64 * 4];
    807
    1
            let norm_data: Vec<u8> = vec![0u8; 64 * 4];
    808
    809
    1
            let data = create_safetensors_bytes(&[
    810
    1
                ("model.embed_tokens.weight", "F32", &[100, 64], &embed_data),
    811
    1
                ("model.norm.weight", "F32", &[64], &norm_data),
    812
    1
            ]);
    813
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    814
    815
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    816
    1
            assert!(result.is_ok());
    817
    1
            let transformer = result.expect("operation failed");
    818
    1
            assert!((transformer.config.eps - 1e-5).abs() < 1e-9);
    819
    1
        }
    820
    821
        #[test]
    822
    1
        fn test_safetensors_to_apr_converter_struct_ext_cov() {
    823
            // Test that SafetensorsToAprConverter is a unit struct
    824
    1
            let _converter = SafetensorsToAprConverter;
    825
            // This just ensures the struct exists and can be instantiated
    826
    1
        }
    827
    828
        #[test]
    829
    1
        fn test_convert_architecture_from_config_ext_cov() {
    830
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    831
    1
            let model_path = temp_dir.path().join("model.safetensors");
    832
    1
            let config_path = temp_dir.path().join("config.json");
    833
    834
    1
            let config = create_config_json(64, 0, 4, 100);
    835
    1
            std::fs::write(&config_path, config).expect("write config");
    836
    837
    1
            let embed_data: Vec<u8> = vec![0u8; 100 * 64 * 4];
    838
    1
            let norm_data: Vec<u8> = vec![0u8; 64 * 4];
    839
    840
    1
            let data = create_safetensors_bytes(&[
    841
    1
                ("model.embed_tokens.weight", "F32", &[100, 64], &embed_data),
    842
    1
                ("model.norm.weight", "F32", &[64], &norm_data),
    843
    1
            ]);
    844
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    845
    846
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    847
    1
            assert!(result.is_ok());
    848
    1
            let transformer = result.expect("operation failed");
    849
    1
            assert_eq!(transformer.config.architecture, "LlamaForCausalLM");
    850
    1
        }
    851
    852
        /// Helper to create all layer tensors for a single transformer layer
    853
    3
        fn create_layer_tensors(
    854
    3
            layer_idx: usize,
    855
    3
            hidden_dim: usize,
    856
    3
            intermediate_dim: usize,
    857
    3
        ) -> Vec<(&'static str, String, Vec<usize>, Vec<u8>)> {
    858
            // Lease the string from Box to get 'static lifetime approximation
    859
            // We'll build it differently - just use the layer_idx in the data
    860
    861
    3
            let prefix = format!("model.layers.{layer_idx}");
    862
    863
            // Calculate tensor sizes
    864
    3
            let attn_norm_size = hidden_dim;
    865
    3
            let q_size = hidden_dim * hidden_dim;
    866
    3
            let k_size = hidden_dim * hidden_dim;
    867
    3
            let v_size = hidden_dim * hidden_dim;
    868
    3
            let o_size = hidden_dim * hidden_dim;
    869
    3
            let ffn_norm_size = hidden_dim;
    870
    3
            let gate_size = hidden_dim * intermediate_dim;
    871
    3
            let up_size = hidden_dim * intermediate_dim;
    872
    3
            let down_size = intermediate_dim * hidden_dim;
    873
    874
    3
            vec![
    875
    3
                (
    876
    3
                    "attn_norm",
    877
    3
                    format!("{prefix}.input_layernorm.weight"),
    878
    3
                    vec![attn_norm_size],
    879
    3
                    vec![0u8; attn_norm_size * 4],
    880
    3
                ),
    881
    3
                (
    882
    3
                    "q_proj",
    883
    3
                    format!("{prefix}.self_attn.q_proj.weight"),
    884
    3
                    vec![hidden_dim, hidden_dim],
    885
    3
                    vec![0u8; q_size * 4],
    886
    3
                ),
    887
    3
                (
    888
    3
                    "k_proj",
    889
    3
                    format!("{prefix}.self_attn.k_proj.weight"),
    890
    3
                    vec![hidden_dim, hidden_dim],
    891
    3
                    vec![0u8; k_size * 4],
    892
    3
                ),
    893
    3
                (
    894
    3
                    "v_proj",
    895
    3
                    format!("{prefix}.self_attn.v_proj.weight"),
    896
    3
                    vec![hidden_dim, hidden_dim],
    897
    3
                    vec![0u8; v_size * 4],
    898
    3
                ),
    899
    3
                (
    900
    3
                    "o_proj",
    901
    3
                    format!("{prefix}.self_attn.o_proj.weight"),
    902
    3
                    vec![hidden_dim, hidden_dim],
    903
    3
                    vec![0u8; o_size * 4],
    904
    3
                ),
    905
    3
                (
    906
    3
                    "ffn_norm",
    907
    3
                    format!("{prefix}.post_attention_layernorm.weight"),
    908
    3
                    vec![ffn_norm_size],
    909
    3
                    vec![0u8; ffn_norm_size * 4],
    910
    3
                ),
    911
    3
                (
    912
    3
                    "gate_proj",
    913
    3
                    format!("{prefix}.mlp.gate_proj.weight"),
    914
    3
                    vec![intermediate_dim, hidden_dim],
    915
    3
                    vec![0u8; gate_size * 4],
    916
    3
                ),
    917
    3
                (
    918
    3
                    "up_proj",
    919
    3
                    format!("{prefix}.mlp.up_proj.weight"),
    920
    3
                    vec![intermediate_dim, hidden_dim],
    921
    3
                    vec![0u8; up_size * 4],
    922
    3
                ),
    923
    3
                (
    924
    3
                    "down_proj",
    925
    3
                    format!("{prefix}.mlp.down_proj.weight"),
    926
    3
                    vec![hidden_dim, intermediate_dim],
    927
    3
                    vec![0u8; down_size * 4],
    928
    3
                ),
    929
            ]
    930
    3
        }
    931
    932
        #[test]
    933
    1
        fn test_convert_with_single_layer_ext_cov() {
    934
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    935
    1
            let model_path = temp_dir.path().join("model.safetensors");
    936
    1
            let config_path = temp_dir.path().join("config.json");
    937
    938
    1
            let hidden_dim = 16;
    939
    1
            let intermediate_dim = 64;
    940
    1
            let vocab_size = 50;
    941
    1
            let num_layers = 1;
    942
    1
            let num_heads = 4;
    943
    944
            // Config
    945
    1
            let config = format!(
    946
    1
                r#"{{
    947
    1
                    "hidden_size": {},
    948
    1
                    "num_hidden_layers": {},
    949
    1
                    "num_attention_heads": {},
    950
    1
                    "vocab_size": {},
    951
    1
                    "intermediate_size": {},
    952
    1
                    "max_position_embeddings": 128,
    953
    1
                    "rms_norm_eps": 1e-6
    954
    1
                }}"#,
    955
                hidden_dim, num_layers, num_heads, vocab_size, intermediate_dim
    956
            );
    957
    1
            std::fs::write(&config_path, config).expect("write config");
    958
    959
            // Build layer tensors
    960
    1
            let layer_tensors = create_layer_tensors(0, hidden_dim, intermediate_dim);
    961
    962
            // Build safetensors with all required tensors
    963
    1
            let embed_data: Vec<u8> = vec![0u8; vocab_size * hidden_dim * 4];
    964
    1
            let norm_data: Vec<u8> = vec![0u8; hidden_dim * 4];
    965
    966
            // Create a comprehensive tensor list
    967
            use serde_json::json;
    968
    1
            let mut tensor_entries = serde_json::Map::new();
    969
    1
            let mut all_data = Vec::new();
    970
    1
            let mut offset = 0usize;
    971
    972
            // Add embed_tokens
    973
    1
            tensor_entries.insert(
    974
    1
                "model.embed_tokens.weight".to_string(),
    975
    1
                json!({
    976
    1
                    "dtype": "F32",
    977
    1
                    "shape": [vocab_size, hidden_dim],
    978
    1
                    "data_offsets": [offset, offset + embed_data.len()]
    979
                }),
    980
            );
    981
    1
            all_data.extend(&embed_data);
    982
    1
            offset += embed_data.len();
    983
    984
            // Add norm
    985
    1
            tensor_entries.insert(
    986
    1
                "model.norm.weight".to_string(),
    987
    1
                json!({
    988
    1
                    "dtype": "F32",
    989
    1
                    "shape": [hidden_dim],
    990
    1
                    "data_offsets": [offset, offset + norm_data.len()]
    991
                }),
    992
            );
    993
    1
            all_data.extend(&norm_data);
    994
    1
            offset += norm_data.len();
    995
    996
            // Add layer tensors
    997
    10
            for (_, name9, shape9, data9) in &layer_tensors {
    998
    9
                tensor_entries.insert(
    999
    9
                    name.clone(),
    1000
    9
                    json!({
    1001
    9
                        "dtype": "F32",
    1002
    9
                        "shape": shape,
    1003
    9
                        "data_offsets": [offset, offset + data.len()]
    1004
    9
                    }),
    1005
    9
                );
    1006
    9
                all_data.extend(data);
    1007
    9
                offset += data.len();
    1008
    9
            }
    1009
    1010
    1
            let json_obj = serde_json::Value::Object(tensor_entries);
    1011
    1
            let json_bytes = json_obj.to_string().into_bytes();
    1012
    1013
    1
            let mut safetensors_data = Vec::new();
    1014
    1
            safetensors_data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1015
    1
            safetensors_data.extend_from_slice(&json_bytes);
    1016
    1
            safetensors_data.extend_from_slice(&all_data);
    1017
    1018
    1
            std::fs::write(&model_path, safetensors_data).expect("write safetensors");
    1019
    1020
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    1021
    1
            assert!(result.is_ok(), "Conversion failed: {:?}"0, result0.err0());
    1022
    1023
    1
            let transformer = result.expect("operation failed");
    1024
    1
            assert_eq!(transformer.config.hidden_dim, hidden_dim);
    1025
    1
            assert_eq!(transformer.config.num_layers, num_layers);
    1026
    1
            assert_eq!(transformer.layers.len(), num_layers);
    1027
    1028
            // Verify layer structure
    1029
    1
            let layer = &transformer.layers[0];
    1030
    1
            assert_eq!(layer.attn_norm_weight.len(), hidden_dim);
    1031
    1
            assert_eq!(layer.qkv_weight.len(), hidden_dim * 3 * hidden_dim);
    1032
    1
            assert_eq!(layer.attn_output_weight.len(), hidden_dim * hidden_dim);
    1033
    1
            assert_eq!(layer.ffn_up_weight.len(), hidden_dim * intermediate_dim);
    1034
    1
            assert_eq!(layer.ffn_down_weight.len(), intermediate_dim * hidden_dim);
    1035
    1
            assert!(layer.ffn_gate_weight.is_some());
    1036
    1
            assert!(layer.ffn_norm_weight.is_some());
    1037
    1
        }
    1038
    1039
        #[test]
    1040
    1
        fn test_convert_with_multiple_layers_ext_cov() {
    1041
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    1042
    1
            let model_path = temp_dir.path().join("model.safetensors");
    1043
    1
            let config_path = temp_dir.path().join("config.json");
    1044
    1045
    1
            let hidden_dim = 8;
    1046
    1
            let intermediate_dim = 32;
    1047
    1
            let vocab_size = 20;
    1048
    1
            let num_layers = 2;
    1049
    1
            let num_heads = 2;
    1050
    1051
    1
            let config = format!(
    1052
    1
                r#"{{
    1053
    1
                    "hidden_size": {},
    1054
    1
                    "num_hidden_layers": {},
    1055
    1
                    "num_attention_heads": {},
    1056
    1
                    "vocab_size": {},
    1057
    1
                    "intermediate_size": {}
    1058
    1
                }}"#,
    1059
                hidden_dim, num_layers, num_heads, vocab_size, intermediate_dim
    1060
            );
    1061
    1
            std::fs::write(&config_path, config).expect("write config");
    1062
    1063
            // Build tensors for multiple layers
    1064
            use serde_json::json;
    1065
    1
            let mut tensor_entries = serde_json::Map::new();
    1066
    1
            let mut all_data = Vec::new();
    1067
    1
            let mut offset = 0usize;
    1068
    1069
            // Add embed_tokens
    1070
    1
            let embed_data: Vec<u8> = vec![0u8; vocab_size * hidden_dim * 4];
    1071
    1
            tensor_entries.insert(
    1072
    1
                "model.embed_tokens.weight".to_string(),
    1073
    1
                json!({
    1074
    1
                    "dtype": "F32",
    1075
    1
                    "shape": [vocab_size, hidden_dim],
    1076
    1
                    "data_offsets": [offset, offset + embed_data.len()]
    1077
                }),
    1078
            );
    1079
    1
            all_data.extend(&embed_data);
    1080
    1
            offset += embed_data.len();
    1081
    1082
            // Add norm
    1083
    1
            let norm_data: Vec<u8> = vec![0u8; hidden_dim * 4];
    1084
    1
            tensor_entries.insert(
    1085
    1
                "model.norm.weight".to_string(),
    1086
    1
                json!({
    1087
    1
                    "dtype": "F32",
    1088
    1
                    "shape": [hidden_dim],
    1089
    1
                    "data_offsets": [offset, offset + norm_data.len()]
    1090
                }),
    1091
            );
    1092
    1
            all_data.extend(&norm_data);
    1093
    1
            offset += norm_data.len();
    1094
    1095
            // Add all layer tensors
    1096
    2
            for layer_idx in 0..num_layers1 {
    1097
    2
                let layer_tensors = create_layer_tensors(layer_idx, hidden_dim, intermediate_dim);
    1098
    20
                for (_, name18, shape18, data18) in &layer_tensors {
    1099
    18
                    tensor_entries.insert(
    1100
    18
                        name.clone(),
    1101
    18
                        json!({
    1102
    18
                            "dtype": "F32",
    1103
    18
                            "shape": shape,
    1104
    18
                            "data_offsets": [offset, offset + data.len()]
    1105
    18
                        }),
    1106
    18
                    );
    1107
    18
                    all_data.extend(data);
    1108
    18
                    offset += data.len();
    1109
    18
                }
    1110
            }
    1111
    1112
    1
            let json_obj = serde_json::Value::Object(tensor_entries);
    1113
    1
            let json_bytes = json_obj.to_string().into_bytes();
    1114
    1115
    1
            let mut safetensors_data = Vec::new();
    1116
    1
            safetensors_data.extend_from_slice(&(json_bytes.len() as u64).to_le_bytes());
    1117
    1
            safetensors_data.extend_from_slice(&json_bytes);
    1118
    1
            safetensors_data.extend_from_slice(&all_data);
    1119
    1120
    1
            std::fs::write(&model_path, safetensors_data).expect("write safetensors");
    1121
    1122
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    1123
    1
            assert!(result.is_ok(), "Conversion failed: {:?}"0, result0.err0());
    1124
    1125
    1
            let transformer = result.expect("operation failed");
    1126
    1
            assert_eq!(transformer.layers.len(), num_layers);
    1127
    1
        }
    1128
    1129
        #[test]
    1130
    1
        fn test_extract_layer_missing_input_layernorm_ext_cov() {
    1131
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    1132
    1
            let model_path = temp_dir.path().join("model.safetensors");
    1133
    1
            let config_path = temp_dir.path().join("config.json");
    1134
    1135
    1
            let config = create_config_json(16, 1, 4, 50);
    1136
    1
            std::fs::write(&config_path, config).expect("write config");
    1137
    1138
            // Missing layer 0 input_layernorm
    1139
    1
            let embed_data: Vec<u8> = vec![0u8; 50 * 16 * 4];
    1140
    1
            let norm_data: Vec<u8> = vec![0u8; 16 * 4];
    1141
    1
            let data = create_safetensors_bytes(&[
    1142
    1
                ("model.embed_tokens.weight", "F32", &[50, 16], &embed_data),
    1143
    1
                ("model.norm.weight", "F32", &[16], &norm_data),
    1144
    1
            ]);
    1145
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    1146
    1147
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    1148
    1
            assert!(result.is_err());
    1149
    1
        }
    1150
    1151
        #[test]
    1152
    1
        fn test_convert_with_num_kv_heads_ext_cov() {
    1153
    1
            let temp_dir = TempDir::new().expect("create temp dir");
    1154
    1
            let model_path = temp_dir.path().join("model.safetensors");
    1155
    1
            let config_path = temp_dir.path().join("config.json");
    1156
    1157
            // Config with GQA (num_key_value_heads < num_attention_heads)
    1158
    1
            let config = r#"{
    1159
    1
                "hidden_size": 64,
    1160
    1
                "num_hidden_layers": 0,
    1161
    1
                "num_attention_heads": 8,
    1162
    1
                "num_key_value_heads": 4,
    1163
    1
                "vocab_size": 100
    1164
    1
            }"#;
    1165
    1
            std::fs::write(&config_path, config).expect("write config");
    1166
    1167
    1
            let embed_data: Vec<u8> = vec![0u8; 100 * 64 * 4];
    1168
    1
            let norm_data: Vec<u8> = vec![0u8; 64 * 4];
    1169
    1
            let data = create_safetensors_bytes(&[
    1170
    1
                ("model.embed_tokens.weight", "F32", &[100, 64], &embed_data),
    1171
    1
                ("model.norm.weight", "F32", &[64], &norm_data),
    1172
    1
            ]);
    1173
    1
            std::fs::write(&model_path, data).expect("write safetensors");
    1174
    1175
    1
            let result = SafetensorsToAprConverter::convert(&model_path);
    1176
    1
            assert!(result.is_ok());
    1177
    1
            let transformer = result.expect("operation failed");
    1178
    1
            assert_eq!(transformer.config.num_heads, 8);
    1179
    1
            assert_eq!(transformer.config.num_kv_heads, 4);
    1180
    1
        }
    1181
    }