settled-reach/server/sr-voice/src/inference.rs

use std::num::NonZeroU32;
use std::path::Path;
use std::time::Instant;

use llama_cpp_2::context::params::LlamaContextParams;
use llama_cpp_2::llama_backend::LlamaBackend;
use llama_cpp_2::llama_batch::LlamaBatch;
use llama_cpp_2::model::params::LlamaModelParams;
use llama_cpp_2::model::{AddBos, LlamaModel, Special};
use llama_cpp_2::sampling::LlamaSampler;

use crate::VoiceError;

/// Configuration for model loading and inference.
pub struct InferenceConfig {
    pub model_path: String,
    pub threads: u32,
    pub ctx_size: u32,
    pub seed: Option<u32>,
}

/// Result of a single generation call.
#[derive(serde::Serialize)]
pub struct GenerationResult {
    pub text: String,
    pub tokens_generated: u32,
    pub generation_time_ms: u64,
    pub tokens_per_sec: f64,
    pub prefill_time_ms: u64,
}

/// Wraps llama.cpp model and context for text generation.
pub struct InferenceEngine {
    backend: LlamaBackend,
    model: LlamaModel,
    ctx_size: u32,
    threads: u32,
}

impl InferenceEngine {
    /// Load a GGUF model from disk.
    ///
    /// Offloads all layers to the GPU via ROCm. The binary is built with
    /// llama-cpp-rs + ROCm support (see Makefile `build-sr-voice` target),
    /// but `LlamaModelParams::default()` sets `n_gpu_layers = 0`, which
    /// runs the entire model on CPU at ~10× lower throughput. Setting
    /// `n_gpu_layers` to a large sentinel value (999) asks llama.cpp to
    /// offload every layer the model has; it clamps to the real count.
    /// For Gemma 2 2B (27 layers) this fully GPU-offloads the model.
    pub fn load(config: &InferenceConfig) -> Result<Self, VoiceError> {
        let backend =
            LlamaBackend::init().map_err(|e| VoiceError::ModelLoadFailed(e.to_string()))?;

        let model_params = LlamaModelParams::default().with_n_gpu_layers(999);
        let model = LlamaModel::load_from_file(
            &backend,
            Path::new(&config.model_path),
            &model_params,
        )
        .map_err(|e| VoiceError::ModelLoadFailed(e.to_string()))?;

        Ok(Self {
            backend,
            model,
            ctx_size: config.ctx_size,
            threads: config.threads,
        })
    }

    /// Generate text from a prompt.
    pub fn generate(
        &self,
        prompt: &str,
        max_tokens: u32,
        temperature: f32,
        top_p: f32,
        seed: Option<u32>,
    ) -> Result<GenerationResult, VoiceError> {
        let ctx_params = LlamaContextParams::default()
            .with_n_ctx(NonZeroU32::new(self.ctx_size))
            .with_n_threads(self.threads as i32)
            .with_n_threads_batch(self.threads as i32);

        let mut ctx = self
            .model
            .new_context(&self.backend, ctx_params)
            .map_err(|e| VoiceError::InferenceFailed(e.to_string()))?;

        // Tokenize the prompt
        let tokens = self
            .model
            .str_to_token(prompt, AddBos::Always)
            .map_err(|e| VoiceError::InferenceFailed(e.to_string()))?;

        if tokens.len() as u32 >= self.ctx_size {
            return Err(VoiceError::InferenceFailed(format!(
                "Prompt ({} tokens) exceeds context size ({})",
                tokens.len(),
                self.ctx_size
            )));
        }

        // Prefill: evaluate the prompt tokens
        let prefill_start = Instant::now();
        let mut batch = LlamaBatch::new(self.ctx_size as usize, 1);
        for (i, &token) in tokens.iter().enumerate() {
            let is_last = i == tokens.len() - 1;
            batch
                .add(token, i as i32, &[0], is_last)
                .map_err(|e| VoiceError::InferenceFailed(e.to_string()))?;
        }
        ctx.decode(&mut batch)
            .map_err(|e| VoiceError::InferenceFailed(e.to_string()))?;
        let prefill_time_ms = prefill_start.elapsed().as_millis() as u64;

        // Generation loop
        let gen_start = Instant::now();
        let mut generated_tokens: u32 = 0;
        let mut output = String::new();
        let mut cur_pos = tokens.len() as i32;

        let mut sampler = LlamaSampler::chain_simple([
            LlamaSampler::temp(temperature),
            LlamaSampler::top_p(top_p, 1),
            LlamaSampler::dist(seed.unwrap_or(1234)),
        ]);

        loop {
            if generated_tokens >= max_tokens {
                break;
            }

            let logits_index = batch.n_tokens() - 1;
            let token = sampler.sample(&ctx, logits_index);

            if self.model.is_eog_token(token) {
                break;
            }

            #[allow(deprecated)]
            let piece = self
                .model
                .token_to_str(token, Special::Tokenize)
                .map_err(|e| VoiceError::InferenceFailed(e.to_string()))?;
            output.push_str(&piece);
            generated_tokens += 1;

            batch.clear();
            batch
                .add(token, cur_pos, &[0], true)
                .map_err(|e| VoiceError::InferenceFailed(e.to_string()))?;
            cur_pos += 1;

            ctx.decode(&mut batch)
                .map_err(|e| VoiceError::InferenceFailed(e.to_string()))?;
        }

        let generation_time_ms = gen_start.elapsed().as_millis() as u64;
        let tokens_per_sec = if generation_time_ms > 0 {
            (generated_tokens as f64 / generation_time_ms as f64) * 1000.0
        } else {
            0.0
        };

        Ok(GenerationResult {
            text: output,
            tokens_generated: generated_tokens,
            generation_time_ms,
            tokens_per_sec,
            prefill_time_ms,
        })
    }
}