car_inference/tasks/

generate.rs

//! Text generation with sampling.

use candle_core::Tensor;
use serde::{Deserialize, Serialize};

use crate::backend::CandleBackend;
use crate::InferenceError;

/// Parameters controlling generation behavior.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerateParams {
    /// Sampling temperature (0.0 = greedy, 1.0 = full distribution).
    #[serde(default = "default_temperature")]
    pub temperature: f64,
    /// Top-p (nucleus) sampling threshold.
    #[serde(default = "default_top_p")]
    pub top_p: f64,
    /// Top-k sampling (0 = disabled).
    #[serde(default)]
    pub top_k: usize,
    /// Maximum tokens to generate.
    #[serde(default = "default_max_tokens")]
    pub max_tokens: usize,
    /// Stop sequences — generation halts when any is produced.
    #[serde(default)]
    pub stop: Vec<String>,
}

fn default_temperature() -> f64 { 0.7 }
fn default_top_p() -> f64 { 0.9 }
fn default_max_tokens() -> usize { 512 }

impl Default for GenerateParams {
    fn default() -> Self {
        Self {
            temperature: default_temperature(),
            top_p: default_top_p(),
            top_k: 0,
            max_tokens: default_max_tokens(),
            stop: Vec::new(),
        }
    }
}

/// A text generation request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerateRequest {
    /// The prompt to complete.
    pub prompt: String,
    /// Optional model override.
    pub model: Option<String>,
    /// Generation parameters.
    #[serde(default)]
    pub params: GenerateParams,
    /// Optional memory context to prepend to the prompt.
    /// When provided, this is injected as a system-level context block
    /// before the user prompt, grounding the model's response.
    #[serde(default)]
    pub context: Option<String>,
}

/// Wrap a raw prompt in Qwen3 chat format if it's not already formatted.
/// Disables thinking mode by default for faster, more direct responses.
/// When context is provided, it is injected into the system message to ground
/// the model's response with memory.
fn apply_chat_template(prompt: &str, context: Option<&str>) -> String {
    if prompt.contains("<|im_start|>") {
        return prompt.to_string();
    }
    match context {
        Some(ctx) => format!(
            "<|im_start|>system\nYou are a helpful assistant. Use the following context to inform your response. /no_think\n\n{ctx}<|im_end|>\n\
             <|im_start|>user\n{prompt}<|im_end|>\n\
             <|im_start|>assistant\n"
        ),
        None => format!(
            "<|im_start|>system\nYou are a helpful assistant. /no_think<|im_end|>\n\
             <|im_start|>user\n{prompt}<|im_end|>\n\
             <|im_start|>assistant\n"
        ),
    }
}

/// Strip Qwen3 thinking tags from output (public for vision module).
pub fn strip_thinking_pub(text: &str) -> String {
    strip_thinking(text)
}

/// Strip Qwen3 thinking tags from output.
fn strip_thinking(text: &str) -> String {
    if let Some(end) = text.find("</think>") {
        text[end + 8..].trim_start().to_string()
    } else if text.starts_with("<think>") {
        // Still in thinking mode, return as-is
        text.to_string()
    } else {
        text.to_string()
    }
}

/// Callback for FLARE-style re-retrieval during generation.
/// Called with partial generation text, returns additional context or None.
pub type RetrievalCallback = Box<dyn Fn(&str) -> Option<String> + Send>;

/// Generate text from a prompt using the loaded model.
/// `retrieval_cb`: optional callback for confidence-triggered re-retrieval (FLARE).
pub async fn generate(
    backend: &mut CandleBackend,
    req: GenerateRequest,
) -> Result<String, InferenceError> {
    // Reset KV cache so each generation starts fresh (prevents cross-call state bleed)
    backend.clear_kv_cache();

    let formatted = apply_chat_template(&req.prompt, req.context.as_deref());
    let tokens = backend.encode(&formatted)?;
    let eos = backend.eos_token_id();
    let eos_alt = backend.token_id("<|im_end|>");
    let params = &req.params;

    if tokens.is_empty() {
        return Ok(String::new());
    }

    // Truncate to model's max context length minus generation headroom.
    // This prevents KV cache overflow on long prompts.
    let max_ctx = backend.context_length().unwrap_or(32768);
    let headroom = params.max_tokens.min(max_ctx / 4);
    let max_prompt = max_ctx.saturating_sub(headroom);
    let tokens = if tokens.len() > max_prompt {
        eprintln!(
            "[car-inference] truncating prompt from {} to {} tokens (context_length={})",
            tokens.len(), max_prompt, max_ctx
        );
        tokens[tokens.len() - max_prompt..].to_vec()
    } else {
        tokens
    };

    let mut generated = Vec::new();

    // Prefill: process all prompt tokens, sample first generated token from prefill logits
    let logits = backend.forward(&tokens, 0)?;
    let mut next_token = sample_token(&logits, params)?;

    for _i in 0..params.max_tokens {
        // Check EOS
        if eos.map_or(false, |id| next_token == id)
            || eos_alt.map_or(false, |id| next_token == id)
        {
            break;
        }

        generated.push(next_token);

        // Check stop sequences
        if !params.stop.is_empty() {
            let text_so_far = backend.decode(&generated)?;
            if params.stop.iter().any(|s| text_so_far.contains(s)) {
                break;
            }
        }

        // Generate next token
        let pos = tokens.len() + generated.len() - 1;
        let logits = backend.forward(&[next_token], pos)?;
        next_token = sample_token(&logits, params)?;
    }

    let text = backend.decode(&generated)?;
    Ok(strip_thinking(&text))
}

/// Generate with FLARE-style confidence-triggered re-retrieval.
///
/// Monitors token logit confidence during generation. When a window of
/// low-confidence tokens is detected, pauses, re-queries memory with the
/// partial generation, and resumes with augmented context.
pub async fn generate_with_retrieval(
    backend: &mut CandleBackend,
    mut req: GenerateRequest,
    retrieval_cb: RetrievalCallback,
) -> Result<String, InferenceError> {
    // First pass: generate normally
    backend.clear_kv_cache();
    let formatted = apply_chat_template(&req.prompt, req.context.as_deref());
    let tokens = backend.encode(&formatted)?;
    let eos = backend.eos_token_id();
    let eos_alt = backend.token_id("<|im_end|>");
    let params = req.params.clone();

    if tokens.is_empty() {
        return Ok(String::new());
    }

    let mut generated = Vec::new();
    let mut low_confidence_count = 0u32;
    let mut retrieval_attempts = 0u32;
    let max_retrievals = 2;
    let confidence_threshold = 0.4f32;
    let low_confidence_window = 3u32;

    let logits = backend.forward(&tokens, 0)?;
    let mut next_token = sample_token(&logits, &params)?;

    for _i in 0..params.max_tokens {
        if eos.map_or(false, |id| next_token == id)
            || eos_alt.map_or(false, |id| next_token == id)
        {
            break;
        }

        generated.push(next_token);

        // Generate next token and check confidence
        let pos = tokens.len() + generated.len() - 1;
        let logits = backend.forward(&[next_token], pos)?;

        // Check max logit probability for confidence
        let logits_f32: Vec<f32> = logits.squeeze(0)
            .unwrap_or(logits.clone())
            .to_dtype(candle_core::DType::F32)
            .map_err(|e| InferenceError::InferenceFailed(format!("dtype: {e}")))?
            .to_vec1()
            .unwrap_or_default();

        if !logits_f32.is_empty() {
            // Compute softmax max probability
            let max_logit = logits_f32.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            let exp_sum: f32 = logits_f32.iter().map(|&v| (v - max_logit).exp()).sum();
            let max_prob = 1.0 / exp_sum; // probability of the top token

            if max_prob < confidence_threshold {
                low_confidence_count += 1;
            } else {
                low_confidence_count = 0;
            }

            // Trigger re-retrieval after sustained low confidence
            if low_confidence_count >= low_confidence_window
                && retrieval_attempts < max_retrievals
            {
                retrieval_attempts += 1;
                low_confidence_count = 0;

                // Use partial generation as re-retrieval query
                let partial = backend.decode(&generated)?;
                if let Some(new_context) = retrieval_cb(&partial) {
                    // Restart generation with augmented context
                    let combined_context = match req.context.take() {
                        Some(old) => format!("{}\n\n{}", old, new_context),
                        None => new_context,
                    };
                    req.context = Some(combined_context);

                    // Re-encode and restart
                    backend.clear_kv_cache();
                    let new_formatted = apply_chat_template(&req.prompt, req.context.as_deref());
                    let new_tokens = backend.encode(&new_formatted)?;
                    generated.clear();

                    let logits = backend.forward(&new_tokens, 0)?;
                    next_token = sample_token(&logits, &params)?;
                    continue;
                }
            }
        }

        next_token = sample_token(&logits, &params)?;
    }

    let text = backend.decode(&generated)?;
    Ok(strip_thinking(&text))
}

/// Sample a token, suppressing specific token IDs (set to -inf before sampling).
pub fn sample_token_suppress(logits: &Tensor, params: &GenerateParams, suppress: &[u32]) -> Result<u32, InferenceError> {
    if suppress.is_empty() {
        return sample_token(logits, params);
    }
    // Clone logits and set suppressed tokens to -inf
    let mut logits_vec: Vec<f32> = logits.squeeze(0)
        .unwrap_or(logits.clone())
        .to_dtype(candle_core::DType::F32)
        .map_err(|e| InferenceError::InferenceFailed(format!("dtype: {e}")))?
        .to_vec1()
        .map_err(|e| InferenceError::InferenceFailed(format!("to_vec: {e}")))?;
    // Handle 2D logits (take last row)
    let dims = logits.dims();
    if dims.len() == 2 {
        let vocab = dims[dims.len() - 1];
        let start = logits_vec.len() - vocab;
        logits_vec = logits_vec[start..].to_vec();
    }
    for &id in suppress {
        if (id as usize) < logits_vec.len() {
            logits_vec[id as usize] = f32::NEG_INFINITY;
        }
    }
    let modified = Tensor::from_vec(logits_vec, logits.squeeze(0).unwrap_or(logits.clone()).shape(), logits.device())
        .map_err(|e| InferenceError::InferenceFailed(format!("from_vec: {e}")))?;
    sample_token(&modified, params)
}

/// Sample a token from logits using temperature + top-p + top-k.
pub fn sample_token(logits: &Tensor, params: &GenerateParams) -> Result<u32, InferenceError> {
    let logits = logits
        .squeeze(0)
        .map_err(|e| InferenceError::InferenceFailed(format!("squeeze: {e}")))?;
    let logits = logits
        .to_dtype(candle_core::DType::F32)
        .map_err(|e| InferenceError::InferenceFailed(format!("dtype: {e}")))?;

    // Get last position's logits
    let dim = logits.dims();
    let logits = if dim.len() == 2 {
        logits
            .get(dim[0] - 1)
            .map_err(|e| InferenceError::InferenceFailed(format!("get last: {e}")))?
    } else {
        logits
    };

    // Greedy decoding
    if params.temperature <= 0.0 {
        let token = logits
            .argmax(0)
            .map_err(|e| InferenceError::InferenceFailed(format!("argmax: {e}")))?
            .to_scalar::<u32>()
            .map_err(|e| InferenceError::InferenceFailed(format!("scalar: {e}")))?;
        return Ok(token);
    }

    // Temperature scaling
    let logits = (&logits / params.temperature)
        .map_err(|e| InferenceError::InferenceFailed(format!("temp scale: {e}")))?;

    let mut logits_vec: Vec<f32> = logits
        .to_vec1()
        .map_err(|e| InferenceError::InferenceFailed(format!("to_vec: {e}")))?;

    // Top-k filtering
    if params.top_k > 0 && params.top_k < logits_vec.len() {
        let mut indexed: Vec<(usize, f32)> = logits_vec.iter().copied().enumerate().collect();
        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        let threshold = indexed[params.top_k].1;
        for v in &mut logits_vec {
            if *v < threshold {
                *v = f32::NEG_INFINITY;
            }
        }
    }

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

    // Top-p (nucleus) filtering
    if params.top_p < 1.0 {
        let mut sorted_indices: Vec<usize> = (0..probs.len()).collect();
        sorted_indices.sort_by(|&a, &b| {
            probs[b].partial_cmp(&probs[a]).unwrap_or(std::cmp::Ordering::Equal)
        });

        let mut cumsum = 0.0f32;
        let mut cutoff_idx = sorted_indices.len();
        for (i, &idx) in sorted_indices.iter().enumerate() {
            cumsum += probs[idx];
            if cumsum > params.top_p as f32 {
                cutoff_idx = i + 1;
                break;
            }
        }

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

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

    // Categorical sample
    let r: f32 = rand_f32();
    let mut cumsum = 0.0f32;
    for (i, &p) in probs.iter().enumerate() {
        cumsum += p;
        if cumsum >= r {
            return Ok(i as u32);
        }
    }

    // Fallback: return highest prob token
    Ok(probs
        .iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(i, _)| i as u32)
        .unwrap_or(0))
}

/// Random float in [0, 1) using the rand crate.
fn rand_f32() -> f32 {
    rand::random::<f32>()
}