## 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/layers/model.rs

    Line
    Count
    Source
    1
    //! Model components for transformer inference
    2
    //!
    3
    //! Extracted from layers/mod.rs (PMAT-802) to reduce module size.
    4
    //! Contains:
    5
    //! - KVCache: Key-Value cache for efficient autoregressive generation
    6
    //! - TransformerBlock: Single transformer layer (attention + FFN)
    7
    //! - Embedding: Token embedding layer
    8
    //! - Model: Full transformer model for inference
    9
    //! - ModelConfig: Configuration for transformer models
    10
    11
    use crate::{
    12
        error::{RealizarError, Result},
    13
        generate::{sample_token, GenerationConfig},
    14
        tensor::Tensor,
    15
    };
    16
    17
    use super::{FeedForward, LayerNorm, Linear, MultiHeadAttention};
    18
    19
    /// Key-Value Cache for efficient transformer inference
    20
    ///
    21
    /// Stores key and value tensors from previous positions to avoid
    22
    /// recomputation during autoregressive generation. Each forward pass
    23
    /// only computes K/V for the new token and appends to the cache.
    24
    ///
    25
    /// # Usage
    26
    ///
    27
    /// 1. Create cache with `KVCache::new(num_layers, max_seq_len, head_dim)`
    28
    /// 2. At each generation step, call `update` with new K/V
    29
    /// 3. Use `get_key`/`get_value` to retrieve cached tensors
    30
    /// 4. Call `clear` to reset for new sequence
    31
    #[derive(Debug, Clone)]
    32
    pub struct KVCache {
    33
        /// Number of transformer layers
    34
        num_layers: usize,
    35
        /// Maximum sequence length
    36
        max_seq_len: usize,
    37
        /// Dimension per head
    38
        head_dim: usize,
    39
        /// Current sequence position
    40
        current_pos: usize,
    41
        /// Cached keys for each layer: `[num_layers][max_seq_len * head_dim]`
    42
        keys: Vec<Vec<f32>>,
    43
        /// Cached values for each layer: `[num_layers][max_seq_len * head_dim]`
    44
        values: Vec<Vec<f32>>,
    45
    }
    46
    47
    impl KVCache {
    48
        /// Create a new KV cache
    49
        ///
    50
        /// # Arguments
    51
        ///
    52
        /// * `num_layers` - Number of transformer layers to cache
    53
        /// * `max_seq_len` - Maximum sequence length to cache
    54
        /// * `head_dim` - Dimension per attention head
    55
        ///
    56
        /// # Errors
    57
        ///
    58
        /// Returns error if any dimension is zero
    59
    20
        pub fn new(num_layers: usize, max_seq_len: usize, head_dim: usize) -> Result<Self> {
    60
    20
            if num_layers == 0 {
    61
    2
                return Err(RealizarError::InvalidShape {
    62
    2
                    reason: "num_layers must be > 0".to_string(),
    63
    2
                });
    64
    18
            }
    65
    18
            if max_seq_len == 0 {
    66
    2
                return Err(RealizarError::InvalidShape {
    67
    2
                    reason: "max_seq_len must be > 0".to_string(),
    68
    2
                });
    69
    16
            }
    70
    16
            if head_dim == 0 {
    71
    2
                return Err(RealizarError::InvalidShape {
    72
    2
                    reason: "head_dim must be > 0".to_string(),
    73
    2
                });
    74
    14
            }
    75
    76
    14
            let cache_size = max_seq_len * head_dim;
    77
    14
            let keys = vec![vec![0.0; cache_size]; num_layers];
    78
    14
            let values = vec![vec![0.0; cache_size]; num_layers];
    79
    80
    14
            Ok(Self {
    81
    14
                num_layers,
    82
    14
                max_seq_len,
    83
    14
                head_dim,
    84
    14
                current_pos: 0,
    85
    14
                keys,
    86
    14
                values,
    87
    14
            })
    88
    20
        }
    89
    90
        /// Update cache with new key/value for a layer
    91
        ///
    92
        /// # Arguments
    93
        ///
    94
        /// * `layer` - Layer index
    95
        /// * `key` - New key tensor `[head_dim]`
    96
        /// * `value` - New value tensor `[head_dim]`
    97
        ///
    98
        /// # Errors
    99
        ///
    100
        /// Returns error if layer is out of bounds, cache is full, or tensor sizes don't match
    101
    13
        pub fn update(&mut self, layer: usize, key: &Tensor<f32>, value: &Tensor<f32>) -> Result<()> {
    102
    13
            if layer >= self.num_layers {
    103
    1
                return Err(RealizarError::InvalidShape {
    104
    1
                    reason: format!(
    105
    1
                        "Layer {} out of bounds (max {})",
    106
    1
                        layer,
    107
    1
                        self.num_layers - 1
    108
    1
                    ),
    109
    1
                });
    110
    12
            }
    111
    12
            if self.current_pos >= self.max_seq_len {
    112
    1
                return Err(RealizarError::InvalidShape {
    113
    1
                    reason: format!(
    114
    1
                        "Cache full at position {} (max {})",
    115
    1
                        self.current_pos, self.max_seq_len
    116
    1
                    ),
    117
    1
                });
    118
    11
            }
    119
    120
    11
            let k_data = key.data();
    121
    11
            let v_data = value.data();
    122
    123
    11
            if k_data.len() != self.head_dim {
    124
    1
                return Err(RealizarError::InvalidShape {
    125
    1
                    reason: format!("Key size {} != head_dim {}", k_data.len(), self.head_dim),
    126
    1
                });
    127
    10
            }
    128
    10
            if v_data.len() != self.head_dim {
    129
    1
                return Err(RealizarError::InvalidShape {
    130
    1
                    reason: format!("Value size {} != head_dim {}", v_data.len(), self.head_dim),
    131
    1
                });
    132
    9
            }
    133
    134
            // Copy key and value into cache at current position
    135
    9
            let offset = self.current_pos * self.head_dim;
    136
    9
            self.keys[layer][offset..offset + self.head_dim].copy_from_slice(k_data);
    137
    9
            self.values[layer][offset..offset + self.head_dim].copy_from_slice(v_data);
    138
    139
    9
            Ok(())
    140
    13
        }
    141
    142
        /// Advance to next position after updating all layers
    143
        ///
    144
        /// Call this after updating all layers for the current position
    145
    8
        pub fn advance(&mut self) {
    146
    8
            if self.current_pos < self.max_seq_len {
    147
    8
                self.current_pos += 1;
    148
    8
            }0
    149
    8
        }
    150
    151
        /// Get cached keys for a layer up to current position
    152
        ///
    153
        /// # Arguments
    154
        ///
    155
        /// * `layer` - Layer index
    156
        ///
    157
        /// # Returns
    158
        ///
    159
        /// Tensor with shape `[current_pos, head_dim]`
    160
        ///
    161
        /// # Errors
    162
        ///
    163
        /// Returns error if layer is out of bounds
    164
    6
        pub fn get_key(&self, layer: usize) -> Result<Tensor<f32>> {
    165
    6
            if layer >= self.num_layers {
    166
    1
                return Err(RealizarError::InvalidShape {
    167
    1
                    reason: format!(
    168
    1
                        "Layer {} out of bounds (max {})",
    169
    1
                        layer,
    170
    1
                        self.num_layers - 1
    171
    1
                    ),
    172
    1
                });
    173
    5
            }
    174
    175
    5
            if self.current_pos == 0 {
    176
                // Return empty tensor with shape [0, head_dim] is invalid
    177
                // Return [1, head_dim] with zeros for consistency
    178
    1
                return Tensor::from_vec(vec![1, self.head_dim], vec![0.0; self.head_dim]);
    179
    4
            }
    180
    181
    4
            let size = self.current_pos * self.head_dim;
    182
    4
            let data = self.keys[layer][..size].to_vec();
    183
    4
            Tensor::from_vec(vec![self.current_pos, self.head_dim], data)
    184
    6
        }
    185
    186
        /// Get cached values for a layer up to current position
    187
        ///
    188
        /// # Arguments
    189
        ///
    190
        /// * `layer` - Layer index
    191
        ///
    192
        /// # Returns
    193
        ///
    194
        /// Tensor with shape `[current_pos, head_dim]`
    195
        ///
    196
        /// # Errors
    197
        ///
    198
        /// Returns error if layer is out of bounds
    199
    4
        pub fn get_value(&self, layer: usize) -> Result<Tensor<f32>> {
    200
    4
            if layer >= self.num_layers {
    201
    1
                return Err(RealizarError::InvalidShape {
    202
    1
                    reason: format!(
    203
    1
                        "Layer {} out of bounds (max {})",
    204
    1
                        layer,
    205
    1
                        self.num_layers - 1
    206
    1
                    ),
    207
    1
                });
    208
    3
            }
    209
    210
    3
            if self.current_pos == 0 {
    211
    1
                return Tensor::from_vec(vec![1, self.head_dim], vec![0.0; self.head_dim]);
    212
    2
            }
    213
    214
    2
            let size = self.current_pos * self.head_dim;
    215
    2
            let data = self.values[layer][..size].to_vec();
    216
    2
            Tensor::from_vec(vec![self.current_pos, self.head_dim], data)
    217
    4
        }
    218
    219
        /// Clear cache and reset position to 0
    220
    1
        pub fn clear(&mut self) {
    221
    1
            self.current_pos = 0;
    222
            // Optionally zero out the cache (not strictly necessary)
    223
    1
            for layer in 0..self.num_layers {
    224
    1
                self.keys[layer].fill(0.0);
    225
    1
                self.values[layer].fill(0.0);
    226
    1
            }
    227
    1
        }
    228
    229
        /// Get current sequence position
    230
        #[must_use]
    231
    5
        pub fn current_pos(&self) -> usize {
    232
    5
            self.current_pos
    233
    5
        }
    234
    235
        /// Get number of layers
    236
        #[must_use]
    237
    4
        pub fn num_layers(&self) -> usize {
    238
    4
            self.num_layers
    239
    4
        }
    240
    241
        /// Get maximum sequence length
    242
        #[must_use]
    243
    2
        pub fn max_seq_len(&self) -> usize {
    244
    2
            self.max_seq_len
    245
    2
        }
    246
    247
        /// Get head dimension
    248
        #[must_use]
    249
    2
        pub fn head_dim(&self) -> usize {
    250
    2
            self.head_dim
    251
    2
        }
    252
    253
        /// Check if cache is full
    254
        #[must_use]
    255
    3
        pub fn is_full(&self) -> bool {
    256
    3
            self.current_pos >= self.max_seq_len
    257
    3
        }
    258
    }
    259
    260
    /// Transformer Block (Pre-norm architecture)
    261
    ///
    262
    /// A single transformer block combining self-attention and feed-forward layers
    263
    /// with residual connections and layer normalization.
    264
    ///
    265
    /// # Architecture
    266
    ///
    267
    /// ```text
    268
    /// Input
    269
    ///   │
    270
    ///   ├──────────────────┐
    271
    ///   ▼                  │
    272
    /// LayerNorm            │
    273
    ///   ▼                  │
    274
    /// Attention            │
    275
    ///   ▼                  │
    276
    ///   + <────────────────┘ (residual)
    277
    ///   │
    278
    ///   ├──────────────────┐
    279
    ///   ▼                  │
    280
    /// LayerNorm            │
    281
    ///   ▼                  │
    282
    /// FFN                  │
    283
    ///   ▼                  │
    284
    ///   + <────────────────┘ (residual)
    285
    ///   │
    286
    /// Output
    287
    /// ```
    288
    ///
    289
    /// This is the pre-norm architecture used in `LLaMA`, GPT-NeoX, and modern transformers.
    290
    #[derive(Debug, Clone)]
    291
    pub struct TransformerBlock {
    292
        /// Layer normalization before attention
    293
        attn_norm: LayerNorm,
    294
        /// Multi-head self-attention layer with Q/K/V/O projections
    295
        attention: MultiHeadAttention,
    296
        /// Layer normalization before FFN
    297
        ffn_norm: LayerNorm,
    298
        /// Feed-forward network
    299
        ffn: FeedForward,
    300
        /// Hidden dimension
    301
        hidden_dim: usize,
    302
        /// Number of attention heads
    303
        num_heads: usize,
    304
    }
    305
    306
    impl TransformerBlock {
    307
        /// Create a new transformer block
    308
        ///
    309
        /// # Arguments
    310
        ///
    311
        /// * `hidden_dim` - Hidden dimension (model dimension)
    312
        /// * `num_heads` - Number of attention heads
    313
        /// * `intermediate_dim` - FFN intermediate dimension
    314
        /// * `eps` - Layer normalization epsilon
    315
        ///
    316
        /// # Errors
    317
        ///
    318
        /// Returns error if:
    319
        /// - `hidden_dim` is zero or not divisible by `num_heads`
    320
        /// - `num_heads` is zero
    321
        /// - `intermediate_dim` is zero
    322
    196
        pub fn new(
    323
    196
            hidden_dim: usize,
    324
    196
            num_heads: usize,
    325
    196
            intermediate_dim: usize,
    326
    196
            eps: f32,
    327
    196
        ) -> Result<Self> {
    328
    196
            if hidden_dim == 0 {
    329
    1
                return Err(RealizarError::InvalidShape {
    330
    1
                    reason: "hidden_dim must be > 0".to_string(),
    331
    1
                });
    332
    195
            }
    333
    195
            if num_heads == 0 {
    334
    1
                return Err(RealizarError::InvalidShape {
    335
    1
                    reason: "num_heads must be > 0".to_string(),
    336
    1
                });
    337
    194
            }
    338
    194
            if !hidden_dim.is_multiple_of(num_heads) {
    339
    1
                return Err(RealizarError::InvalidShape {
    340
    1
                    reason: format!(
    341
    1
                        "hidden_dim {hidden_dim} must be divisible by num_heads {num_heads}"
    342
    1
                    ),
    343
    1
                });
    344
    193
            }
    345
    346
    193
            let attn_norm = LayerNorm::new(hidden_dim, eps)?0;
    347
            // Use standard MHA with Q/K/V/O projections
    348
    193
            let attention = MultiHeadAttention::mha(hidden_dim, num_heads)?0;
    349
    193
            let ffn_norm = LayerNorm::new(hidden_dim, eps)?0;
    350
    193
            let ffn192 = FeedForward::new(hidden_dim, intermediate_dim)?1;
    351
    352
    192
            Ok(Self {
    353
    192
                attn_norm,
    354
    192
                attention,
    355
    192
                ffn_norm,
    356
    192
                ffn,
    357
    192
                hidden_dim,
    358
    192
                num_heads,
    359
    192
            })
    360
    196
        }
    361
    362
        /// Forward pass through the transformer block
    363
        ///
    364
        /// # Arguments
    365
        ///
    366
        /// * `input` - Input tensor `[seq_len, hidden_dim]`
    367
        ///
    368
        /// # Returns
    369
        ///
    370
        /// Output tensor `[seq_len, hidden_dim]`
    371
        ///
    372
        /// # Errors
    373
        ///
    374
        /// Returns error if input shape is invalid
    375
        ///
    376
        /// # Note
    377
        ///
    378
        /// This simplified implementation uses the same input for Q, K, V (self-attention).
    379
        /// Production models would compute Q, K, V projections separately.
    380
    1.58k
        pub fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
    381
    1.58k
            let shape = input.shape();
    382
    383
    1.58k
            if shape.is_empty() {
    384
    0
                return Err(RealizarError::InvalidShape {
    385
    0
                    reason: "Input tensor must have at least 1 dimension".to_string(),
    386
    0
                });
    387
    1.58k
            }
    388
    389
    1.58k
            let last_dim = shape[shape.len() - 1];
    390
    1.58k
            if last_dim != self.hidden_dim {
    391
    1
                return Err(RealizarError::InvalidShape {
    392
    1
                    reason: format!(
    393
    1
                        "Expected last dimension {}, got {}",
    394
    1
                        self.hidden_dim, last_dim
    395
    1
                    ),
    396
    1
                });
    397
    1.58k
            }
    398
    399
            // Pre-norm attention block
    400
    1.58k
            let normed = self.attn_norm.forward(input)?0;
    401
    402
            // Self-attention with proper Q/K/V/O projections via MultiHeadAttention
    403
    1.58k
            let attn_out = self.attention.forward(&normed)?0;
    404
    405
            // Residual connection
    406
    1.58k
            let mut residual1 = Vec::with_capacity(input.data().len());
    407
    3.72M
            for (i, &val) in input.data()1.58k.iter1.58k().enumerate1.58k() {
    408
    3.72M
                residual1.push(val + attn_out.data()[i]);
    409
    3.72M
            }
    410
    1.58k
            let after_attn = Tensor::from_vec(shape.to_vec(), residual1)?0;
    411
    412
            // Pre-norm FFN block
    413
    1.58k
            let normed2 = self.ffn_norm.forward(&after_attn)?0;
    414
    1.58k
            let ffn_out = self.ffn.forward(&normed2)?0;
    415
    416
            // Residual connection
    417
    1.58k
            let mut residual2 = Vec::with_capacity(after_attn.data().len());
    418
    3.72M
            for (i, &val) in after_attn.data()1.58k.iter1.58k().enumerate1.58k() {
    419
    3.72M
                residual2.push(val + ffn_out.data()[i]);
    420
    3.72M
            }
    421
    422
    1.58k
            Tensor::from_vec(shape.to_vec(), residual2)
    423
    1.58k
        }
    424
    425
        /// Get hidden dimension
    426
        #[must_use]
    427
    4
        pub fn hidden_dim(&self) -> usize {
    428
    4
            self.hidden_dim
    429
    4
        }
    430
    431
        /// Get mutable reference to attention layer normalization
    432
    1
        pub fn attn_norm_mut(&mut self) -> &mut LayerNorm {
    433
    1
            &mut self.attn_norm
    434
    1
        }
    435
    436
        /// Get mutable reference to multi-head attention
    437
    1
        pub fn attention_mut(&mut self) -> &mut MultiHeadAttention {
    438
    1
            &mut self.attention
    439
    1
        }
    440
    441
        /// Get number of attention heads
    442
        #[must_use]
    443
    0
        pub fn num_heads(&self) -> usize {
    444
    0
            self.num_heads
    445
    0
        }
    446
    447
        /// Get mutable reference to FFN layer normalization
    448
    1
        pub fn ffn_norm_mut(&mut self) -> &mut LayerNorm {
    449
    1
            &mut self.ffn_norm
    450
    1
        }
    451
    452
        /// Get mutable reference to FFN
    453
    1
        pub fn ffn_mut(&mut self) -> &mut FeedForward {
    454
    1
            &mut self.ffn
    455
    1
        }
    456
    }
    457
    458
    /// Embedding layer for converting token IDs to vectors
    459
    ///
    460
    /// Maps discrete token IDs to continuous vector representations.
    461
    /// This is the first layer in a transformer model.
    462
    #[derive(Debug, Clone)]
    463
    pub struct Embedding {
    464
        /// Vocabulary size
    465
        vocab_size: usize,
    466
        /// Embedding dimension
    467
        embed_dim: usize,
    468
        /// Embedding weights: `[vocab_size, embed_dim]`
    469
        weights: Vec<f32>,
    470
    }
    471
    472
    impl Embedding {
    473
        /// Create a new embedding layer
    474
        ///
    475
        /// # Arguments
    476
        ///
    477
        /// * `vocab_size` - Size of vocabulary
    478
        /// * `embed_dim` - Dimension of embedding vectors
    479
        ///
    480
        /// # Errors
    481
        ///
    482
        /// Returns error if `vocab_size` or `embed_dim` is zero
    483
    180
        pub fn new(vocab_size: usize, embed_dim: usize) -> Result<Self> {
    484
    180
            if vocab_size == 0 {
    485
    3
                return Err(RealizarError::InvalidShape {
    486
    3
                    reason: "vocab_size must be > 0".to_string(),
    487
    3
                });
    488
    177
            }
    489
    177
            if embed_dim == 0 {
    490
    2
                return Err(RealizarError::InvalidShape {
    491
    2
                    reason: "embed_dim must be > 0".to_string(),
    492
    2
                });
    493
    175
            }
    494
    495
    175
            let weights = vec![0.0; vocab_size * embed_dim];
    496
    497
    175
            Ok(Self {
    498
    175
                vocab_size,
    499
    175
                embed_dim,
    500
    175
                weights,
    501
    175
            })
    502
    180
        }
    503
    504
        /// Look up embeddings for token IDs
    505
        ///
    506
        /// # Arguments
    507
        ///
    508
        /// * `token_ids` - Slice of token IDs
    509
        ///
    510
        /// # Returns
    511
        ///
    512
        /// Tensor with shape `[seq_len, embed_dim]`
    513
        ///
    514
        /// # Errors
    515
        ///
    516
        /// Returns error if any token ID is out of bounds
    517
    1.38k
        pub fn forward(&self, token_ids: &[usize]) -> Result<Tensor<f32>> {
    518
    1.38k
            if token_ids.is_empty() {
    519
    2
                return Err(RealizarError::InvalidShape {
    520
    2
                    reason: "Token IDs cannot be empty".to_string(),
    521
    2
                });
    522
    1.38k
            }
    523
    524
    1.38k
            let seq_len = token_ids.len();
    525
    1.38k
            let mut output = Vec::with_capacity(seq_len * self.embed_dim);
    526
    527
    114k
            for &token_id113k in token_ids {
    528
    113k
                if token_id >= self.vocab_size {
    529
    3
                    return Err(RealizarError::InvalidShape {
    530
    3
                        reason: format!(
    531
    3
                            "Token ID {token_id} out of bounds (vocab_size={})",
    532
    3
                            self.vocab_size
    533
    3
                        ),
    534
    3
                    });
    535
    113k
                }
    536
    537
    113k
                let offset = token_id * self.embed_dim;
    538
    113k
                output.extend_from_slice(&self.weights[offset..offset + self.embed_dim]);
    539
            }
    540
    541
    1.37k
            Tensor::from_vec(vec![seq_len, self.embed_dim], output)
    542
    1.38k
        }
    543
    544
        /// Get vocabulary size
    545
        #[must_use]
    546
    4
        pub fn vocab_size(&self) -> usize {
    547
    4
            self.vocab_size
    548
    4
        }
    549
    550
        /// Get embedding dimension
    551
        #[must_use]
    552
    4
        pub fn embed_dim(&self) -> usize {
    553
    4
            self.embed_dim
    554
    4
        }
    555
    556
        /// Get mutable access to weights for loading
    557
    6
        pub fn weights_mut(&mut self) -> &mut [f32] {
    558
    6
            &mut self.weights
    559
    6
        }
    560
    }
    561
    562
    /// Transformer Language Model
    563
    ///
    564
    /// Complete transformer model for language modeling:
    565
    /// - Token embedding
    566
    /// - Stack of transformer blocks
    567
    /// - Final layer normalization
    568
    /// - Output projection (LM head)
    569
    ///
    570
    /// # Architecture
    571
    ///
    572
    /// ```text
    573
    /// Token IDs → Embedding → [TransformerBlock × N] → LayerNorm → Linear → Logits
    574
    /// ```
    575
    #[derive(Debug, Clone)]
    576
    pub struct Model {
    577
        /// Token embedding layer
    578
        embedding: Embedding,
    579
        /// Stack of transformer blocks
    580
        blocks: Vec<TransformerBlock>,
    581
        /// Final layer normalization
    582
        final_norm: LayerNorm,
    583
        /// Output projection (LM head)
    584
        lm_head: Linear,
    585
        /// Model configuration
    586
        config: ModelConfig,
    587
    }
    588
    589
    /// Configuration for the transformer model
    590
    #[derive(Debug, Clone)]
    591
    pub struct ModelConfig {
    592
        /// Vocabulary size
    593
        pub vocab_size: usize,
    594
        /// Hidden dimension
    595
        pub hidden_dim: usize,
    596
        /// Number of attention heads
    597
        pub num_heads: usize,
    598
        /// Number of transformer blocks
    599
        pub num_layers: usize,
    600
        /// FFN intermediate dimension
    601
        pub intermediate_dim: usize,
    602
        /// Layer normalization epsilon
    603
        pub eps: f32,
    604
    }
    605
    606
    impl Model {
    607
        /// Create a new transformer model
    608
        ///
    609
        /// # Arguments
    610
        ///
    611
        /// * `config` - Model configuration
    612
        ///
    613
        /// # Errors
    614
        ///
    615
        /// Returns error if configuration is invalid
    616
    163
        pub fn new(config: ModelConfig) -> Result<Self> {
    617
    163
            let embedding = Embedding::new(config.vocab_size, config.hidden_dim)?0;
    618
    619
    163
            let mut blocks = Vec::with_capacity(config.num_layers);
    620
    163
            for _ in 0..config.num_layers {
    621
    183
                blocks.push(TransformerBlock::new(
    622
    183
                    config.hidden_dim,
    623
    183
                    config.num_heads,
    624
    183
                    config.intermediate_dim,
    625
    183
                    config.eps,
    626
    0
                )?);
    627
            }
    628
    629
    163
            let final_norm = LayerNorm::new(config.hidden_dim, config.eps)?0;
    630
    163
            let lm_head = Linear::new(config.hidden_dim, config.vocab_size)?0;
    631
    632
    163
            Ok(Self {
    633
    163
                embedding,
    634
    163
                blocks,
    635
    163
                final_norm,
    636
    163
                lm_head,
    637
    163
                config,
    638
    163
            })
    639
    163
        }
    640
    641
        /// Forward pass through the model
    642
        ///
    643
        /// # Arguments
    644
        ///
    645
        /// * `token_ids` - Input token IDs
    646
        ///
    647
        /// # Returns
    648
        ///
    649
        /// Logits tensor with shape `[seq_len, vocab_size]`
    650
        ///
    651
        /// # Errors
    652
        ///
    653
        /// Returns error if input is invalid
    654
    1.35k
        pub fn forward(&self, token_ids: &[usize]) -> Result<Tensor<f32>> {
    655
            // Embed tokens
    656
    1.35k
            let mut hidden = self.embedding.forward(token_ids)?0;
    657
    658
            // Pass through transformer blocks
    659
    2.92k
            for block1.57k in &self.blocks {
    660
    1.57k
                hidden = block.forward(&hidden)?0;
    661
            }
    662
    663
            // Final layer norm
    664
    1.35k
            hidden = self.final_norm.forward(&hidden)?0;
    665
    666
            // Project to vocabulary
    667
    1.35k
            self.lm_head.forward(&hidden)
    668
    1.35k
        }
    669
    670
        /// Get model configuration
    671
        #[must_use]
    672
    6
        pub fn config(&self) -> &ModelConfig {
    673
    6
            &self.config
    674
    6
        }
    675
    676
        /// Get mutable reference to embedding layer
    677
    1
        pub fn embedding_mut(&mut self) -> &mut Embedding {
    678
    1
            &mut self.embedding
    679
    1
        }
    680
    681
        /// Get mutable reference to transformer blocks
    682
    1
        pub fn blocks_mut(&mut self) -> &mut [TransformerBlock] {
    683
    1
            &mut self.blocks
    684
    1
        }
    685
    686
        /// Get mutable reference to final layer norm
    687
    1
        pub fn final_norm_mut(&mut self) -> &mut LayerNorm {
    688
    1
            &mut self.final_norm
    689
    1
        }
    690
    691
        /// Get mutable reference to LM head
    692
    1
        pub fn lm_head_mut(&mut self) -> &mut Linear {
    693
    1
            &mut self.lm_head
    694
    1
        }
    695
    696
        /// Get number of parameters in the model (approximate)
    697
        #[must_use]
    698
    1
        pub fn num_parameters(&self) -> usize {
    699
    1
            let embed_params = self.config.vocab_size * self.config.hidden_dim;
    700
    1
            let block_params = self.config.num_layers
    701
    1
                * (
    702
    1
                    // Attention (Q, K, V, O projections would be here in full impl)
    703
    1
                    // For now just count layer norms and FFN
    704
    1
                    2 * self.config.hidden_dim  // Layer norm weights
    705
    1
                    + self.config.hidden_dim * self.config.intermediate_dim  // fc1
    706
    1
                    + self.config.intermediate_dim * self.config.hidden_dim
    707
    1
                    // fc2
    708
    1
                );
    709
    1
            let head_params = self.config.hidden_dim * self.config.vocab_size;
    710
    711
    1
            embed_params + block_params + head_params
    712
    1
        }
    713
    714
        /// Generate tokens autoregressively
    715
        ///
    716
        /// # Arguments
    717
        ///
    718
        /// * `prompt` - Initial token IDs
    719
        /// * `config` - Generation configuration
    720
        ///
    721
        /// # Returns
    722
        ///
    723
        /// Vector of generated token IDs (including prompt)
    724
        ///
    725
        /// # Errors
    726
        ///
    727
        /// Returns error if generation fails
    728
        ///
    729
        /// # Example
    730
        ///
    731
        /// ```rust,ignore
    732
        /// let generated = model.generate(&[1, 2, 3], &GenerationConfig::greedy())?;
    733
        /// ```
    734
    52
        pub fn generate(&self, prompt: &[usize], config: &GenerationConfig) -> Result<Vec<usize>> {
    735
    52
            if prompt.is_empty() {
    736
    1
                return Err(RealizarError::InvalidShape {
    737
    1
                    reason: "Prompt cannot be empty".to_string(),
    738
    1
                });
    739
    51
            }
    740
    741
    51
            let mut tokens = prompt.to_vec();
    742
    51
            let mut rng_state = config.seed.unwrap_or(42);
    743
    744
    51
            for _ in 0..config.max_tokens {
    745
                // Forward pass
    746
    1.33k
                let logits = self.forward(&tokens)?0;
    747
    748
                // Get logits for last position
    749
    1.33k
                let seq_len = tokens.len();
    750
    1.33k
                let vocab_size = self.config.vocab_size;
    751
    1.33k
                let last_logits_start = (seq_len - 1) * vocab_size;
    752
    1.33k
                let last_logits = &logits.data()[last_logits_start..last_logits_start + vocab_size];
    753
    754
    1.33k
                let last_logits_tensor = Tensor::from_vec(vec![vocab_size], last_logits.to_vec())?0;
    755
    756
                // Simple LCG for random number generation
    757
    1.33k
                rng_state = rng_state
    758
    1.33k
                    .wrapping_mul(6_364_136_223_846_793_005)
    759
    1.33k
                    .wrapping_add(1);
    760
                #[allow(clippy::cast_precision_loss)]
    761
    1.33k
                let rng_value = (rng_state >> 33) as f32 / (1u64 << 31) as f32;
    762
    763
                // Sample next token
    764
    1.33k
                let next_token1.33k = sample_token(&last_logits_tensor, config, rng_value)?1;
    765
    766
                // Check for EOS
    767
    1.33k
                if let Some(eos_id100) = config.eos_token_id {
    768
    100
                    if next_token == eos_id {
    769
    0
                        break;
    770
    100
                    }
    771
    1.23k
                }
    772
    773
    1.33k
                tokens.push(next_token);
    774
            }
    775
    776
    50
            Ok(tokens)
    777
    52
        }
    778
    }