## 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/generate/sampler.rs

    Line
    Count
    Source
    1
    //! Advanced Sampling Strategies (PMAT-802)
    2
    //!
    3
    //! Extracted from generate/mod.rs - Advanced sampling algorithms.
    4
    //!
    5
    //! ## Contents
    6
    //! - Stop sequence detection
    7
    //! - Repetition, presence/frequency penalties
    8
    //! - Min-p, Mirostat, TFS, typical sampling
    9
    //! - DRY, XTC, Eta sampling
    10
    //! - Token healing
    11
    12
    use crate::tensor::Tensor;
    13
    use crate::error::Result;
    14
    use super::{
    15
        GenerationConfig, apply_temperature, sample_greedy, sample_token,
    16
    };
    17
    use std::collections::HashMap;
    18
    19
    // ==================== Advanced Sampling Features ====================
    20
    21
    use serde::{Deserialize, Serialize};
    22
    23
    /// Stop sequence detector for generation termination
    24
    ///
    25
    /// Detects when generated text matches stop sequences and signals termination.
    26
    /// Supports both token ID sequences and string patterns.
    27
    #[derive(Debug, Clone, Default)]
    28
    pub struct StopSequenceDetector {
    29
        /// Token ID sequences to stop on
    30
        token_sequences: Vec<Vec<usize>>,
    31
        /// String patterns to stop on
    32
        string_patterns: Vec<String>,
    33
        /// Buffer for partial matches (token-based)
    34
        token_buffer: Vec<usize>,
    35
        /// Maximum sequence length to track
    36
        max_seq_len: usize,
    37
    }
    38
    39
    impl StopSequenceDetector {
    40
        /// Create new stop sequence detector
    41
    10
        pub fn new() -> Self {
    42
    10
            Self {
    43
    10
                token_sequences: Vec::new(),
    44
    10
                string_patterns: Vec::new(),
    45
    10
                token_buffer: Vec::new(),
    46
    10
                max_seq_len: 0,
    47
    10
            }
    48
    10
        }
    49
    50
        /// Add a token ID sequence as stop condition
    51
        #[must_use]
    52
    6
        pub fn with_token_sequence(mut self, sequence: Vec<usize>) -> Self {
    53
    6
            if !sequence.is_empty() {
    54
    6
                self.max_seq_len = self.max_seq_len.max(sequence.len());
    55
    6
                self.token_sequences.push(sequence);
    56
    6
            }0
    57
    6
            self
    58
    6
        }
    59
    60
        /// Add a string pattern as stop condition
    61
        #[must_use]
    62
    3
        pub fn with_string_pattern(mut self, pattern: impl Into<String>) -> Self {
    63
    3
            let pattern = pattern.into();
    64
    3
            if !pattern.is_empty() {
    65
    3
                self.string_patterns.push(pattern);
    66
    3
            }0
    67
    3
            self
    68
    3
        }
    69
    70
        /// Add multiple stop sequences from strings
    71
        #[must_use]
    72
    2
        pub fn with_stop_strings(mut self, stops: Vec<String>) -> Self {
    73
    5
            for stop3 in stops {
    74
    3
                if !stop.is_empty() {
    75
    3
                    self.string_patterns.push(stop);
    76
    3
                }0
    77
            }
    78
    2
            self
    79
    2
        }
    80
    81
        /// Check if a new token triggers a stop condition
    82
        ///
    83
        /// Returns true if generation should stop.
    84
    109
        pub fn check_token(&mut self, token_id: usize) -> bool {
    85
            // Add to buffer
    86
    109
            self.token_buffer.push(token_id);
    87
    88
            // Trim buffer to max sequence length
    89
    109
            if self.token_buffer.len() > self.max_seq_len && self.max_seq_len > 098 {
    90
    98
                self.token_buffer.remove(0);
    91
    98
            }11
    92
    93
            // Check token sequences
    94
    216
            for seq109 in &self.token_sequences {
    95
    109
                if self.token_buffer.ends_with(seq) {
    96
    2
                    return true;
    97
    107
                }
    98
            }
    99
    100
    107
            false
    101
    109
        }
    102
    103
        /// Check if generated text contains a stop string
    104
        ///
    105
        /// Returns Some(position) if stop found, None otherwise.
    106
    6
        pub fn check_text(&self, text: &str) -> Option<usize> {
    107
    10
            for pattern8 in &self.string_patterns {
    108
    8
                if let Some(pos4) = text.find(pattern) {
    109
    4
                    return Some(pos);
    110
    4
                }
    111
            }
    112
    2
            None
    113
    6
        }
    114
    115
        /// Reset detector state
    116
    1
        pub fn reset(&mut self) {
    117
    1
            self.token_buffer.clear();
    118
    1
        }
    119
    120
        /// Check if detector has any stop conditions configured
    121
    4
        pub fn has_conditions(&self) -> bool {
    122
    4
            !self.token_sequences.is_empty() || !self.string_patterns.is_empty()2
    123
    4
        }
    124
    }
    125
    126
    /// Repetition penalty configuration
    127
    ///
    128
    /// Penalizes tokens that have appeared in the context to reduce repetition.
    129
    /// Higher values = stronger penalty (1.0 = no penalty).
    130
    #[derive(Debug, Clone, Serialize, Deserialize)]
    131
    pub struct RepetitionPenaltyConfig {
    132
        /// Penalty multiplier for repeated tokens (1.0 = no penalty, >1.0 = penalty)
    133
        pub penalty: f32,
    134
        /// Number of recent tokens to consider (0 = all)
    135
        pub window_size: usize,
    136
    }
    137
    138
    impl Default for RepetitionPenaltyConfig {
    139
    1
        fn default() -> Self {
    140
    1
            Self {
    141
    1
                penalty: 1.0, // No penalty by default
    142
    1
                window_size: 64,
    143
    1
            }
    144
    1
        }
    145
    }
    146
    147
    impl RepetitionPenaltyConfig {
    148
        /// Create with specified penalty
    149
    9
        pub fn new(penalty: f32) -> Self {
    150
    9
            Self {
    151
    9
                penalty,
    152
    9
                window_size: 64,
    153
    9
            }
    154
    9
        }
    155
    156
        /// Set window size for context
    157
        #[must_use]
    158
    2
        pub fn with_window(mut self, window_size: usize) -> Self {
    159
    2
            self.window_size = window_size;
    160
    2
            self
    161
    2
        }
    162
    163
        /// Check if penalty is enabled
    164
    6
        pub fn is_enabled(&self) -> bool {
    165
    6
            (self.penalty - 1.0).abs() > 1e-6
    166
    6
        }
    167
    }
    168
    169
    /// Apply repetition penalty to logits
    170
    ///
    171
    /// Divides logits of tokens that appear in context by the penalty factor.
    172
    ///
    173
    /// # Arguments
    174
    ///
    175
    /// * `logits` - Raw logits from model
    176
    /// * `context_tokens` - List of previously generated token IDs
    177
    /// * `config` - Repetition penalty configuration
    178
    ///
    179
    /// # Returns
    180
    ///
    181
    /// Logits with repetition penalty applied
    182
    4
    pub fn apply_repetition_penalty(
    183
    4
        logits: &Tensor<f32>,
    184
    4
        context_tokens: &[usize],
    185
    4
        config: &RepetitionPenaltyConfig,
    186
    4
    ) -> Tensor<f32> {
    187
    4
        if !config.is_enabled() || context_tokens3.is_empty3() {
    188
    1
            return logits.clone();
    189
    3
        }
    190
    191
    3
        let data = logits.data();
    192
    3
        let mut penalized = data.to_vec();
    193
    3
        let vocab_size = data.len();
    194
    195
        // Get relevant context window
    196
    3
        let window_start = if config.window_size > 0 && context_tokens.len() > config.window_size {
    197
    1
            context_tokens.len() - config.window_size
    198
        } else {
    199
    2
            0
    200
        };
    201
    3
        let relevant_tokens = &context_tokens[window_start..];
    202
    203
        // Apply penalty to each token in context
    204
    11
        for &token_id8 in relevant_tokens {
    205
    8
            if token_id < vocab_size {
    206
    8
                let logit = penalized[token_id];
    207
                // For positive logits, divide by penalty
    208
                // For negative logits, multiply by penalty
    209
    8
                penalized[token_id] = if logit > 0.0 {
    210
    7
                    logit / config.penalty
    211
                } else {
    212
    1
                    logit * config.penalty
    213
                };
    214
    0
            }
    215
        }
    216
    217
    3
        Tensor::from_vec(logits.shape().to_vec(), penalized)
    218
    3
            .expect("Shape should match original logits")
    219
    4
    }
    220
    221
    /// Presence and frequency penalty configuration (OpenAI-style)
    222
    ///
    223
    /// - Presence penalty: Constant penalty for tokens that appear at least once
    224
    /// - Frequency penalty: Penalty proportional to token frequency
    225
    #[derive(Debug, Clone, Serialize, Deserialize)]
    226
    pub struct PresenceFrequencyPenalty {
    227
        /// Presence penalty (penalty if token appeared at all)
    228
        pub presence_penalty: f32,
    229
        /// Frequency penalty (penalty per occurrence)
    230
        pub frequency_penalty: f32,
    231
    }
    232
    233
    impl Default for PresenceFrequencyPenalty {
    234
    1
        fn default() -> Self {
    235
    1
            Self {
    236
    1
                presence_penalty: 0.0,
    237
    1
                frequency_penalty: 0.0,
    238
    1
            }
    239
    1
        }
    240
    }
    241
    242
    impl PresenceFrequencyPenalty {
    243
        /// Create new penalty config
    244
    8
        pub fn new(presence: f32, frequency: f32) -> Self {
    245
    8
            Self {
    246
    8
                presence_penalty: presence,
    247
    8
                frequency_penalty: frequency,
    248
    8
            }
    249
    8
        }
    250
    251
        /// Check if any penalty is enabled
    252
    6
        pub fn is_enabled(&self) -> bool {
    253
    6
            self.presence_penalty.abs() > 1e-6 || self.frequency_penalty.abs() > 1e-62
    254
    6
        }
    255
    }
    256
    257
    /// Apply presence and frequency penalties to logits
    258
    ///
    259
    /// Formula: logit -= presence_penalty * (1 if token in context else 0)
    260
    /// Formula: logit -= frequency_penalty * count(token in context)
    261
    ///
    262
    /// # Arguments
    263
    ///
    264
    /// * `logits` - Raw logits from model
    265
    /// * `context_tokens` - List of previously generated token IDs
    266
    /// * `config` - Presence/frequency penalty configuration
    267
    ///
    268
    /// # Returns
    269
    ///
    270
    /// Logits with penalties applied
    271
    4
    pub fn apply_presence_frequency_penalty(
    272
    4
        logits: &Tensor<f32>,
    273
    4
        context_tokens: &[usize],
    274
    4
        config: &PresenceFrequencyPenalty,
    275
    4
    ) -> Tensor<f32> {
    276
    4
        if !config.is_enabled() || context_tokens.is_empty() {
    277
    0
            return logits.clone();
    278
    4
        }
    279
    280
    4
        let data = logits.data();
    281
    4
        let mut penalized = data.to_vec();
    282
    4
        let vocab_size = data.len();
    283
    284
        // Count token frequencies
    285
    4
        let mut token_counts: HashMap<usize, usize> = HashMap::new();
    286
    17
        for &token_id13 in context_tokens {
    287
    13
            if token_id < vocab_size {
    288
    13
                *token_counts.entry(token_id).or_insert(0) += 1;
    289
    13
            }0
    290
        }
    291
    292
        // Apply penalties
    293
    12
        for (token_id8, count8) in token_counts {
    294
    8
            let presence = if count > 0 { 1.0 } else { 0.00 };
    295
    8
            penalized[token_id] -= config.presence_penalty * presence;
    296
    8
            penalized[token_id] -= config.frequency_penalty * (count as f32);
    297
        }
    298
    299
    4
        Tensor::from_vec(logits.shape().to_vec(), penalized)
    300
    4
            .expect("Shape should match original logits")
    301
    4
    }
    302
    303
    /// Logit bias configuration
    304
    ///
    305
    /// Allows adjusting specific token probabilities before sampling.
    306
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    307
    pub struct LogitBias {
    308
        /// Map of token ID to bias value (added to logit)
    309
        biases: HashMap<usize, f32>,
    310
    }
    311
    312
    impl LogitBias {
    313
        /// Create empty logit bias
    314
    5
        pub fn new() -> Self {
    315
    5
            Self {
    316
    5
                biases: HashMap::new(),
    317
    5
            }
    318
    5
        }
    319
    320
        /// Add bias for a specific token
    321
        #[must_use]
    322
    8
        pub fn with_bias(mut self, token_id: usize, bias: f32) -> Self {
    323
    8
            self.biases.insert(token_id, bias);
    324
    8
            self
    325
    8
        }
    326
    327
        /// Add multiple biases from a map
    328
        #[must_use]
    329
    0
        pub fn with_biases(mut self, biases: HashMap<usize, f32>) -> Self {
    330
    0
            self.biases.extend(biases);
    331
    0
            self
    332
    0
        }
    333
    334
        /// Check if any biases are configured
    335
    5
        pub fn is_empty(&self) -> bool {
    336
    5
            self.biases.is_empty()
    337
    5
        }
    338
    339
        /// Get bias for a token (0.0 if not set)
    340
    3
        pub fn get(&self, token_id: usize) -> f32 {
    341
    3
            self.biases.get(&token_id).copied().unwrap_or(0.0)
    342
    3
        }
    343
    }
    344
    345
    /// Apply logit bias to logits
    346
    ///
    347
    /// # Arguments
    348
    ///
    349
    /// * `logits` - Raw logits from model
    350
    /// * `bias` - Logit bias configuration
    351
    ///
    352
    /// # Returns
    353
    ///
    354
    /// Logits with biases applied
    355
    3
    pub fn apply_logit_bias(logits: &Tensor<f32>, bias: &LogitBias) -> Tensor<f32> {
    356
    3
        if bias.is_empty() {
    357
    0
            return logits.clone();
    358
    3
        }
    359
    360
    3
        let data = logits.data();
    361
    3
        let mut biased = data.to_vec();
    362
    3
        let vocab_size = data.len();
    363
    364
    8
        for (&token_id5, &bias_value5) in &bias.biases {
    365
    5
            if token_id < vocab_size {
    366
    4
                biased[token_id] += bias_value;
    367
    4
            }1
    368
        }
    369
    370
    3
        Tensor::from_vec(logits.shape().to_vec(), biased).expect("Shape should match original logits")
    371
    3
    }
    372
    373
    // ===== Prompt Caching =====
    374
    375
    /// Prompt cache entry
    376
    #[derive(Debug, Clone)]
    377
    pub struct PromptCacheEntry {
    378
        /// Token sequence
    379
        pub tokens: Vec<usize>,
    380
        /// Cached KV state (simplified - in practice would be actual KV tensors)
    381
        pub kv_hash: u64,
    382
        /// Number of times this entry has been hit
    383
        pub hit_count: usize,
    384
        /// Last access timestamp
    385
        pub last_access: std::time::Instant,
    386
    }
    387
    388
    /// Prompt cache for efficient prefix reuse
    389
    ///
    390
    /// Caches prompt prefixes to avoid recomputation when generating multiple
    391
    /// completions with the same prefix.
    392
    #[derive(Debug)]
    393
    pub struct PromptCache {
    394
        /// Cache entries keyed by token sequence hash
    395
        entries: std::collections::HashMap<u64, PromptCacheEntry>,
    396
        /// Maximum cache size
    397
        max_entries: usize,
    398
    }
    399
    400
    impl Default for PromptCache {
    401
    1
        fn default() -> Self {
    402
    1
            Self::new(100)
    403
    1
        }
    404
    }
    405
    406
    impl PromptCache {
    407
        /// Create new prompt cache
    408
    8
        pub fn new(max_entries: usize) -> Self {
    409
    8
            Self {
    410
    8
                entries: std::collections::HashMap::new(),
    411
    8
                max_entries,
    412
    8
            }
    413
    8
        }
    414
    415
        /// Compute hash for token sequence
    416
    18
        fn hash_tokens(tokens: &[usize]) -> u64 {
    417
            use std::hash::{Hash, Hasher};
    418
    18
            let mut hasher = std::collections::hash_map::DefaultHasher::new();
    419
    18
            tokens.hash(&mut hasher);
    420
    18
            hasher.finish()
    421
    18
        }
    422
    423
        /// Find longest matching prefix in cache
    424
    5
        pub fn find_prefix(&mut self, tokens: &[usize]) -> Option<(usize, u64)> {
    425
            // Try progressively shorter prefixes
    426
    8
            for len in (1..=tokens.len()5).rev5() {
    427
    8
                let prefix = &tokens[..len];
    428
    8
                let hash = Self::hash_tokens(prefix);
    429
    8
                if let Some(entry4) = self.entries.get_mut(&hash) {
    430
    4
                    entry.hit_count += 1;
    431
    4
                    entry.last_access = std::time::Instant::now();
    432
    4
                    return Some((len, entry.kv_hash));
    433
    4
                }
    434
            }
    435
    1
            None
    436
    5
        }
    437
    438
        /// Add entry to cache
    439
    10
        pub fn add(&mut self, tokens: Vec<usize>, kv_hash: u64) {
    440
            // Evict if at capacity
    441
    10
            if self.entries.len() >= self.max_entries {
    442
    1
                self.evict_lru();
    443
    9
            }
    444
    445
    10
            let hash = Self::hash_tokens(&tokens);
    446
    10
            self.entries.insert(
    447
    10
                hash,
    448
    10
                PromptCacheEntry {
    449
    10
                    tokens,
    450
    10
                    kv_hash,
    451
    10
                    hit_count: 0,
    452
    10
                    last_access: std::time::Instant::now(),
    453
    10
                },
    454
            );
    455
    10
        }
    456
    457
        /// Evict least recently used entry
    458
    1
        fn evict_lru(&mut self) {
    459
    1
            if let Some((&key, _)) = self.entries.iter().min_by_key(|(_, v)| v.last_access) {
    460
    1
                self.entries.remove(&key);
    461
    1
            }0
    462
    1
        }
    463
    464
        /// Get cache size
    465
    5
        pub fn len(&self) -> usize {
    466
    5
            self.entries.len()
    467
    5
        }
    468
    469
        /// Check if cache is empty
    470
    3
        pub fn is_empty(&self) -> bool {
    471
    3
            self.entries.is_empty()
    472
    3
        }
    473
    474
        /// Clear all entries
    475
    1
        pub fn clear(&mut self) {
    476
    1
            self.entries.clear();
    477
    1
        }
    478
    479
        /// Get cache statistics
    480
    1
        pub fn stats(&self) -> PromptCacheStats {
    481
    1
            let total_hits: usize = self.entries.values().map(|e| e.hit_count).sum();
    482
    1
            PromptCacheStats {
    483
    1
                entries: self.entries.len(),
    484
    1
                total_hits,
    485
    1
                max_entries: self.max_entries,
    486
    1
            }
    487
    1
        }
    488
    }
    489
    490
    /// Prompt cache statistics
    491
    #[derive(Debug, Clone)]
    492
    pub struct PromptCacheStats {
    493
        /// Number of entries in cache
    494
        pub entries: usize,
    495
        /// Total cache hits
    496
        pub total_hits: usize,
    497
        /// Maximum cache size
    498
        pub max_entries: usize,
    499
    }
    500
    501
    /// Beam search state for a single hypothesis
    502
    #[derive(Debug, Clone)]
    503
    pub struct BeamHypothesis {
    504
        /// Token sequence generated so far
    505
        pub tokens: Vec<usize>,
    506
        /// Cumulative log probability
    507
        pub score: f32,
    508
        /// Whether this hypothesis has finished (hit EOS)
    509
        pub finished: bool,
    510
    }
    511
    512
    impl BeamHypothesis {
    513
        /// Create a new hypothesis starting with given tokens
    514
    12
        pub fn new(tokens: Vec<usize>, score: f32) -> Self {
    515
    12
            Self {
    516
    12
                tokens,
    517
    12
                score,
    518
    12
                finished: false,
    519
    12
            }
    520
    12
        }
    521
    522
        /// Extend hypothesis with a new token
    523
        #[must_use]
    524
    6
        pub fn extend(&self, token: usize, log_prob: f32, is_eos: bool) -> Self {
    525
    6
            let mut new_tokens = self.tokens.clone();
    526
    6
            new_tokens.push(token);
    527
    6
            Self {
    528
    6
                tokens: new_tokens,
    529
    6
                score: self.score + log_prob,
    530
    6
                finished: is_eos,
    531
    6
            }
    532
    6
        }
    533
    534
        /// Get length-normalized score
    535
    8
        pub fn normalized_score(&self, length_penalty: f32) -> f32 {
    536
    8
            let len = self.tokens.len() as f32;
    537
    8
            self.score / len.powf(length_penalty)
    538
    8
        }
    539
    }
    540
    541
    /// Beam search configuration
    542
    #[derive(Debug, Clone)]
    543
    pub struct BeamSearchConfig {
    544
        /// Number of beams (hypotheses) to keep
    545
        pub num_beams: usize,
    546
        /// Length penalty (>1.0 favors longer sequences, <1.0 favors shorter)
    547
        pub length_penalty: f32,
    548
        /// Early stopping: stop when num_beams hypotheses are finished
    549
        pub early_stopping: bool,
    550
        /// Number of beams to return
    551
        pub num_return: usize,
    552
    }
    553
    554
    impl Default for BeamSearchConfig {
    555
    9
        fn default() -> Self {
    556
    9
            Self {
    557
    9
                num_beams: 4,
    558
    9
                length_penalty: 1.0,
    559
    9
                early_stopping: true,
    560
    9
                num_return: 1,
    561
    9
            }
    562
    9
        }
    563
    }
    564
    565
    impl BeamSearchConfig {
    566
        /// Create new beam search config
    567
    8
        pub fn new(num_beams: usize) -> Self {
    568
    8
            Self {
    569
    8
                num_beams,
    570
    8
                ..Default::default()
    571
    8
            }
    572
    8
        }
    573
    574
        /// Set length penalty
    575
        #[must_use]
    576
    2
        pub fn with_length_penalty(mut self, penalty: f32) -> Self {
    577
    2
            self.length_penalty = penalty;
    578
    2
            self
    579
    2
        }
    580
    581
        /// Set early stopping
    582
        #[must_use]
    583
    5
        pub fn with_early_stopping(mut self, early: bool) -> Self {
    584
    5
            self.early_stopping = early;
    585
    5
            self
    586
    5
        }
    587
    588
        /// Set number of sequences to return
    589
        #[must_use]
    590
    2
        pub fn with_num_return(mut self, n: usize) -> Self {
    591
    2
            self.num_return = n;
    592
    2
            self
    593
    2
        }
    594
    }
    595
    596
    /// Beam search state manager
    597
    #[derive(Debug, Clone)]
    598
    pub struct BeamSearchState {
    599
        /// Current hypotheses
    600
        pub hypotheses: Vec<BeamHypothesis>,
    601
        /// Finished hypotheses
    602
        pub finished: Vec<BeamHypothesis>,
    603
        /// Configuration
    604
        pub config: BeamSearchConfig,
    605
    }
    606
    607
    impl BeamSearchState {
    608
        /// Create new beam search state
    609
    6
        pub fn new(config: BeamSearchConfig, initial_tokens: Vec<usize>) -> Self {
    610
    6
            let hypotheses = vec![BeamHypothesis::new(initial_tokens, 0.0)];
    611
    6
            Self {
    612
    6
                hypotheses,
    613
    6
                finished: Vec::new(),
    614
    6
                config,
    615
    6
            }
    616
    6
        }
    617
    618
        /// Process a step with log probabilities for each hypothesis
    619
        ///
    620
        /// # Arguments
    621
        ///
    622
        /// * `log_probs_per_hyp` - Log probabilities for each token, for each hypothesis
    623
        /// * `eos_token` - Optional end-of-sequence token ID
    624
    1
        pub fn step(&mut self, log_probs_per_hyp: &[Vec<f32>], eos_token: Option<usize>) {
    625
    1
            let mut candidates: Vec<BeamHypothesis> = Vec::new();
    626
    627
    1
            for (hyp_idx, hyp) in self.hypotheses.iter().enumerate() {
    628
    1
                if hyp.finished {
    629
    0
                    candidates.push(hyp.clone());
    630
    0
                    continue;
    631
    1
                }
    632
    633
    1
                let log_probs = &log_probs_per_hyp[hyp_idx];
    634
    635
                // Get top-k tokens for this hypothesis (k = num_beams * 2 for safety)
    636
    1
                let mut indexed: Vec<(usize, f32)> = log_probs
    637
    1
                    .iter()
    638
    1
                    .enumerate()
    639
    5
                    .map1(|(i, &lp)| (i, lp))
    640
    1
                    .collect();
    641
    4
                indexed1.sort_by1(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    642
    643
    4
                for &(token, log_prob) in indexed.iter()1.take1(self.config.num_beams * 21) {
    644
    4
                    let is_eos = eos_token == Some(token);
    645
    4
                    let new_hyp = hyp.extend(token, log_prob, is_eos);
    646
    647
    4
                    if is_eos {
    648
    0
                        self.finished.push(new_hyp);
    649
    4
                    } else {
    650
    4
                        candidates.push(new_hyp);
    651
    4
                    }
    652
                }
    653
            }
    654
    655
            // Select top num_beams hypotheses by normalized score
    656
    3
            candidates1.sort_by1(|a, b| {
    657
    3
                let score_a = a.normalized_score(self.config.length_penalty);
    658
    3
                let score_b = b.normalized_score(self.config.length_penalty);
    659
    3
                score_b
    660
    3
                    .partial_cmp(&score_a)
    661
    3
                    .unwrap_or(std::cmp::Ordering::Equal)
    662
    3
            });
    663
    664
    1
            self.hypotheses = candidates.into_iter().take(self.config.num_beams).collect();
    665
    1
        }
    666
    667
        /// Check if search should stop
    668
    4
        pub fn should_stop(&self) -> bool {
    669
    4
            if self.config.early_stopping && self.finished2.len() >= self.config.num_beams {
    670
    1
                return true;
    671
    3
            }
    672
    3
            self.hypotheses.is_empty() || self.hypotheses.iter()2.all2(|h| h.finished)
    673
    4
        }
    674
    675
        /// Get best completed hypotheses
    676
    0
        pub fn best_hypotheses(&self) -> Vec<BeamHypothesis> {
    677
    0
            let mut all: Vec<_> = self
    678
    0
                .finished
    679
    0
                .iter()
    680
    0
                .chain(self.hypotheses.iter())
    681
    0
                .cloned()
    682
    0
                .collect();
    683
    0
            all.sort_by(|a, b| {
    684
    0
                let score_a = a.normalized_score(self.config.length_penalty);
    685
    0
                let score_b = b.normalized_score(self.config.length_penalty);
    686
    0
                score_b
    687
    0
                    .partial_cmp(&score_a)
    688
    0
                    .unwrap_or(std::cmp::Ordering::Equal)
    689
    0
            });
    690
    0
            all.into_iter().take(self.config.num_return).collect()
    691
    0
        }
    692
    }
    693
    694
    /// Streaming generation callback type
    695
    ///
    696
    /// The callback receives:
    697
    /// - token_id: The generated token ID
    698
    /// - token_text: Optional decoded text for the token
    699
    /// - is_final: Whether this is the last token
    700
    ///
    701
    /// Returns `true` to continue, `false` to stop generation
    702
    pub type StreamCallback = Box<dyn FnMut(usize, Option<&str>, bool) -> bool + Send>;
    703
    704
    /// Streaming generation state
    705
    #[derive(Debug)]
    706
    pub struct StreamingGenerator {
    707
        /// Tokens generated so far
    708
        pub tokens: Vec<usize>,
    709
        /// Generated text so far
    710
        pub text: String,
    711
        /// Whether generation is complete
    712
        pub finished: bool,
    713
        /// Total tokens generated
    714
        pub total_tokens: usize,
    715
    }
    716
    717
    impl StreamingGenerator {
    718
        /// Create new streaming generator
    719
    8
        pub fn new() -> Self {
    720
    8
            Self {
    721
    8
                tokens: Vec::new(),
    722
    8
                text: String::new(),
    723
    8
                finished: false,
    724
    8
                total_tokens: 0,
    725
    8
            }
    726
    8
        }
    727
    728
        /// Add a generated token
    729
    15
        pub fn add_token(&mut self, token_id: usize, token_text: Option<&str>) {
    730
    15
            self.tokens.push(token_id);
    731
    15
            if let Some(text9) = token_text {
    732
    9
                self.text.push_str(text);
    733
    9
            }6
    734
    15
            self.total_tokens += 1;
    735
    15
        }
    736
    737
        /// Mark generation as finished
    738
    1
        pub fn finish(&mut self) {
    739
    1
            self.finished = true;
    740
    1
        }
    741
    742
        /// Get current token count
    743
    4
        pub fn token_count(&self) -> usize {
    744
    4
            self.total_tokens
    745
    4
        }
    746
    }
    747
    748
    impl Default for StreamingGenerator {
    749
    1
        fn default() -> Self {
    750
    1
            Self::new()
    751
    1
        }
    752
    }
    753
    754
    /// Extended generation configuration with advanced sampling options
    755
    #[derive(Debug, Clone, Default)]
    756
    pub struct AdvancedGenerationConfig {
    757
        /// Base generation config
    758
        pub base: GenerationConfig,
    759
        /// Stop sequence detector
    760
        pub stop_detector: Option<StopSequenceDetector>,
    761
        /// Repetition penalty config
    762
        pub repetition_penalty: Option<RepetitionPenaltyConfig>,
    763
        /// Presence/frequency penalties
    764
        pub presence_frequency: Option<PresenceFrequencyPenalty>,
    765
        /// Logit bias
    766
        pub logit_bias: Option<LogitBias>,
    767
    }
    768
    769
    impl AdvancedGenerationConfig {
    770
        /// Create with base config
    771
    2
        pub fn new(base: GenerationConfig) -> Self {
    772
    2
            Self {
    773
    2
                base,
    774
    2
                ..Default::default()
    775
    2
            }
    776
    2
        }
    777
    778
        /// Add stop sequences
    779
        #[must_use]
    780
    1
        pub fn with_stop_sequences(mut self, stops: Vec<String>) -> Self {
    781
    1
            self.stop_detector = Some(StopSequenceDetector::new().with_stop_strings(stops));
    782
    1
            self
    783
    1
        }
    784
    785
        /// Add repetition penalty
    786
        #[must_use]
    787
    2
        pub fn with_repetition_penalty(mut self, penalty: f32) -> Self {
    788
    2
            self.repetition_penalty = Some(RepetitionPenaltyConfig::new(penalty));
    789
    2
            self
    790
    2
        }
    791
    792
        /// Add presence/frequency penalties
    793
        #[must_use]
    794
    2
        pub fn with_presence_frequency(mut self, presence: f32, frequency: f32) -> Self {
    795
    2
            self.presence_frequency = Some(PresenceFrequencyPenalty::new(presence, frequency));
    796
    2
            self
    797
    2
        }
    798
    799
        /// Add logit bias
    800
        #[must_use]
    801
    2
        pub fn with_logit_bias(mut self, bias: LogitBias) -> Self {
    802
    2
            self.logit_bias = Some(bias);
    803
    2
            self
    804
    2
        }
    805
    }
    806
    807
    /// Apply all configured penalties and biases to logits
    808
    ///
    809
    /// # Arguments
    810
    ///
    811
    /// * `logits` - Raw logits from model
    812
    /// * `context_tokens` - Previously generated tokens
    813
    /// * `config` - Advanced generation configuration
    814
    ///
    815
    /// # Returns
    816
    ///
    817
    /// Logits with all penalties applied
    818
    2
    pub fn apply_all_penalties(
    819
    2
        logits: &Tensor<f32>,
    820
    2
        context_tokens: &[usize],
    821
    2
        config: &AdvancedGenerationConfig,
    822
    2
    ) -> Tensor<f32> {
    823
    2
        let mut result = logits.clone();
    824
    825
        // Apply repetition penalty
    826
    2
        if let Some(ref rep_config1) = config.repetition_penalty {
    827
    1
            result = apply_repetition_penalty(&result, context_tokens, rep_config);
    828
    1
        }
    829
    830
        // Apply presence/frequency penalty
    831
    2
        if let Some(ref pf_config1) = config.presence_frequency {
    832
    1
            result = apply_presence_frequency_penalty(&result, context_tokens, pf_config);
    833
    1
        }
    834
    835
        // Apply logit bias
    836
    2
        if let Some(ref bias1) = config.logit_bias {
    837
    1
            result = apply_logit_bias(&result, bias);
    838
    1
        }
    839
    840
    2
        result
    841
    2
    }
    842
    843
    // ============================================================================
    844
    // Dynamic Temperature (temp_ext) - Entropy-based temperature adjustment
    845
    // ============================================================================
    846
    847
    /// Configuration for dynamic temperature (temp_ext)
    848
    ///
    849
    /// Adjusts temperature based on the entropy of the probability distribution.
    850
    /// When entropy is low (confident), uses higher temperature to increase diversity.
    851
    /// When entropy is high (uncertain), uses lower temperature to focus on likely tokens.
    852
    ///
    853
    /// Reference: llama.cpp `llama_sampler_init_temp_ext`
    854
    #[derive(Debug, Clone, Serialize, Deserialize)]
    855
    pub struct DynTempConfig {
    856
        /// Base temperature
    857
        pub temp: f32,
    858
        /// Range around base temperature (min = temp - delta, max = temp + delta)
    859
        pub delta: f32,
    860
        /// Exponent for entropy mapping (higher = more aggressive adjustment)
    861
        pub exponent: f32,
    862
    }
    863
    864
    impl Default for DynTempConfig {
    865
    1
        fn default() -> Self {
    866
    1
            Self {
    867
    1
                temp: 1.0,
    868
    1
                delta: 0.0,
    869
    1
                exponent: 1.0,
    870
    1
            }
    871
    1
        }
    872
    }
    873
    874
    impl DynTempConfig {
    875
        /// Create a new dynamic temperature config
    876
    7
        pub fn new(temp: f32, delta: f32, exponent: f32) -> Self {
    877
    7
            Self {
    878
    7
                temp,
    879
    7
                delta,
    880
    7
                exponent,
    881
    7
            }
    882
    7
        }
    883
    884
        /// Create with just temperature (no dynamic adjustment)
    885
    2
        pub fn static_temp(temp: f32) -> Self {
    886
    2
            Self {
    887
    2
                temp,
    888
    2
                delta: 0.0,
    889
    2
                exponent: 1.0,
    890
    2
            }
    891
    2
        }
    892
    }
    893
    894
    /// Apply dynamic temperature based on entropy
    895
    ///
    896
    /// The algorithm:
    897
    /// 1. Calculate max possible entropy: -log(1/n)
    898
    /// 2. Calculate actual entropy: -sum(p * log(p))
    899
    /// 3. Normalize entropy to [0, 1]
    900
    /// 4. Map to temperature: min_temp + (max_temp - min_temp) * pow(norm_entropy, exponent)
    901
    /// 5. Apply calculated temperature to logits
    902
    ///
    903
    /// # Arguments
    904
    ///
    905
    /// * `logits` - Raw logits from model
    906
    /// * `config` - Dynamic temperature configuration
    907
    ///
    908
    /// # Returns
    909
    ///
    910
    /// Logits with dynamic temperature applied
    911
    6
    pub fn apply_dynamic_temperature(logits: &Tensor<f32>, config: &DynTempConfig) -> Tensor<f32> {
    912
        // If no delta, just apply static temperature
    913
    6
        if config.delta <= 0.0 {
    914
    1
            return apply_temperature(logits, config.temp).unwrap_or_else(|_| logits0.clone0());
    915
    5
        }
    916
    917
    5
        let data = logits.data();
    918
    5
        if data.len() <= 1 {
    919
    1
            return logits.clone();
    920
    4
        }
    921
    922
        // Calculate softmax probabilities
    923
    4
        let max_logit = data.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    924
    20
        let exp_sum4: f324 = data4.iter4().map4(|x| (x - max_logit).exp()).sum4();
    925
    4
        let probs: Vec<f32> = data
    926
    4
            .iter()
    927
    20
            .map4(|x| (x - max_logit).exp() / exp_sum)
    928
    4
            .collect();
    929
    930
        // Calculate maximum possible entropy: -log(1/n) = log(n)
    931
    4
        let max_entropy = (data.len() as f32).ln();
    932
    933
        // Calculate actual entropy: -sum(p * log(p))
    934
    4
        let entropy: f32 = probs
    935
    4
            .iter()
    936
    20
            .filter4(|&&p| p > 0.0)
    937
    20
            .map4(|&p| -p * p.ln())
    938
    4
            .sum();
    939
    940
        // Normalize entropy to [0, 1]
    941
    4
        let normalized_entropy = if max_entropy > 0.0 {
    942
    4
            (entropy / max_entropy).clamp(0.0, 1.0)
    943
        } else {
    944
    0
            0.0
    945
        };
    946
    947
        // Calculate dynamic temperature
    948
    4
        let min_temp = (config.temp - config.delta).max(0.0);
    949
    4
        let max_temp = config.temp + config.delta;
    950
    4
        let dyn_temp = min_temp + (max_temp - min_temp) * normalized_entropy.powf(config.exponent);
    951
    952
        // Apply calculated temperature
    953
    4
        apply_temperature(logits, dyn_temp).unwrap_or_else(|_| logits0.clone0())
    954
    6
    }
    955
    956
    // ============================================================================
    957
    // Infill/FIM Sampler - Fill-in-the-Middle for code completion
    958
    // ============================================================================
    959
    960
    /// Configuration for infill/FIM (Fill-in-the-Middle) sampling
    961
    ///
    962
    /// Used for code completion where the model generates text to fill a gap.
    963
    /// Handles EOG (End-of-Generation) tokens specially to determine when to stop.
    964
    ///
    965
    /// Reference: llama.cpp `llama_sampler_init_infill`
    966
    #[derive(Debug, Clone, Serialize, Deserialize)]
    967
    pub struct InfillConfig {
    968
        /// EOG (End-of-Generation) token IDs
    969
        pub eog_tokens: Vec<usize>,
    970
        /// Ratio threshold: if 3*p_eog*n > p_txt, force EOG
    971
        pub eog_ratio_threshold: f32,
    972
    }
    973
    974
    impl Default for InfillConfig {
    975
    2
        fn default() -> Self {
    976
    2
            Self {
    977
    2
                eog_tokens: vec![],
    978
    2
                eog_ratio_threshold: 3.0,
    979
    2
            }
    980
    2
        }
    981
    }
    982
    983
    impl InfillConfig {
    984
        /// Create a new infill config with EOG tokens
    985
    7
        pub fn new(eog_tokens: Vec<usize>) -> Self {
    986
    7
            Self {
    987
    7
                eog_tokens,
    988
    7
                eog_ratio_threshold: 3.0,
    989
    7
            }
    990
    7
        }
    991
    992
        /// Set the EOG ratio threshold
    993
        #[must_use]
    994
    1
        pub fn with_threshold(mut self, threshold: f32) -> Self {
    995
    1
            self.eog_ratio_threshold = threshold;
    996
    1
            self
    997
    1
        }
    998
    }
    999
    1000
    /// Result of infill sampling
    1001
    #[derive(Debug, Clone)]
    1002
    pub struct InfillResult {
    1003
        /// Modified logits (with non-EOG tokens potentially zeroed)
    1004
        pub logits: Tensor<f32>,
    1005
        /// Whether to force EOG token
    1006
        pub force_eog: bool,
    1007
        /// Probability sum of text tokens
    1008
        pub p_txt: f32,
    1009
        /// Probability sum of EOG tokens
    1010
        pub p_eog: f32,
    1011
    }
    1012
    1013
    /// Apply infill sampling logic
    1014
    ///
    1015
    /// This determines if the model should stop generating (emit EOG) based on
    1016
    /// the relative probabilities of EOG vs text tokens.
    1017
    ///
    1018
    /// # Arguments
    1019
    ///
    1020
    /// * `logits` - Raw logits from model
    1021
    /// * `config` - Infill configuration
    1022
    ///
    1023
    /// # Returns
    1024
    ///
    1025
    /// `InfillResult` with modified logits and EOG decision
    1026
    5
    pub fn apply_infill_sampling(logits: &Tensor<f32>, config: &InfillConfig) -> InfillResult {
    1027
    5
        let data = logits.data();
    1028
    5
        if data.is_empty() || config.eog_tokens.is_empty() {
    1029
    1
            return InfillResult {
    1030
    1
                logits: logits.clone(),
    1031
    1
                force_eog: false,
    1032
    1
                p_txt: 1.0,
    1033
    1
                p_eog: 0.0,
    1034
    1
            };
    1035
    4
        }
    1036
    1037
        // Calculate softmax probabilities
    1038
    4
        let max_logit = data.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    1039
    20
        let exp_sum4: f324 = data4.iter4().map4(|x| (x - max_logit).exp()).sum4();
    1040
    4
        let probs: Vec<f32> = data
    1041
    4
            .iter()
    1042
    20
            .map4(|x| (x - max_logit).exp() / exp_sum)
    1043
    4
            .collect();
    1044
    1045
        // Calculate p_eog and p_txt
    1046
    4
        let mut p_eog: f32 = 0.0;
    1047
    4
        let mut p_txt: f32 = 0.0;
    1048
    1049
    20
        for (i, &p) in probs.iter()4.enumerate4() {
    1050
    20
            if config.eog_tokens.contains(&i) {
    1051
    5
                p_eog += p;
    1052
    15
            } else {
    1053
    15
                p_txt += p;
    1054
    15
            }
    1055
        }
    1056
    1057
        // Check if we should force EOG
    1058
        // Condition: 3 * p_eog * n > p_txt
    1059
    4
        let n = data.len() as f32;
    1060
    4
        let force_eog = config.eog_ratio_threshold * p_eog * n > p_txt;
    1061
    1062
    4
        if force_eog {
    1063
            // Keep only EOG tokens
    1064
    3
            let mut new_data = vec![f32::NEG_INFINITY; data.len()];
    1065
    3
            let mut eog_sum = 0.0;
    1066
    1067
    7
            for &eog_id4 in &config.eog_tokens {
    1068
    4
                if eog_id < data.len() {
    1069
    4
                    new_data[eog_id] = data[eog_id];
    1070
    4
                    eog_sum += probs[eog_id];
    1071
    4
                }0
    1072
            }
    1073
    1074
            // Renormalize EOG tokens
    1075
    3
            if eog_sum > 0.0 {
    1076
    7
                for &eog_id4 in &config.eog_tokens {
    1077
    4
                    if eog_id < data.len() && new_data[eog_id] > f32::NEG_INFINITY {
    1078
    4
                        // Convert back to logit scale
    1079
    4
                        let normalized_p = probs[eog_id] / eog_sum;
    1080
    4
                        new_data[eog_id] = normalized_p.ln();
    1081
    4
                    }0
    1082
                }
    1083
    0
            }
    1084
    1085
            InfillResult {
    1086
    3
                logits: Tensor::from_vec(logits.shape().to_vec(), new_data)
    1087
    3
                    .unwrap_or_else(|_| logits0.clone0()),
    1088
                force_eog: true,
    1089
    3
                p_txt,
    1090
    3
                p_eog,
    1091
            }
    1092
        } else {
    1093
    1
            InfillResult {
    1094
    1
                logits: logits.clone(),
    1095
    1
                force_eog: false,
    1096
    1
                p_txt,
    1097
    1
                p_eog,
    1098
    1
            }
    1099
        }
    1100
    5
    }
    1101
    1102
    // ============================================================================
    1103
    // Sampler Chain - Composable sampler pipeline
    1104
    // ============================================================================
    1105
    1106
    /// Trait for samplers that can be chained together
    1107
    pub trait Sampler: Send + Sync {
    1108
        /// Get the sampler name
    1109
        fn name(&self) -> &'static str;
    1110
    1111
        /// Apply the sampler to logits (in-place modification)
    1112
        fn apply(&self, logits: &mut Tensor<f32>, context: &SamplerContext);
    1113
    1114
        /// Clone the sampler (for use in chains)
    1115
        fn clone_box(&self) -> Box<dyn Sampler>;
    1116
    }
    1117
    1118
    /// Context passed to samplers during application
    1119
    #[derive(Debug, Clone, Default)]
    1120
    pub struct SamplerContext {
    1121
        /// Previously generated tokens
    1122
        pub tokens: Vec<usize>,
    1123
        /// Random value for stochastic samplers [0, 1)
    1124
        pub rng_value: f32,
    1125
        /// Current generation step
    1126
        pub step: usize,
    1127
    }
    1128
    1129
    impl SamplerContext {
    1130
        /// Create a new sampler context
    1131
    6
        pub fn new() -> Self {
    1132
    6
            Self::default()
    1133
    6
        }
    1134
    1135
        /// Set tokens
    1136
        #[must_use]
    1137
    1
        pub fn with_tokens(mut self, tokens: Vec<usize>) -> Self {
    1138
    1
            self.tokens = tokens;
    1139
    1
            self
    1140
    1
        }
    1141
    1142
        /// Set RNG value
    1143
        #[must_use]
    1144
    1
        pub fn with_rng(mut self, rng_value: f32) -> Self {
    1145
    1
            self.rng_value = rng_value;
    1146
    1
            self
    1147
    1
        }
    1148
    1149
        /// Set step
    1150
        #[must_use]
    1151
    1
        pub fn with_step(mut self, step: usize) -> Self {
    1152
    1
            self.step = step;
    1153
    1
            self
    1154
    1
        }
    1155
    }
    1156
    1157
    /// A chain of samplers applied in sequence
    1158
    pub struct SamplerChain {
    1159
        samplers: Vec<Box<dyn Sampler>>,
    1160
    }
    1161
    1162
    impl Default for SamplerChain {
    1163
    1
        fn default() -> Self {
    1164
    1
            Self::new()
    1165
    1
        }
    1166
    }
    1167
    1168
    impl SamplerChain {
    1169
        /// Create a new empty sampler chain
    1170
    9
        pub fn new() -> Self {
    1171
    9
            Self { samplers: vec![] }
    1172
    9
        }
    1173
    1174
        /// Add a sampler to the chain (builder pattern)
    1175
        #[must_use]
    1176
    11
        pub fn with_sampler<S: Sampler + 'static>(mut self, sampler: S) -> Self {
    1177
    11
            self.samplers.push(Box::new(sampler));
    1178
    11
            self
    1179
    11
        }
    1180
    1181
        /// Push a boxed sampler to the chain
    1182
    1
        pub fn push(&mut self, sampler: Box<dyn Sampler>) {
    1183
    1
            self.samplers.push(sampler);
    1184
    1
        }
    1185
    1186
        /// Get the number of samplers in the chain
    1187
    5
        pub fn len(&self) -> usize {
    1188
    5
            self.samplers.len()
    1189
    5
        }
    1190
    1191
        /// Check if the chain is empty
    1192
    2
        pub fn is_empty(&self) -> bool {
    1193
    2
            self.samplers.is_empty()
    1194
    2
        }
    1195
    1196
        /// Get sampler names in order
    1197
    3
        pub fn names(&self) -> Vec<&'static str> {
    1198
    6
            self.samplers.iter()3.map3(|s| s.name()).collect3()
    1199
    3
        }
    1200
    1201
        /// Apply all samplers in sequence
    1202
    3
        pub fn apply(&self, logits: &mut Tensor<f32>, context: &SamplerContext) {
    1203
    8
            for sampler5 in &self.samplers {
    1204
    5
                sampler.apply(logits, context);
    1205
    5
            }
    1206
    3
        }
    1207
    1208
        /// Sample a token after applying all samplers
    1209
        ///
    1210
        /// # Errors
    1211
        ///
    1212
        /// Returns error if sampling fails
    1213
    2
        pub fn sample(&self, logits: &Tensor<f32>, context: &SamplerContext) -> Result<usize> {
    1214
    2
            let mut modified = logits.clone();
    1215
    2
            self.apply(&mut modified, context);
    1216
    2
            sample_greedy(&modified)
    1217
    2
        }
    1218
    }
    1219
    1220
    impl Clone for SamplerChain {
    1221
    1
        fn clone(&self) -> Self {
    1222
            Self {
    1223
    2
                samplers: self.samplers.iter()1.map1(|s| s.clone_box()).collect1(),
    1224
            }
    1225
    1
        }
    1226
    }
    1227
    1228
    // Concrete sampler implementations for the chain
    1229
    1230
    /// Temperature sampler
    1231
    #[derive(Debug, Clone)]
    1232
    pub struct TemperatureSampler {
    1233
        /// Temperature value (1.0 = no change)
    1234
        pub temp: f32,
    1235
    }
    1236
    1237
    impl TemperatureSampler {
    1238
        /// Create a new temperature sampler
    1239
    8
        pub fn new(temp: f32) -> Self {
    1240
    8
            Self { temp }
    1241
    8
        }
    1242
    }
    1243
    1244
    impl Sampler for TemperatureSampler {
    1245
    4
        fn name(&self) -> &'static str {
    1246
    4
            "temperature"
    1247
    4
        }
    1248
    1249
    3
        fn apply(&self, logits: &mut Tensor<f32>, _context: &SamplerContext) {
    1250
    3
            if let Ok(result) = apply_temperature(logits, self.temp) {
    1251
    3
                *logits = result;
    1252
    3
            }0
    1253
    3
        }
    1254
    1255
    1
        fn clone_box(&self) -> Box<dyn Sampler> {
    1256
    1
            Box::new(self.clone())
    1257
    1
        }
    1258
    }
    1259
    1260
    /// Dynamic temperature sampler
    1261
    #[derive(Debug, Clone)]
    1262
    pub struct DynTempSampler {
    1263
        /// Dynamic temperature configuration
    1264
        pub config: DynTempConfig,
    1265
    }
    1266
    1267
    impl DynTempSampler {
    1268
        /// Create a new dynamic temperature sampler
    1269
    1
        pub fn new(config: DynTempConfig) -> Self {
    1270
    1
            Self { config }
    1271
    1
        }
    1272
    }
    1273
    1274
    impl Sampler for DynTempSampler {
    1275
    1
        fn name(&self) -> &'static str {
    1276
    1
            "dyn_temp"
    1277
    1
        }
    1278
    1279
    0
        fn apply(&self, logits: &mut Tensor<f32>, _context: &SamplerContext) {
    1280
    0
            *logits = apply_dynamic_temperature(logits, &self.config);
    1281
    0
        }
    1282
    1283
    0
        fn clone_box(&self) -> Box<dyn Sampler> {
    1284
    0
            Box::new(self.clone())
    1285
    0
        }
    1286
    }
    1287
    1288
    /// Top-K sampler
    1289
    #[derive(Debug, Clone)]
    1290
    pub struct TopKSampler {
    1291
        /// Number of top tokens to consider
    1292
        pub k: usize,
    1293
    }
    1294
    1295
    impl TopKSampler {
    1296
        /// Create a new top-k sampler
    1297
    5
        pub fn new(k: usize) -> Self {
    1298
    5
            Self { k }
    1299
    5
        }
    1300
    }
    1301
    1302
    impl Sampler for TopKSampler {
    1303
    3
        fn name(&self) -> &'static str {
    1304
    3
            "top_k"
    1305
    3
        }
    1306
    1307
    2
        fn apply(&self, logits: &mut Tensor<f32>, _context: &SamplerContext) {
    1308
            // Apply top-k by zeroing out tokens outside top-k
    1309
    2
            let data = logits.data();
    1310
    2
            let mut indexed: Vec<(usize, f32)> = data.iter().copied().enumerate().collect();
    1311
    54
            indexed2.sort_by2(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    1312
    1313
    2
            let mut new_data = vec![f32::NEG_INFINITY; data.len()];
    1314
    12
            for (idx, logit) in indexed.iter()2.take2(self.k2) {
    1315
    12
                new_data[*idx] = *logit;
    1316
    12
            }
    1317
    1318
    2
            if let Ok(result) = Tensor::from_vec(logits.shape().to_vec(), new_data) {
    1319
    2
                *logits = result;
    1320
    2
            }0
    1321
    2
        }
    1322
    1323
    1
        fn clone_box(&self) -> Box<dyn Sampler> {
    1324
    1
            Box::new(self.clone())
    1325
    1
        }
    1326
    }
    1327
    1328
    /// Top-P (nucleus) sampler
    1329
    #[derive(Debug, Clone)]
    1330
    pub struct TopPSampler {
    1331
        /// Cumulative probability threshold (0.0 to 1.0)
    1332
        pub p: f32,
    1333
    }
    1334
    1335
    impl TopPSampler {
    1336
        /// Create a new top-p sampler
    1337
    4
        pub fn new(p: f32) -> Self {
    1338
    4
            Self { p }
    1339
    4
        }
    1340
    }
    1341
    1342
    impl Sampler for TopPSampler {
    1343
    2
        fn name(&self) -> &'static str {
    1344
    2
            "top_p"
    1345
    2
        }
    1346
    1347
    2
        fn apply(&self, logits: &mut Tensor<f32>, _context: &SamplerContext) {
    1348
    2
            let data = logits.data();
    1349
    1350
            // Calculate softmax
    1351
    2
            let max_logit = data.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    1352
    15
            let exp_sum2: f322 = data2.iter2().map2(|x| (x - max_logit).exp()).sum2();
    1353
    2
            let mut indexed: Vec<(usize, f32, f32)> = data
    1354
    2
                .iter()
    1355
    2
                .enumerate()
    1356
    15
                .map2(|(i, &logit)| (i, logit, (logit - max_logit).exp() / exp_sum))
    1357
    2
                .collect();
    1358
    1359
            // Sort by probability descending
    1360
    50
            indexed2.sort_by2(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
    1361
    1362
            // Find cutoff
    1363
    2
            let mut cumsum = 0.0;
    1364
    2
            let mut cutoff_idx = indexed.len();
    1365
    4
            for (i, (_, _, prob)) in indexed.iter()2.enumerate2() {
    1366
    4
                cumsum += prob;
    1367
    4
                if cumsum >= self.p {
    1368
    2
                    cutoff_idx = i + 1;
    1369
    2
                    break;
    1370
    2
                }
    1371
            }
    1372
    1373
            // Zero out tokens below cutoff
    1374
    2
            let mut new_data = vec![f32::NEG_INFINITY; data.len()];
    1375
    4
            for (idx, logit, _) in indexed.iter()2.take2(cutoff_idx2) {
    1376
    4
                new_data[*idx] = *logit;
    1377
    4
            }
    1378
    1379
    2
            if let Ok(result) = Tensor::from_vec(logits.shape().to_vec(), new_data) {
    1380
    2
                *logits = result;
    1381
    2
            }0
    1382
    2
        }
    1383
    1384
    0
        fn clone_box(&self) -> Box<dyn Sampler> {
    1385
    0
            Box::new(self.clone())
    1386
    0
        }
    1387
    }
    1388
    1389
    /// Repetition penalty sampler
    1390
    #[derive(Debug, Clone)]
    1391
    pub struct RepetitionPenaltySampler {
    1392
        /// Repetition penalty configuration
    1393
        pub config: RepetitionPenaltyConfig,
    1394
    }
    1395
    1396
    impl RepetitionPenaltySampler {
    1397
        /// Create a new repetition penalty sampler
    1398
    1
        pub fn new(config: RepetitionPenaltyConfig) -> Self {
    1399
    1
            Self { config }
    1400
    1
        }
    1401
    }
    1402
    1403
    impl Sampler for RepetitionPenaltySampler {
    1404
    1
        fn name(&self) -> &'static str {
    1405
    1
            "repetition_penalty"
    1406
    1
        }
    1407
    1408
    0
        fn apply(&self, logits: &mut Tensor<f32>, context: &SamplerContext) {
    1409
    0
            *logits = apply_repetition_penalty(logits, &context.tokens, &self.config);
    1410
    0
        }
    1411
    1412
    0
        fn clone_box(&self) -> Box<dyn Sampler> {
    1413
    0
            Box::new(self.clone())
    1414
    0
        }
    1415
    }
    1416
    1417
    /// Infill sampler
    1418
    #[derive(Debug, Clone)]
    1419
    pub struct InfillSampler {
    1420
        /// Infill/FIM configuration
    1421
        pub config: InfillConfig,
    1422
    }
    1423
    1424
    impl InfillSampler {
    1425
        /// Create a new infill sampler
    1426
    1
        pub fn new(config: InfillConfig) -> Self {
    1427
    1
            Self { config }
    1428
    1
        }
    1429
    }
    1430
    1431
    impl Sampler for InfillSampler {
    1432
    1
        fn name(&self) -> &'static str {
    1433
    1
            "infill"
    1434
    1
        }
    1435
    1436
    0
        fn apply(&self, logits: &mut Tensor<f32>, _context: &SamplerContext) {
    1437
    0
            let result = apply_infill_sampling(logits, &self.config);
    1438
    0
            *logits = result.logits;
    1439
    0
        }
    1440
    1441
    0
        fn clone_box(&self) -> Box<dyn Sampler> {
    1442
    0
            Box::new(self.clone())
    1443
    0
        }
    1444
    }
    1445
    1446
    // =============================================================================
    1447
    // LogitProcessor Trait (RLZR-GEN-001)
    1448
    // =============================================================================
    1449
    //
    1450
    // Composable logit processing for text generation pipelines.
    1451
    // Based on HuggingFace Transformers LogitsProcessor pattern.
    1452
    //
    1453
    // References:
    1454
    // - Holtzman et al. (2020) "The Curious Case of Neural Text Degeneration"
    1455
    // - Wolf et al. (2020) "Transformers: State-of-the-Art NLP"
    1456
    // =============================================================================
    1457
    1458
    /// Context available during logit processing
    1459
    ///
    1460
    /// Provides information about the current generation state to processors.
    1461
    #[derive(Debug, Clone)]
    1462
    pub struct LogitProcessorContext<'a> {
    1463
        /// Previously generated tokens (including initial prompt)
    1464
        pub tokens: &'a [u32],
    1465
        /// Current generation step (0-indexed, after initial tokens)
    1466
        pub step: usize,
    1467
        /// Vocabulary size
    1468
        pub n_vocab: usize,
    1469
    }
    1470
    1471
    impl<'a> LogitProcessorContext<'a> {
    1472
        /// Create a new context
    1473
        #[must_use]
    1474
    21
        pub fn new(tokens: &'a [u32], step: usize, n_vocab: usize) -> Self {
    1475
    21
            Self {
    1476
    21
                tokens,
    1477
    21
                step,
    1478
    21
                n_vocab,
    1479
    21
            }
    1480
    21
        }
    1481
    }
    1482
    1483
    /// Logit processor trait for composable pre-sampling transforms
    1484
    ///
    1485
    /// Processors are applied in order before sampling. They can:
    1486
    /// - Set logits to -inf to suppress tokens
    1487
    /// - Add penalties (repetition, length)
    1488
    /// - Scale logits (temperature)
    1489
    ///
    1490
    /// # Example
    1491
    ///
    1492
    /// ```rust,ignore
    1493
    /// use realizar::generate::{LogitProcessor, LogitProcessorContext};
    1494
    ///
    1495
    /// struct MyProcessor;
    1496
    ///
    1497
    /// impl LogitProcessor for MyProcessor {
    1498
    ///     fn process(&self, logits: &mut [f32], ctx: &LogitProcessorContext) {
    1499
    ///         // Suppress token 0
    1500
    ///         logits[0] = f32::NEG_INFINITY;
    1501
    ///     }
    1502
    /// }
    1503
    /// ```
    1504
    pub trait LogitProcessor: Send + Sync {
    1505
        /// Process logits in-place before sampling
    1506
        ///
    1507
        /// # Arguments
    1508
        ///
    1509
        /// * `logits` - Mutable slice of logits to modify
    1510
        /// * `ctx` - Context with token history and generation state
    1511
        fn process(&self, logits: &mut [f32], ctx: &LogitProcessorContext);
    1512
    1513
        /// Human-readable name for debugging and tracing
    1514
    0
        fn name(&self) -> &'static str {
    1515
    0
            "unnamed"
    1516
    0
        }
    1517
    }
    1518
    1519
    /// Suppress specific tokens by setting their logits to -inf
    1520
    ///
    1521
    /// Use this to prevent certain tokens from being generated, such as:
    1522
    /// - Special tokens (SOT, PREV, SOLM in Whisper)
    1523
    /// - Profanity or sensitive content
    1524
    /// - Invalid tokens for the current context
    1525
    #[derive(Debug, Clone)]
    1526
    pub struct TokenSuppressor {
    1527
        /// Token IDs to suppress
    1528
        suppress_ids: Vec<u32>,
    1529
    }
    1530
    1531
    impl TokenSuppressor {
    1532
        /// Create a new token suppressor
    1533
        ///
    1534
        /// # Arguments
    1535
        ///
    1536
        /// * `suppress_ids` - Token IDs to suppress (set to -inf)
    1537
        #[must_use]
    1538
    9
        pub fn new(suppress_ids: Vec<u32>) -> Self {
    1539
    9
            Self { suppress_ids }
    1540
    9
        }
    1541
    1542
        /// Create from a slice of token IDs
    1543
        #[must_use]
    1544
    0
        pub fn from_slice(suppress_ids: &[u32]) -> Self {
    1545
    0
            Self {
    1546
    0
                suppress_ids: suppress_ids.to_vec(),
    1547
    0
            }
    1548
    0
        }
    1549
    }
    1550
    1551
    impl LogitProcessor for TokenSuppressor {
    1552
    9
        fn process(&self, logits: &mut [f32], _ctx: &LogitProcessorContext) {
    1553
    21
            for &token_id12 in &self.suppress_ids {
    1554
    12
                if (token_id as usize) < logits.len() {
    1555
    10
                    logits[token_id as usize] = f32::NEG_INFINITY;
    1556
    10
                }2
    1557
            }
    1558
    9
        }
    1559
    1560
    2
        fn name(&self) -> &'static str {
    1561
    2
            "token_suppressor"
    1562
    2
        }
    1563
    }
    1564
    1565
    /// Penalize repeated tokens to reduce repetitive generation
    1566
    ///
    1567
    /// Applies a penalty to tokens that have appeared in the recent context.
    1568
    /// Penalty > 1.0 reduces probability, < 1.0 increases it.
    1569
    ///
    1570
    /// Based on: Keskar et al. (2019) "CTRL: A Conditional Transformer Language Model"
    1571
    #[derive(Debug, Clone)]
    1572
    pub struct RepetitionPenalty {
    1573
        /// Penalty multiplier (> 1.0 to penalize, < 1.0 to encourage)
    1574
        penalty: f32,
    1575
        /// Look-back window size (0 = entire history)
    1576
        window: usize,
    1577
    }
    1578
    1579
    impl RepetitionPenalty {
    1580
        /// Create a new repetition penalty processor
    1581
        ///
    1582
        /// # Arguments
    1583
        ///
    1584
        /// * `penalty` - Penalty multiplier (typical: 1.0-2.0)
    1585
        /// * `window` - Look-back window (0 = use all tokens)
    1586
        #[must_use]
    1587
    1
        pub fn new(penalty: f32, window: usize) -> Self {
    1588
    1
            Self { penalty, window }
    1589
    1
        }
    1590
    1591
        /// Create with default window (entire history)
    1592
        #[must_use]
    1593
    4
        pub fn with_penalty(penalty: f32) -> Self {
    1594
    4
            Self { penalty, window: 0 }
    1595
    4
        }
    1596
    }
    1597
    1598
    impl LogitProcessor for RepetitionPenalty {
    1599
    3
        fn process(&self, logits: &mut [f32], ctx: &LogitProcessorContext) {
    1600
            // Determine which tokens to consider
    1601
    3
            let tokens = if self.window > 0 && ctx.tokens1.len() > self.window {
    1602
    1
                &ctx.tokens[ctx.tokens.len() - self.window..]
    1603
            } else {
    1604
    2
                ctx.tokens
    1605
            };
    1606
    1607
            // Apply penalty to tokens that have appeared
    1608
    9
            for &token_id6 in tokens {
    1609
    6
                if (token_id as usize) < logits.len() {
    1610
    6
                    let logit = logits[token_id as usize];
    1611
                    // Apply penalty: divide positive logits, multiply negative logits
    1612
    6
                    logits[token_id as usize] = if logit > 0.0 {
    1613
    5
                        logit / self.penalty
    1614
                    } else {
    1615
    1
                        logit * self.penalty
    1616
                    };
    1617
    0
                }
    1618
            }
    1619
    3
        }
    1620
    1621
    1
        fn name(&self) -> &'static str {
    1622
    1
            "repetition_penalty"
    1623
    1
        }
    1624
    }
    1625
    1626
    /// Scale logits by temperature
    1627
    ///
    1628
    /// Temperature > 1.0 increases randomness (flatter distribution)
    1629
    /// Temperature < 1.0 decreases randomness (sharper distribution)
    1630
    /// Temperature = 1.0 has no effect
    1631
    #[derive(Debug, Clone)]
    1632
    pub struct TemperatureScaler {
    1633
        /// Temperature value (must be > 0)
    1634
        temperature: f32,
    1635
    }
    1636
    1637
    impl TemperatureScaler {
    1638
        /// Create a new temperature scaler
    1639
        ///
    1640
        /// # Arguments
    1641
        ///
    1642
        /// * `temperature` - Temperature value (> 0)
    1643
        ///
    1644
        /// # Panics
    1645
        ///
    1646
        /// Panics if temperature <= 0
    1647
        #[must_use]
    1648
    5
        pub fn new(temperature: f32) -> Self {
    1649
    5
            assert!(temperature > 0.0, "Temperature must be positive"1);
    1650
    4
            Self { temperature }
    1651
    4
        }
    1652
    }
    1653
    1654
    impl LogitProcessor for TemperatureScaler {
    1655
    3
        fn process(&self, logits: &mut [f32], _ctx: &LogitProcessorContext) {
    1656
    3
            if (self.temperature - 1.0).abs() > 1e-6 {
    1657
    6
                for logit in logits2.iter_mut2() {
    1658
    6
                    *logit /= self.temperature;
    1659
    6
                }
    1660
    1
            }
    1661
    3
        }
    1662
    1663
    1
        fn name(&self) -> &'static str {
    1664
    1
            "temperature_scaler"
    1665
    1
        }
    1666
    }
    1667
    1668
    /// Chain of logit processors applied in order
    1669
    ///
    1670
    /// Allows composing multiple processors into a single processing step.
    1671
    #[derive(Default)]
    1672
    pub struct LogitProcessorChain {
    1673
        processors: Vec<Box<dyn LogitProcessor>>,
    1674
    }
    1675
    1676
    impl LogitProcessorChain {
    1677
        /// Create an empty processor chain
    1678
        #[must_use]
    1679
    9
        pub fn new() -> Self {
    1680
    9
            Self {
    1681
    9
                processors: Vec::new(),
    1682
    9
            }
    1683
    9
        }
    1684
    1685
        /// Add a processor to the chain (builder pattern)
    1686
        #[must_use]
    1687
    10
        pub fn with_processor<P: LogitProcessor + 'static>(mut self, processor: P) -> Self {
    1688
    10
            self.processors.push(Box::new(processor));
    1689
    10
            self
    1690
    10
        }
    1691
    1692
        /// Add a boxed processor to the chain (builder pattern)
    1693
        #[must_use]
    1694
    0
        pub fn with_boxed_processor(mut self, processor: Box<dyn LogitProcessor>) -> Self {
    1695
    0
            self.processors.push(processor);
    1696
    0
            self
    1697
    0
        }
    1698
    1699
        /// Process logits through all processors in order
    1700
    13
        pub fn process(&self, logits: &mut [f32], ctx: &LogitProcessorContext) {
    1701
    21
            for processor8 in &self.processors {
    1702
    8
                processor.process(logits, ctx);
    1703
    8
            }
    1704
    13
        }
    1705
    1706
        /// Get the number of processors in the chain
    1707
        #[must_use]
    1708
    2
        pub fn len(&self) -> usize {
    1709
    2
            self.processors.len()
    1710
    2
        }
    1711
    1712
        /// Check if the chain is empty
    1713
        #[must_use]
    1714
    2
        pub fn is_empty(&self) -> bool {
    1715
    2
            self.processors.is_empty()
    1716
    2
        }
    1717
    1718
        /// Get processor names for debugging
    1719
        #[must_use]
    1720
    1
        pub fn processor_names(&self) -> Vec<&str> {
    1721
    3
            self.processors.iter()1.map1(|p| p.name()).collect1()
    1722
    1
        }
    1723
    }
    1724
    1725
    impl LogitProcessor for LogitProcessorChain {
    1726
    1
        fn process(&self, logits: &mut [f32], ctx: &LogitProcessorContext) {
    1727
    1
            LogitProcessorChain::process(self, logits, ctx);
    1728
    1
        }
    1729
    1730
    1
        fn name(&self) -> &'static str {
    1731
    1
            "processor_chain"
    1732
    1
        }
    1733
    }
    1734
    1735
    /// Model trait for generation pipeline
    1736
    ///
    1737
    /// Implement this trait to use your model with GenerationPipeline.
    1738
    pub trait GenerativeModel {
    1739
        /// Forward pass producing logits for next token
    1740
        ///
    1741
        /// # Arguments
    1742
        ///
    1743
        /// * `tokens` - Current token sequence
    1744
        ///
    1745
        /// # Returns
    1746
        ///
    1747
        /// Logits for vocabulary (shape: [vocab_size])
    1748
        fn forward(&mut self, tokens: &[u32]) -> Result<Vec<f32>>;
    1749
    1750
        /// Get vocabulary size
    1751
        fn vocab_size(&self) -> usize;
    1752
    1753
        /// Reset any cached state (e.g., KV cache)
    1754
    0
        fn reset(&mut self) {}
    1755
    }
    1756
    1757
    /// Generation pipeline with processor chain
    1758
    ///
    1759
    /// Orchestrates the generation loop with:
    1760
    /// 1. Model forward pass
    1761
    /// 2. Logit processing
    1762
    /// 3. Token sampling
    1763
    /// 4. EOS detection
    1764
    ///
    1765
    /// # Example
    1766
    ///
    1767
    /// ```rust,ignore
    1768
    /// use realizar::generate::{GenerationPipeline, TokenSuppressor, GenerationConfig};
    1769
    ///
    1770
    /// let pipeline = GenerationPipeline::new(model)
    1771
    ///     .add_processor(TokenSuppressor::new(vec![0, 1, 2]))
    1772
    ///     .with_config(GenerationConfig::greedy().with_eos_token_id(50256));
    1773
    ///
    1774
    /// let tokens = pipeline.generate(&[1, 2, 3])?;
    1775
    /// ```
    1776
    pub struct GenerationPipeline<M: GenerativeModel> {
    1777
        model: M,
    1778
        processors: LogitProcessorChain,
    1779
        config: GenerationConfig,
    1780
    }
    1781
    1782
    impl<M: GenerativeModel> GenerationPipeline<M> {
    1783
        /// Create a new generation pipeline
    1784
        #[must_use]
    1785
    4
        pub fn new(model: M) -> Self {
    1786
    4
            Self {
    1787
    4
                model,
    1788
    4
                processors: LogitProcessorChain::new(),
    1789
    4
                config: GenerationConfig::default(),
    1790
    4
            }
    1791
    4
        }
    1792
    1793
        /// Add a logit processor to the pipeline
    1794
        #[must_use]
    1795
    2
        pub fn add_processor<P: LogitProcessor + 'static>(mut self, processor: P) -> Self {
    1796
    2
            self.processors = self.processors.with_processor(processor);
    1797
    2
            self
    1798
    2
        }
    1799
    1800
        /// Set generation configuration
    1801
        #[must_use]
    1802
    4
        pub fn with_config(mut self, config: GenerationConfig) -> Self {
    1803
    4
            self.config = config;
    1804
    4
            self
    1805
    4
        }
    1806
    1807
        /// Generate tokens starting from initial sequence
    1808
        ///
    1809
        /// # Arguments
    1810
        ///
    1811
        /// * `initial_tokens` - Starting token sequence (prompt)
    1812
        ///
    1813
        /// # Returns
    1814
        ///
    1815
        /// Generated token sequence (including initial tokens)
    1816
    4
        pub fn generate(&mut self, initial_tokens: &[u32]) -> Result<Vec<u32>> {
    1817
    4
            let mut tokens = initial_tokens.to_vec();
    1818
    4
            let n_vocab = self.model.vocab_size();
    1819
    4
            let eos_token = self.config.eos_token_id;
    1820
    1821
            // Simple PRNG for sampling (deterministic with seed)
    1822
    4
            let mut rng_state = self.config.seed.unwrap_or(42);
    1823
    1824
    11
            for step in 0..self.config.max_tokens4 {
    1825
                // Forward pass
    1826
    11
                let mut logits = self.model.forward(&tokens)?0;
    1827
    1828
                // Apply logit processors
    1829
    11
                let ctx = LogitProcessorContext::new(&tokens, step, n_vocab);
    1830
    11
                self.processors.process(&mut logits, &ctx);
    1831
    1832
                // Sample next token
    1833
    11
                let logits_tensor = Tensor::from_vec(vec![logits.len()], logits)?0;
    1834
    1835
                // Simple LCG for RNG
    1836
    11
                rng_state = rng_state
    1837
    11
                    .wrapping_mul(6_364_136_223_846_793_005)
    1838
    11
                    .wrapping_add(1);
    1839
    11
                let rng_value = (rng_state >> 33) as f32 / (1u64 << 31) as f32;
    1840
    1841
    11
                let next_token = sample_token(&logits_tensor, &self.config, rng_value)?0 as u32;
    1842
    1843
    11
                tokens.push(next_token);
    1844
    1845
                // Check for EOS
    1846
    11
                if let Some(eos7) = eos_token {
    1847
    7
                    if next_token == eos as u32 {
    1848
    2
                        break;
    1849
    5
                    }
    1850
    4
                }
    1851
            }
    1852
    1853
    4
            Ok(tokens)
    1854
    4
        }
    1855
    1856
        /// Get reference to the model
    1857
        #[must_use]
    1858
    0
        pub fn model(&self) -> &M {
    1859
    0
            &self.model
    1860
    0
        }
    1861
    1862
        /// Get mutable reference to the model
    1863
    0
        pub fn model_mut(&mut self) -> &mut M {
    1864
    0
            &mut self.model
    1865
    0
        }
    1866
    1867
        /// Get reference to the processor chain
    1868
        #[must_use]
    1869
    0
        pub fn processors(&self) -> &LogitProcessorChain {
    1870
    0
            &self.processors
    1871
    0
        }
    1872
    1873
        /// Get reference to the config
    1874
        #[must_use]
    1875
    0
        pub fn config(&self) -> &GenerationConfig {
    1876
    0
            &self.config
    1877
    0
        }
    1878
    }
    1879
    1880
    // Tests extracted to tests.rs (PMAT-802)