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

    Line
    Count
    Source
    1
    //! Text generation and sampling strategies
    2
    //!
    3
    //! This module provides the generation loop for autoregressive text generation
    4
    //! and various sampling strategies for token selection.
    5
    //!
    6
    //! # Sampling Strategies
    7
    //!
    8
    //! - **Greedy**: Always select the most probable token
    9
    //! - **Top-k**: Sample from the k most probable tokens
    10
    //! - **Top-p (nucleus)**: Sample from tokens with cumulative probability ≤ p
    11
    //! - **Temperature**: Scale logits before softmax to control randomness
    12
    13
    use crate::{
    14
        error::{RealizarError, Result},
    15
        layers::softmax,
    16
        tensor::Tensor,
    17
    };
    18
    19
    // Submodules
    20
    mod algorithms;
    21
    mod sampler;
    22
    23
    // Re-exports from algorithms (unique sampling algorithms)
    24
    pub use algorithms::{
    25
        sample_min_p, MirostatState, sample_mirostat, sample_tfs, sample_typical,
    26
        DryConfig, apply_dry_penalty, XtcConfig, apply_xtc, EtaConfig, sample_eta,
    27
        TokenHealingConfig, TokenHealingResult, analyze_token_healing, CfgConfig, apply_cfg,
    28
    };
    29
    30
    // Re-exports from sampler (advanced sampling infrastructure)
    31
    pub use sampler::{
    32
        StopSequenceDetector, RepetitionPenaltyConfig, apply_repetition_penalty,
    33
        PresenceFrequencyPenalty, apply_presence_frequency_penalty, LogitBias, apply_logit_bias,
    34
        PromptCache, PromptCacheEntry, PromptCacheStats, BeamHypothesis, BeamSearchConfig,
    35
        BeamSearchState, StreamingGenerator, AdvancedGenerationConfig, apply_all_penalties,
    36
        DynTempConfig, apply_dynamic_temperature, InfillConfig, InfillResult, apply_infill_sampling,
    37
        SamplerContext, SamplerChain, Sampler, TemperatureSampler, DynTempSampler, TopKSampler,
    38
        TopPSampler, RepetitionPenaltySampler, InfillSampler, LogitProcessorContext,
    39
        LogitProcessor, TokenSuppressor, RepetitionPenalty, TemperatureScaler,
    40
        LogitProcessorChain, GenerativeModel, GenerationPipeline,
    41
    };
    42
    43
    /// Sample from a probability distribution using a random value
    44
    ///
    45
    /// # Arguments
    46
    ///
    47
    /// * `probs` - Probabilities (must sum to 1)
    48
    /// * `indices` - Corresponding indices for each probability
    49
    /// * `rng_value` - Random value in [0, 1)
    50
    ///
    51
    /// # Returns
    52
    ///
    53
    /// Selected index
    54
    322
    pub(crate) fn sample_from_distribution(probs: &[f32], indices: &[usize], rng_value: f32) -> usize {
    55
    322
        let mut cumsum = 0.0;
    56
    11.8k
        for (i, &prob) in probs322.iter322().enumerate322() {
    57
    11.8k
            cumsum += prob;
    58
    11.8k
            if rng_value < cumsum {
    59
    322
                return indices[i];
    60
    11.5k
            }
    61
        }
    62
        // Fallback to last token
    63
    0
        indices[indices.len() - 1]
    64
    322
    }
    65
    66
    /// Convert logits to softmax probabilities for a subset
    67
    ///
    68
    /// # Arguments
    69
    ///
    70
    /// * `indexed` - Pairs of (index, logit) sorted by logit descending
    71
    ///
    72
    /// # Returns
    73
    ///
    74
    /// Probabilities for the subset
    75
    19
    pub(crate) fn logits_to_probs(indexed: &[(usize, f32)]) -> Vec<f32> {
    76
    19
        let max_logit = indexed[0].1;
    77
    96
        let exp_vals19: Vec<f32>19 = indexed19.iter19().map19(|(_, l)| (l - max_logit).exp()).collect19();
    78
    19
        let sum_exp: f32 = exp_vals.iter().sum();
    79
    96
        exp_vals.iter()19.map19(|e| e / sum_exp).collect19()
    80
    19
    }
    81
    82
    /// Build nucleus for top-p sampling
    83
    ///
    84
    /// # Arguments
    85
    ///
    86
    /// * `indexed` - Pairs of (index, prob) sorted by prob descending
    87
    /// * `p` - Cumulative probability threshold
    88
    ///
    89
    /// # Returns
    90
    ///
    91
    /// Nucleus of (index, prob) pairs with cumulative probability >= p
    92
    271
    pub(crate) fn build_nucleus(indexed: &[(usize, f32)], p: f32) -> Vec<(usize, f32)> {
    93
    271
        let mut cumsum = 0.0;
    94
    271
        let mut nucleus = Vec::new();
    95
    24.3k
        for &(idx, prob) in indexed {
    96
    24.3k
            nucleus.push((idx, prob));
    97
    24.3k
            cumsum += prob;
    98
    24.3k
            if cumsum >= p {
    99
    271
                break;
    100
    24.0k
            }
    101
        }
    102
    271
        nucleus
    103
    271
    }
    104
    105
    /// Sampling strategy for token selection
    106
    #[derive(Debug, Clone, Copy, PartialEq)]
    107
    pub enum SamplingStrategy {
    108
        /// Always select the most probable token
    109
        Greedy,
    110
        /// Sample from the k most probable tokens
    111
        TopK {
    112
            /// Number of top tokens to consider
    113
            k: usize,
    114
        },
    115
        /// Sample from tokens with cumulative probability ≤ p
    116
        TopP {
    117
            /// Cumulative probability threshold
    118
            p: f32,
    119
        },
    120
    }
    121
    122
    /// Configuration for text generation
    123
    #[derive(Debug, Clone)]
    124
    pub struct GenerationConfig {
    125
        /// Maximum number of tokens to generate
    126
        pub max_tokens: usize,
    127
        /// Sampling strategy
    128
        pub strategy: SamplingStrategy,
    129
        /// Temperature for scaling logits (1.0 = no scaling)
    130
        pub temperature: f32,
    131
        /// Token ID for end-of-sequence
    132
        pub eos_token_id: Option<usize>,
    133
        /// Random seed for reproducibility
    134
        pub seed: Option<u64>,
    135
    }
    136
    137
    impl Default for GenerationConfig {
    138
    59
        fn default() -> Self {
    139
    59
            Self {
    140
    59
                max_tokens: 100,
    141
    59
                strategy: SamplingStrategy::Greedy,
    142
    59
                temperature: 1.0,
    143
    59
                eos_token_id: None,
    144
    59
                seed: None,
    145
    59
            }
    146
    59
        }
    147
    }
    148
    149
    impl GenerationConfig {
    150
        /// Create a new generation config with greedy sampling
    151
        #[must_use]
    152
    16
        pub fn greedy() -> Self {
    153
    16
            Self {
    154
    16
                strategy: SamplingStrategy::Greedy,
    155
    16
                ..Default::default()
    156
    16
            }
    157
    16
        }
    158
    159
        /// Create a new generation config with top-k sampling
    160
        #[must_use]
    161
    3
        pub fn top_k(k: usize) -> Self {
    162
    3
            Self {
    163
    3
                strategy: SamplingStrategy::TopK { k },
    164
    3
                ..Default::default()
    165
    3
            }
    166
    3
        }
    167
    168
        /// Create a new generation config with top-p (nucleus) sampling
    169
        #[must_use]
    170
    2
        pub fn top_p(p: f32) -> Self {
    171
    2
            Self {
    172
    2
                strategy: SamplingStrategy::TopP { p },
    173
    2
                ..Default::default()
    174
    2
            }
    175
    2
        }
    176
    177
        /// Set temperature
    178
        #[must_use]
    179
    31
        pub fn with_temperature(mut self, temperature: f32) -> Self {
    180
    31
            self.temperature = temperature;
    181
    31
            self
    182
    31
        }
    183
    184
        /// Set maximum tokens
    185
        #[must_use]
    186
    41
        pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
    187
    41
            self.max_tokens = max_tokens;
    188
    41
            self
    189
    41
        }
    190
    191
        /// Set end-of-sequence token ID
    192
        #[must_use]
    193
    4
        pub fn with_eos_token_id(mut self, eos_token_id: usize) -> Self {
    194
    4
            self.eos_token_id = Some(eos_token_id);
    195
    4
            self
    196
    4
        }
    197
    198
        /// Set random seed
    199
        #[must_use]
    200
    13
        pub fn with_seed(mut self, seed: u64) -> Self {
    201
    13
            self.seed = Some(seed);
    202
    13
            self
    203
    13
        }
    204
    }
    205
    206
    /// Apply temperature scaling to logits
    207
    ///
    208
    /// # Arguments
    209
    ///
    210
    /// * `logits` - Raw logits from model
    211
    /// * `temperature` - Temperature value (> 0)
    212
    ///
    213
    /// # Returns
    214
    ///
    215
    /// Scaled logits
    216
    ///
    217
    /// # Errors
    218
    ///
    219
    /// Returns error if temperature is not positive
    220
    1.36k
    pub fn apply_temperature(logits: &Tensor<f32>, temperature: f32) -> Result<Tensor<f32>> {
    221
    1.36k
        if temperature <= 0.0 {
    222
    3
            return Err(RealizarError::InvalidShape {
    223
    3
                reason: "Temperature must be positive".to_string(),
    224
    3
            });
    225
    1.35k
        }
    226
    227
    1.35k
        if (temperature - 1.0).abs() < 1e-6 {
    228
            // No scaling needed
    229
    540
            return Ok(logits.clone());
    230
    819
        }
    231
    232
    819
        let data = logits.data();
    233
    80.8k
        let scaled819: Vec<f32>819 = data819.iter819().map819(|&x| x / temperature).collect819();
    234
    819
        Tensor::from_vec(logits.shape().to_vec(), scaled)
    235
    1.36k
    }
    236
    237
    /// Greedy sampling: select the token with highest probability
    238
    ///
    239
    /// # Arguments
    240
    ///
    241
    /// * `logits` - Logits for the vocabulary
    242
    ///
    243
    /// # Returns
    244
    ///
    245
    /// Index of the selected token
    246
    ///
    247
    /// # Errors
    248
    ///
    249
    /// Returns error if logits are empty
    250
    1.07k
    pub fn sample_greedy(logits: &Tensor<f32>) -> Result<usize> {
    251
    1.07k
        let data = logits.data();
    252
    1.07k
        if data.is_empty() {
    253
    0
            return Err(RealizarError::InvalidShape {
    254
    0
                reason: "Logits cannot be empty".to_string(),
    255
    0
            });
    256
    1.07k
        }
    257
    258
    1.07k
        let mut max_idx = 0;
    259
    1.07k
        let mut max_val = data[0];
    260
    301k
        for (i, &val) in data1.07k.iter1.07k().enumerate1.07k().skip1.07k(1) {
    261
    301k
            if val > max_val {
    262
    28
                max_val = val;
    263
    28
                max_idx = i;
    264
    301k
            }
    265
        }
    266
    267
    1.07k
        Ok(max_idx)
    268
    1.07k
    }
    269
    270
    /// Top-k sampling: sample from the k most probable tokens
    271
    ///
    272
    /// # Arguments
    273
    ///
    274
    /// * `logits` - Logits for the vocabulary
    275
    /// * `k` - Number of top tokens to consider
    276
    /// * `rng_value` - Random value in [0, 1) for sampling
    277
    ///
    278
    /// # Returns
    279
    ///
    280
    /// Index of the selected token
    281
    ///
    282
    /// # Errors
    283
    ///
    284
    /// Returns error if k is 0 or logits are empty
    285
    20
    pub fn sample_top_k(logits: &Tensor<f32>, k: usize, rng_value: f32) -> Result<usize> {
    286
    20
        let data = logits.data();
    287
    20
        if data.is_empty() {
    288
    0
            return Err(RealizarError::InvalidShape {
    289
    0
                reason: "Logits cannot be empty".to_string(),
    290
    0
            });
    291
    20
        }
    292
    20
        if k == 0 {
    293
    1
            return Err(RealizarError::InvalidShape {
    294
    1
                reason: "k must be > 0".to_string(),
    295
    1
            });
    296
    19
        }
    297
    298
        // Create (index, logit) pairs and sort by logit descending
    299
    19
        let mut indexed: Vec<(usize, f32)> = data.iter().copied().enumerate().collect();
    300
    1.16k
        indexed19.sort_by19(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    301
    302
        // Take top k
    303
    19
        let top_k: Vec<(usize, f32)> = indexed.into_iter().take(k.min(data.len())).collect();
    304
    305
        // Convert to probabilities and sample
    306
    19
        let probs = logits_to_probs(&top_k);
    307
    19
        let indices: Vec<usize> = top_k.iter().map(|(idx, _)| *idx).collect();
    308
    19
        Ok(sample_from_distribution(&probs, &indices, rng_value))
    309
    20
    }
    310
    311
    /// Top-p (nucleus) sampling: sample from tokens with cumulative probability ≤ p
    312
    ///
    313
    /// # Arguments
    314
    ///
    315
    /// * `logits` - Logits for the vocabulary
    316
    /// * `p` - Cumulative probability threshold
    317
    /// * `rng_value` - Random value in [0, 1) for sampling
    318
    ///
    319
    /// # Returns
    320
    ///
    321
    /// Index of the selected token
    322
    ///
    323
    /// # Errors
    324
    ///
    325
    /// Returns error if p is not in (0, 1] or logits are empty
    326
    274
    pub fn sample_top_p(logits: &Tensor<f32>, p: f32, rng_value: f32) -> Result<usize> {
    327
    274
        let data = logits.data();
    328
    274
        if data.is_empty() {
    329
    0
            return Err(RealizarError::InvalidShape {
    330
    0
                reason: "Logits cannot be empty".to_string(),
    331
    0
            });
    332
    274
        }
    333
    274
        if p <= 0.0 || p > 1.0272 {
    334
    3
            return Err(RealizarError::InvalidShape {
    335
    3
                reason: "p must be in (0, 1]".to_string(),
    336
    3
            });
    337
    271
        }
    338
    339
        // Convert logits to probabilities
    340
    271
        let probs_tensor = softmax(logits)?0;
    341
    271
        let probs = probs_tensor.data();
    342
    343
        // Create (index, prob) pairs and sort by prob descending
    344
    271
        let mut indexed: Vec<(usize, f32)> = probs.iter().copied().enumerate().collect();
    345
    26.4k
        indexed271.sort_by271(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    346
    347
        // Build nucleus (cumulative probability <= p)
    348
    271
        let nucleus = build_nucleus(&indexed, p);
    349
    350
        // Renormalize and sample
    351
    271
        let nucleus_sum: f32 = nucleus.iter().map(|(_, prob)| prob).sum();
    352
    24.3k
        let normalized_probs271: Vec<f32>271 = nucleus.iter()271.map271(|(_, prob)| prob / nucleus_sum).collect271();
    353
    271
        let indices: Vec<usize> = nucleus.iter().map(|(idx, _)| *idx).collect();
    354
    355
    271
        Ok(sample_from_distribution(
    356
    271
            &normalized_probs,
    357
    271
            &indices,
    358
    271
            rng_value,
    359
    271
        ))
    360
    274
    }
    361
    362
    /// Sample a token based on the sampling strategy
    363
    ///
    364
    /// # Arguments
    365
    ///
    366
    /// * `logits` - Logits for the vocabulary
    367
    /// * `config` - Generation configuration
    368
    /// * `rng_value` - Random value in [0, 1) for sampling (ignored for greedy)
    369
    ///
    370
    /// # Returns
    371
    ///
    372
    /// Index of the selected token
    373
    ///
    374
    /// # Errors
    375
    ///
    376
    /// Returns error if temperature is invalid or sampling fails
    377
    1.34k
    pub fn sample_token(
    378
    1.34k
        logits: &Tensor<f32>,
    379
    1.34k
        config: &GenerationConfig,
    380
    1.34k
        rng_value: f32,
    381
    1.34k
    ) -> Result<usize> {
    382
        // Apply temperature
    383
    1.34k
        let scaled_logits1.34k = apply_temperature(logits, config.temperature)?1;
    384
    385
    1.34k
        match config.strategy {
    386
    1.06k
            SamplingStrategy::Greedy => sample_greedy(&scaled_logits),
    387
    15
            SamplingStrategy::TopK { k } => sample_top_k(&scaled_logits, k, rng_value),
    388
    268
            SamplingStrategy::TopP { p } => sample_top_p(&scaled_logits, p, rng_value),
    389
        }
    390
    1.34k
    }
    391
    392
    393
    // Tests extracted to tests.rs (PMAT-802)
    394
    #[cfg(test)]
    395
    #[path = "tests.rs"]
    396
    mod generate_tests;