rustytorch/crates/production/rtx-serving-api/src/sampling_strategies.rs

//! Advanced sampling strategies for LLM generation
//!
//! Provides comprehensive sampling algorithms including:
//! - Nucleus sampling (top-p) with dynamic threshold adjustment
//! - Top-k sampling with adaptive k selection
//! - Temperature scaling with per-token adjustment
//! - Repetition penalty with context-aware scoring
//! - Presence penalty with semantic similarity detection
//! - Custom sampling strategies with pluggable algorithms

use anyhow::{Result, anyhow};
use rand::SeedableRng;
use rand::distributions::{Distribution, WeightedIndex};
use rand::rngs::StdRng;
use serde::{Deserialize, Serialize};
use std::{cmp::Ordering, collections::HashMap};

/// Sampling configuration
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SamplingConfig {
    pub strategy: SamplingStrategy,
    pub temperature: f32,
    pub top_k: Option<usize>,
    pub top_p: Option<f32>,
    pub repetition_penalty: f32,
    pub presence_penalty: f32,
    pub frequency_penalty: f32,
    pub length_penalty: f32,
    pub diversity_penalty: f32,
    pub typical_p: Option<f32>,
    pub eta_cutoff: Option<f32>,
    pub epsilon_cutoff: Option<f32>,
    pub min_p: Option<f32>,
}

impl Default for SamplingConfig {
    fn default() -> Self {
        Self {
            strategy: SamplingStrategy::TopP,
            temperature: 1.0,
            top_k: Some(50),
            top_p: Some(0.95),
            repetition_penalty: 1.1,
            presence_penalty: 0.0,
            frequency_penalty: 0.0,
            length_penalty: 1.0,
            diversity_penalty: 0.0,
            typical_p: None,
            eta_cutoff: None,
            epsilon_cutoff: None,
            min_p: None,
        }
    }
}

/// Available sampling strategies
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum SamplingStrategy {
    Greedy,
    TopK,
    TopP,
    Temperature,
    Typical,
    Combined,
    Adaptive,
    Custom,
}

/// Token with probability and metadata
#[derive(Debug, Clone, PartialEq)]
pub struct TokenCandidate {
    pub token_id: u32,
    pub token_text: String,
    pub log_prob: f32,
    pub prob: f32,
    pub rank: usize,
    pub is_repeated: bool,
    pub frequency_count: u32,
    pub semantic_score: Option<f32>,
}

impl TokenCandidate {
    /// Create new token candidate
    #[must_use]
    pub fn new(token_id: u32, token_text: String, log_prob: f32) -> Self {
        Self {
            token_id,
            token_text,
            log_prob,
            prob: log_prob.exp(),
            rank: 0,
            is_repeated: false,
            frequency_count: 0,
            semantic_score: None,
        }
    }
}

impl PartialOrd for TokenCandidate {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        // Sort by probability (descending)
        other.prob.partial_cmp(&self.prob)
    }
}

/// Context for repetition and frequency tracking
#[derive(Debug, Clone, Default)]
pub struct SamplingContext {
    pub generated_tokens: Vec<u32>,
    pub token_frequencies: HashMap<u32, u32>,
    pub recent_ngrams: HashMap<Vec<u32>, u32>,
    pub semantic_embeddings: HashMap<u32, Vec<f32>>,
    pub position: usize,
}

impl SamplingContext {
    /// Add generated token to context
    pub fn add_token(&mut self, token_id: u32) {
        self.generated_tokens.push(token_id);
        *self.token_frequencies.entry(token_id).or_insert(0) += 1;
        self.position += 1;

        // Update n-grams for repetition detection
        self.update_ngrams(token_id);
    }

    /// Update n-gram tracking
    fn update_ngrams(&mut self, token_id: u32) {
        let _ = token_id;
        // Track 2-grams, 3-grams, and 4-grams.
        // `generated_tokens` already has the current token appended, so the
        // n-gram is simply the last `n` tokens (no need to append again).
        for n in 2..=4 {
            if self.generated_tokens.len() >= n {
                let ngram = self.generated_tokens[(self.generated_tokens.len() - n)..].to_vec();

                *self.recent_ngrams.entry(ngram).or_insert(0) += 1;
            }
        }
    }

    /// Check if token would create repetition
    #[must_use]
    pub fn would_repeat(&self, token_id: u32, n: usize) -> bool {
        if self.generated_tokens.len() < n - 1 {
            return false;
        }

        let potential_ngram = self.generated_tokens[(self.generated_tokens.len() - n + 1)..]
            .iter()
            .copied()
            .chain(std::iter::once(token_id))
            .collect::<Vec<u32>>();

        self.recent_ngrams
            .get(&potential_ngram)
            .copied()
            .unwrap_or(0)
            > 0
    }

    /// Get frequency penalty for token
    #[must_use]
    pub fn get_frequency_penalty(&self, token_id: u32, penalty: f32) -> f32 {
        let count = self.token_frequencies.get(&token_id).copied().unwrap_or(0);
        if count > 0 {
            penalty * count as f32
        } else {
            0.0
        }
    }

    /// Get presence penalty for token
    #[must_use]
    pub fn get_presence_penalty(&self, token_id: u32, penalty: f32) -> f32 {
        if self.token_frequencies.contains_key(&token_id) {
            penalty
        } else {
            0.0
        }
    }
}

/// Advanced sampler with multiple strategies
pub struct AdvancedSampler {
    config: SamplingConfig,
    rng: StdRng,
}

impl Default for AdvancedSampler {
    fn default() -> Self {
        Self::new(SamplingConfig::default())
    }
}

impl AdvancedSampler {
    /// Create new sampler with configuration
    #[must_use]
    pub fn new(config: SamplingConfig) -> Self {
        Self {
            config,
            rng: StdRng::from_entropy(),
        }
    }

    /// Sample next token from logits
    pub fn sample(
        &mut self,
        logits: &[f32],
        token_texts: &[String],
        context: &mut SamplingContext,
    ) -> Result<TokenCandidate> {
        if logits.is_empty() {
            return Err(anyhow!("Empty logits"));
        }

        // Create token candidates
        let mut candidates = self.create_candidates(logits, token_texts, context)?;

        // Apply sampling strategy
        let selected = match self.config.strategy {
            SamplingStrategy::Greedy => self.greedy_sample(&candidates)?,
            SamplingStrategy::TopK => self.top_k_sample(&mut candidates, context)?,
            SamplingStrategy::TopP => self.nucleus_sample(&mut candidates, context)?,
            SamplingStrategy::Temperature => self.temperature_sample(&mut candidates, context)?,
            SamplingStrategy::Typical => self.typical_sampling(&mut candidates, context)?,
            SamplingStrategy::Combined => self.combined_sampling(&mut candidates, context)?,
            SamplingStrategy::Adaptive => self.adaptive_sampling(&mut candidates, context)?,
            SamplingStrategy::Custom => self.custom_sampling(&mut candidates, context)?,
        };

        // Update context
        context.add_token(selected.token_id);

        Ok(selected)
    }

    /// Create token candidates from logits
    fn create_candidates(
        &self,
        logits: &[f32],
        token_texts: &[String],
        context: &SamplingContext,
    ) -> Result<Vec<TokenCandidate>> {
        let mut candidates = Vec::new();

        for (i, (&logit, token_text)) in logits.iter().zip(token_texts).enumerate() {
            let mut candidate = TokenCandidate::new(i as u32, token_text.clone(), logit);

            // Apply penalties
            candidate = self.apply_penalties(candidate, context);

            candidates.push(candidate);
        }

        // Sort by probability and assign ranks
        candidates.sort_by(|a, b| b.prob.total_cmp(&a.prob));
        for (rank, candidate) in candidates.iter_mut().enumerate() {
            candidate.rank = rank;
        }

        Ok(candidates)
    }

    /// Apply various penalties to token candidate
    fn apply_penalties(
        &self,
        mut candidate: TokenCandidate,
        context: &SamplingContext,
    ) -> TokenCandidate {
        let mut adjusted_log_prob = candidate.log_prob;

        // Repetition penalty
        if context.token_frequencies.contains_key(&candidate.token_id) {
            let penalty = if self.config.repetition_penalty > 1.0 {
                -self.config.repetition_penalty.ln()
            } else {
                self.config.repetition_penalty.ln()
            };
            adjusted_log_prob += penalty;
            candidate.is_repeated = true;
        }

        // Frequency penalty
        let freq_penalty =
            context.get_frequency_penalty(candidate.token_id, self.config.frequency_penalty);
        adjusted_log_prob -= freq_penalty;

        // Presence penalty
        let presence_penalty =
            context.get_presence_penalty(candidate.token_id, self.config.presence_penalty);
        adjusted_log_prob -= presence_penalty;

        candidate.log_prob = adjusted_log_prob;
        candidate.prob = adjusted_log_prob.exp();
        candidate.frequency_count = context
            .token_frequencies
            .get(&candidate.token_id)
            .copied()
            .unwrap_or(0);

        candidate
    }

    /// Greedy sampling - select highest probability token
    fn greedy_sample(&self, candidates: &[TokenCandidate]) -> Result<TokenCandidate> {
        candidates
            .first()
            .ok_or_else(|| anyhow!("No candidates available"))
            .cloned()
    }

    /// Top-k sampling
    fn top_k_sample(
        &mut self,
        candidates: &mut [TokenCandidate],
        _context: &SamplingContext,
    ) -> Result<TokenCandidate> {
        let k = self.config.top_k.unwrap_or(candidates.len());
        let k = k.min(candidates.len());

        // Take top-k candidates
        let top_candidates = &candidates[..k];

        // Apply temperature scaling
        let scaled_probs = self.apply_temperature(top_candidates);

        // Sample from distribution
        self.sample_from_distribution(&scaled_probs, top_candidates)
    }

    /// Nucleus (top-p) sampling
    fn nucleus_sample(
        &mut self,
        candidates: &mut [TokenCandidate],
        _context: &SamplingContext,
    ) -> Result<TokenCandidate> {
        let p = self.config.top_p.unwrap_or(1.0);

        // Find nucleus (cumulative probability >= p)
        let mut cumulative_prob = 0.0;
        let mut nucleus_size = 0;

        for (i, candidate) in candidates.iter().enumerate() {
            cumulative_prob += candidate.prob;
            nucleus_size = i + 1;

            if cumulative_prob >= p {
                break;
            }
        }

        let nucleus = &candidates[..nucleus_size];

        // Apply temperature scaling
        let scaled_probs = self.apply_temperature(nucleus);

        // Sample from nucleus
        self.sample_from_distribution(&scaled_probs, nucleus)
    }

    /// Temperature sampling
    fn temperature_sample(
        &mut self,
        candidates: &mut [TokenCandidate],
        _context: &SamplingContext,
    ) -> Result<TokenCandidate> {
        let scaled_probs = self.apply_temperature(candidates);
        self.sample_from_distribution(&scaled_probs, candidates)
    }

    /// Typical sampling (entropy-based)
    fn typical_sampling(
        &mut self,
        candidates: &mut [TokenCandidate],
        _context: &SamplingContext,
    ) -> Result<TokenCandidate> {
        let typical_p = self.config.typical_p.unwrap_or(0.95);

        // Calculate entropy
        let entropy = -candidates.iter().map(|c| c.prob * c.prob.ln()).sum::<f32>();

        // Filter candidates by typical probability
        let mut typical_candidates = Vec::new();
        for candidate in candidates.iter() {
            let surprisal = -candidate.prob.ln();
            let typical_score = (surprisal - entropy).abs();

            if typical_score <= typical_p {
                typical_candidates.push((candidate.clone(), typical_score));
            }
        }

        if typical_candidates.is_empty() {
            return self.greedy_sample(candidates);
        }

        // Sort by typical score (lower is more typical)
        typical_candidates.sort_by(|a, b| a.1.total_cmp(&b.1));

        let candidates_only: Vec<_> = typical_candidates.iter().map(|(c, _)| c.clone()).collect();
        let scaled_probs = self.apply_temperature(&candidates_only);

        self.sample_from_distribution(&scaled_probs, &candidates_only)
    }

    /// Combined sampling strategy
    fn combined_sampling(
        &mut self,
        candidates: &mut [TokenCandidate],
        _context: &SamplingContext,
    ) -> Result<TokenCandidate> {
        // Apply multiple filters in sequence

        // 1. Top-k filtering
        let k = self.config.top_k.unwrap_or(candidates.len());
        let k = k.min(candidates.len());
        let mut filtered = candidates[..k].to_vec();

        // 2. Top-p filtering
        if let Some(p) = self.config.top_p {
            let mut cumulative_prob = 0.0;
            let mut nucleus_size = 0;

            for (i, candidate) in filtered.iter().enumerate() {
                cumulative_prob += candidate.prob;
                nucleus_size = i + 1;

                if cumulative_prob >= p {
                    break;
                }
            }

            filtered.truncate(nucleus_size);
        }

        // 3. Min-p filtering
        if let Some(min_p) = self.config.min_p {
            let max_prob = filtered.first().map_or(0.0, |c| c.prob);
            let threshold = min_p * max_prob;
            filtered.retain(|c| c.prob >= threshold);
        }

        if filtered.is_empty() {
            return self.greedy_sample(candidates);
        }

        // Apply temperature and sample
        let scaled_probs = self.apply_temperature(&filtered);
        self.sample_from_distribution(&scaled_probs, &filtered)
    }

    /// Adaptive sampling based on context
    fn adaptive_sampling(
        &mut self,
        candidates: &mut [TokenCandidate],
        context: &SamplingContext,
    ) -> Result<TokenCandidate> {
        // Adapt parameters based on generation context
        let mut adaptive_config = self.config.clone();

        // Reduce temperature as sequence gets longer (more conservative)
        let length_factor = 1.0 - (context.position as f32 / 1000.0).min(0.5);
        adaptive_config.temperature *= length_factor;

        // Increase repetition penalty if repetition detected
        let repetition_rate = context
            .recent_ngrams
            .values()
            .filter(|&&count| count > 1)
            .count() as f32
            / context.recent_ngrams.len().max(1) as f32;
        if repetition_rate > 0.3 {
            adaptive_config.repetition_penalty *= 1.5;
        }

        // Temporarily update config
        let original_config = self.config.clone();
        self.config = adaptive_config;

        let result = self.combined_sampling(candidates, context);

        // Restore original config
        self.config = original_config;

        result
    }

    /// Custom sampling strategy (placeholder for extensibility)
    fn custom_sampling(
        &mut self,
        candidates: &mut [TokenCandidate],
        context: &SamplingContext,
    ) -> Result<TokenCandidate> {
        // Default to combined sampling for now
        // This can be extended to support custom algorithms
        self.combined_sampling(candidates, context)
    }

    /// Apply temperature scaling to probabilities
    fn apply_temperature(&self, candidates: &[TokenCandidate]) -> Vec<f32> {
        let temperature = self.config.temperature.max(0.01); // Prevent division by zero

        let scaled_logits: Vec<f32> = candidates
            .iter()
            .map(|c| c.log_prob / temperature)
            .collect();

        // Compute softmax
        let max_logit = scaled_logits
            .iter()
            .fold(f32::NEG_INFINITY, |a, &b| a.max(b));
        let exp_logits: Vec<f32> = scaled_logits
            .iter()
            .map(|&logit| (logit - max_logit).exp())
            .collect();

        let sum_exp: f32 = exp_logits.iter().sum();

        exp_logits
            .iter()
            .map(|&exp_logit| exp_logit / sum_exp)
            .collect()
    }

    /// Sample from probability distribution
    fn sample_from_distribution(
        &mut self,
        probabilities: &[f32],
        candidates: &[TokenCandidate],
    ) -> Result<TokenCandidate> {
        if probabilities.is_empty() || candidates.is_empty() {
            return Err(anyhow!("Empty probabilities or candidates"));
        }

        // Create weighted distribution
        let dist = WeightedIndex::new(probabilities)
            .map_err(|e| anyhow!("Failed to create weighted distribution: {e}"))?;

        let index = dist.sample(&mut self.rng);

        Ok(candidates[index].clone())
    }

    /// Update sampling configuration
    pub fn update_config(&mut self, config: SamplingConfig) {
        self.config = config;
    }

    /// Get current configuration
    #[must_use]
    pub fn get_config(&self) -> &SamplingConfig {
        &self.config
    }
}

/// Batch sampler for efficient parallel sampling
pub struct BatchSampler {
    samplers: Vec<AdvancedSampler>,
}

impl BatchSampler {
    /// Create new batch sampler
    #[must_use]
    pub fn new(batch_size: usize, config: SamplingConfig) -> Self {
        let samplers = (0..batch_size)
            .map(|_| AdvancedSampler::new(config.clone()))
            .collect();

        Self { samplers }
    }

    /// Sample from multiple sequences in parallel
    pub fn batch_sample(
        &mut self,
        batch_logits: &[Vec<f32>],
        batch_token_texts: &[Vec<String>],
        batch_contexts: &mut [SamplingContext],
    ) -> Result<Vec<TokenCandidate>> {
        if batch_logits.len() != batch_contexts.len()
            || batch_logits.len() != batch_token_texts.len()
            || batch_logits.len() != self.samplers.len()
        {
            return Err(anyhow!("Batch size mismatch"));
        }

        let mut results = Vec::new();

        for (i, ((logits, token_texts), context)) in batch_logits
            .iter()
            .zip(batch_token_texts)
            .zip(batch_contexts.iter_mut())
            .enumerate()
        {
            let sampled = self.samplers[i].sample(logits, token_texts, context)?;
            results.push(sampled);
        }

        Ok(results)
    }
}

/// Sampling statistics for analysis and tuning
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SamplingStats {
    pub total_samples: u64,
    pub strategy_usage: HashMap<SamplingStrategy, u64>,
    pub average_entropy: f32,
    pub repetition_rate: f32,
    pub diversity_score: f32,
    pub temperature_usage: Vec<f32>,
    pub top_k_usage: Vec<usize>,
    pub top_p_usage: Vec<f32>,
}

impl Default for SamplingStats {
    fn default() -> Self {
        Self {
            total_samples: 0,
            strategy_usage: HashMap::new(),
            average_entropy: 0.0,
            repetition_rate: 0.0,
            diversity_score: 0.0,
            temperature_usage: Vec::new(),
            top_k_usage: Vec::new(),
            top_p_usage: Vec::new(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_token_candidate_creation() {
        let candidate = TokenCandidate::new(1, "hello".to_string(), 0.5);
        assert_eq!(candidate.token_id, 1);
        assert_eq!(candidate.token_text, "hello");
        assert_eq!(candidate.log_prob, 0.5);
        assert!((candidate.prob - 0.5_f32.exp()).abs() < 1e-6);
    }

    #[test]
    fn test_sampling_context() {
        let mut context = SamplingContext::default();

        context.add_token(1);
        context.add_token(2);
        context.add_token(1); // Repeat

        assert_eq!(context.generated_tokens, vec![1, 2, 1]);
        assert_eq!(context.token_frequencies[&1], 2);
        assert_eq!(context.token_frequencies[&2], 1);

        // Test repetition detection
        assert!(context.would_repeat(2, 2)); // Would create [2, 1, 2] bigram
    }

    #[test]
    fn test_sampling_config_default() {
        let config = SamplingConfig::default();
        assert_eq!(config.strategy, SamplingStrategy::TopP);
        assert_eq!(config.temperature, 1.0);
        assert_eq!(config.repetition_penalty, 1.1);
    }

    #[test]
    fn test_advanced_sampler_creation() {
        let config = SamplingConfig::default();
        let sampler = AdvancedSampler::new(config.clone());
        assert_eq!(*sampler.get_config(), config);
    }

    #[test]
    fn test_greedy_sampling() {
        let mut sampler = AdvancedSampler::new(SamplingConfig {
            strategy: SamplingStrategy::Greedy,
            ..Default::default()
        });

        let logits = vec![0.1, 0.8, 0.3]; // Middle token has highest probability
        let token_texts = vec!["a".to_string(), "b".to_string(), "c".to_string()];
        let mut context = SamplingContext::default();

        let result = sampler.sample(&logits, &token_texts, &mut context).unwrap();
        assert_eq!(result.token_id, 1); // Highest probability token
        assert_eq!(result.token_text, "b");
    }

    #[test]
    fn test_temperature_scaling() {
        let sampler = AdvancedSampler::new(SamplingConfig {
            temperature: 2.0,
            ..Default::default()
        });

        let candidates = vec![
            TokenCandidate::new(0, "a".to_string(), 0.8),
            TokenCandidate::new(1, "b".to_string(), 0.2),
        ];

        let scaled_probs = sampler.apply_temperature(&candidates);

        // Baseline softmax with no temperature scaling (temperature = 1.0),
        // to compare against the effect of the higher temperature above.
        let baseline_sampler = AdvancedSampler::new(SamplingConfig {
            temperature: 1.0,
            ..Default::default()
        });
        let baseline_probs = baseline_sampler.apply_temperature(&candidates);

        // With higher temperature, probabilities should be more uniform:
        // the top candidate's probability shrinks and the runner-up's grows.
        assert!(scaled_probs[0] < baseline_probs[0]);
        assert!(scaled_probs[1] > baseline_probs[1]);
    }

    #[test]
    fn test_batch_sampler() {
        let config = SamplingConfig::default();
        let mut batch_sampler = BatchSampler::new(2, config);

        let batch_logits = vec![vec![0.1, 0.8, 0.3], vec![0.5, 0.2, 0.7]];

        let batch_token_texts = vec![
            vec!["a".to_string(), "b".to_string(), "c".to_string()],
            vec!["x".to_string(), "y".to_string(), "z".to_string()],
        ];

        let mut batch_contexts = vec![SamplingContext::default(), SamplingContext::default()];

        let results = batch_sampler
            .batch_sample(&batch_logits, &batch_token_texts, &mut batch_contexts)
            .unwrap();

        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_penalty_application() {
        let sampler = AdvancedSampler::new(SamplingConfig {
            repetition_penalty: 2.0,
            frequency_penalty: 0.5,
            presence_penalty: 0.3,
            ..Default::default()
        });

        let mut context = SamplingContext::default();
        context.add_token(1); // Add token to create repetition

        let candidate = TokenCandidate::new(1, "test".to_string(), 0.5);
        let penalized = sampler.apply_penalties(candidate, &context);

        // Should have lower probability due to penalties
        assert!(penalized.prob < 0.5_f32.exp());
        assert!(penalized.is_repeated);
    }

    #[test]
    fn test_nucleus_sampling_nucleus_selection() {
        // Test that nucleus sampling correctly selects the nucleus
        let candidates = vec![
            TokenCandidate {
                prob: 0.6,
                ..TokenCandidate::new(0, "a".to_string(), 0.6_f32.ln())
            },
            TokenCandidate {
                prob: 0.3,
                ..TokenCandidate::new(1, "b".to_string(), 0.3_f32.ln())
            },
            TokenCandidate {
                prob: 0.08,
                ..TokenCandidate::new(2, "c".to_string(), 0.08_f32.ln())
            },
            TokenCandidate {
                prob: 0.02,
                ..TokenCandidate::new(3, "d".to_string(), 0.02_f32.ln())
            },
        ];

        // With p=0.9, nucleus should include first 3 tokens (0.6 + 0.3 + 0.08 = 0.98 > 0.9)
        let p = 0.9;
        let mut cumulative_prob = 0.0;
        let mut nucleus_size = 0;

        for (i, candidate) in candidates.iter().enumerate() {
            cumulative_prob += candidate.prob;
            nucleus_size = i + 1;

            // Use a small epsilon so that floating-point rounding in the
            // cumulative sum doesn't cause the nucleus to be cut off one
            // token earlier than the intended threshold crossing.
            if cumulative_prob >= p + 1e-6 {
                break;
            }
        }

        assert_eq!(nucleus_size, 3);
    }
}