## 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/apr/mod.rs

    Line
    Count
    Source
    1
    //! Aprender .apr format support for realizar (APR v2 only)
    2
    //!
    3
    //! This module provides loading and inference for models in Aprender's native
    4
    //! .apr v2 format (Magic: `APR\0` = 0x41505232).
    5
    //!
    6
    //! ## Format Structure (APR v2, 64-byte header)
    7
    //!
    8
    //! ```text
    9
    //! ┌─────────────────────────────────────────────────────────────┐
    10
    //! │ Header (64 bytes)                                           │
    11
    //! │   - Magic: "APR\0" (4 bytes)                                 │
    12
    //! │   - Version: major.minor (2 bytes)                          │
    13
    //! │   - Flags (2 bytes)                                         │
    14
    //! │   - Tensor count (4 bytes)                                  │
    15
    //! │   - Metadata offset (8 bytes)                               │
    16
    //! │   - Metadata size (4 bytes)                                 │
    17
    //! │   - Tensor index offset (8 bytes)                           │
    18
    //! │   - Data offset (8 bytes)                                   │
    19
    //! │   - Checksum (4 bytes)                                      │
    20
    //! │   - Reserved (20 bytes)                                     │
    21
    //! ├─────────────────────────────────────────────────────────────┤
    22
    //! │ JSON Metadata (padded to 64-byte boundary)                  │
    23
    //! ├─────────────────────────────────────────────────────────────┤
    24
    //! │ Tensor Index (sorted by name)                               │
    25
    //! ├─────────────────────────────────────────────────────────────┤
    26
    //! │ Tensor Data (each tensor 64-byte aligned)                   │
    27
    //! └─────────────────────────────────────────────────────────────┘
    28
    //! ```
    29
    //!
    30
    //! ## Example
    31
    //!
    32
    //! ```rust,ignore
    33
    //! use realizar::apr::AprV2Model;
    34
    //!
    35
    //! let model = AprV2Model::load("model.apr")?;
    36
    //! println!("Tensors: {}", model.tensor_count());
    37
    //! ```
    38
    39
    use std::collections::HashMap;
    40
    use std::fs::{self, File};
    41
    use std::path::{Path, PathBuf};
    42
    43
    use serde::{Deserialize, Serialize};
    44
    45
    use crate::error::{RealizarError, Result};
    46
    47
    // PMAT-802: Extracted modules
    48
    #[cfg(feature = "cuda")]
    49
    mod cuda;
    50
    mod helpers;
    51
    mod tokenizer;
    52
    53
    #[cfg(feature = "cuda")]
    54
    pub use cuda::AprV2ModelCuda;
    55
    pub use helpers::{is_apr_file, detect_format, simd_dot};
    56
    use helpers::{rms_norm, matmul, simple_attention};
    57
    #[cfg(feature = "cuda")]
    58
    use helpers::transpose_matrix;
    59
    pub use tokenizer::{BpeTokenizer, SimpleTokenizer, byte_to_bpe_char};
    60
    use tokenizer::bpe_encode;
    61
    62
    // ============================================================================
    63
    // Memory-mapped model data (Heijunka - Level Loading)
    64
    // ============================================================================
    65
    //
    66
    // References:
    67
    // - Didona et al. (2022): mmap vs read() achieves 2.3x throughput for sequential access
    68
    // - Chu (2011): LMDB design - let kernel manage pages, don't fight the VM subsystem
    69
    // - Vahalia (1996): SIGBUS behavior on truncated mmap
    70
    //
    71
    // This abstraction allows models to be loaded via:
    72
    // 1. Memory mapping (mmap) - zero-copy, kernel manages pages, no zram pressure
    73
    // 2. Heap allocation (Vec<u8>) - required for compressed files after decompression
    74
    75
    /// Model data storage abstraction for zero-copy access.
    76
    ///
    77
    /// # Memory Management
    78
    ///
    79
    /// When using `Mmap` variant:
    80
    /// - Data is not copied into userspace heap
    81
    /// - Kernel demand-pages from disk on access
    82
    /// - After GPU transfer, call `release_cpu_pages()` to advise kernel
    83
    /// - Pages backed by file (not zram) when evicted
    84
    ///
    85
    /// When using `Heap` variant:
    86
    /// - Used for compressed files (must decompress to Vec<u8>)
    87
    /// - Standard heap allocation behavior
    88
    /// - May be compressed to zram when idle
    89
    #[derive(Debug)]
    90
    pub enum ModelData {
    91
        /// Memory-mapped file (zero-copy, kernel-managed paging)
    92
        #[cfg(not(target_arch = "wasm32"))]
    93
        Mmap {
    94
            /// Memory-mapped region
    95
            mmap: memmap2::Mmap,
    96
            /// Original file path (for diagnostics)
    97
            path: PathBuf,
    98
        },
    99
        /// Heap-allocated data (for compressed files or WASM)
    100
        Heap(Vec<u8>),
    101
    }
    102
    103
    impl ModelData {
    104
        /// Open a file with memory mapping.
    105
        ///
    106
        /// # Safety
    107
        ///
    108
        /// Uses `memmap2::Mmap` which requires:
    109
        /// - File must not be truncated while mapped (SIGBUS on Unix)
    110
        /// - File must not be modified while mapped (undefined behavior)
    111
        ///
    112
        /// # References
    113
        ///
    114
        /// - Vahalia (1996): SIGBUS from truncated mmap
    115
        /// - memmap2 crate safety documentation
    116
        #[cfg(not(target_arch = "wasm32"))]
    117
        #[allow(unsafe_code)]
    118
    8
        pub fn open_mmap(path: impl AsRef<Path>) -> Result<Self> {
    119
    8
            let path_ref = path.as_ref();
    120
    8
            let file7 = File::open(path_ref).map_err(|e| RealizarError::IoError {
    121
    1
                message: format!("Failed to open file '{}': {e}", path_ref.display()),
    122
    1
            })?;
    123
    124
            // SAFETY: File is opened read-only. We document the single-writer
    125
            // assumption. Callers should validate checksums before trusting data.
    126
            // SIGBUS can occur if file is truncated externally - this is documented.
    127
    7
            let mmap = unsafe {
    128
    7
                memmap2::MmapOptions::new()
    129
    7
                    .map(&file)
    130
    7
                    .map_err(|e| RealizarError::IoError {
    131
    0
                        message: format!("Failed to mmap file '{}': {e}", path_ref.display()),
    132
    0
                    })?
    133
            };
    134
    135
    7
            Ok(Self::Mmap {
    136
    7
                mmap,
    137
    7
                path: path_ref.to_path_buf(),
    138
    7
            })
    139
    8
        }
    140
    141
        /// Create from heap-allocated data (for compressed files).
    142
        #[must_use]
    143
    162
        pub fn from_vec(data: Vec<u8>) -> Self {
    144
    162
            Self::Heap(data)
    145
    162
        }
    146
    147
        /// Get the data as a byte slice.
    148
        #[must_use]
    149
    183
        pub fn as_slice(&self) -> &[u8] {
    150
    183
            match self {
    151
                #[cfg(not(target_arch = "wasm32"))]
    152
    7
                Self::Mmap { mmap, .. } => mmap,
    153
    176
                Self::Heap(data) => data,
    154
            }
    155
    183
        }
    156
    157
        /// Get data length.
    158
        #[must_use]
    159
    8
        pub fn len(&self) -> usize {
    160
    8
            self.as_slice().len()
    161
    8
        }
    162
    163
        /// Check if data is empty.
    164
        #[must_use]
    165
    8
        pub fn is_empty(&self) -> bool {
    166
    8
            self.as_slice().is_empty()
    167
    8
        }
    168
    169
        /// Release CPU pages after GPU transfer (Unix only).
    170
        ///
    171
        /// Calls `madvise(MADV_DONTNEED)` to tell the kernel these pages
    172
        /// are no longer needed. The kernel will:
    173
        /// - Drop pages immediately (not compress to zram)
    174
        /// - Re-fault from disk if accessed again
    175
        ///
    176
        /// # When to Call
    177
        ///
    178
        /// After `cuMemcpy()` completes for all tensors.
    179
        ///
    180
        /// # Safety
    181
        ///
    182
        /// Uses `unchecked_advise` because `MADV_DONTNEED` is in the
    183
        /// `UncheckedAdvice` enum. This is safe for read-only mmaps where
    184
        /// data can be re-faulted from the backing file.
    185
        ///
    186
        /// # References
    187
        ///
    188
        /// - Didona et al. (2022): madvise for memory management
    189
        #[cfg(all(unix, not(target_arch = "wasm32")))]
    190
        #[allow(unsafe_code)]
    191
    3
        pub fn release_cpu_pages(&self) -> Result<()> {
    192
    3
            match self {
    193
    2
                Self::Mmap { mmap, path } => {
    194
                    // SAFETY: We opened the file read-only, so MADV_DONTNEED is safe -
    195
                    // the kernel will re-fault pages from the backing file if accessed.
    196
                    unsafe {
    197
    2
                        mmap.unchecked_advise(memmap2::UncheckedAdvice::DontNeed)
    198
    2
                            .map_err(|e| RealizarError::IoError {
    199
    0
                                message: format!(
    200
    0
                                    "madvise(MADV_DONTNEED) failed for '{}': {e}",
    201
    0
                                    path.display()
    202
                                ),
    203
    0
                            })
    204
                    }
    205
                },
    206
                Self::Heap(_) => {
    207
                    // No-op for heap data - kernel manages via normal VM pressure
    208
    1
                    Ok(())
    209
                },
    210
            }
    211
    3
        }
    212
    213
        /// Advise sequential access pattern (Unix only).
    214
        ///
    215
        /// Call before linear scan through model data.
    216
        #[cfg(all(unix, not(target_arch = "wasm32")))]
    217
    4
        pub fn advise_sequential(&self) -> Result<()> {
    218
    4
            match self {
    219
    3
                Self::Mmap { mmap, path } => {
    220
    3
                    mmap.advise(memmap2::Advice::Sequential)
    221
    3
                        .map_err(|e| RealizarError::IoError {
    222
    0
                            message: format!(
    223
    0
                                "madvise(MADV_SEQUENTIAL) failed for '{}': {e}",
    224
    0
                                path.display()
    225
                            ),
    226
    0
                        })
    227
                },
    228
    1
                Self::Heap(_) => Ok(()),
    229
            }
    230
    4
        }
    231
    232
        /// Check if this is memory-mapped data.
    233
        #[must_use]
    234
    7
        pub fn is_mmap(&self) -> bool {
    235
    7
            match self {
    236
                #[cfg(not(target_arch = "wasm32"))]
    237
    3
                Self::Mmap { .. } => true,
    238
    4
                Self::Heap(_) => false,
    239
            }
    240
    7
        }
    241
    }
    242
    243
    /// Magic number: "APR" followed by version byte
    244
    /// - Legacy: APR\0 (0x41, 0x50, 0x52, 0x00)
    245
    /// - v1: APR1 (0x41, 0x50, 0x52, 0x31)
    246
    /// - v2: APR2 (0x41, 0x50, 0x52, 0x32)
    247
    pub const MAGIC_PREFIX: [u8; 3] = [0x41, 0x50, 0x52]; // "APR"
    248
    249
    /// Legacy magic for compatibility
    250
    pub const MAGIC: [u8; 4] = [0x41, 0x50, 0x52, 0x00];
    251
    252
    /// Header size in bytes (64-byte aligned)
    253
    pub const HEADER_SIZE: usize = 64;
    254
    255
    /// Tensor alignment in bytes
    256
    pub const ALIGNMENT: usize = 64;
    257
    258
    // ============================================================================
    259
    // Dequantization helpers for quantized tensor formats
    260
    // ============================================================================
    261
    262
    /// Convert F16 (IEEE 754 half-precision) to F32
    263
    #[inline]
    264
    67
    pub(crate) fn f16_to_f32(bits: u16) -> f32 {
    265
    67
        let sign = u32::from((bits >> 15) & 1);
    266
    67
        let exp = u32::from((bits >> 10) & 0x1F);
    267
    67
        let mant = u32::from(bits & 0x3FF);
    268
    269
    67
        if exp == 0 {
    270
    14
            if mant == 0 {
    271
                // Zero
    272
    8
                f32::from_bits(sign << 31)
    273
            } else {
    274
                // Subnormal - convert to normalized f32
    275
    6
                let mut m = mant;
    276
    6
                let mut e = 0i32;
    277
    48
                while (m & 0x400) == 0 {
    278
    42
                    m <<= 1;
    279
    42
                    e -= 1;
    280
    42
                }
    281
    6
                m &= 0x3FF;
    282
    6
                let f32_exp = (127 - 15 + 1 + e) as u32;
    283
    6
                f32::from_bits((sign << 31) | (f32_exp << 23) | (m << 13))
    284
            }
    285
    53
        } else if exp == 31 {
    286
            // Inf or NaN
    287
    7
            if mant == 0 {
    288
    4
                f32::from_bits((sign << 31) | (0xFF << 23))
    289
            } else {
    290
    3
                f32::from_bits((sign << 31) | (0xFF << 23) | (mant << 13))
    291
            }
    292
        } else {
    293
            // Normal number
    294
    46
            let f32_exp = (exp as i32 - 15 + 127) as u32;
    295
    46
            f32::from_bits((sign << 31) | (f32_exp << 23) | (mant << 13))
    296
        }
    297
    67
    }
    298
    299
    /// Dequantize F16 data to F32
    300
    9
    pub(crate) fn dequantize_f16(bytes: &[u8], num_elements: usize) -> Vec<f32> {
    301
    9
        let mut result = Vec::with_capacity(num_elements);
    302
    15
        for chunk in bytes9.chunks_exact9(2) {
    303
    15
            let bits = u16::from_le_bytes([chunk[0], chunk[1]]);
    304
    15
            result.push(f16_to_f32(bits));
    305
    15
        }
    306
    9
        result.truncate(num_elements);
    307
    9
        result
    308
    9
    }
    309
    310
    /// Dequantize Q8_0 format (GGUF compatible)
    311
    /// Q8_0: blocks of 32 elements, each block has 2-byte f16 scale + 32 bytes of int8 quants
    312
    9
    pub(crate) fn dequantize_q8_0(bytes: &[u8], num_elements: usize) -> Vec<f32> {
    313
        const BLOCK_SIZE: usize = 32;
    314
        const BLOCK_BYTES: usize = 2 + 32; // f16 scale + 32 int8 values
    315
    316
    9
        let mut result = Vec::with_capacity(num_elements);
    317
    9
        let mut offset = 0;
    318
    319
    21
        while result.len() < num_elements && offset + BLOCK_BYTES13 <= bytes.len() {
    320
            // Read scale (f16)
    321
    12
            let scale_bits = u16::from_le_bytes([bytes[offset], bytes[offset + 1]]);
    322
    12
            let scale = f16_to_f32(scale_bits);
    323
    12
            offset += 2;
    324
    325
            // Read 32 int8 values
    326
    374
            for i363 in 0..32 {
    327
    363
                if result.len() >= num_elements {
    328
    1
                    break;
    329
    362
                }
    330
    362
                let v = f32::from(bytes[offset + i] as i8);
    331
    362
                result.push(v * scale);
    332
            }
    333
    12
            offset += 32;
    334
        }
    335
    336
    9
        result.truncate(num_elements);
    337
    9
        result
    338
    9
    }
    339
    340
    /// Dequantize Q4_K format (GGUF K-quants)
    341
    /// Q4_K: super blocks of 256 elements
    342
    /// Each super block: d (f16) + dmin (f16) + scales (12 bytes) + qs (128 bytes) = 144 bytes
    343
    7
    pub(crate) fn dequantize_q4_k(bytes: &[u8], num_elements: usize) -> Vec<f32> {
    344
        const SUPER_BLOCK_SIZE: usize = 256;
    345
        const SUPER_BLOCK_BYTES: usize = 2 + 2 + 12 + 128; // 144 bytes
    346
    347
    7
        let mut result = Vec::with_capacity(num_elements);
    348
    7
        let mut offset = 0;
    349
    350
    13
        while result.len() < num_elements && offset + SUPER_BLOCK_BYTES7 <= bytes.len() {
    351
            // Read d (f16 scale) and dmin (f16 min)
    352
    6
            let d = f16_to_f32(u16::from_le_bytes([bytes[offset], bytes[offset + 1]]));
    353
    6
            let dmin = f16_to_f32(u16::from_le_bytes([bytes[offset + 2], bytes[offset + 3]]));
    354
    6
            offset += 4;
    355
    356
            // Read scales (12 bytes = 8 6-bit scale values packed)
    357
    6
            let scales_bytes = &bytes[offset..offset + 12];
    358
    6
            let mut scales = [0u8; 8];
    359
    6
            let mut mins = [0u8; 8];
    360
    361
            // Unpack 6-bit scales and mins from 12 bytes
    362
    30
            for i24 in 0..4 {
    363
    24
                scales[i] = scales_bytes[i] & 0x3F;
    364
    24
                scales[i + 4] = scales_bytes[i + 4] & 0x3F;
    365
    24
                mins[i] = (scales_bytes[i] >> 6) | ((scales_bytes[i + 8] & 0x0F) << 2);
    366
    24
                mins[i + 4] = (scales_bytes[i + 4] >> 6) | ((scales_bytes[i + 8] >> 4) << 2);
    367
    24
            }
    368
    6
            offset += 12;
    369
    370
            // Read 128 bytes = 256 4-bit quantized values
    371
    6
            let qs = &bytes[offset..offset + 128];
    372
    6
            offset += 128;
    373
    374
            // Dequantize: each sub-block has 32 elements (8 sub-blocks total)
    375
    54
            for j48 in 0..8 {
    376
    48
                let scale = d * f32::from(scales[j]);
    377
    48
                let min_val = dmin * f32::from(mins[j]);
    378
    379
    619
                for l584 in 0..16 {
    380
    584
                    if result.len() >= num_elements {
    381
    13
                        break;
    382
    571
                    }
    383
    571
                    let q_byte = qs[j * 16 + l];
    384
    571
                    let q0 = (q_byte & 0x0F) as f32;
    385
    571
                    let q1 = (q_byte >> 4) as f32;
    386
    571
                    result.push(q0 * scale - min_val);
    387
    571
                    if result.len() < num_elements {
    388
    570
                        result.push(q1 * scale - min_val);
    389
    570
                    }1
    390
                }
    391
            }
    392
        }
    393
    394
    7
        result.truncate(num_elements);
    395
    7
        result
    396
    7
    }
    397
    398
    /// Dequantize Q6_K format (GGUF K-quants)
    399
    /// Q6_K: super blocks of 256 elements
    400
    /// Each super block: ql (128 bytes) + qh (64 bytes) + scales (16 bytes) + d (f16) = 210 bytes
    401
    7
    pub(crate) fn dequantize_q6_k(bytes: &[u8], num_elements: usize) -> Vec<f32> {
    402
        const SUPER_BLOCK_SIZE: usize = 256;
    403
        const SUPER_BLOCK_BYTES: usize = 128 + 64 + 16 + 2; // 210 bytes
    404
    405
    7
        let mut result = Vec::with_capacity(num_elements);
    406
    7
        let mut offset = 0;
    407
    408
    13
        while result.len() < num_elements && offset + SUPER_BLOCK_BYTES7 <= bytes.len() {
    409
            // Read ql (128 bytes = low 4 bits of 256 6-bit values)
    410
    6
            let ql = &bytes[offset..offset + 128];
    411
    6
            offset += 128;
    412
    413
            // Read qh (64 bytes = high 2 bits of 256 6-bit values)
    414
    6
            let qh = &bytes[offset..offset + 64];
    415
    6
            offset += 64;
    416
    417
            // Read scales (16 bytes = 16 int8 scales)
    418
    6
            let scales = &bytes[offset..offset + 16];
    419
    6
            offset += 16;
    420
    421
            // Read d (f16)
    422
    6
            let d = f16_to_f32(u16::from_le_bytes([bytes[offset], bytes[offset + 1]]));
    423
    6
            offset += 2;
    424
    425
            // Dequantize 16 sub-blocks of 16 elements each
    426
    102
            for j96 in 0..16 {
    427
    96
                let scale = d * f32::from(scales[j] as i8);
    428
    429
    650
                for l581 in 0..8 {
    430
    581
                    if result.len() >= num_elements {
    431
    27
                        break;
    432
    554
                    }
    433
    554
                    let idx = j * 8 + l;
    434
    554
                    let ql_byte = ql[idx];
    435
    554
                    let qh_byte = qh[idx / 2];
    436
    437
                    // Extract two 6-bit values
    438
    554
                    let qh_shift = (l % 2) * 4;
    439
    554
                    let q0 = ((ql_byte & 0x0F) | ((qh_byte >> qh_shift) & 0x03) << 4) as i8 - 32;
    440
    554
                    let q1 = ((ql_byte >> 4) | (((qh_byte >> qh_shift) >> 2) & 0x03) << 4) as i8 - 32;
    441
    442
    554
                    result.push(f32::from(q0) * scale);
    443
    554
                    if result.len() < num_elements {
    444
    553
                        result.push(f32::from(q1) * scale);
    445
    553
                    }1
    446
                }
    447
            }
    448
        }
    449
    450
    7
        result.truncate(num_elements);
    451
    7
        result
    452
    7
    }
    453
    454
    // ============================================================================
    455
    // Quantization type mapping for GPU kernels
    456
    // ============================================================================
    457
    458
    /// Map APR dtype string to GGML quantization type ID.
    459
    ///
    460
    /// These IDs are used by `load_quantized_weights_with_type()` to select
    461
    /// the correct GPU dequantization kernel (Q4K GEMV, Q6K GEMV, etc.).
    462
    #[inline]
    463
    43
    pub(crate) fn dtype_to_ggml_qtype(dtype: &str) -> Option<u32> {
    464
    43
        match dtype {
    465
    43
            "Q4_K" | "q4_k"40 => Some(12)5, // GGML_TYPE_Q4_K
    466
    38
            "Q5_K" | "q5_k"35 => Some(13)4, // GGML_TYPE_Q5_K
    467
    34
            "Q6_K" | "q6_k"31 => Some(14)4, // GGML_TYPE_Q6_K
    468
    30
            "Q8_0" | "q8_0"26 => Some(8)5,  // GGML_TYPE_Q8_0
    469
    25
            "Q4_0" | "q4_0"21 => Some(2)5,  // GGML_TYPE_Q4_0
    470
    20
            "Q4_1" | "q4_1"17 => Some(3)4,  // GGML_TYPE_Q4_1
    471
    16
            "Q5_0" | "q5_0"13 => Some(6)4,  // GGML_TYPE_Q5_0
    472
    12
            _ => None,                   // F32/F16 are not quantized
    473
        }
    474
    43
    }
    475
    476
    /// Check if dtype is a quantized format that can use GPU dequant kernels.
    477
    #[inline]
    478
    18
    pub(crate) fn is_quantized_dtype(dtype: &str) -> bool {
    479
    18
        dtype_to_ggml_qtype(dtype).is_some()
    480
    18
    }
    481
    482
    /// APR v2 feature flags
    483
    #[derive(Debug, Clone, Copy, Default)]
    484
    pub struct AprFlags(u16);
    485
    486
    impl AprFlags {
    487
        /// LZ4 compression enabled
    488
        pub const LZ4_COMPRESSED: u16 = 0x0001;
    489
        /// Zstandard compression enabled
    490
        pub const ZSTD_COMPRESSED: u16 = 0x0002;
    491
        /// Model is encrypted
    492
        pub const ENCRYPTED: u16 = 0x0004;
    493
        /// Model has cryptographic signature
    494
        pub const SIGNED: u16 = 0x0008;
    495
        /// Model is sharded across multiple files
    496
        pub const SHARDED: u16 = 0x0010;
    497
        /// Weights are quantized (int8/int4)
    498
        pub const QUANTIZED: u16 = 0x0020;
    499
        /// Model includes embedded vocabulary
    500
        pub const HAS_VOCAB: u16 = 0x0200;
    501
    502
        /// Create flags from raw bits
    503
        #[must_use]
    504
    194
        pub const fn new(bits: u16) -> Self {
    505
    194
            Self(bits)
    506
    194
        }
    507
    508
        /// Check if model uses compression (LZ4 or Zstd)
    509
        #[must_use]
    510
    166
        pub const fn is_compressed(&self) -> bool {
    511
    166
            self.0 & (Self::LZ4_COMPRESSED | Self::ZSTD_COMPRESSED) != 0
    512
    166
        }
    513
    514
        /// Check if model uses LZ4 compression
    515
        #[must_use]
    516
    10
        pub const fn is_lz4(&self) -> bool {
    517
    10
            self.0 & Self::LZ4_COMPRESSED != 0
    518
    10
        }
    519
    520
        /// Check if model uses ZSTD compression
    521
        #[must_use]
    522
    8
        pub const fn is_zstd(&self) -> bool {
    523
    8
            self.0 & Self::ZSTD_COMPRESSED != 0
    524
    8
        }
    525
    526
        /// Check if model is encrypted
    527
        #[must_use]
    528
    161
        pub const fn is_encrypted(&self) -> bool {
    529
    161
            self.0 & Self::ENCRYPTED != 0
    530
    161
        }
    531
    532
        /// Check if weights are quantized
    533
        #[must_use]
    534
    10
        pub const fn is_quantized(&self) -> bool {
    535
    10
            self.0 & Self::QUANTIZED != 0
    536
    10
        }
    537
    538
        /// Check if model includes embedded vocabulary
    539
        #[must_use]
    540
    4
        pub const fn has_vocab(&self) -> bool {
    541
    4
            self.0 & Self::HAS_VOCAB != 0
    542
    4
        }
    543
    }
    544
    545
    /// APR v2 file header (64 bytes)
    546
    #[derive(Debug, Clone)]
    547
    pub struct AprHeader {
    548
        /// Magic number ("APR\0")
    549
        pub magic: [u8; 4],
    550
        /// Format version (major, minor)
    551
        pub version: (u8, u8),
    552
        /// Feature flags
    553
        pub flags: AprFlags,
    554
        /// Number of tensors
    555
        pub tensor_count: u32,
    556
        /// Offset to metadata section
    557
        pub metadata_offset: u64,
    558
        /// Size of metadata section
    559
        pub metadata_size: u32,
    560
        /// Offset to tensor index
    561
        pub tensor_index_offset: u64,
    562
        /// Offset to tensor data
    563
        pub data_offset: u64,
    564
        /// Header checksum (CRC32)
    565
        pub checksum: u32,
    566
    }
    567
    568
    impl AprHeader {
    569
        /// Parse header from bytes
    570
    187
        pub fn from_bytes(data: &[u8]) -> Result<Self> {
    571
    187
            if data.len() < HEADER_SIZE {
    572
    5
                return Err(RealizarError::FormatError {
    573
    5
                    reason: format!(
    574
    5
                        ".apr header too small: {} bytes (need {})",
    575
    5
                        data.len(),
    576
    5
                        HEADER_SIZE
    577
    5
                    ),
    578
    5
                });
    579
    182
            }
    580
    581
            // Check magic - first 3 bytes must be "APR", 4th byte is version
    582
    182
            let magic: [u8; 4] = data[0..4]
    583
    182
                .try_into()
    584
    182
                .map_err(|_| RealizarError::FormatError {
    585
    0
                    reason: "Failed to read magic bytes".to_string(),
    586
    0
                })?;
    587
    588
            // Validate magic prefix (APR)
    589
    182
            if magic[0..3] != MAGIC_PREFIX {
    590
    7
                return Err(RealizarError::FormatError {
    591
    7
                    reason: format!(
    592
    7
                        "Invalid .apr magic: expected APR {:?}, got {:?}",
    593
    7
                        MAGIC_PREFIX, &magic[0..3]
    594
    7
                    ),
    595
    7
                });
    596
    175
            }
    597
    598
            // Validate version byte (0, '1', or '2')
    599
    175
            let version_byte = magic[3];
    600
    175
            if version_byte != 0 && version_byte != b'1'0 && version_byte != b'2'0 {
    601
    0
                return Err(RealizarError::FormatError {
    602
    0
                    reason: format!(
    603
    0
                        "Invalid .apr version byte: expected 0, '1', or '2', got {}",
    604
    0
                        version_byte
    605
    0
                    ),
    606
    0
                });
    607
    175
            }
    608
    609
            // APR v1 (magic "APR1") has different header layout - not supported for inference
    610
            // APR v1 is used by Whisper models but has inline tensor index format
    611
    175
            if version_byte == b'1' {
    612
    0
                return Err(RealizarError::UnsupportedOperation {
    613
    0
                    operation: "load_apr_v1".to_string(),
    614
    0
                    reason: "APR v1 format not supported for inference. \
    615
    0
                            Use 'apr convert model.apr -o model_v2.apr --format apr2' \
    616
    0
                            to convert to APR v2 format, or use the GGUF version.".to_string(),
    617
    0
                });
    618
    175
            }
    619
    620
    175
            let version = (data[4], data[5]);
    621
    175
            let flags = AprFlags::new(u16::from_le_bytes([data[6], data[7]]));
    622
    175
            let tensor_count = u32::from_le_bytes([data[8], data[9], data[10], data[11]]);
    623
    175
            let metadata_offset = u64::from_le_bytes([
    624
    175
                data[12], data[13], data[14], data[15], data[16], data[17], data[18], data[19],
    625
    175
            ]);
    626
    175
            let metadata_size = u32::from_le_bytes([data[20], data[21], data[22], data[23]]);
    627
    175
            let tensor_index_offset = u64::from_le_bytes([
    628
    175
                data[24], data[25], data[26], data[27], data[28], data[29], data[30], data[31],
    629
    175
            ]);
    630
    175
            let data_offset = u64::from_le_bytes([
    631
    175
                data[32], data[33], data[34], data[35], data[36], data[37], data[38], data[39],
    632
    175
            ]);
    633
    175
            let checksum = u32::from_le_bytes([data[40], data[41], data[42], data[43]]);
    634
    635
    175
            Ok(Self {
    636
    175
                magic,
    637
    175
                version,
    638
    175
                flags,
    639
    175
                tensor_count,
    640
    175
                metadata_offset,
    641
    175
                metadata_size,
    642
    175
                tensor_index_offset,
    643
    175
                data_offset,
    644
    175
                checksum,
    645
    175
            })
    646
    187
        }
    647
    }
    648
    649
    /// Tensor entry in the index
    650
    #[derive(Debug, Clone, Serialize, Deserialize)]
    651
    pub struct TensorEntry {
    652
        /// Tensor name (e.g., "model.layers.0.attention.wq")
    653
        pub name: String,
    654
        /// Data type (e.g., "F32", "F16", "BF16", "I8")
    655
        pub dtype: String,
    656
        /// Tensor dimensions
    657
        pub shape: Vec<usize>,
    658
        /// Byte offset from data section start
    659
        pub offset: u64,
    660
        /// Size in bytes
    661
        pub size: u64,
    662
    }
    663
    664
    impl TensorEntry {
    665
        /// Parse tensor entry from binary format (aprender v2 format)
    666
        ///
    667
        /// Binary format:
    668
        /// - name_len (2 bytes LE) + name bytes
    669
        /// - dtype (1 byte)
    670
        /// - ndim (1 byte) + dims (8 bytes LE each, up to 8)
    671
        /// - offset (8 bytes LE)
    672
        /// - size (8 bytes LE)
    673
    202
        pub fn from_binary(data: &[u8]) -> Result<(Self, usize)> {
    674
    202
            if data.len() < 4 {
    675
    1
                return Err(RealizarError::FormatError {
    676
    1
                    reason: "Tensor entry too short".to_string(),
    677
    1
                });
    678
    201
            }
    679
    680
    201
            let mut pos = 0;
    681
    682
            // Name
    683
    201
            let name_len = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize;
    684
    201
            pos += 2;
    685
    686
    201
            if data.len() < pos + name_len + 2 {
    687
    104
                return Err(RealizarError::FormatError {
    688
    104
                    reason: "Tensor entry truncated at name".to_string(),
    689
    104
                });
    690
    97
            }
    691
    692
    97
            let name = String::from_utf8_lossy(&data[pos..pos + name_len]).to_string();
    693
    97
            pos += name_len;
    694
    695
            // Dtype (1 byte)
    696
    97
            let dtype_byte = data[pos];
    697
    97
            pos += 1;
    698
    97
            let dtype = match dtype_byte {
    699
    80
                0 => "F32",
    700
    3
                1 => "F16",
    701
    3
                2 => "BF16",
    702
    2
                3 => "I8",
    703
    1
                4 => "I16",
    704
    1
                5 => "I32",
    705
    1
                6 => "I64",
    706
    1
                7 => "U8",
    707
    1
                8 => "Q4_K",  // GGUF Q4_K_M quantization (4.5 bits/element)
    708
    1
                9 => "Q6_K",  // GGUF Q6_K quantization (6.5 bits/element)
    709
    2
                10 => "Q8_0", // GGUF Q8_0 quantization (8 bits/element)
    710
    1
                _ => "F32",
    711
            }
    712
    97
            .to_string();
    713
    714
            // Shape: ndim (1 byte) + dims
    715
    97
            let ndim = data[pos] as usize;
    716
    97
            pos += 1;
    717
    718
    97
            if data.len() < pos + ndim * 8 + 16 {
    719
    2
                return Err(RealizarError::FormatError {
    720
    2
                    reason: "Tensor entry truncated at shape".to_string(),
    721
    2
                });
    722
    95
            }
    723
    724
    95
            let mut shape = Vec::with_capacity(ndim);
    725
    166
            for _ in 0..ndim95 {
    726
    166
                let dim = u64::from_le_bytes([
    727
    166
                    data[pos],
    728
    166
                    data[pos + 1],
    729
    166
                    data[pos + 2],
    730
    166
                    data[pos + 3],
    731
    166
                    data[pos + 4],
    732
    166
                    data[pos + 5],
    733
    166
                    data[pos + 6],
    734
    166
                    data[pos + 7],
    735
    166
                ]) as usize;
    736
    166
                pos += 8;
    737
    166
                shape.push(dim);
    738
    166
            }
    739
    740
            // Offset and size
    741
    95
            let offset = u64::from_le_bytes([
    742
    95
                data[pos],
    743
    95
                data[pos + 1],
    744
    95
                data[pos + 2],
    745
    95
                data[pos + 3],
    746
    95
                data[pos + 4],
    747
    95
                data[pos + 5],
    748
    95
                data[pos + 6],
    749
    95
                data[pos + 7],
    750
    95
            ]);
    751
    95
            pos += 8;
    752
    753
    95
            let size = u64::from_le_bytes([
    754
    95
                data[pos],
    755
    95
                data[pos + 1],
    756
    95
                data[pos + 2],
    757
    95
                data[pos + 3],
    758
    95
                data[pos + 4],
    759
    95
                data[pos + 5],
    760
    95
                data[pos + 6],
    761
    95
                data[pos + 7],
    762
    95
            ]);
    763
    95
            pos += 8;
    764
    765
    95
            Ok((
    766
    95
                Self {
    767
    95
                    name,
    768
    95
                    dtype,
    769
    95
                    shape,
    770
    95
                    offset,
    771
    95
                    size,
    772
    95
                },
    773
    95
                pos,
    774
    95
            ))
    775
    202
        }
    776
    777
        /// Calculate element count from shape
    778
    16
        pub fn element_count(&self) -> usize {
    779
    16
            self.shape.iter().product()
    780
    16
        }
    781
    }
    782
    783
    /// Model metadata from .apr file
    784
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    785
    pub struct AprMetadata {
    786
        /// Model type (e.g., "transformer_lm", "whisper", "llama")
    787
        #[serde(default)]
    788
        pub model_type: Option<String>,
    789
        /// Human-readable model name
    790
        #[serde(default)]
    791
        pub name: Option<String>,
    792
        /// Model architecture family
    793
        #[serde(default)]
    794
        pub architecture: Option<String>,
    795
        /// Hidden dimension size
    796
        #[serde(default)]
    797
        pub hidden_size: Option<usize>,
    798
        /// Number of transformer layers
    799
        #[serde(default)]
    800
        pub num_layers: Option<usize>,
    801
        /// Number of attention heads
    802
        #[serde(default)]
    803
        pub num_heads: Option<usize>,
    804
        /// Number of key-value heads (for GQA, defaults to num_heads)
    805
        #[serde(default)]
    806
        pub num_kv_heads: Option<usize>,
    807
        /// Vocabulary size
    808
        #[serde(default)]
    809
        pub vocab_size: Option<usize>,
    810
        /// FFN intermediate dimension
    811
        #[serde(default)]
    812
        pub intermediate_size: Option<usize>,
    813
        /// Maximum context/sequence length
    814
        #[serde(default)]
    815
        pub max_position_embeddings: Option<usize>,
    816
        /// RoPE theta for position encoding
    817
        #[serde(default)]
    818
        pub rope_theta: Option<f32>,
    819
        /// RoPE type: 0=NORM (adjacent pairs), 2=NEOX (split halves)
    820
        /// CORRECTNESS-011: Qwen2.5 models require rope_type=2 (NEOX style)
    821
        #[serde(default)]
    822
        pub rope_type: Option<u32>,
    823
        /// Layer norm epsilon
    824
        #[serde(default)]
    825
        pub rms_norm_eps: Option<f32>,
    826
        /// Additional metadata fields
    827
        #[serde(flatten)]
    828
        pub extra: HashMap<String, serde_json::Value>,
    829
    }
    830
    831
    impl AprMetadata {
    832
        /// Check if this model has transformer configuration
    833
        #[must_use]
    834
    18
        pub fn is_transformer(&self) -> bool {
    835
    18
            self.hidden_size.is_some()
    836
    14
                && self.num_layers.is_some()
    837
    9
                && self.num_heads.is_some()
    838
    8
                && self.vocab_size.is_some()
    839
    18
        }
    840
    841
        /// Extract embedded tokenizer vocabulary from APR metadata (GH-156)
    842
        ///
    843
        /// APR files can contain embedded tokenizer data in the metadata section.
    844
        /// This is the preferred way to decode tokens - no sibling files needed.
    845
        ///
    846
        /// # Returns
    847
        /// - `Some(vocab)` if tokenizer.vocabulary is present in metadata
    848
        /// - `None` if no embedded tokenizer
    849
        #[must_use]
    850
    0
        pub fn get_embedded_vocabulary(&self) -> Option<Vec<String>> {
    851
    0
            let vocab_value = self.extra.get("tokenizer.vocabulary")?;
    852
    0
            let vocab_array = vocab_value.as_array()?;
    853
    854
    0
            let vocab: Vec<String> = vocab_array
    855
    0
                .iter()
    856
    0
                .filter_map(|v| v.as_str().map(String::from))
    857
    0
                .collect();
    858
    859
    0
            if vocab.is_empty() {
    860
    0
                None
    861
            } else {
    862
    0
                Some(vocab)
    863
            }
    864
    0
        }
    865
    866
        /// Get embedded BOS token ID from APR metadata
    867
        #[must_use]
    868
    0
        pub fn get_embedded_bos_token_id(&self) -> Option<u32> {
    869
    0
            self.extra
    870
    0
                .get("tokenizer.bos_token_id")
    871
    0
                .and_then(|v| v.as_u64())
    872
    0
                .map(|v| v as u32)
    873
    0
        }
    874
    875
        /// Get embedded EOS token ID from APR metadata
    876
        #[must_use]
    877
    0
        pub fn get_embedded_eos_token_id(&self) -> Option<u32> {
    878
    0
            self.extra
    879
    0
                .get("tokenizer.eos_token_id")
    880
    0
                .and_then(|v| v.as_u64())
    881
    0
                .map(|v| v as u32)
    882
    0
        }
    883
    }
    884
    885
    /// APR v2 model for realizar inference
    886
    ///
    887
    /// # Memory Management
    888
    ///
    889
    /// Uses memory-mapped I/O for uncompressed files to avoid zram pressure.
    890
    /// After loading tensors to GPU, call `release_cpu_pages()` to advise
    891
    /// the kernel that pages can be dropped (re-faulted from disk if needed).
    892
    ///
    893
    /// # References
    894
    ///
    895
    /// - Didona et al. (2022): mmap vs read() performance
    896
    /// - See docs/model-loading.md for full design rationale
    897
    #[derive(Debug)]
    898
    pub struct AprV2Model {
    899
        /// Header information
    900
        header: AprHeader,
    901
        /// Model metadata
    902
        metadata: AprMetadata,
    903
        /// Tensor index
    904
        tensors: Vec<TensorEntry>,
    905
        /// Raw file data (mmap for uncompressed, heap for compressed)
    906
        data: ModelData,
    907
    }
    908
    909
    impl AprV2Model {
    910
        /// Load a model from a .apr file using memory mapping.
    911
        ///
    912
        /// # Memory Efficiency
    913
        ///
    914
        /// For uncompressed files, uses `mmap()` for zero-copy access.
    915
        /// The kernel manages pages via demand paging - only accessed
    916
        /// pages are loaded into RAM. After GPU transfer, call
    917
        /// `release_cpu_pages()` to advise the kernel to drop pages.
    918
        ///
    919
        /// For compressed files, falls back to heap allocation after
    920
        /// decompression (mmap not possible for decompressed data).
    921
        ///
    922
        /// # References
    923
        ///
    924
        /// - Didona et al. (2022): mmap achieves 2.3x throughput vs read()
    925
        /// - See docs/model-loading.md for design rationale
    926
        #[cfg(not(target_arch = "wasm32"))]
    927
    5
        pub fn load<P: AsRef<Path>>(path: P) -> Result<Self> {
    928
            use std::io::Read;
    929
    930
    5
            let path_ref = path.as_ref();
    931
    932
            // Read just the header first to check for compression
    933
    5
            let mut file4 = File::open(path_ref).map_err(|e| RealizarError::IoError {
    934
    1
                message: format!("Failed to open .apr file: {e}"),
    935
    1
            })?;
    936
    937
    4
            let mut header_buf = [0u8; HEADER_SIZE];
    938
    4
            file.read_exact(&mut header_buf)
    939
    4
                .map_err(|e| RealizarError::IoError {
    940
    1
                    message: format!("Failed to read .apr header: {e}"),
    941
    1
                })?;
    942
    943
    3
            let header2 = AprHeader::from_bytes(&header_buf)?1;
    944
    945
            // Check for unsupported features
    946
    2
            if header.flags.is_encrypted() {
    947
    0
                return Err(RealizarError::FormatError {
    948
    0
                    reason: "Encrypted .apr files not yet supported".to_string(),
    949
    0
                });
    950
    2
            }
    951
    952
            // Choose loading strategy based on compression
    953
    2
            let data = if header.flags.is_compressed() {
    954
                // Compressed: must read entire file into heap, then decompress
    955
    0
                drop(file); // Close file handle
    956
    0
                let raw_data = std::fs::read(path_ref).map_err(|e| RealizarError::IoError {
    957
    0
                    message: format!("Failed to read compressed .apr file: {e}"),
    958
    0
                })?;
    959
    0
                let decompressed = Self::decompress_apr_data(&header, raw_data)?;
    960
    0
                ModelData::from_vec(decompressed)
    961
            } else {
    962
                // Uncompressed: use mmap for zero-copy access
    963
    2
                drop(file); // Close file handle before mmap
    964
    2
                ModelData::open_mmap(path_ref)?0
    965
            };
    966
    967
            // Advise sequential access pattern for parsing
    968
            #[cfg(unix)]
    969
    2
            let _ = data.advise_sequential();
    970
    971
    2
            Self::from_model_data(header, data)
    972
    5
        }
    973
    974
        /// Load a model from a .apr file (WASM fallback).
    975
        #[cfg(target_arch = "wasm32")]
    976
        pub fn load<P: AsRef<Path>>(path: P) -> Result<Self> {
    977
            let raw_data = std::fs::read(path.as_ref()).map_err(|e| RealizarError::IoError {
    978
                message: format!("Failed to read .apr file: {e}"),
    979
            })?;
    980
            Self::from_bytes(raw_data)
    981
        }
    982
    983
        /// Load a model from bytes (heap-allocated).
    984
        ///
    985
        /// Use this for:
    986
        /// - Compressed files after decompression
    987
        /// - Data received over network
    988
        /// - WASM environments (no mmap support)
    989
        ///
    990
        /// For file-based loading with mmap support, use `load()` instead.
    991
    153
        pub fn from_bytes(data: Vec<u8>) -> Result<Self> {
    992
            // Parse header
    993
    153
            let header151 = AprHeader::from_bytes(&data)?2;
    994
    995
            // Check for unsupported features
    996
    151
            if header.flags.is_encrypted() {
    997
    2
                return Err(RealizarError::FormatError {
    998
    2
                    reason: "Encrypted .apr files not yet supported".to_string(),
    999
    2
                });
    1000
    149
            }
    1001
    1002
            // Decompress data if needed (GH-35)
    1003
    149
            let data148 = if header.flags.is_compressed() {
    1004
    1
                Self::decompress_apr_data(&header, data)?
    1005
            } else {
    1006
    148
                data
    1007
            };
    1008
    1009
    148
            Self::from_model_data(header, ModelData::from_vec(data))
    1010
    153
        }
    1011
    1012
        /// Internal: construct model from header and ModelData.
    1013
    150
        fn from_model_data(header: AprHeader, data: ModelData) -> Result<Self> {
    1014
    150
            let data_slice = data.as_slice();
    1015
    1016
            // Parse metadata
    1017
    150
            let metadata_start = header.metadata_offset as usize;
    1018
    150
            let metadata_end = metadata_start + header.metadata_size as usize;
    1019
    1020
    150
            if data_slice.len() < metadata_end {
    1021
    1
                return Err(RealizarError::FormatError {
    1022
    1
                    reason: format!(
    1023
    1
                        ".apr file truncated: metadata extends to {} but file is {} bytes",
    1024
    1
                        metadata_end,
    1025
    1
                        data_slice.len()
    1026
    1
                    ),
    1027
    1
                });
    1028
    149
            }
    1029
    1030
    149
            let metadata: AprMetadata = if header.metadata_size > 0 {
    1031
    145
                serde_json::from_slice(&data_slice[metadata_start..metadata_end]).unwrap_or_default()
    1032
            } else {
    1033
    4
                AprMetadata::default()
    1034
            };
    1035
    1036
            // Parse tensor index (binary format from aprender v2)
    1037
    149
            let index_start = header.tensor_index_offset as usize;
    1038
    149
            let index_end = header.data_offset as usize;
    1039
    1040
    149
            let mut tensors = Vec::with_capacity(header.tensor_count as usize);
    1041
    149
            if index_start < index_end && index_end145 <= data_slice.len() {
    1042
    145
                let index_data = &data_slice[index_start..index_end];
    1043
    145
                let mut pos = 0;
    1044
    1045
    220
                while pos < index_data.len() && tensors179.len() < header.tensor_count as usize {
    1046
    179
                    match TensorEntry::from_binary(&index_data[pos..]) {
    1047
    75
                        Ok((entry, consumed)) => {
    1048
    75
                            tensors.push(entry);
    1049
    75
                            pos += consumed;
    1050
    75
                        },
    1051
    104
                        Err(_) => break, // Stop on parse error
    1052
                    }
    1053
                }
    1054
    4
            }
    1055
    1056
    149
            Ok(Self {
    1057
    149
                header,
    1058
    149
                metadata,
    1059
    149
                tensors,
    1060
    149
                data,
    1061
    149
            })
    1062
    150
        }
    1063
    1064
        /// Decompress APR data based on compression flags (GH-35)
    1065
        ///
    1066
        /// The compressed format stores: header (64 bytes, uncompressed) + compressed payload.
    1067
        /// We decompress the payload and reconstruct the full data vector.
    1068
        #[allow(unreachable_patterns)] // Pattern varies based on apr-compression feature
    1069
    1
        fn decompress_apr_data(header: &AprHeader, data: Vec<u8>) -> Result<Vec<u8>> {
    1070
            #[cfg(feature = "apr-compression")]
    1071
            let compressed_payload = &data[HEADER_SIZE..];
    1072
    1073
            #[cfg(feature = "apr-compression")]
    1074
            {
    1075
                let decompressed = if header.flags.is_lz4() {
    1076
                    lz4_flex::decompress_size_prepended(compressed_payload).map_err(|e| {
    1077
                        RealizarError::FormatError {
    1078
                            reason: format!("LZ4 decompression failed: {e}"),
    1079
                        }
    1080
                    })?
    1081
                } else if header.flags.is_zstd() {
    1082
                    zstd::decode_all(compressed_payload).map_err(|e| RealizarError::FormatError {
    1083
                        reason: format!("ZSTD decompression failed: {e}"),
    1084
                    })?
    1085
                } else {
    1086
                    // Unknown compression - should not happen
    1087
                    return Err(RealizarError::FormatError {
    1088
                        reason: "Unknown compression algorithm in APR flags".to_string(),
    1089
                    });
    1090
                };
    1091
    1092
                // Reconstruct full data: header + decompressed payload
    1093
                let mut result = Vec::with_capacity(HEADER_SIZE + decompressed.len());
    1094
                result.extend_from_slice(&data[..HEADER_SIZE]);
    1095
                result.extend_from_slice(&decompressed);
    1096
                Ok(result)
    1097
            }
    1098
    1099
            #[cfg(not(feature = "apr-compression"))]
    1100
            {
    1101
    1
                let _ = (header, &data); // Suppress unused warnings
    1102
    1
                Err(RealizarError::FormatError {
    1103
    1
                    reason: "Compressed .apr files require 'apr-compression' feature".to_string(),
    1104
    1
                })
    1105
            }
    1106
    1
        }
    1107
    1108
        /// Get number of tensors
    1109
        #[must_use]
    1110
    6
        pub fn tensor_count(&self) -> u32 {
    1111
    6
            self.header.tensor_count
    1112
    6
        }
    1113
    1114
        /// Get tensor names
    1115
        #[must_use]
    1116
    3
        pub fn tensor_names(&self) -> Vec<&str> {
    1117
    3
            self.tensors.iter().map(|t| t.name.as_str()).collect()
    1118
    3
        }
    1119
    1120
        /// Get metadata
    1121
        #[must_use]
    1122
    7
        pub fn metadata(&self) -> &AprMetadata {
    1123
    7
            &self.metadata
    1124
    7
        }
    1125
    1126
        /// Get tensor by name
    1127
        #[must_use]
    1128
    33
        pub fn get_tensor(&self, name: &str) -> Option<&TensorEntry> {
    1129
    51
            self.tensors.iter()33.find33(|t| t.name == name)
    1130
    33
        }
    1131
    1132
        /// Get tensor data as f32 slice
    1133
    10
        pub fn get_tensor_f32(&self, name: &str) -> Result<Vec<f32>> {
    1134
    10
            let entry7 = self
    1135
    10
                .get_tensor(name)
    1136
    10
                .ok_or_else(|| RealizarError::FormatError {
    1137
    3
                    reason: format!("Tensor not found: {name}"),
    1138
    3
                })?;
    1139
    1140
    7
            let start = (self.header.data_offset + entry.offset) as usize;
    1141
    7
            let end = start + entry.size as usize;
    1142
    7
            let data_slice = self.data.as_slice();
    1143
    1144
    7
            if end > data_slice.len() {
    1145
    0
                return Err(RealizarError::FormatError {
    1146
    0
                    reason: format!("Tensor data out of bounds: {name}"),
    1147
    0
                });
    1148
    7
            }
    1149
    1150
    7
            let bytes = &data_slice[start..end];
    1151
    1152
            // Calculate total number of elements from shape
    1153
    7
            let num_elements: usize = entry.shape.iter().product();
    1154
    1155
            // Parse based on dtype
    1156
    7
            match entry.dtype.as_str() {
    1157
    7
                "F32" | "f32"3 => {
    1158
    4
                    let floats: Vec<f32> = bytes
    1159
    4
                        .chunks_exact(4)
    1160
    40
                        .map4(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
    1161
    4
                        .collect();
    1162
    4
                    Ok(floats)
    1163
                },
    1164
    3
                "F16" | "f16"2 => Ok(1dequantize_f161(bytes, num_elements)),
    1165
    2
                "Q8_0" | "q8_0"1 => Ok(1dequantize_q8_01(bytes, num_elements)),
    1166
    1
                "Q4_K" | "q4_k" => Ok(0dequantize_q4_k0(bytes, num_elements)),
    1167
    1
                "Q6_K" | "q6_k" => Ok(0dequantize_q6_k0(bytes, num_elements)),
    1168
    1
                dtype => Err(RealizarError::FormatError {
    1169
    1
                    reason: format!("Unsupported tensor dtype: {dtype}"),
    1170
    1
                }),
    1171
            }
    1172
    10
        }
    1173
    1174
        /// Get raw tensor bytes
    1175
    4
        pub fn get_tensor_bytes(&self, name: &str) -> Result<&[u8]> {
    1176
    4
            let entry2 = self
    1177
    4
                .get_tensor(name)
    1178
    4
                .ok_or_else(|| RealizarError::FormatError {
    1179
    2
                    reason: format!("Tensor not found: {name}"),
    1180
    2
                })?;
    1181
    1182
    2
            let start = (self.header.data_offset + entry.offset) as usize;
    1183
    2
            let end = start + entry.size as usize;
    1184
    2
            let data_slice = self.data.as_slice();
    1185
    1186
    2
            if end > data_slice.len() {
    1187
    0
                return Err(RealizarError::FormatError {
    1188
    0
                    reason: format!("Tensor data out of bounds: {name}"),
    1189
    0
                });
    1190
    2
            }
    1191
    1192
    2
            Ok(&data_slice[start..end])
    1193
    4
        }
    1194
    1195
        /// Release CPU pages after GPU transfer (Unix only).
    1196
        ///
    1197
        /// Advises the kernel that the mapped pages are no longer needed.
    1198
        /// The kernel will drop pages immediately (not compress to zram)
    1199
        /// and re-fault from disk if accessed again.
    1200
        ///
    1201
        /// # When to Call
    1202
        ///
    1203
        /// After all tensor data has been copied to GPU via `cuMemcpy()`.
    1204
        /// This is the key method for reducing zram pressure.
    1205
        ///
    1206
        /// # Example
    1207
        ///
    1208
        /// ```rust,ignore
    1209
        /// let model = AprV2Model::load("model.apr")?;
    1210
        /// for name in model.tensor_names() {
    1211
        ///     let bytes = model.get_tensor_bytes(&name)?;
    1212
        ///     cuda::memcpy_htod(gpu_ptr, bytes);
    1213
        /// }
    1214
        /// // Free CPU pages now that data is on GPU
    1215
        /// model.release_cpu_pages()?;
    1216
        /// ```
    1217
        #[cfg(all(unix, not(target_arch = "wasm32")))]
    1218
    1
        pub fn release_cpu_pages(&self) -> Result<()> {
    1219
    1
            self.data.release_cpu_pages()
    1220
    1
        }
    1221
    1222
        /// Check if model is using memory-mapped I/O.
    1223
        ///
    1224
        /// Returns `true` if the model was loaded via mmap (uncompressed file).
    1225
        /// Returns `false` if the model is heap-allocated (compressed file or WASM).
    1226
        #[must_use]
    1227
    3
        pub fn is_mmap(&self) -> bool {
    1228
    3
            self.data.is_mmap()
    1229
    3
        }
    1230
    1231
        /// Estimate total parameters
    1232
        #[must_use]
    1233
    3
        pub fn estimated_parameters(&self) -> usize {
    1234
    3
            self.tensors
    1235
    3
                .iter()
    1236
    3
                .map(|t| t.shape.iter().product::<usize>())
    1237
    3
                .sum()
    1238
    3
        }
    1239
    1240
        /// Run inference on input features (for simple models)
    1241
        ///
    1242
        /// For transformer models, use `forward()` instead.
    1243
        ///
    1244
        /// # Arguments
    1245
        ///
    1246
        /// * `features` - Input feature vector
    1247
        ///
    1248
        /// # Returns
    1249
        ///
    1250
        /// Output vector
    1251
        ///
    1252
        /// # Errors
    1253
        ///
    1254
        /// Returns error if model has no tensors
    1255
    5
        pub fn predict(&self, features: &[f32]) -> Result<Vec<f32>> {
    1256
    5
            if self.tensors.is_empty() && self.header.tensor_count == 00 {
    1257
    0
                let sum: f32 = features.iter().sum();
    1258
    0
                return Ok(vec![sum]);
    1259
    5
            }
    1260
    1261
            // Linear model: y = Wx + b (if we have weights)
    1262
    5
            if let Some(weight1) = self.get_tensor("weight") {
    1263
    1
                let weights = self.get_tensor_f32("weight")?0;
    1264
    1
                let bias = self.get_tensor_f32("bias").unwrap_or_default();
    1265
    1266
    1
                let output_dim = weight.shape.first().copied().unwrap_or(1);
    1267
    1
                let input_dim = weight.shape.get(1).copied().unwrap_or(features.len());
    1268
    1269
    1
                let mut output = vec![0.0; output_dim];
    1270
    2
                for (i, out) in output.iter_mut()1.enumerate1() {
    1271
    6
                    for (j, &feat) in features2.iter2().take2(input_dim2).enumerate2() {
    1272
    6
                        *out += weights.get(i * input_dim + j).copied().unwrap_or(0.0) * feat;
    1273
    6
                    }
    1274
    2
                    *out += bias.get(i).copied().unwrap_or(0.0);
    1275
                }
    1276
    1
                return Ok(output);
    1277
    4
            }
    1278
    1279
    4
            let sum: f32 = features.iter().sum();
    1280
    4
            Ok(vec![sum])
    1281
    5
        }
    1282
    1283
        /// Run transformer forward pass on token IDs
    1284
        ///
    1285
        /// Returns logits for the next token prediction.
    1286
        ///
    1287
        /// # Arguments
    1288
        ///
    1289
        /// * `token_ids` - Input token IDs
    1290
        ///
    1291
        /// # Returns
    1292
        ///
    1293
        /// Logits vector of size `vocab_size`
    1294
        ///
    1295
        /// # Errors
    1296
        ///
    1297
        /// Returns error if model is not a transformer or tensors are missing
    1298
    6
        pub fn forward(&self, token_ids: &[u32]) -> Result<Vec<f32>> {
    1299
    6
            if token_ids.is_empty() {
    1300
    2
                return Err(RealizarError::InvalidShape {
    1301
    2
                    reason: "Token sequence cannot be empty".to_string(),
    1302
    2
                });
    1303
    4
            }
    1304
    1305
    4
            if !self.metadata.is_transformer() {
    1306
    3
                return Err(RealizarError::FormatError {
    1307
    3
                    reason: "Model is not a transformer (missing config)".to_string(),
    1308
    3
                });
    1309
    1
            }
    1310
    1311
    1
            let hidden_dim = self.metadata.hidden_size.unwrap_or(0);
    1312
    1
            let num_layers = self.metadata.num_layers.unwrap_or(0);
    1313
    1
            let num_heads = self.metadata.num_heads.unwrap_or(1);
    1314
    1
            let num_kv_heads = self.metadata.num_kv_heads.unwrap_or(num_heads);
    1315
    1
            let vocab_size = self.metadata.vocab_size.unwrap_or(0);
    1316
    1
            let intermediate_dim = self.metadata.intermediate_size.unwrap_or(hidden_dim * 4);
    1317
    1
            let eps = self.metadata.rms_norm_eps.unwrap_or(1e-6);
    1318
    1319
            // 1. Token embedding lookup
    1320
    1
            let embed_name0 = self.find_tensor_name(&[
    1321
    1
                "model.embed_tokens.weight",
    1322
    1
                "embed_tokens.weight", // SafeTensors (no model. prefix)
    1323
    1
                "transformer.wte.weight",
    1324
    1
                "embeddings.word_embeddings.weight",
    1325
    1
                "tok_embeddings.weight",
    1326
    1
                "token_embd.weight", // GGUF naming convention
    1327
    1
            ])?;
    1328
    1329
    0
            let embeddings = self.get_tensor_f32(&embed_name)?;
    1330
    0
            let mut hidden = Vec::with_capacity(token_ids.len() * hidden_dim);
    1331
    1332
    0
            for &token_id in token_ids {
    1333
    0
                let offset = (token_id as usize) * hidden_dim;
    1334
    0
                if offset + hidden_dim <= embeddings.len() {
    1335
    0
                    hidden.extend_from_slice(&embeddings[offset..offset + hidden_dim]);
    1336
    0
                } else {
    1337
    0
                    hidden.extend(std::iter::repeat_n(0.0, hidden_dim));
    1338
    0
                }
    1339
            }
    1340
    1341
            // 2. Process through transformer layers
    1342
    0
            for layer_idx in 0..num_layers {
    1343
                // Try common naming patterns (HuggingFace, SafeTensors, GPT-2, LLaMA, GGUF)
    1344
    0
                let attn_norm_name = self.find_tensor_name(&[
    1345
    0
                    &format!("model.layers.{layer_idx}.input_layernorm.weight"),
    1346
    0
                    &format!("layers.{layer_idx}.input_layernorm.weight"), // SafeTensors
    1347
    0
                    &format!("transformer.h.{layer_idx}.ln_1.weight"),
    1348
    0
                    &format!("layers.{layer_idx}.attention_norm.weight"),
    1349
    0
                    &format!("blk.{layer_idx}.attn_norm.weight"), // GGUF naming
    1350
    0
                ])?;
    1351
    1352
    0
                let q_name = self.find_tensor_name(&[
    1353
    0
                    &format!("model.layers.{layer_idx}.self_attn.q_proj.weight"),
    1354
    0
                    &format!("layers.{layer_idx}.self_attn.q_proj.weight"), // SafeTensors
    1355
    0
                    &format!("transformer.h.{layer_idx}.attn.q_proj.weight"),
    1356
    0
                    &format!("layers.{layer_idx}.attention.wq.weight"),
    1357
    0
                    &format!("blk.{layer_idx}.attn_q.weight"), // GGUF naming
    1358
    0
                ])?;
    1359
    1360
    0
                let k_name = self.find_tensor_name(&[
    1361
    0
                    &format!("model.layers.{layer_idx}.self_attn.k_proj.weight"),
    1362
    0
                    &format!("layers.{layer_idx}.self_attn.k_proj.weight"), // SafeTensors
    1363
    0
                    &format!("transformer.h.{layer_idx}.attn.k_proj.weight"),
    1364
    0
                    &format!("layers.{layer_idx}.attention.wk.weight"),
    1365
    0
                    &format!("blk.{layer_idx}.attn_k.weight"), // GGUF naming
    1366
    0
                ])?;
    1367
    1368
    0
                let v_name = self.find_tensor_name(&[
    1369
    0
                    &format!("model.layers.{layer_idx}.self_attn.v_proj.weight"),
    1370
    0
                    &format!("layers.{layer_idx}.self_attn.v_proj.weight"), // SafeTensors
    1371
    0
                    &format!("transformer.h.{layer_idx}.attn.v_proj.weight"),
    1372
    0
                    &format!("layers.{layer_idx}.attention.wv.weight"),
    1373
    0
                    &format!("blk.{layer_idx}.attn_v.weight"), // GGUF naming
    1374
    0
                ])?;
    1375
    1376
    0
                let o_name = self.find_tensor_name(&[
    1377
    0
                    &format!("model.layers.{layer_idx}.self_attn.o_proj.weight"),
    1378
    0
                    &format!("layers.{layer_idx}.self_attn.o_proj.weight"), // SafeTensors
    1379
    0
                    &format!("transformer.h.{layer_idx}.attn.out_proj.weight"),
    1380
    0
                    &format!("layers.{layer_idx}.attention.wo.weight"),
    1381
    0
                    &format!("blk.{layer_idx}.attn_output.weight"), // GGUF naming
    1382
    0
                ])?;
    1383
    1384
                // Load tensors
    1385
    0
                let norm_weight = self.get_tensor_f32(&attn_norm_name)?;
    1386
    0
                let q_weight = self.get_tensor_f32(&q_name)?;
    1387
    0
                let k_weight = self.get_tensor_f32(&k_name)?;
    1388
    0
                let v_weight = self.get_tensor_f32(&v_name)?;
    1389
    0
                let o_weight = self.get_tensor_f32(&o_name)?;
    1390
    1391
                // RMSNorm
    1392
    0
                let normed = rms_norm(&hidden, &norm_weight, eps);
    1393
    1394
                // Attention: Q, K, V projections
    1395
    0
                let seq_len = token_ids.len();
    1396
    0
                let head_dim = hidden_dim / num_heads;
    1397
    1398
    0
                let q = matmul(&normed, &q_weight, seq_len, hidden_dim, hidden_dim);
    1399
    0
                let k = matmul(
    1400
    0
                    &normed,
    1401
    0
                    &k_weight,
    1402
    0
                    seq_len,
    1403
    0
                    hidden_dim,
    1404
    0
                    num_kv_heads * head_dim,
    1405
                );
    1406
    0
                let v = matmul(
    1407
    0
                    &normed,
    1408
    0
                    &v_weight,
    1409
    0
                    seq_len,
    1410
    0
                    hidden_dim,
    1411
    0
                    num_kv_heads * head_dim,
    1412
                );
    1413
    1414
                // Simplified attention (no RoPE for now, full attention)
    1415
    0
                let attn_out = simple_attention(&q, &k, &v, seq_len, num_heads, num_kv_heads, head_dim);
    1416
    1417
                // Output projection
    1418
    0
                let attn_proj = matmul(&attn_out, &o_weight, seq_len, hidden_dim, hidden_dim);
    1419
    1420
                // Residual connection
    1421
    0
                for (h, &a) in hidden.iter_mut().zip(attn_proj.iter()) {
    1422
    0
                    *h += a;
    1423
    0
                }
    1424
    1425
                // FFN
    1426
    0
                let ffn_norm_name = self.find_tensor_name(&[
    1427
    0
                    &format!("model.layers.{layer_idx}.post_attention_layernorm.weight"),
    1428
    0
                    &format!("layers.{layer_idx}.post_attention_layernorm.weight"), // SafeTensors
    1429
    0
                    &format!("transformer.h.{layer_idx}.ln_2.weight"),
    1430
    0
                    &format!("layers.{layer_idx}.ffn_norm.weight"),
    1431
    0
                    &format!("blk.{layer_idx}.ffn_norm.weight"), // GGUF naming
    1432
    0
                ])?;
    1433
    1434
    0
                let gate_name = self.find_tensor_name(&[
    1435
    0
                    &format!("model.layers.{layer_idx}.mlp.gate_proj.weight"),
    1436
    0
                    &format!("layers.{layer_idx}.mlp.gate_proj.weight"), // SafeTensors
    1437
    0
                    &format!("transformer.h.{layer_idx}.mlp.gate_proj.weight"),
    1438
    0
                    &format!("layers.{layer_idx}.feed_forward.w1.weight"),
    1439
    0
                    &format!("blk.{layer_idx}.ffn_gate.weight"), // GGUF naming
    1440
    0
                ])?;
    1441
    1442
    0
                let up_name = self.find_tensor_name(&[
    1443
    0
                    &format!("model.layers.{layer_idx}.mlp.up_proj.weight"),
    1444
    0
                    &format!("layers.{layer_idx}.mlp.up_proj.weight"), // SafeTensors
    1445
    0
                    &format!("transformer.h.{layer_idx}.mlp.up_proj.weight"),
    1446
    0
                    &format!("layers.{layer_idx}.feed_forward.w3.weight"),
    1447
    0
                    &format!("blk.{layer_idx}.ffn_up.weight"), // GGUF naming
    1448
    0
                ])?;
    1449
    1450
    0
                let down_name = self.find_tensor_name(&[
    1451
    0
                    &format!("model.layers.{layer_idx}.mlp.down_proj.weight"),
    1452
    0
                    &format!("layers.{layer_idx}.mlp.down_proj.weight"), // SafeTensors
    1453
    0
                    &format!("transformer.h.{layer_idx}.mlp.down_proj.weight"),
    1454
    0
                    &format!("layers.{layer_idx}.feed_forward.w2.weight"),
    1455
    0
                    &format!("blk.{layer_idx}.ffn_down.weight"), // GGUF naming
    1456
    0
                ])?;
    1457
    1458
    0
                let ffn_norm = self.get_tensor_f32(&ffn_norm_name)?;
    1459
    0
                let gate = self.get_tensor_f32(&gate_name)?;
    1460
    0
                let up = self.get_tensor_f32(&up_name)?;
    1461
    0
                let down = self.get_tensor_f32(&down_name)?;
    1462
    1463
    0
                let normed = rms_norm(&hidden, &ffn_norm, eps);
    1464
    0
                let gate_out = matmul(&normed, &gate, seq_len, hidden_dim, intermediate_dim);
    1465
    0
                let up_out = matmul(&normed, &up, seq_len, hidden_dim, intermediate_dim);
    1466
    1467
                // SiLU activation and element-wise multiply
    1468
    0
                let mut ffn_hidden = Vec::with_capacity(seq_len * intermediate_dim);
    1469
    0
                for (g, u) in gate_out.iter().zip(up_out.iter()) {
    1470
    0
                    let silu = g * (1.0 / (1.0 + (-g).exp()));
    1471
    0
                    ffn_hidden.push(silu * u);
    1472
    0
                }
    1473
    1474
    0
                let ffn_out = matmul(&ffn_hidden, &down, seq_len, intermediate_dim, hidden_dim);
    1475
    1476
                // Residual
    1477
    0
                for (h, &f) in hidden.iter_mut().zip(ffn_out.iter()) {
    1478
    0
                    *h += f;
    1479
    0
                }
    1480
            }
    1481
    1482
            // 3. Final layer norm
    1483
    0
            let final_norm_name = self.find_tensor_name(&[
    1484
    0
                "model.norm.weight",
    1485
    0
                "norm.weight", // SafeTensors
    1486
    0
                "transformer.ln_f.weight",
    1487
    0
                "output_norm.weight", // GGUF naming
    1488
    0
            ])?;
    1489
    0
            let final_norm = self.get_tensor_f32(&final_norm_name)?;
    1490
    0
            let hidden = rms_norm(&hidden, &final_norm, eps);
    1491
    1492
            // 4. LM head (last token only for generation)
    1493
    0
            let lm_head_name = self.find_tensor_name(&[
    1494
    0
                "lm_head.weight",
    1495
    0
                "output.weight",
    1496
    0
                "model.embed_tokens.weight", // Tied embeddings
    1497
    0
                "embed_tokens.weight",       // SafeTensors tied embeddings
    1498
    0
            ])?;
    1499
    0
            let lm_head = self.get_tensor_f32(&lm_head_name)?;
    1500
    1501
            // Get hidden state for last token
    1502
    0
            let last_hidden = &hidden[hidden.len() - hidden_dim..];
    1503
    1504
            // Project to vocab
    1505
    0
            let mut logits = vec![0.0; vocab_size];
    1506
    0
            for (i, logit) in logits.iter_mut().enumerate() {
    1507
    0
                for (j, &h) in last_hidden.iter().enumerate() {
    1508
    0
                    *logit += h * lm_head.get(i * hidden_dim + j).copied().unwrap_or(0.0);
    1509
    0
                }
    1510
            }
    1511
    1512
    0
            Ok(logits)
    1513
    6
        }
    1514
    1515
        /// Autoregressive text generation.
    1516
        ///
    1517
        /// Generates tokens one at a time using greedy decoding (argmax sampling).
    1518
        ///
    1519
        /// # Arguments
    1520
        ///
    1521
        /// * `input_tokens` - Initial token sequence (prompt)
    1522
        /// * `max_new_tokens` - Maximum number of new tokens to generate
    1523
        /// * `eos_token_id` - End-of-sequence token ID (stops generation early)
    1524
        ///
    1525
        /// # Returns
    1526
        ///
    1527
        /// Complete token sequence including input and generated tokens
    1528
        ///
    1529
        /// # Errors
    1530
        ///
    1531
        /// Returns error if model is not a transformer or forward pass fails
    1532
    3
        pub fn generate(
    1533
    3
            &self,
    1534
    3
            input_tokens: &[u32],
    1535
    3
            max_new_tokens: usize,
    1536
    3
            eos_token_id: Option<u32>,
    1537
    3
        ) -> Result<Vec<u32>> {
    1538
    3
            if input_tokens.is_empty() {
    1539
    1
                return Err(RealizarError::InvalidShape {
    1540
    1
                    reason: "Input tokens cannot be empty".to_string(),
    1541
    1
                });
    1542
    2
            }
    1543
    1544
    2
            let mut tokens = input_tokens.to_vec();
    1545
    2
            let vocab_size = self.metadata.vocab_size.unwrap_or(0);
    1546
    1547
    2
            for _ in 0..max_new_tokens {
    1548
                // Forward pass to get logits for next token
    1549
    1
                let logits0 = self.forward(&tokens)?;
    1550
    1551
                // Greedy sampling: pick token with highest logit
    1552
    0
                let next_token = logits
    1553
    0
                    .iter()
    1554
    0
                    .enumerate()
    1555
    0
                    .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
    1556
    0
                    .map_or(0, |(idx, _)| idx as u32);
    1557
    1558
                // Check for EOS
    1559
    0
                if let Some(eos) = eos_token_id {
    1560
    0
                    if next_token == eos {
    1561
    0
                        break;
    1562
    0
                    }
    1563
    0
                }
    1564
    1565
                // Sanity check: don't append invalid tokens
    1566
    0
                if (next_token as usize) >= vocab_size && vocab_size > 0 {
    1567
    0
                    break;
    1568
    0
                }
    1569
    1570
    0
                tokens.push(next_token);
    1571
            }
    1572
    1573
    1
            Ok(tokens)
    1574
    3
        }
    1575
    1576
        /// Find first matching tensor name from candidates
    1577
    1
        fn find_tensor_name(&self, candidates: &[&str]) -> Result<String> {
    1578
    7
            for &name6 in candidates {
    1579
    6
                if self.get_tensor(name).is_some() {
    1580
    0
                    return Ok(name.to_string());
    1581
    6
                }
    1582
            }
    1583
    1
            Err(RealizarError::FormatError {
    1584
    1
                reason: format!("No matching tensor found. Tried: {:?}", candidates),
    1585
    1
            })
    1586
    1
        }
    1587
    1588
        /// Load tokenizer from sibling tokenizer.json file
    1589
        ///
    1590
        /// Looks for tokenizer.json in the same directory as the model file.
    1591
        /// Returns (vocab, bos_token_id, eos_token_id) if found.
    1592
    5
        pub fn load_tokenizer_from_sibling(
    1593
    5
            model_path: &Path,
    1594
    5
        ) -> Option<(Vec<String>, Option<u32>, Option<u32>)> {
    1595
    5
            let tokenizer_path = model_path.with_file_name("tokenizer.json");
    1596
    5
            if !tokenizer_path.exists() {
    1597
    1
                return None;
    1598
    4
            }
    1599
    1600
    4
            let content = fs::read_to_string(&tokenizer_path).ok()?0;
    1601
    4
            let json3: serde_json::Value3 = serde_json::from_str(&content).ok()?1;
    1602
    1603
            // Extract vocabulary from model.vocab
    1604
    3
            let vocab_obj1 = json.get("model")?1.get2("vocab")?1;
    1605
    1
            let vocab_map = vocab_obj.as_object()?0;
    1606
    1607
            // Build vocab vector (sorted by ID)
    1608
    1
            let mut vocab_vec: Vec<(String, u32)> = vocab_map
    1609
    1
                .iter()
    1610
    4
                .filter_map1(|(token, id)| Some((token.clone(), id.as_u64()?0 as u32)))
    1611
    1
                .collect();
    1612
    1
            vocab_vec.sort_by_key(|(_, id)| *id);
    1613
    1614
    1
            let vocab: Vec<String> = vocab_vec.into_iter().map(|(token, _)| token).collect();
    1615
    1616
            // Extract special tokens
    1617
    1
            let mut bos_id = None;
    1618
    1
            let mut eos_id = None;
    1619
    1620
    1
            if let Some(added_tokens) = json.get("added_tokens").and_then(|v| v.as_array()) {
    1621
    3
                for token2 in added_tokens {
    1622
    2
                    let content = token.get("content").and_then(|v| v.as_str());
    1623
    2
                    let id = token
    1624
    2
                        .get("id")
    1625
    2
                        .and_then(serde_json::Value::as_u64)
    1626
    2
                        .map(|v| v as u32);
    1627
    1628
    2
                    if let (Some(content), Some(id)) = (content, id) {
    1629
    2
                        if content == "<|endoftext|>" || content == "</s>" || content == "<eos>"1 {
    1630
    1
                            eos_id = Some(id);
    1631
    1
                        }
    1632
    2
                        if content == "<s>" || content == "<bos>"1 {
    1633
    1
                            bos_id = Some(id);
    1634
    1
                        }
    1635
    0
                    }
    1636
                }
    1637
    0
            }
    1638
    1639
    1
            Some((vocab, bos_id, eos_id))
    1640
    5
        }
    1641
    1642
        /// Decode token IDs to text using vocabulary
    1643
        ///
    1644
        /// If vocab is not available, returns formatted token IDs.
    1645
    22
        pub fn decode_tokens(vocab: &[String], token_ids: &[u32]) -> String {
    1646
    22
            let mut result = String::new();
    1647
    70
            for &id48 in token_ids {
    1648
    48
                if let Some(token39) = vocab.get(id as usize) {
    1649
    39
                    // Handle byte-level BPE encoding (Ġ = space prefix)
    1650
    39
                    let decoded = token
    1651
    39
                        .replace("Ġ", " ")
    1652
    39
                        .replace("Ċ", "\n")
    1653
    39
                        .replace("ĉ", "\t");
    1654
    39
                    result.push_str(&decoded);
    1655
    39
                } else {
    1656
    9
                    result.push_str(&format!("[{}]", id));
    1657
    9
                }
    1658
            }
    1659
    22
            result
    1660
    22
        }
    1661
    1662
        /// Encode text to token IDs using BPE tokenization
    1663
        ///
    1664
        /// Loads vocab and merges from tokenizer.json, then performs BPE encoding.
    1665
        /// Returns None if tokenizer not found or encoding fails.
    1666
    3
        pub fn encode_text(model_path: &Path, text: &str) -> Option<Vec<u32>> {
    1667
    3
            let tokenizer_path = model_path.with_file_name("tokenizer.json");
    1668
    3
            if !tokenizer_path.exists() {
    1669
    1
                return None;
    1670
    2
            }
    1671
    1672
    2
            let content = fs::read_to_string(&tokenizer_path).ok()?0;
    1673
    2
            let json1: serde_json::Value1 = serde_json::from_str(&content).ok()?1;
    1674
    1675
            // Extract vocabulary (token -> id)
    1676
    1
            let vocab_obj = json.get("model")?0.get("vocab")?0;
    1677
    1
            let vocab_map = vocab_obj.as_object()?0;
    1678
    1
            let token_to_id: HashMap<String, u32> = vocab_map
    1679
    1
                .iter()
    1680
    4
                .filter_map1(|(token, id)| Some((token.clone(), id.as_u64()?0 as u32)))
    1681
    1
                .collect();
    1682
    1683
            // Extract merges (pair rules for BPE)
    1684
    1
            let merges = json.get("model")?0.get("merges")?0.as_array()?0;
    1685
    1686
    1
            let merge_rules: Vec<(String, String)> = merges
    1687
    1
                .iter()
    1688
    1
                .filter_map(|m| {0
    1689
    0
                    let s = m.as_str()?;
    1690
    0
                    let parts: Vec<&str> = s.splitn(2, ' ').collect();
    1691
    0
                    if parts.len() == 2 {
    1692
    0
                        Some((parts[0].to_string(), parts[1].to_string()))
    1693
                    } else {
    1694
    0
                        None
    1695
                    }
    1696
    0
                })
    1697
    1
                .collect();
    1698
    1699
            // BPE encoding: convert text to byte-level tokens, then apply merges
    1700
    1
            let tokens = bpe_encode(text, &token_to_id, &merge_rules);
    1701
    1
            Some(tokens)
    1702
    3
        }
    1703
    1704
        /// Load tokenizer from embedded APR metadata (GH-156)
    1705
        ///
    1706
        /// APR files can contain embedded tokenizer data - this is the preferred
    1707
        /// way to decode tokens since it doesn't require sibling files.
    1708
        ///
    1709
        /// Returns a simple decode-only tokenizer (no BPE encoding support).
    1710
    0
        pub fn load_embedded_tokenizer(&self) -> Option<SimpleTokenizer> {
    1711
    0
            let vocab = self.metadata.get_embedded_vocabulary()?;
    1712
    0
            let bos_id = self.metadata.get_embedded_bos_token_id();
    1713
    0
            let eos_id = self.metadata.get_embedded_eos_token_id();
    1714
    1715
    0
            Some(SimpleTokenizer {
    1716
    0
                id_to_token: vocab,
    1717
    0
                bos_token_id: bos_id,
    1718
    0
                eos_token_id: eos_id,
    1719
    0
            })
    1720
    0
        }
    1721
    1722
        /// Load a full tokenizer struct from sibling tokenizer.json
    1723
        ///
    1724
        /// Returns a BpeTokenizer that can be reused for multiple encode/decode calls.
    1725
        /// For decode-only operations, prefer `load_embedded_tokenizer()` first.
    1726
    3
        pub fn load_tokenizer(model_path: &Path) -> Option<BpeTokenizer> {
    1727
    3
            let tokenizer_path = model_path.with_file_name("tokenizer.json");
    1728
    3
            if !tokenizer_path.exists() {
    1729
    1
                return None;
    1730
    2
            }
    1731
    1732
    2
            let content = fs::read_to_string(&tokenizer_path).ok()?0;
    1733
    2
            let json1: serde_json::Value1 = serde_json::from_str(&content).ok()?1;
    1734
    1735
            // Extract vocabulary
    1736
    1
            let vocab_obj = json.get("model")?0.get("vocab")?0;
    1737
    1
            let vocab_map = vocab_obj.as_object()?0;
    1738
    1739
    1
            let mut token_to_id: HashMap<String, u32> = HashMap::new();
    1740
    1
            let mut id_to_token: Vec<String> = Vec::new();
    1741
    1742
    1
            let mut vocab_vec: Vec<(String, u32)> = vocab_map
    1743
    1
                .iter()
    1744
    5
                .filter_map1(|(token, id)| Some((token.clone(), id.as_u64()?0 as u32)))
    1745
    1
                .collect();
    1746
    1
            vocab_vec.sort_by_key(|(_, id)| *id);
    1747
    1748
    6
            for (token5, id5) in vocab_vec {
    1749
    5
                token_to_id.insert(token.clone(), id);
    1750
                // Pad id_to_token if needed
    1751
    10
                while id_to_token.len() <= id as usize {
    1752
    5
                    id_to_token.push(String::new());
    1753
    5
                }
    1754
    5
                id_to_token[id as usize] = token;
    1755
            }
    1756
    1757
            // Extract merges
    1758
    1
            let merges = json.get("model")?0.get("merges")?0.as_array()?0;
    1759
    1
            let merge_rules: Vec<(String, String)> = merges
    1760
    1
                .iter()
    1761
    1
                .filter_map(|m| {
    1762
    1
                    let s = m.as_str()?0;
    1763
    1
                    let parts: Vec<&str> = s.splitn(2, ' ').collect();
    1764
    1
                    if parts.len() == 2 {
    1765
    1
                        Some((parts[0].to_string(), parts[1].to_string()))
    1766
                    } else {
    1767
    0
                        None
    1768
                    }
    1769
    1
                })
    1770
    1
                .collect();
    1771
    1772
            // Extract special tokens
    1773
    1
            let mut bos_id = None;
    1774
    1
            let mut eos_id = None;
    1775
    1776
    1
            if let Some(added_tokens) = json.get("added_tokens").and_then(|v| v.as_array()) {
    1777
    3
                for token2 in added_tokens {
    1778
    2
                    let content = token.get("content").and_then(|v| v.as_str());
    1779
    2
                    let id = token
    1780
    2
                        .get("id")
    1781
    2
                        .and_then(serde_json::Value::as_u64)
    1782
    2
                        .map(|v| v as u32);
    1783
    1784
    2
                    if let (Some(content), Some(id)) = (content, id) {
    1785
    2
                        if content == "<|endoftext|>" || content == "</s>"1 || content == "<eos>"1 {
    1786
    1
                            eos_id = Some(id);
    1787
    1
                        }
    1788
    2
                        if content == "<s>" || content == "<bos>" {
    1789
    1
                            bos_id = Some(id);
    1790
    1
                        }
    1791
    0
                    }
    1792
                }
    1793
    0
            }
    1794
    1795
    1
            Some(BpeTokenizer {
    1796
    1
                token_to_id,
    1797
    1
                id_to_token,
    1798
    1
                merge_rules,
    1799
    1
                bos_id,
    1800
    1
                eos_id,
    1801
    1
            })
    1802
    3
        }
    1803
    }
    1804
    1805
    /// Legacy type alias for APR v2 model
    1806
    pub type AprModel = AprV2Model;
    1807
    /// Legacy type alias (model types are now in metadata)
    1808
    pub type AprModelType = ();
    1809
    1810
    // =============================================================================
    1811
    1812
    use memmap2::Mmap;
    1813
    1814
    /// Memory-mapped APR model for fast loading and GPU inference
    1815
    ///
    1816
    /// Similar to MappedGGUFModel, this provides zero-copy access to APR tensor data.
    1817
    /// The file is memory-mapped for fast startup (~36x faster than full file read).
    1818
    #[derive(Debug)]
    1819
    pub struct MappedAprModel {
    1820
        /// APR header
    1821
        pub header: AprHeader,
    1822
        /// Model metadata
    1823
        pub metadata: AprMetadata,
    1824
        /// Tensor index
    1825
        pub tensors: Vec<TensorEntry>,
    1826
        /// Memory-mapped file data
    1827
        mmap: Mmap,
    1828
    }
    1829
    1830
    impl MappedAprModel {
    1831
        /// Load an APR model with memory mapping for fast startup
    1832
        ///
    1833
        /// # Arguments
    1834
        /// * `path` - Path to the .apr file
    1835
        ///
    1836
        /// # Errors
    1837
        /// Returns error if file cannot be opened or has invalid format.
    1838
    0
        pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Self> {
    1839
    0
            let file = File::open(path.as_ref()).map_err(|e| RealizarError::IoError {
    1840
    0
                message: format!("Failed to open .apr file: {e}"),
    1841
    0
            })?;
    1842
    1843
            // SAFETY: File is opened read-only, callers validate format before trusting data
    1844
    0
            let mmap = unsafe {
    1845
    0
                Mmap::map(&file).map_err(|e| RealizarError::IoError {
    1846
    0
                    message: format!("Failed to mmap .apr file: {e}"),
    1847
    0
                })?
    1848
            };
    1849
    1850
    0
            Self::from_mmap(mmap)
    1851
    0
        }
    1852
    1853
        /// Create from existing memory map
    1854
    0
        fn from_mmap(mmap: Mmap) -> Result<Self> {
    1855
    0
            let data = &mmap[..];
    1856
    1857
            // Parse header
    1858
    0
            let header = AprHeader::from_bytes(data)?;
    1859
    1860
            // Validate magic
    1861
    0
            if header.magic != MAGIC {
    1862
    0
                return Err(RealizarError::FormatError {
    1863
    0
                    reason: "Invalid APR magic bytes".to_string(),
    1864
    0
                });
    1865
    0
            }
    1866
    1867
            // Parse metadata
    1868
    0
            let metadata_start = header.metadata_offset as usize;
    1869
    0
            let metadata_end = metadata_start + header.metadata_size as usize;
    1870
    1871
    0
            if data.len() < metadata_end {
    1872
    0
                return Err(RealizarError::FormatError {
    1873
    0
                    reason: "APR file truncated: metadata extends past EOF".to_string(),
    1874
    0
                });
    1875
    0
            }
    1876
    1877
    0
            let metadata: AprMetadata = if header.metadata_size > 0 {
    1878
    0
                serde_json::from_slice(&data[metadata_start..metadata_end]).unwrap_or_default()
    1879
            } else {
    1880
    0
                AprMetadata::default()
    1881
            };
    1882
    1883
            // Parse tensor index
    1884
    0
            let index_start = header.tensor_index_offset as usize;
    1885
    0
            let index_end = header.data_offset as usize;
    1886
    1887
    0
            let mut tensors = Vec::with_capacity(header.tensor_count as usize);
    1888
    0
            if index_start < index_end && index_end <= data.len() {
    1889
    0
                let index_data = &data[index_start..index_end];
    1890
    0
                let mut pos = 0;
    1891
    1892
    0
                while pos < index_data.len() && tensors.len() < header.tensor_count as usize {
    1893
    0
                    match TensorEntry::from_binary(&index_data[pos..]) {
    1894
    0
                        Ok((entry, consumed)) => {
    1895
    0
                            tensors.push(entry);
    1896
    0
                            pos += consumed;
    1897
    0
                        },
    1898
    0
                        Err(_) => break,
    1899
                    }
    1900
                }
    1901
    0
            }
    1902
    1903
    0
            Ok(Self {
    1904
    0
                header,
    1905
    0
                metadata,
    1906
    0
                tensors,
    1907
    0
                mmap,
    1908
    0
            })
    1909
    0
        }
    1910
    1911
        /// Get raw file data (for tensor access)
    1912
        #[must_use]
    1913
    0
        pub fn data(&self) -> &[u8] {
    1914
    0
            &self.mmap[..]
    1915
    0
        }
    1916
    1917
        /// Get file size in bytes
    1918
        #[must_use]
    1919
    0
        pub fn file_size(&self) -> usize {
    1920
    0
            self.mmap.len()
    1921
    0
        }
    1922
    1923
        /// Get tensor count
    1924
        #[must_use]
    1925
    0
        pub fn tensor_count(&self) -> usize {
    1926
    0
            self.tensors.len()
    1927
    0
        }
    1928
    1929
        /// Get data offset (start of tensor data section)
    1930
        #[must_use]
    1931
    0
        pub fn data_offset(&self) -> u64 {
    1932
    0
            self.header.data_offset
    1933
    0
        }
    1934
    1935
        /// Find tensor by name
    1936
        #[must_use]
    1937
    0
        pub fn find_tensor(&self, name: &str) -> Option<&TensorEntry> {
    1938
    0
            self.tensors.iter().find(|t| t.name == name)
    1939
    0
        }
    1940
    1941
        /// Get raw tensor data by name
    1942
    0
        pub fn get_tensor_data(&self, name: &str) -> Result<&[u8]> {
    1943
    0
            let tensor = self
    1944
    0
                .find_tensor(name)
    1945
    0
                .ok_or_else(|| RealizarError::FormatError {
    1946
    0
                    reason: format!("Tensor not found: {name}"),
    1947
    0
                })?;
    1948
    1949
    0
            let start = self.header.data_offset as usize + tensor.offset as usize;
    1950
    0
            let end = start + tensor.size as usize;
    1951
    1952
    0
            if end > self.mmap.len() {
    1953
    0
                return Err(RealizarError::FormatError {
    1954
    0
                    reason: format!("Tensor {name} extends past EOF"),
    1955
    0
                });
    1956
    0
            }
    1957
    1958
    0
            Ok(&self.mmap[start..end])
    1959
    0
        }
    1960
    1961
        /// Convert APR dtype string to GGML qtype
    1962
        #[must_use]
    1963
    0
        pub fn dtype_to_qtype(dtype: &str) -> u32 {
    1964
    0
            match dtype {
    1965
    0
                "F32" => 0,
    1966
    0
                "F16" => 1,
    1967
    0
                "Q4_0" => 2,
    1968
    0
                "Q4_1" => 3,
    1969
    0
                "Q5_0" => 6,
    1970
    0
                "Q5_1" => 7,
    1971
    0
                "Q8_0" => 8,
    1972
    0
                "Q8_1" => 9,
    1973
    0
                "Q2_K" => 10,
    1974
    0
                "Q3_K" => 11,
    1975
    0
                "Q4_K" => 12,
    1976
    0
                "Q5_K" => 13,
    1977
    0
                "Q6_K" => 14,
    1978
    0
                "IQ2_XXS" => 16,
    1979
    0
                "IQ2_XS" => 17,
    1980
    0
                "BF16" => 30,
    1981
    0
                _ => 0, // Default to F32
    1982
            }
    1983
    0
        }
    1984
    }
    1985
    1986
    // Tests extracted to tests.rs (PMAT-802)
    1987
    #[cfg(test)]
    1988
    #[path = "tests.rs"]
    1989
    mod apr_tests;