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

    Line
    Count
    Source
    1
    //! Advanced Sampling Algorithms (PMAT-802)
    2
    //!
    3
    //! Unique sampling algorithms not in sampler.rs:
    4
    //! - Min-p sampling
    5
    //! - Mirostat (v1/v2) adaptive sampling
    6
    //! - Tail-Free Sampling (TFS)
    7
    //! - Typical sampling
    8
    //! - DRY (Don't Repeat Yourself) penalty
    9
    //! - XTC (Exclude Top Choices) sampling
    10
    //! - Eta sampling
    11
    //! - Token healing
    12
    //! - Classifier-Free Guidance (CFG)
    13
    14
    use crate::error::{RealizarError, Result};
    15
    use crate::layers::softmax;
    16
    use crate::tensor::Tensor;
    17
    use serde::{Deserialize, Serialize};
    18
    use super::{sample_greedy, sample_from_distribution};
    19
    20
    /// Sample using min-p (minimum probability) sampling.
    21
    ///
    22
    /// Filters tokens with probability below `min_p * max_prob` threshold.
    23
    6
    pub fn sample_min_p(logits: &Tensor<f32>, min_p: f32, rng_value: f32) -> Result<usize> {
    24
    6
        let data = logits.data();
    25
    6
        if data.is_empty() {
    26
    0
            return Err(RealizarError::InvalidShape {
    27
    0
                reason: "Logits cannot be empty".to_string(),
    28
    0
            });
    29
    6
        }
    30
    6
        if !(0.0..=1.0).contains(&min_p) {
    31
    0
            return Err(RealizarError::InvalidShape {
    32
    0
                reason: "min_p must be in [0, 1]".to_string(),
    33
    0
            });
    34
    6
        }
    35
    36
        // Convert to probabilities
    37
    6
        let probs_tensor = softmax(logits)?0;
    38
    6
        let probs = probs_tensor.data();
    39
    40
        // Find max probability
    41
    6
        let max_prob = probs.iter().copied().fold(0.0_f32, f32::max);
    42
    6
        let threshold = min_p * max_prob;
    43
    44
        // Keep tokens above threshold
    45
    6
        let mut candidates: Vec<(usize, f32)> = probs
    46
    6
            .iter()
    47
    6
            .copied()
    48
    6
            .enumerate()
    49
    19
            .filter6(|(_, p)| *p >= threshold)
    50
    6
            .collect();
    51
    52
        // Sort by probability descending
    53
    6
        candidates.sort_by(|a, b| b.15.partial_cmp5(&a.15).unwrap_or5(std::cmp::Ordering::Equal5));
    54
    55
    6
        if candidates.is_empty() {
    56
            // Fallback to argmax
    57
    0
            return sample_greedy(logits);
    58
    6
        }
    59
    60
        // Renormalize and sample
    61
    6
        let sum: f32 = candidates.iter().map(|(_, p)| p).sum();
    62
    10
        let normalized6: Vec<f32>6 = candidates.iter()6.map6(|(_, p)| p / sum).collect6();
    63
    6
        let indices: Vec<usize> = candidates.iter().map(|(idx, _)| *idx).collect();
    64
    65
    6
        Ok(sample_from_distribution(&normalized, &indices, rng_value))
    66
    6
    }
    67
    68
    /// Mirostat sampling state for adaptive perplexity targeting
    69
    ///
    70
    /// Implements Mirostat 2.0 algorithm from the paper:
    71
    /// "Mirostat: A Neural Text Decoding Algorithm that Directly Controls Perplexity"
    72
    #[derive(Debug, Clone)]
    73
    pub struct MirostatState {
    74
        /// Target surprise value (tau)
    75
        pub tau: f32,
    76
        /// Learning rate (eta)
    77
        pub eta: f32,
    78
        /// Current surprise estimate (mu)
    79
        pub mu: f32,
    80
    }
    81
    82
    impl Default for MirostatState {
    83
    4
        fn default() -> Self {
    84
    4
            Self {
    85
    4
                tau: 5.0, // Default target surprise
    86
    4
                eta: 0.1, // Learning rate
    87
    4
                mu: 10.0, // Initial mu = 2 * tau
    88
    4
            }
    89
    4
        }
    90
    }
    91
    92
    impl MirostatState {
    93
        /// Create new Mirostat state with specified tau
    94
    3
        pub fn new(tau: f32) -> Self {
    95
    3
            Self {
    96
    3
                tau,
    97
    3
                eta: 0.1,
    98
    3
                mu: 2.0 * tau,
    99
    3
            }
    100
    3
        }
    101
    102
        /// Set learning rate
    103
        #[must_use]
    104
    2
        pub fn with_eta(mut self, eta: f32) -> Self {
    105
    2
            self.eta = eta;
    106
    2
            self
    107
    2
        }
    108
    109
        /// Update mu based on observed surprise
    110
    15
        pub fn update(&mut self, observed_surprise: f32) {
    111
    15
            self.mu -= self.eta * (observed_surprise - self.tau);
    112
    15
        }
    113
    }
    114
    115
    /// Mirostat 2.0 sampling: adaptive sampling to target perplexity
    116
    ///
    117
    /// # Arguments
    118
    ///
    119
    /// * `logits` - Logits for the vocabulary
    120
    /// * `state` - Mirostat state (will be updated)
    121
    /// * `rng_value` - Random value in [0, 1) for sampling
    122
    ///
    123
    /// # Returns
    124
    ///
    125
    /// Index of the selected token
    126
    ///
    127
    /// # Errors
    128
    ///
    129
    /// Returns error if logits are empty
    130
    13
    pub fn sample_mirostat(
    131
    13
        logits: &Tensor<f32>,
    132
    13
        state: &mut MirostatState,
    133
    13
        rng_value: f32,
    134
    13
    ) -> Result<usize> {
    135
    13
        let data = logits.data();
    136
    13
        if data.is_empty() {
    137
    0
            return Err(RealizarError::InvalidShape {
    138
    0
                reason: "Logits cannot be empty".to_string(),
    139
    0
            });
    140
    13
        }
    141
    142
        // Convert to probabilities
    143
    13
        let probs_tensor = softmax(logits)?0;
    144
    13
        let probs = probs_tensor.data();
    145
    146
        // Sort by probability descending
    147
    13
        let mut indexed: Vec<(usize, f32)> = probs.iter().copied().enumerate().collect();
    148
    49
        indexed13.sort_by13(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    149
    150
        // Save top candidate for fallback
    151
    13
        let top_candidate = indexed[0];
    152
    153
        // Calculate surprise values and find cutoff
    154
    13
        let mut candidates = Vec::new();
    155
    59
        for (idx58, prob58) in indexed {
    156
    58
            let surprise = -prob.ln();
    157
    58
            if surprise > state.mu {
    158
    12
                break;
    159
    46
            }
    160
    46
            candidates.push((idx, prob));
    161
        }
    162
    163
        // Ensure at least one candidate
    164
    13
        if candidates.is_empty() {
    165
    0
            candidates.push(top_candidate);
    166
    13
        }
    167
    168
        // Renormalize and sample
    169
    13
        let sum: f32 = candidates.iter().map(|(_, p)| p).sum();
    170
    46
        let normalized13: Vec<f32>13 = candidates.iter()13.map13(|(_, p)| p / sum).collect13();
    171
    13
        let indices: Vec<usize> = candidates.iter().map(|(idx, _)| *idx).collect();
    172
    173
    13
        let selected = sample_from_distribution(&normalized, &indices, rng_value);
    174
    15
        let selected_idx13 = indices.iter()13.position13(|&i| i == selected).unwrap_or13(0);
    175
    13
        let selected_prob = candidates[selected_idx].1;
    176
    177
        // Update mu based on observed surprise
    178
    13
        let observed_surprise = -selected_prob.ln();
    179
    13
        state.update(observed_surprise);
    180
    181
    13
        Ok(selected)
    182
    13
    }
    183
    184
    /// Tail-Free Sampling (TFS): Filter tokens based on probability second derivatives
    185
    ///
    186
    /// TFS analyzes the "tail" of the probability distribution and removes tokens
    187
    /// in the low-probability tail. It computes second derivatives to find where
    188
    /// the distribution starts to flatten out.
    189
    ///
    190
    /// # Arguments
    191
    ///
    192
    /// * `logits` - Logits for the vocabulary
    193
    /// * `z` - TFS parameter (0.0 to 1.0, higher = more tokens kept)
    194
    /// * `rng_value` - Random value in [0, 1) for sampling
    195
    ///
    196
    /// # Returns
    197
    ///
    198
    /// Index of the selected token
    199
    ///
    200
    /// # Errors
    201
    ///
    202
    /// Returns error if logits are empty
    203
    6
    pub fn sample_tfs(logits: &Tensor<f32>, z: f32, rng_value: f32) -> Result<usize> {
    204
    6
        let data = logits.data();
    205
    6
        if data.is_empty() {
    206
    0
            return Err(crate::error::RealizarError::InvalidShape {
    207
    0
                reason: "Logits cannot be empty".to_string(),
    208
    0
            });
    209
    6
        }
    210
    211
        // Convert to probabilities
    212
    6
        let max_logit = data.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    213
    23
        let exp_logits6: Vec<f32>6 = data6.iter6().map6(|&x| (x - max_logit).exp()).collect6();
    214
    6
        let sum: f32 = exp_logits.iter().sum();
    215
    23
        let probs6: Vec<f32>6 = exp_logits.iter()6.map6(|&x| x / sum).collect6();
    216
    217
        // Sort by probability descending
    218
    23
        let mut indexed6: Vec<(usize, f32)>6 = probs.iter()6.enumerate6().map6(|(i, &p)| (i, p)).collect6();
    219
    17
        indexed6.sort_by6(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    220
    221
    6
        if indexed.len() < 3 {
    222
            // Not enough tokens for second derivative, use greedy
    223
    2
            return Ok(indexed[0].0);
    224
    4
        }
    225
    226
        // Compute first derivatives (differences between consecutive probabilities)
    227
    4
        let first_derivatives: Vec<f32> = indexed
    228
    4
            .windows(2)
    229
    16
            .map4(|w| (w[0].1 - w[1].1).abs())
    230
    4
            .collect();
    231
    232
        // Compute second derivatives
    233
    4
        let second_derivatives: Vec<f32> = first_derivatives
    234
    4
            .windows(2)
    235
    12
            .map4(|w| (w[0] - w[1]).abs())
    236
    4
            .collect();
    237
    238
        // Normalize second derivatives
    239
    4
        let sum_second: f32 = second_derivatives.iter().sum();
    240
    4
        let normalized: Vec<f32> = if sum_second > 1e-9 {
    241
    9
            second_derivatives.iter()3.map3(|&x| x / sum_second).collect3()
    242
        } else {
    243
    1
            vec![1.0 / second_derivatives.len() as f32; second_derivatives.len()]
    244
        };
    245
    246
        // Find cumulative sum and cutoff point
    247
    4
        let mut cumsum = 0.0;
    248
    4
        let mut cutoff_idx = indexed.len();
    249
    9
        for (i, &val) in normalized.iter()4.enumerate4() {
    250
    9
            cumsum += val;
    251
    9
            if cumsum > z {
    252
    3
                cutoff_idx = i + 2; // +2 because second derivative is 2 steps behind
    253
    3
                break;
    254
    6
            }
    255
        }
    256
    257
        // Keep tokens up to cutoff
    258
    4
        let kept: Vec<(usize, f32)> = indexed.into_iter().take(cutoff_idx.max(1)).collect();
    259
    260
        // Renormalize and sample
    261
    4
        let sum_kept: f32 = kept.iter().map(|(_, p)| p).sum();
    262
    14
        let normalized_kept4: Vec<f32>4 = kept.iter()4.map4(|(_, p)| p / sum_kept).collect4();
    263
    4
        let indices: Vec<usize> = kept.iter().map(|(idx, _)| *idx).collect();
    264
    265
    4
        Ok(sample_from_distribution(
    266
    4
            &normalized_kept,
    267
    4
            &indices,
    268
    4
            rng_value,
    269
    4
        ))
    270
    6
    }
    271
    272
    /// Locally Typical Sampling: Sample based on local typicality
    273
    ///
    274
    /// Typical sampling selects tokens whose information content is close to
    275
    /// the expected information content (entropy) of the distribution.
    276
    /// This tends to produce more "typical" text.
    277
    ///
    278
    /// Reference: Meister et al. (2022) "Locally Typical Sampling"
    279
    ///
    280
    /// # Arguments
    281
    ///
    282
    /// * `logits` - Logits for the vocabulary
    283
    /// * `p` - Cumulative probability mass to keep (0.0 to 1.0)
    284
    /// * `rng_value` - Random value in [0, 1) for sampling
    285
    ///
    286
    /// # Returns
    287
    ///
    288
    /// Index of the selected token
    289
    ///
    290
    /// # Errors
    291
    ///
    292
    /// Returns error if logits are empty
    293
    6
    pub fn sample_typical(logits: &Tensor<f32>, p: f32, rng_value: f32) -> Result<usize> {
    294
    6
        let data = logits.data();
    295
    6
        if data.is_empty() {
    296
    0
            return Err(crate::error::RealizarError::InvalidShape {
    297
    0
                reason: "Logits cannot be empty".to_string(),
    298
    0
            });
    299
    6
        }
    300
    301
        // Convert to probabilities
    302
    6
        let max_logit = data.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    303
    22
        let exp_logits6: Vec<f32>6 = data6.iter6().map6(|&x| (x - max_logit).exp()).collect6();
    304
    6
        let sum: f32 = exp_logits.iter().sum();
    305
    22
        let probs6: Vec<f32>6 = exp_logits.iter()6.map6(|&x| x / sum).collect6();
    306
    307
        // Compute entropy (expected information content)
    308
    6
        let entropy: f32 = -probs
    309
    6
            .iter()
    310
    22
            .filter6(|&&p| p > 1e-10)
    311
    22
            .map6(|&p| p * p.ln())
    312
    6
            .sum::<f32>();
    313
    314
        // Compute information content for each token: -log(p)
    315
        // Then compute deviation from entropy: |info - entropy|
    316
    6
        let mut indexed: Vec<(usize, f32, f32)> = probs
    317
    6
            .iter()
    318
    6
            .enumerate()
    319
    22
            .filter6(|(_, &prob)| prob > 1e-10)
    320
    22
            .map6(|(i, &prob)| {
    321
    22
                let info = -prob.ln();
    322
    22
                let deviation = (info - entropy).abs();
    323
    22
                (i, prob, deviation)
    324
    22
            })
    325
    6
            .collect();
    326
    327
        // Sort by deviation (most typical first)
    328
    17
        indexed6.sort_by6(|a, b| a.2.partial_cmp(&b.2).unwrap_or(std::cmp::Ordering::Equal));
    329
    330
        // Keep tokens until cumulative probability exceeds p
    331
    6
        let mut cumsum = 0.0;
    332
    6
        let mut kept: Vec<(usize, f32)> = Vec::new();
    333
    18
        for (idx, prob, _) in indexed {
    334
    18
            kept.push((idx, prob));
    335
    18
            cumsum += prob;
    336
    18
            if cumsum >= p {
    337
    6
                break;
    338
    12
            }
    339
        }
    340
    341
        // Ensure at least one token
    342
    6
        if kept.is_empty() {
    343
    0
            let max_idx = probs
    344
    0
                .iter()
    345
    0
                .enumerate()
    346
    0
                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
    347
    0
                .map_or(0, |(i, _)| i);
    348
    0
            return Ok(max_idx);
    349
    6
        }
    350
    351
        // Renormalize and sample
    352
    6
        let sum_kept: f32 = kept.iter().map(|(_, p)| p).sum();
    353
    18
        let normalized6: Vec<f32>6 = kept.iter()6.map6(|(_, p)| p / sum_kept).collect6();
    354
    6
        let indices: Vec<usize> = kept.iter().map(|(idx, _)| *idx).collect();
    355
    356
    6
        Ok(sample_from_distribution(&normalized, &indices, rng_value))
    357
    6
    }
    358
    359
    /// DRY (Don't Repeat Yourself) sampling configuration
    360
    ///
    361
    /// DRY sampling penalizes n-gram repetitions to prevent the model from
    362
    /// generating repetitive sequences.
    363
    #[derive(Debug, Clone, Serialize, Deserialize)]
    364
    pub struct DryConfig {
    365
        /// Multiplier for the penalty (higher = stronger penalty)
    366
        pub multiplier: f32,
    367
        /// Base value for exponential penalty growth
    368
        pub base: f32,
    369
        /// Minimum n-gram length to consider
    370
        pub allowed_length: usize,
    371
        /// Maximum sequence length to check for repetitions
    372
        pub penalty_last_n: usize,
    373
    }
    374
    375
    impl Default for DryConfig {
    376
    4
        fn default() -> Self {
    377
    4
            Self {
    378
    4
                multiplier: 0.8,
    379
    4
                base: 1.75,
    380
    4
                allowed_length: 2,
    381
    4
                penalty_last_n: 256,
    382
    4
            }
    383
    4
        }
    384
    }
    385
    386
    impl DryConfig {
    387
        /// Create new DRY config with specified multiplier
    388
    3
        pub fn new(multiplier: f32) -> Self {
    389
    3
            Self {
    390
    3
                multiplier,
    391
    3
                ..Default::default()
    392
    3
            }
    393
    3
        }
    394
    395
        /// Set the base for exponential penalty
    396
        #[must_use]
    397
    1
        pub fn with_base(mut self, base: f32) -> Self {
    398
    1
            self.base = base;
    399
    1
            self
    400
    1
        }
    401
    402
        /// Set minimum n-gram length
    403
        #[must_use]
    404
    1
        pub fn with_allowed_length(mut self, len: usize) -> Self {
    405
    1
            self.allowed_length = len;
    406
    1
            self
    407
    1
        }
    408
    409
        /// Set penalty window size
    410
        #[must_use]
    411
    1
        pub fn with_penalty_last_n(mut self, n: usize) -> Self {
    412
    1
            self.penalty_last_n = n;
    413
    1
            self
    414
    1
        }
    415
    416
        /// Check if DRY is enabled
    417
    7
        pub fn is_enabled(&self) -> bool {
    418
    7
            self.multiplier > 0.0
    419
    7
        }
    420
    }
    421
    422
    /// Apply DRY (Don't Repeat Yourself) penalty to logits
    423
    ///
    424
    /// Penalizes tokens that would extend n-gram repetitions in the context.
    425
    ///
    426
    /// # Arguments
    427
    ///
    428
    /// * `logits` - Raw logits from model
    429
    /// * `context_tokens` - List of previously generated token IDs
    430
    /// * `config` - DRY configuration
    431
    ///
    432
    /// # Returns
    433
    ///
    434
    /// Logits with DRY penalty applied
    435
    4
    pub fn apply_dry_penalty(
    436
    4
        logits: &Tensor<f32>,
    437
    4
        context_tokens: &[usize],
    438
    4
        config: &DryConfig,
    439
    4
    ) -> Tensor<f32> {
    440
    4
        if !config.is_enabled() || context_tokens3.len() < config.allowed_length {
    441
    2
            return logits.clone();
    442
    2
        }
    443
    444
    2
        let data = logits.data();
    445
    2
        let mut penalized = data.to_vec();
    446
    447
        // Get relevant context window
    448
    2
        let window_start = if context_tokens.len() > config.penalty_last_n {
    449
    1
            context_tokens.len() - config.penalty_last_n
    450
        } else {
    451
    1
            0
    452
        };
    453
    2
        let context = &context_tokens[window_start..];
    454
    455
        // For each possible next token, check if it would extend a repetition
    456
    10
        for (token_id, logit) in penalized.iter_mut()2.enumerate2() {
    457
    10
            let match_len = find_ngram_match_length(context, token_id, config.allowed_length);
    458
    459
    10
            if match_len >= config.allowed_length {
    460
    1
                // Apply exponential penalty based on match length
    461
    1
                let penalty =
    462
    1
                    config.multiplier * config.base.powi((match_len - config.allowed_length) as i32);
    463
    1
                *logit -= penalty;
    464
    9
            }
    465
        }
    466
    467
    2
        Tensor::from_vec(logits.shape().to_vec(), penalized)
    468
    2
            .expect("Shape should match original logits")
    469
    4
    }
    470
    471
    /// Find the length of the longest n-gram that would be repeated if we add this token
    472
    10
    fn find_ngram_match_length(context: &[usize], next_token: usize, min_len: usize) -> usize {
    473
    10
        if context.len() < min_len {
    474
    0
            return 0;
    475
    10
        }
    476
    477
    10
        let mut max_match = 0;
    478
    479
        // Build the sequence ending with the potential next token
    480
        // Then search for earlier occurrences
    481
    25
        for end_pos in min_len10..=context10.len10() {
    482
    25
            let search_start = context.len() - end_pos;
    483
    25
            let suffix = &context[search_start..];
    484
    485
            // Look for this suffix earlier in the context
    486
    25
            for start20 in 0..(context.len() - end_pos) {
    487
    20
                let potential_end = start + end_pos;
    488
    20
                if potential_end >= context.len() {
    489
    0
                    continue;
    490
    20
                }
    491
    492
                // Check if suffix matches
    493
    20
                if context[start..potential_end] == *suffix {
    494
                    // Check if the next token after this match equals our candidate
    495
    5
                    if potential_end < context.len() && context[potential_end] == next_token {
    496
    1
                        max_match = max_match.max(end_pos + 1);
    497
    4
                    }
    498
    15
                }
    499
            }
    500
        }
    501
    502
    10
        max_match
    503
    10
    }
    504
    505
    // ===== XTC (Exclude Top Choices) Sampling =====
    506
    507
    /// XTC (Exclude Top Choices) sampling configuration
    508
    ///
    509
    /// XTC removes the most likely tokens with some probability, forcing the model
    510
    /// to explore alternative completions. This can increase creativity and diversity.
    511
    #[derive(Debug, Clone, Serialize, Deserialize)]
    512
    pub struct XtcConfig {
    513
        /// Probability of excluding top tokens (0.0 = disabled, 1.0 = always exclude)
    514
        pub probability: f32,
    515
        /// Threshold for excluding tokens (tokens with prob >= threshold may be excluded)
    516
        pub threshold: f32,
    517
        /// Minimum number of tokens to keep after exclusion
    518
        pub min_keep: usize,
    519
    }
    520
    521
    impl Default for XtcConfig {
    522
    6
        fn default() -> Self {
    523
    6
            Self {
    524
    6
                probability: 0.0,
    525
    6
                threshold: 0.5,
    526
    6
                min_keep: 1,
    527
    6
            }
    528
    6
        }
    529
    }
    530
    531
    impl XtcConfig {
    532
        /// Create new XTC config with specified probability
    533
    4
        pub fn new(probability: f32) -> Self {
    534
    4
            Self {
    535
    4
                probability,
    536
    4
                ..Default::default()
    537
    4
            }
    538
    4
        }
    539
    540
        /// Set threshold
    541
        #[must_use]
    542
    3
        pub fn with_threshold(mut self, threshold: f32) -> Self {
    543
    3
            self.threshold = threshold;
    544
    3
            self
    545
    3
        }
    546
    547
        /// Set minimum tokens to keep
    548
        #[must_use]
    549
    2
        pub fn with_min_keep(mut self, min_keep: usize) -> Self {
    550
    2
            self.min_keep = min_keep;
    551
    2
            self
    552
    2
        }
    553
    554
        /// Check if XTC is enabled
    555
    6
        pub fn is_enabled(&self) -> bool {
    556
    6
            self.probability > 0.0
    557
    6
        }
    558
    }
    559
    560
    /// Apply XTC (Exclude Top Choices) sampling
    561
    ///
    562
    /// XTC randomly excludes top tokens to increase diversity.
    563
    ///
    564
    /// # Arguments
    565
    ///
    566
    /// * `logits` - Raw logits from the model
    567
    /// * `config` - XTC configuration
    568
    /// * `rng_value` - Random value [0, 1) for stochastic exclusion decision
    569
    ///
    570
    /// # Returns
    571
    ///
    572
    /// Modified logits with top choices potentially excluded
    573
    4
    pub fn apply_xtc(logits: &Tensor<f32>, config: &XtcConfig, rng_value: f32) -> Tensor<f32> {
    574
    4
        if !config.is_enabled() || rng_value >= config.probability3 {
    575
    2
            return logits.clone();
    576
    2
        }
    577
    578
    2
        let data = logits.data();
    579
    2
        if data.len() <= config.min_keep {
    580
    0
            return logits.clone();
    581
    2
        }
    582
    583
        // Convert to probabilities
    584
    2
        let max_logit = data.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    585
    8
        let exp_logits2: Vec<f32>2 = data2.iter2().map2(|&x| (x - max_logit).exp()).collect2();
    586
    2
        let sum: f32 = exp_logits.iter().sum();
    587
    8
        let probs2: Vec<f32>2 = exp_logits.iter()2.map2(|&x| x / sum).collect2();
    588
    589
        // Find tokens above threshold
    590
    2
        let mut excluded_count = 0;
    591
    2
        let mut modified = data.to_vec();
    592
    593
        // Sort by probability descending to find top tokens
    594
    8
        let mut indexed2: Vec<(usize, f32)>2 = probs.iter()2.enumerate2().map2(|(i, &p)| (i, p)).collect2();
    595
    6
        indexed2.sort_by2(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    596
    597
        // Exclude top tokens above threshold, respecting min_keep
    598
    10
        for (idx8, prob8) in &indexed {
    599
    8
            if *prob >= config.threshold && data3.len() - excluded_count > config.min_keep {
    600
    2
                modified[*idx] = f32::NEG_INFINITY;
    601
    2
                excluded_count += 1;
    602
    6
            }
    603
        }
    604
    605
    2
        Tensor::from_vec(logits.shape().to_vec(), modified).expect("Shape should match original logits")
    606
    4
    }
    607
    608
    // ===== Eta Sampling =====
    609
    610
    /// Eta Sampling (entropy-based truncation)
    611
    ///
    612
    /// Eta sampling dynamically adjusts the truncation threshold based on the
    613
    /// entropy of the probability distribution. Higher entropy = more tokens kept.
    614
    #[derive(Debug, Clone, Serialize, Deserialize)]
    615
    pub struct EtaConfig {
    616
        /// Eta parameter (controls sensitivity to entropy)
    617
        pub eta: f32,
    618
        /// Minimum probability to keep (absolute floor)
    619
        pub min_p: f32,
    620
    }
    621
    622
    impl Default for EtaConfig {
    623
    6
        fn default() -> Self {
    624
    6
            Self {
    625
    6
                eta: 0.3,
    626
    6
                min_p: 0.0001,
    627
    6
            }
    628
    6
        }
    629
    }
    630
    631
    impl EtaConfig {
    632
        /// Create new Eta config
    633
    2
        pub fn new(eta: f32) -> Self {
    634
    2
            Self {
    635
    2
                eta,
    636
    2
                ..Default::default()
    637
    2
            }
    638
    2
        }
    639
    640
        /// Set minimum probability
    641
        #[must_use]
    642
    1
        pub fn with_min_p(mut self, min_p: f32) -> Self {
    643
    1
            self.min_p = min_p;
    644
    1
            self
    645
    1
        }
    646
    647
        /// Check if eta sampling is enabled
    648
    2
        pub fn is_enabled(&self) -> bool {
    649
    2
            self.eta > 0.0
    650
    2
        }
    651
    }
    652
    653
    /// Apply Eta sampling
    654
    ///
    655
    /// # Arguments
    656
    ///
    657
    /// * `logits` - Raw logits from the model
    658
    /// * `config` - Eta configuration
    659
    /// * `rng_value` - Random value [0, 1) for sampling
    660
    ///
    661
    /// # Returns
    662
    ///
    663
    /// Index of the selected token
    664
    ///
    665
    /// # Errors
    666
    ///
    667
    /// Returns error if logits are empty
    668
    3
    pub fn sample_eta(logits: &Tensor<f32>, config: &EtaConfig, rng_value: f32) -> Result<usize> {
    669
    3
        let data = logits.data();
    670
    3
        if data.is_empty() {
    671
    0
            return Err(crate::error::RealizarError::InvalidShape {
    672
    0
                reason: "Logits cannot be empty".to_string(),
    673
    0
            });
    674
    3
        }
    675
    676
        // Convert to probabilities
    677
    3
        let max_logit = data.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    678
    10
        let exp_logits3: Vec<f32>3 = data3.iter3().map3(|&x| (x - max_logit).exp()).collect3();
    679
    3
        let sum: f32 = exp_logits.iter().sum();
    680
    10
        let probs3: Vec<f32>3 = exp_logits.iter()3.map3(|&x| x / sum).collect3();
    681
    682
        // Compute entropy
    683
    3
        let entropy: f32 = -probs
    684
    3
            .iter()
    685
    10
            .filter3(|&&p| p > 1e-10)
    686
    10
            .map3(|&p| p * p.ln())
    687
    3
            .sum::<f32>();
    688
    689
        // Compute dynamic threshold: eta * exp(-entropy)
    690
    3
        let threshold = (config.eta * (-entropy).exp()).max(config.min_p);
    691
    692
        // Keep tokens above threshold
    693
    3
        let mut indexed: Vec<(usize, f32)> = probs
    694
    3
            .iter()
    695
    3
            .enumerate()
    696
    10
            .filter3(|(_, &p)| p >= threshold)
    697
    8
            .map3(|(i, &p)| (i, p))
    698
    3
            .collect();
    699
    700
        // Ensure at least one token
    701
    3
        if indexed.is_empty() {
    702
    0
            let max_idx = probs
    703
    0
                .iter()
    704
    0
                .enumerate()
    705
    0
                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
    706
    0
                .map_or(0, |(i, _)| i);
    707
    0
            return Ok(max_idx);
    708
    3
        }
    709
    710
        // Sort by probability descending
    711
    5
        indexed3.sort_by3(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    712
    713
        // Renormalize and sample
    714
    3
        let sum_kept: f32 = indexed.iter().map(|(_, p)| p).sum();
    715
    8
        let normalized3: Vec<f32>3 = indexed.iter()3.map3(|(_, p)| p / sum_kept).collect3();
    716
    3
        let indices: Vec<usize> = indexed.iter().map(|(idx, _)| *idx).collect();
    717
    718
    3
        Ok(sample_from_distribution(&normalized, &indices, rng_value))
    719
    3
    }
    720
    721
    // ===== Token Healing =====
    722
    723
    /// Token Healing configuration
    724
    ///
    725
    /// Token healing fixes broken token boundaries by backing up and re-tokenizing
    726
    /// when a partial token is detected at the prompt boundary.
    727
    #[derive(Debug, Clone, Default)]
    728
    pub struct TokenHealingConfig {
    729
        /// Enable token healing
    730
        pub enabled: bool,
    731
        /// Maximum characters to back up
    732
        pub max_backup_chars: usize,
    733
    }
    734
    735
    impl TokenHealingConfig {
    736
        /// Create new token healing config
    737
    1
        pub fn new(enabled: bool) -> Self {
    738
    1
            Self {
    739
    1
                enabled,
    740
    1
                max_backup_chars: 10,
    741
    1
            }
    742
    1
        }
    743
    744
        /// Set max backup characters
    745
        #[must_use]
    746
    1
        pub fn with_max_backup(mut self, chars: usize) -> Self {
    747
    1
            self.max_backup_chars = chars;
    748
    1
            self
    749
    1
        }
    750
    }
    751
    752
    /// Token healing result
    753
    #[derive(Debug, Clone)]
    754
    pub struct TokenHealingResult {
    755
        /// Adjusted prompt tokens (may be shorter than original)
    756
        pub adjusted_tokens: Vec<usize>,
    757
        /// Prefix constraint for first generated token
    758
        pub prefix_constraint: Option<String>,
    759
        /// Number of tokens removed from end
    760
        pub tokens_removed: usize,
    761
    }
    762
    763
    /// Analyze prompt for token healing
    764
    ///
    765
    /// Detects if the last token is a partial token that should be healed.
    766
    /// This is a simplified implementation - full implementation requires tokenizer access.
    767
    ///
    768
    /// # Arguments
    769
    ///
    770
    /// * `prompt_tokens` - Original prompt tokens
    771
    /// * `last_token_text` - Text of the last token (if available)
    772
    ///
    773
    /// # Returns
    774
    ///
    775
    /// Token healing result with adjusted tokens
    776
    4
    pub fn analyze_token_healing(
    777
    4
        prompt_tokens: &[usize],
    778
    4
        last_token_text: Option<&str>,
    779
    4
    ) -> TokenHealingResult {
    780
        // Simple heuristic: if last token is a partial word (no space, single char),
    781
        // we might want to heal it
    782
    4
        let should_heal = last_token_text.is_some_and(|text| {
    783
    4
            !text.is_empty()
    784
    4
                && !text.starts_with(' ')
    785
    3
                && text.len() <= 3
    786
    2
                && text.chars().all(char::is_alphanumeric)
    787
    4
        });
    788
    789
    4
        if should_heal && !prompt_tokens.is_empty()2 {
    790
    1
            TokenHealingResult {
    791
    1
                adjusted_tokens: prompt_tokens[..prompt_tokens.len() - 1].to_vec(),
    792
    1
                prefix_constraint: last_token_text.map(String::from),
    793
    1
                tokens_removed: 1,
    794
    1
            }
    795
        } else {
    796
    3
            TokenHealingResult {
    797
    3
                adjusted_tokens: prompt_tokens.to_vec(),
    798
    3
                prefix_constraint: None,
    799
    3
                tokens_removed: 0,
    800
    3
            }
    801
        }
    802
    4
    }
    803
    804
    // ===== Classifier-Free Guidance (CFG) =====
    805
    806
    /// Classifier-Free Guidance configuration
    807
    ///
    808
    /// CFG improves generation quality by comparing conditional and unconditional
    809
    /// logits, amplifying the difference to steer generation.
    810
    #[derive(Debug, Clone, Serialize, Deserialize)]
    811
    pub struct CfgConfig {
    812
        /// Guidance scale (1.0 = no guidance, higher = stronger guidance)
    813
        pub scale: f32,
    814
        /// Negative prompt tokens (for unconditional generation)
    815
        pub negative_prompt_tokens: Vec<usize>,
    816
    }
    817
    818
    impl Default for CfgConfig {
    819
    2
        fn default() -> Self {
    820
    2
            Self {
    821
    2
                scale: 1.0,
    822
    2
                negative_prompt_tokens: Vec::new(),
    823
    2
            }
    824
    2
        }
    825
    }
    826
    827
    impl CfgConfig {
    828
        /// Create new CFG config with specified scale
    829
    1
        pub fn new(scale: f32) -> Self {
    830
    1
            Self {
    831
    1
                scale,
    832
    1
                ..Default::default()
    833
    1
            }
    834
    1
        }
    835
    836
        /// Set negative prompt tokens
    837
        #[must_use]
    838
    1
        pub fn with_negative_prompt(mut self, tokens: Vec<usize>) -> Self {
    839
    1
            self.negative_prompt_tokens = tokens;
    840
    1
            self
    841
    1
        }
    842
    843
        /// Check if CFG is enabled
    844
    2
        pub fn is_enabled(&self) -> bool {
    845
    2
            self.scale > 1.0
    846
    2
        }
    847
    }
    848
    849
    /// Apply Classifier-Free Guidance
    850
    ///
    851
    /// Combines conditional and unconditional logits using the CFG formula:
    852
    /// output = unconditional + scale * (conditional - unconditional)
    853
    ///
    854
    /// # Arguments
    855
    ///
    856
    /// * `conditional_logits` - Logits from the model with the prompt
    857
    /// * `unconditional_logits` - Logits from the model with negative/empty prompt
    858
    /// * `scale` - Guidance scale
    859
    ///
    860
    /// # Returns
    861
    ///
    862
    /// Guided logits
    863
    ///
    864
    /// # Errors
    865
    ///
    866
    /// Returns error if conditional and unconditional logits have different shapes
    867
    4
    pub fn apply_cfg(
    868
    4
        conditional_logits: &Tensor<f32>,
    869
    4
        unconditional_logits: &Tensor<f32>,
    870
    4
        scale: f32,
    871
    4
    ) -> Result<Tensor<f32>> {
    872
    4
        if conditional_logits.shape() != unconditional_logits.shape() {
    873
    1
            return Err(crate::error::RealizarError::ShapeMismatch {
    874
    1
                expected: conditional_logits.shape().to_vec(),
    875
    1
                actual: unconditional_logits.shape().to_vec(),
    876
    1
            });
    877
    3
        }
    878
    879
    3
        let cond = conditional_logits.data();
    880
    3
        let uncond = unconditional_logits.data();
    881
    882
        // CFG formula: uncond + scale * (cond - uncond)
    883
    3
        let guided: Vec<f32> = cond
    884
    3
            .iter()
    885
    3
            .zip(uncond.iter())
    886
    11
            .map3(|(&c, &u)| u + scale * (c - u))
    887
    3
            .collect();
    888
    889
    3
        Tensor::from_vec(conditional_logits.shape().to_vec(), guided)
    890
    4
    }
    891