xserv/crates/xserv-model/src/sampling.rs

use half::bf16;
use rand::Rng;
use xserv_tensor::{DType, Device, Tensor};

#[derive(Clone)]
pub struct SamplingParams {
    pub temperature: f32,
    pub top_k: usize,
    pub top_p: f32,
}

impl Default for SamplingParams {
    fn default() -> Self {
        Self { temperature: 0.0, top_k: 0, top_p: 1.0 }
    }
}

/// Sample a token from logits with shape [seq_len, vocab_size].
/// Uses the last position's logits. Handles both F32 and BF16 dtypes.
pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
    assert_eq!(logits.ndim(), 2);
    let vocab_size = logits.shape()[1];
    let seq_len = logits.shape()[0];
    let logits_cpu = logits.to_device(Device::Cpu);

    // Extract last row as f32
    let last_row: Vec<f32> = match logits.dtype() {
        DType::F32 => {
            let data = logits_cpu.as_slice::<f32>();
            data[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec()
        }
        DType::BF16 => {
            let data = logits_cpu.as_slice::<bf16>();
            data[(seq_len - 1) * vocab_size..seq_len * vocab_size]
                .iter()
                .map(|v| v.to_f32())
                .collect()
        }
        _ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
    };

    // Greedy
    if params.temperature == 0.0 {
        return argmax(&last_row);
    }

    // Apply temperature
    let mut logits_f32: Vec<f32> = last_row.iter().map(|v| v / params.temperature).collect();

    // Top-k filtering
    if params.top_k > 0 && params.top_k < vocab_size {
        let mut indices: Vec<usize> = (0..vocab_size).collect();
        indices.select_nth_unstable_by(params.top_k, |&a, &b| {
            logits_f32[b].partial_cmp(&logits_f32[a]).unwrap()
        });
        // Everything after top_k should be masked
        for &i in &indices[params.top_k..] {
            logits_f32[i] = f32::NEG_INFINITY;
        }
    }

    // Top-p (nucleus) filtering
    if params.top_p < 1.0 {
        // Sort indices by descending logit value
        let mut indices: Vec<usize> = (0..vocab_size).collect();
        indices.sort_unstable_by(|&a, &b| logits_f32[b].partial_cmp(&logits_f32[a]).unwrap());

        // Compute softmax probabilities for the sorted order
        let max_val = logits_f32[indices[0]];
        let sorted_probs: Vec<f32> = indices
            .iter()
            .map(|&i| (logits_f32[i] - max_val).exp())
            .collect();
        let sum: f32 = sorted_probs.iter().sum();
        let sorted_probs: Vec<f32> = sorted_probs.iter().map(|v| v / sum).collect();

        // Cumulative sum, find cutoff
        let mut cumsum = 0.0f32;
        let mut cutoff = indices.len();
        for (rank, &prob) in sorted_probs.iter().enumerate() {
            cumsum += prob;
            if cumsum > params.top_p {
                cutoff = rank + 1; // keep at least this many
                break;
            }
        }

        // Mask everything beyond cutoff
        for &i in &indices[cutoff..] {
            logits_f32[i] = f32::NEG_INFINITY;
        }
    }

    // Softmax
    let max_val = logits_f32.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = logits_f32.iter().map(|v| (v - max_val).exp()).collect();
    let sum: f32 = exps.iter().sum();
    let probs: Vec<f32> = exps.iter().map(|v| v / sum).collect();

    // Weighted random sampling
    let mut rng = rand::thread_rng();
    let r: f32 = rng.r#gen();
    let mut cumsum = 0.0f32;
    for (i, &p) in probs.iter().enumerate() {
        cumsum += p;
        if cumsum > r {
            return i as u32;
        }
    }

    // Fallback (rounding edge case)
    (vocab_size - 1) as u32
}

fn argmax(data: &[f32]) -> u32 {
    data.iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
        .map(|(i, _)| i as u32)
        .unwrap()
}