MiniCPM5-1B-rust/src/inference/generation.rs

use burn::tensor::backend::Backend;
use burn::tensor::{Float, Int, Tensor};
use rand::Rng;

use crate::config::EosTokenId;
use crate::model::{LlamaForCausalLM, LlamaKVCache};

pub struct GenerationConfig {
    pub max_new_tokens: Option<usize>,
    pub temperature: f32,
    pub top_p: f32,
}

impl Default for GenerationConfig {
    fn default() -> Self {
        Self {
            max_new_tokens: None,
            temperature: 1.0,
            top_p: 1.0,
        }
    }
}

pub fn generate_with_cache<B: Backend>(
    model: &LlamaForCausalLM<B>,
    input_ids: &[u32],
    config: &GenerationConfig,
    eos_token_id: &EosTokenId,
    device: &B::Device,
) -> Vec<u32> {
    let mut output_ids = input_ids.to_vec();
    let mut cache: Option<LlamaKVCache<B>> = None;

    // 第一步：输入完整 prompt，建立初始 cache
    {
        let input_ints: Vec<i64> = input_ids.iter().map(|&x| x as i64).collect();
        let input_tensor =
            Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
                .unsqueeze::<2>();

        let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
        cache = Some(new_cache);

        let next_token = sample_last(&logits, config.temperature, config.top_p);
        output_ids.push(next_token);

        if eos_token_id.contains(next_token) {
            return output_ids;
        }
    }

    // 后续步骤：每次只输入 1 个 token，使用 cache
    let mut count = 1;
    loop {
        if let Some(max) = config.max_new_tokens {
            if count >= max {
                break;
            }
        }

        let last_token = output_ids[output_ids.len() - 1];
        let input_ints: Vec<i64> = vec![last_token as i64];
        let input_tensor =
            Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
                .unsqueeze::<2>();

        let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
        cache = Some(new_cache);

        let vocab_size = model.config.vocab_size;
        let next_token_logits = logits
            .clone()
            .slice([0..1, 0..1, 0..vocab_size])
            .reshape([vocab_size]);

        let next_token = sample(&next_token_logits, config.temperature, config.top_p);
        output_ids.push(next_token);

        if eos_token_id.contains(next_token) {
            break;
        }

        count += 1;
    }

    output_ids
}

fn sample_last<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 {
    let shape = logits.shape();
    let dims: [usize; 3] = shape.dims();
    let seq_len = dims[1];
    let vocab_size = dims[2];

    let last_logits = logits
        .clone()
        .slice([0..1, seq_len - 1..seq_len, 0..vocab_size])
        .reshape([vocab_size]);

    sample(&last_logits, temperature, top_p)
}

fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p: f32) -> u32 {
    let data = logits.clone().to_data();
    let values: Vec<f32> = data.as_slice().unwrap().to_vec();

    // temperature = 0 时退化为 greedy
    if temperature <= 0.001 {
        let mut max_idx = 0;
        let mut max_val = f32::NEG_INFINITY;
        for (i, &val) in values.iter().enumerate() {
            if val > max_val {
                max_val = val;
                max_idx = i;
            }
        }
        return max_idx as u32;
    }

    // 应用 temperature
    let scaled: Vec<f32> = values.iter().map(|&v| v / temperature).collect();

    // softmax
    let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let exp_sum: f32 = scaled.iter().map(|&v| (v - max_val).exp()).sum();
    let mut probs: Vec<f32> = scaled.iter().map(|&v| (v - max_val).exp() / exp_sum).collect();

    // top_p 过滤
    if top_p < 1.0 {
        let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();
        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

        let mut cumsum = 0.0;
        let mut cutoff = indexed.len();
        for (i, &(_, p)) in indexed.iter().enumerate() {
            cumsum += p;
            if cumsum > top_p {
                cutoff = i + 1;
                break;
            }
        }

        let keep: std::collections::HashSet<usize> = indexed[..cutoff].iter().map(|(i, _)| *i).collect();
        for (i, p) in probs.iter_mut().enumerate() {
            if !keep.contains(&i) {
                *p = 0.0;
            }
        }

        let sum: f32 = probs.iter().sum();
        if sum > 0.0 {
            for p in probs.iter_mut() {
                *p /= sum;
            }
        }
    }

    // 多项式采样
    let mut rng = rand::thread_rng();
    let r: f32 = rng.gen();
    let mut cumsum = 0.0;
    for (i, &p) in probs.iter().enumerate() {
        cumsum += p;
        if r < cumsum {
            return i as u32;
        }
    }
    probs.len() as u32 - 1
}