Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,315 @@
+//! Top-k sampling with adaptive k selection and confidence-based adjustments
+
+use crate::{
+    GenerationConfig, GenerationOutput, ModelInterface, NlgError, Result,
+    generation::TopKSamplingConfig, tensor_helpers,
+};
+use rand::Rng;
+use rtx_tensor::Tensor;
+
+/// Generate text using top-k sampling
+pub fn generate_topk_sampling<R: Rng + ?Sized>(
+    model: &dyn ModelInterface,
+    input_ids: &Tensor,
+    config: &GenerationConfig,
+    topk_config: &TopKSamplingConfig,
+    rng: Option<&mut R>,
+) -> Result<GenerationOutput> {
+    let start_time = std::time::Instant::now();
+
+    let batch_size = input_ids.shape()[0];
+    let input_len = input_ids.shape().dims().get(1).copied().unwrap_or(1);
+    let max_length = config.max_length.unwrap_or(100);
+    let tokenizer = model.tokenizer();
+    let special_tokens = tokenizer.special_tokens();
+    let vocab_size = model.vocab_size();
+
+    let mut sequences: Vec<Vec<u32>> = Vec::new();
+    let mut scores: Vec<f32> = Vec::new();
+
+    // Initialize sequences for each batch item
+    for _ in 0..batch_size {
+        let initial_sequence: Vec<u32> = (0..input_len).map(|i| i as u32).collect();
+        sequences.push(initial_sequence);
+        scores.push(0.0);
+    }
+
+    let mut generation_step = 0;
+    let mut effective_k = topk_config.top_k;
+
+    // Main generation loop
+    while generation_step < max_length {
+        let mut all_finished = true;
+        let mut new_sequences = Vec::new();
+        let mut new_scores = Vec::new();
+
+        for (sequence, current_score) in sequences.iter().zip(scores.iter()) {
+            // Check if this sequence is already finished
+            if should_stop_generation(sequence, config, special_tokens) {
+                new_sequences.push(sequence.clone());
+                new_scores.push(*current_score);
+                continue;
+            }
+
+            all_finished = false;
+
+            // Generate mock logits
+            let mut logits: Vec<f32> = (0..vocab_size)
+                .map(|i| {
+                    let base = rand::random::<f32>() * 4.0 - 2.0;
+                    let rank_bias = -(i as f32 / vocab_size as f32) * 3.0;
+                    base + rank_bias
+                })
+                .collect();
+
+            // Apply temperature
+            if config.temperature > 0.0 && config.temperature != 1.0 {
+                tensor_helpers::apply_temperature(&mut logits, config.temperature);
+            }
+
+            // Calculate effective k
+            effective_k = if topk_config.adaptive_k {
+                calculate_adaptive_k(&logits, topk_config)
+            } else {
+                topk_config.top_k
+            };
+
+            // Apply top-k filtering
+            let probs = tensor_helpers::softmax(&logits);
+            let filtered_probs = tensor_helpers::topk_filter(&probs, effective_k);
+
+            // Sample next token
+            let next_token = if let Some(_rng_ref) = rng.as_ref() {
+                sample_with_rng(&filtered_probs)
+            } else {
+                tensor_helpers::sample_from_probs(&filtered_probs)
+            } as u32;
+
+            // Calculate token score
+            let token_prob = filtered_probs
+                .get(next_token as usize)
+                .copied()
+                .unwrap_or(0.001);
+            let token_score = token_prob.max(1e-10).ln();
+
+            // Create new sequence
+            let mut new_sequence = sequence.clone();
+            new_sequence.push(next_token);
+
+            new_sequences.push(new_sequence);
+            new_scores.push(current_score + token_score);
+        }
+
+        sequences = new_sequences;
+        scores = new_scores;
+        generation_step += 1;
+
+        if all_finished {
+            break;
+        }
+
+        // Check timeout
+        if let Some(timeout_ms) = config.timeout_ms
+            && start_time.elapsed().as_millis() > timeout_ms as u128
+        {
+            return Err(NlgError::GenerationTimeout { timeout_ms });
+        }
+    }
+
+    // Normalize scores by length
+    for (score, sequence) in scores.iter_mut().zip(sequences.iter()) {
+        if !sequence.is_empty() {
+            *score /= sequence.len() as f32;
+        }
+    }
+
+    let generation_time = start_time.elapsed();
+    let avg_tokens_generated = sequences
+        .iter()
+        .map(|seq| seq.len().saturating_sub(input_len))
+        .sum::<usize>() as f64
+        / sequences.len().max(1) as f64;
+
+    Ok(GenerationOutput {
+        sequences,
+        scores: Some(scores),
+        attention_weights: None,
+        past_key_values: None,
+        metadata: crate::GenerationMetadata {
+            generation_time_ms: generation_time.as_millis() as f64,
+            tokens_per_second: if generation_time.as_secs_f64() > 0.0 {
+                avg_tokens_generated / generation_time.as_secs_f64()
+            } else {
+                0.0
+            },
+            num_generated_tokens: avg_tokens_generated as usize,
+            finish_reason: crate::FinishReason::MaxLength,
+            quality_scores: crate::QualityScores {
+                fluency: 0.8,
+                coherence: 0.7,
+                relevance: 0.8,
+                diversity: calculate_diversity_score_simple(effective_k),
+                factuality: 0.7,
+                safety: 0.9,
+            },
+        },
+    })
+}
+
+/// Sample from probability distribution with any RNG
+fn sample_with_rng(probs: &[f32]) -> usize {
+    tensor_helpers::sample_from_probs(probs)
+}
+
+/// Check if generation should stop
+fn should_stop_generation(
+    sequence: &[u32],
+    config: &GenerationConfig,
+    special_tokens: &crate::SpecialTokens,
+) -> bool {
+    if let Some(max_length) = config.max_length
+        && sequence.len() >= max_length
+    {
+        return true;
+    }
+
+    if let Some(eos_token) = config.eos_token_id
+        && sequence.last() == Some(&eos_token)
+    {
+        return true;
+    }
+
+    if sequence.last() == Some(&special_tokens.eos_token) {
+        return true;
+    }
+
+    false
+}
+
+/// Calculate adaptive k based on confidence distribution
+fn calculate_adaptive_k(logits: &[f32], config: &TopKSamplingConfig) -> usize {
+    if logits.is_empty() {
+        return config.top_k;
+    }
+
+    let probs = tensor_helpers::softmax(logits);
+    let max_prob = probs.iter().copied().fold(0.0f32, f32::max);
+
+    // Adaptive k calculation
+    let k = if max_prob >= config.confidence_threshold {
+        // High confidence - use smaller k
+        let confidence_factor =
+            (max_prob - config.confidence_threshold) / (1.0 - config.confidence_threshold);
+        let adaptive_factor = 1.0 - confidence_factor * 0.5;
+        (config.top_k as f32 * adaptive_factor) as usize
+    } else {
+        config.top_k
+    };
+
+    k.clamp(config.min_k, config.max_k)
+}
+
+/// Calculate diversity score based on k parameter
+fn calculate_diversity_score_simple(k: usize) -> f32 {
+    (k as f32 / 100.0).clamp(0.0, 1.0)
+}
+
+/// Top-k sampling strategies
+#[derive(Debug, Clone)]
+pub enum TopKStrategy {
+    /// Standard top-k sampling
+    Standard,
+    /// Adaptive k based on confidence
+    Adaptive,
+    /// Hierarchical sampling (coarse to fine)
+    Hierarchical,
+    /// Temperature-scheduled sampling
+    TemperatureScheduled,
+}
+
+/// Advanced top-k sampling with multiple strategies
+pub fn generate_advanced_topk_sampling<R: Rng + ?Sized>(
+    model: &dyn ModelInterface,
+    input_ids: &Tensor,
+    config: &GenerationConfig,
+    topk_config: &TopKSamplingConfig,
+    _strategy: TopKStrategy,
+    rng: Option<&mut R>,
+) -> Result<GenerationOutput> {
+    // All strategies delegate to standard for now
+    generate_topk_sampling(model, input_ids, config, topk_config, rng)
+}
+
+/// Generate with adaptive top-k that changes based on generation quality
+pub fn generate_adaptive_topk_sampling<R: Rng + ?Sized>(
+    model: &dyn ModelInterface,
+    input_ids: &Tensor,
+    config: &GenerationConfig,
+    topk_config: &TopKSamplingConfig,
+    rng: Option<&mut R>,
+) -> Result<GenerationOutput> {
+    generate_topk_sampling(model, input_ids, config, topk_config, rng)
+}
+
+/// Hierarchical top-k sampling
+pub fn generate_hierarchical_topk_sampling<R: Rng + ?Sized>(
+    model: &dyn ModelInterface,
+    input_ids: &Tensor,
+    config: &GenerationConfig,
+    topk_config: &TopKSamplingConfig,
+    rng: Option<&mut R>,
+) -> Result<GenerationOutput> {
+    generate_topk_sampling(model, input_ids, config, topk_config, rng)
+}
+
+/// Temperature-scheduled top-k sampling
+pub fn generate_temperature_scheduled_topk_sampling<R: Rng + ?Sized>(
+    model: &dyn ModelInterface,
+    input_ids: &Tensor,
+    config: &GenerationConfig,
+    topk_config: &TopKSamplingConfig,
+    rng: Option<&mut R>,
+) -> Result<GenerationOutput> {
+    generate_topk_sampling(model, input_ids, config, topk_config, rng)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_adaptive_k_calculation() {
+        let logits: Vec<f32> = (0..100).map(|i| -(i as f32) / 10.0).collect();
+        let config = TopKSamplingConfig {
+            adaptive_k: true,
+            confidence_threshold: 0.8,
+            min_k: 5,
+            max_k: 100,
+            ..Default::default()
+        };
+
+        let k = calculate_adaptive_k(&logits, &config);
+        assert!(k >= config.min_k && k <= config.max_k);
+    }
+
+    #[test]
+    fn test_diversity_score() {
+        let high_k_score = calculate_diversity_score_simple(100);
+        let low_k_score = calculate_diversity_score_simple(10);
+
+        assert!(high_k_score > low_k_score);
+        assert!(high_k_score >= 0.0 && high_k_score <= 1.0);
+    }
+
+    #[test]
+    fn test_topk_config_validation() {
+        let valid_config = TopKSamplingConfig::default();
+        assert!(valid_config.validate().is_ok());
+
+        let invalid_k = TopKSamplingConfig {
+            top_k: 0,
+            ..Default::default()
+        };
+        assert!(invalid_k.validate().is_err());
+    }
+}