runtime/
sampler.rs

1//! Greedy sampler implementation
2
3use ndarray::{ArrayViewD, Axis};
4use std::f32;
5
6/// Greedy sampler that selects the token with the highest probability
7pub struct GreedySampler;
8
9impl GreedySampler {
10    /// Create a new greedy sampler
11    pub fn new() -> Self {
12        Self
13    }
14    
15    /// Sample the next token using greedy decoding
16    /// 
17    /// # Arguments
18    /// * `logits` - The logits from the model (unnormalized log probabilities)
19    /// 
20    /// # Returns
21    /// The token ID with the highest probability
22    pub fn sample(&self, logits: ArrayViewD<f32>) -> usize {
23        // For greedy sampling, we simply select the token with the highest logit
24        // In a real implementation, we would:
25        // 1. Apply temperature scaling if needed
26        // 2. Apply top-k or top-p filtering if needed
27        // 3. Select the token with the highest probability
28        
29        // Find the index of the maximum value
30        let mut max_idx = 0;
31        let mut max_val = f32::NEG_INFINITY;
32        
33        // Iterate through all elements to find the maximum
34        for (i, &val) in logits.iter().enumerate() {
35            if val > max_val {
36                max_val = val;
37                max_idx = i;
38            }
39        }
40        
41        max_idx
42    }
43    
44    /// Sample multiple tokens using greedy decoding
45    /// 
46    /// # Arguments
47    /// * `logits` - The logits from the model (batch_size x vocab_size)
48    /// 
49    /// # Returns
50    /// A vector of token IDs, one for each batch item
51    pub fn sample_batch(&self, logits: ArrayViewD<f32>) -> Vec<usize> {
52        // For batch sampling, we sample each batch item independently
53        let batch_size = logits.shape()[0];
54        
55        let mut samples = Vec::with_capacity(batch_size);
56        
57        for batch_idx in 0..batch_size {
58            // Extract logits for this batch item
59            let batch_logits = logits.index_axis(Axis(0), batch_idx);
60            
61            // Sample using the regular sample method
62            let token_id = self.sample(batch_logits.view());
63            samples.push(token_id);
64        }
65        
66        samples
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use ndarray::ArrayD;
74    
75    #[test]
76    fn test_greedy_sampler_single() {
77        let sampler = GreedySampler::new();
78        
79        // Create test logits where the highest value is at index 3
80        let logits = ArrayD::from_shape_vec(vec![5], vec![0.1, 0.2, 0.3, 0.9, 0.4]).unwrap();
81        let logits_view = logits.view();
82        
83        let token_id = sampler.sample(logits_view);
84        assert_eq!(token_id, 3);
85    }
86    
87    #[test]
88    fn test_greedy_sampler_batch() {
89        let sampler = GreedySampler::new();
90        
91        // Create test logits for a batch of 2
92        let logits = ArrayD::from_shape_vec(vec![2, 3], vec![
93            0.1, 0.8, 0.3,  // First batch item - max at index 1
94            0.6, 0.2, 0.9,  // Second batch item - max at index 2
95        ]).unwrap();
96        let logits_view = logits.view();
97        
98        let token_ids = sampler.sample_batch(logits_view);
99        assert_eq!(token_ids, vec![1, 2]);
100    }
101}
runtime/sampler.rs

runtime/
sampler.rs
