fuel-examples 0.20.0

// SPDX-License-Identifier: MIT OR Apache-2.0
#[cfg(feature = "mkl")]
extern crate intel_mkl_src;

#[cfg(feature = "accelerate")]
extern crate accelerate_src;

use anyhow::Error as E;
use clap::{Parser, ValueEnum};
use std::io::Write;
use tokenizers::Tokenizer;

use fuel::lazy_gemma3::{Gemma3Config, GemmaActivation};
use fuel::lazy_quantized_gemma3::QuantizedGemma3Model;

use fuel_examples::token_output_stream::TokenOutputStream;

const DEFAULT_PROMPT: &str = "Write a function to calculate fibonacci num";

#[derive(Clone, Debug, Copy, PartialEq, Eq, ValueEnum)]
enum Which {
    #[value(name = "gemma3-4b-it")]
    Gemma3_4bIt,
}

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// GGUF file to load, typically a .gguf file generated by quantization
    #[arg(long)]
    model: Option<String>,

    /// The initial prompt, use 'interactive' for entering multiple prompts in an interactive way
    /// and 'chat' for an interactive model where history of previous prompts and generated tokens
    /// is preserved.
    #[arg(long)]
    prompt: Option<String>,

    /// The length of the sample to generate (in tokens).
    #[arg(short = 'n', long, default_value_t = 1000)]
    sample_len: usize,

    /// The tokenizer config in json format.
    #[arg(long)]
    tokenizer: Option<String>,

    /// The temperature used to generate samples, use 0 for greedy sampling.
    #[arg(long, default_value_t = 0.8)]
    temperature: f64,

    /// Nucleus sampling probability cutoff.
    #[arg(long)]
    top_p: Option<f64>,

    /// Only sample among the top K samples.
    #[arg(long)]
    top_k: Option<usize>,

    /// The seed to use when generating random samples.
    #[arg(long, default_value_t = 299792458)]
    seed: u64,

    /// Enable tracing (generates a trace-timestamp.json file).
    #[arg(long)]
    tracing: bool,

    /// Process prompt elements separately.
    #[arg(long)]
    split_prompt: bool,

    /// Run on CPU rather than GPU even if a GPU is available.
    #[arg(long)]
    cpu: bool,

    /// Penalty to be applied for repeating tokens, 1. means no penalty.
    #[arg(long, default_value_t = 1.1)]
    repeat_penalty: f32,

    /// The context size to consider for the repeat penalty.
    #[arg(long, default_value_t = 64)]
    repeat_last_n: usize,

    /// The model size to use.
    #[arg(long, default_value = "gemma3-4b-it")]
    which: Which,
}

impl Args {
    fn tokenizer(&self) -> anyhow::Result<Tokenizer> {
        let tokenizer_path = match &self.tokenizer {
            Some(config) => std::path::PathBuf::from(config),
            None => {
                let api = hf_hub::api::sync::Api::new()?;
                let repo = "google/gemma-3-4b-it";
                println!("DEBUG: Downloading tokenizer from {repo}");
                let api = api.model(repo.to_string());
                api.get("tokenizer.json")?
            }
        };
        println!("DEBUG: Loading tokenizer from {tokenizer_path:?}");
        let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg)?;

        Ok(tokenizer)
    }

    fn model(&self) -> anyhow::Result<std::path::PathBuf> {
        let model_path = match &self.model {
            Some(config) => std::path::PathBuf::from(config),
            None => {
                let (repo, filename) = match self.which {
                    Which::Gemma3_4bIt => (
                        "google/gemma-3-4b-it-qat-q4_0-gguf",
                        "gemma-3-4b-it-q4_0.gguf",
                    ),
                };
                let api = hf_hub::api::sync::Api::new()?;
                api.repo(hf_hub::Repo::with_revision(
                    repo.to_string(),
                    hf_hub::RepoType::Model,
                    "main".to_string(),
                ))
                .get(filename)?
            }
        };
        Ok(model_path)
    }
}

fn format_size(size_in_bytes: usize) -> String {
    if size_in_bytes < 1_000 {
        format!("{size_in_bytes}B")
    } else if size_in_bytes < 1_000_000 {
        format!("{:.2}KB", size_in_bytes as f64 / 1e3)
    } else if size_in_bytes < 1_000_000_000 {
        format!("{:.2}MB", size_in_bytes as f64 / 1e6)
    } else {
        format!("{:.2}GB", size_in_bytes as f64 / 1e9)
    }
}

#[derive(Debug)]
enum Prompt {
    Interactive,
    Chat,
    One(String),
}

/// Derive a Gemma3Config from GGUF metadata + tensor-shape inspection.
///
/// Mirrors the eager loader's metadata read for the architecture-side fields
/// (head counts, key length, RMS norm eps, RoPE bases, sliding window) and
/// recovers vocab_size / intermediate_size from the token_embd and ffn_down
/// tensor shapes. Fields that Gemma 3 does not exercise (softcaps,
/// attention_bias) are filled with the architecture's documented defaults.
fn gemma3_config_from_gguf(
    content: &fuel::quantized::gguf_file::Content,
) -> anyhow::Result<Gemma3Config> {
    // Probe architecture prefix (gemma3 / gemma2 / gemma) — matches the
    // eager loader's behaviour.
    let prefix = ["gemma3", "gemma2", "gemma", "gemma-embedding"]
        .iter()
        .find(|p| {
            content
                .metadata
                .contains_key(&format!("{}.attention.head_count", p))
        })
        .copied()
        .unwrap_or("gemma3");

    let md_get = |s: &str| -> anyhow::Result<&fuel::quantized::gguf_file::Value> {
        let key = format!("{prefix}.{s}");
        content
            .metadata
            .get(&key)
            .ok_or_else(|| E::msg(format!("cannot find {key} in metadata")))
    };

    let num_attention_heads = md_get("attention.head_count")?.to_u32()? as usize;
    let num_key_value_heads = md_get("attention.head_count_kv")?.to_u32()? as usize;
    let num_hidden_layers = md_get("block_count")?.to_u32()? as usize;
    let hidden_size = md_get("embedding_length")?.to_u32()? as usize;
    let head_dim = md_get("attention.key_length")?.to_u32()? as usize;
    let rms_norm_eps = md_get("attention.layer_norm_rms_epsilon")?.to_f32()? as f64;
    let sliding_window = md_get("attention.sliding_window")?.to_u32()? as usize;

    let sliding_window_pattern = md_get("attention.sliding_window_type")
        .and_then(|m| Ok(m.to_u32()? as usize))
        .unwrap_or(6);

    let rope_theta = md_get("rope.freq_base")
        .and_then(|m| Ok(m.to_f32()? as f64))
        .unwrap_or(1_000_000.0);
    let rope_local_base_freq = md_get("rope.local_freq_base")
        .and_then(|m| Ok(m.to_f32()? as f64))
        .unwrap_or(10_000.0);

    // Recover vocab_size and intermediate_size from tensor shapes:
    //   token_embd.weight is [vocab_size, hidden_size]
    //   blk.0.ffn_down.weight is [hidden_size, intermediate_size]
    let tok = content
        .tensor_infos
        .get("token_embd.weight")
        .ok_or_else(|| E::msg("gguf missing token_embd.weight"))?;
    let vocab_size = tok.shape.dims()[0];

    let ffn_down = content
        .tensor_infos
        .get("blk.0.ffn_down.weight")
        .ok_or_else(|| E::msg("gguf missing blk.0.ffn_down.weight"))?;
    // ffn_down shape: [hidden_size, intermediate_size] in GGUF [out, in] order.
    let intermediate_size = ffn_down.shape.dims()[1];

    // Gemma 3 dropped the attention/final logit softcaps (which Gemma 2 had).
    // The eager forward in fuel-transformers reflects this — no softcap is
    // applied. Bias projections are absent in Gemma-3 GGUFs (attention_bias=false).
    Ok(Gemma3Config {
        vocab_size,
        hidden_size,
        intermediate_size,
        num_hidden_layers,
        num_attention_heads,
        num_key_value_heads,
        head_dim,
        rms_norm_eps,
        rope_theta,
        rope_local_base_freq,
        max_position_embeddings: 8192,
        sliding_window,
        sliding_window_pattern,
        attention_bias: false,
        hidden_activation: GemmaActivation::GeluPytorchTanh,
        attn_logit_softcapping: None,
        final_logit_softcapping: None,
    })
}

fn apply_repeat_penalty(logits: &mut [f32], penalty: f32, ctx: &[u32]) {
    let mut seen = std::collections::HashSet::new();
    for &t in ctx {
        if !seen.insert(t) {
            continue;
        }
        let i = t as usize;
        if i < logits.len() {
            let v = logits[i];
            logits[i] = if v >= 0.0 { v / penalty } else { v * penalty };
        }
    }
}

fn sample(logits: &[f32], temp: f32, top_k: Option<usize>, top_p: Option<f32>, seed: u64) -> u32 {
    if temp <= 0.0 {
        let mut bi = 0usize;
        let mut b = logits[0];
        for (i, &v) in logits.iter().enumerate().skip(1) {
            if v > b {
                b = v;
                bi = i;
            }
        }
        return bi as u32;
    }
    let m = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let inv = 1.0 / temp.max(1e-6);
    let mut probs: Vec<f32> = logits.iter().map(|&x| ((x - m) * inv).exp()).collect();
    let s: f32 = probs.iter().sum();
    for p in &mut probs {
        *p /= s.max(1e-30);
    }
    let mut idx: Vec<usize> = (0..probs.len()).collect();
    idx.sort_unstable_by(|&a, &b| probs[b].partial_cmp(&probs[a]).unwrap());
    let mut keep = vec![true; probs.len()];
    if let Some(k) = top_k {
        for &i in idx.iter().skip(k) {
            keep[i] = false;
        }
    }
    if let Some(pc) = top_p {
        let mut c = 0.0;
        let mut allow = true;
        for &i in &idx {
            if !keep[i] {
                continue;
            }
            if !allow {
                keep[i] = false;
                continue;
            }
            c += probs[i];
            if c >= pc {
                allow = false;
            }
        }
    }
    let mut f: Vec<f32> = probs
        .iter()
        .enumerate()
        .map(|(i, p)| if keep[i] { *p } else { 0.0 })
        .collect();
    let ss: f32 = f.iter().sum();
    if ss > 0.0 {
        for v in &mut f {
            *v /= ss;
        }
    } else {
        return 0;
    }
    let mut st = seed
        .wrapping_mul(6364136223846793005)
        .wrapping_add(1442695040888963407);
    st ^= st >> 33;
    st = st.wrapping_mul(0xff51_afd7_ed55_8ccd);
    st ^= st >> 33;
    let r = (st as f32) / (u64::MAX as f32);
    let mut c = 0.0;
    for (i, p) in f.iter().enumerate() {
        c += *p;
        if r <= c {
            return i as u32;
        }
    }
    (f.len() - 1) as u32
}

fn main() -> anyhow::Result<()> {
    use tracing_chrome::ChromeLayerBuilder;
    use tracing_subscriber::prelude::*;

    let args = Args::parse();
    let _guard = if args.tracing {
        let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
        tracing_subscriber::registry().with(chrome_layer).init();
        Some(guard)
    } else {
        None
    };

    println!(
        "avx: {}, neon: {}, f16c: {}",
        fuel::utils::with_avx(),
        fuel::utils::with_neon(),
        fuel::utils::with_f16c()
    );
    println!(
        "temp: {:.2} repeat-penalty: {:.2} repeat-last-n: {}",
        args.temperature, args.repeat_penalty, args.repeat_last_n
    );

    let model_path = args.model()?;
    let start = std::time::Instant::now();
    let _device = fuel_examples::device(args.cpu)?;

    // Open the GGUF once to read metadata + compute on-disk size, then drop
    // the mmap before letting the lazy loader re-open it. (The lazy
    // `QuantizedGemma3Model::from_gguf` opens its own mmap internally.)
    let cfg = {
        let mmaped = fuel::quantized::gguf_mmap::MmapedContent::from_path(&model_path)
            .map_err(|e| e.with_path(&model_path))?;
        let header_done = start.elapsed();
        let (_mmap, model_content) = mmaped.into_parts();
        let mut total_size_in_bytes = 0;
        for tensor in model_content.tensor_infos.values() {
            let elem_count = tensor.shape.elem_count();
            total_size_in_bytes +=
                elem_count * tensor.ggml_dtype.type_size() / tensor.ggml_dtype.block_size();
        }
        println!(
            "mmapped {:?} tensors ({}); header in {:.2}s",
            model_content.tensor_infos.len(),
            format_size(total_size_in_bytes),
            header_done.as_secs_f32(),
        );
        gemma3_config_from_gguf(&model_content)?
    };

    let model = QuantizedGemma3Model::from_gguf(&model_path, &cfg)
        .map_err(|e| E::msg(format!("load gguf: {e}")))?;
    println!("model built in {:.2}s total", start.elapsed().as_secs_f32());

    let tokenizer = args.tokenizer()?;

    let mut tos = TokenOutputStream::new(tokenizer);
    println!(
        "DEBUG: Tokenizer vocabulary size: {}",
        tos.tokenizer().get_vocab(true).len()
    );

    let prompt = match args.prompt.as_deref() {
        Some("chat") => Prompt::Chat,
        Some("interactive") => Prompt::Interactive,
        Some(s) => Prompt::One(s.to_string()),
        None => Prompt::One(DEFAULT_PROMPT.to_string()),
    };

    let mut pre_prompt_tokens = vec![];
    for _ in 0.. {
        let prompt_str = match &prompt {
            Prompt::One(prompt) => prompt.clone(),
            Prompt::Interactive | Prompt::Chat => {
                print!("> ");
                std::io::stdout().flush()?;
                let mut prompt = String::new();
                std::io::stdin().read_line(&mut prompt)?;
                if prompt.ends_with('\n') {
                    prompt.pop();
                    if prompt.ends_with('\r') {
                        prompt.pop();
                    }
                }
                // Format for Gemma 3 chat/instruction format
                format!("<start_of_turn> user\n{prompt}<end_of_turn>\n<start_of_turn> model\n")
            }
        };
        print!("{}", prompt_str);

        let tokens = tos
            .tokenizer()
            .encode(prompt_str, true)
            .map_err(anyhow::Error::msg)?;
        let prompt_tokens = [&pre_prompt_tokens, tokens.get_ids()].concat();

        let to_sample = args.sample_len.saturating_sub(1);
        let max_seq_len = 8192; // Gemma 3 context length
        let prompt_tokens = if prompt_tokens.len() + to_sample > max_seq_len - 10 {
            let to_remove = prompt_tokens.len() + to_sample + 10 - max_seq_len;
            prompt_tokens[prompt_tokens.len().saturating_sub(to_remove)..].to_vec()
        } else {
            prompt_tokens
        };
        let mut all_tokens = vec![];
        let vocab_size = cfg.vocab_size;

        // Pick a single logit row [vocab_size] (the last seq position) from
        // a forward call whose logits have shape [1, seq, vocab_size].
        let last_logits = |logits_vec: Vec<f32>, seq: usize| -> Vec<f32> {
            let last_off = (seq - 1) * vocab_size;
            logits_vec[last_off..last_off + vocab_size].to_vec()
        };

        let temp = args.temperature as f32;
        let top_p = args.top_p.map(|p| p as f32);
        let top_k = args.top_k;

        let start_prompt_processing = std::time::Instant::now();
        let mut next_token = if !args.split_prompt {
            let seq = prompt_tokens.len();
            let logits = model
                .forward(&prompt_tokens, 0)
                .map_err(|e| E::msg(format!("forward: {e}")))?;
            let logits_v = logits.realize_f32();
            let last = last_logits(logits_v, seq);
            sample(&last, temp, top_k, top_p, args.seed)
        } else {
            let mut next_token = 0;
            for (pos, token) in prompt_tokens.iter().enumerate() {
                let logits = model
                    .forward(&[*token], pos)
                    .map_err(|e| E::msg(format!("forward: {e}")))?;
                let logits_v = logits.realize_f32();
                let last = last_logits(logits_v, 1);
                next_token = sample(
                    &last,
                    temp,
                    top_k,
                    top_p,
                    args.seed.wrapping_add(pos as u64),
                );
            }
            next_token
        };
        let prompt_dt = start_prompt_processing.elapsed();
        all_tokens.push(next_token);
        if let Some(t) = tos.next_token(next_token)? {
            print!("{t}");
            std::io::stdout().flush()?;
        }

        // For Gemma 3, use the correct end of sequence token
        let eos_token = *tos
            .tokenizer()
            .get_vocab(true)
            .get("<end_of_turn>")
            .unwrap();

        let start_post_prompt = std::time::Instant::now();
        let mut sampled = 0;
        for index in 0..to_sample {
            let logits = model
                .forward(&[next_token], prompt_tokens.len() + index)
                .map_err(|e| E::msg(format!("forward: {e}")))?;
            let logits_v = logits.realize_f32();
            let mut last = last_logits(logits_v, 1);
            if args.repeat_penalty != 1. {
                let start_at = all_tokens.len().saturating_sub(args.repeat_last_n);
                apply_repeat_penalty(&mut last, args.repeat_penalty, &all_tokens[start_at..]);
            }
            next_token = sample(
                &last,
                temp,
                top_k,
                top_p,
                args.seed.wrapping_add((index + 1) as u64),
            );
            all_tokens.push(next_token);
            if let Some(t) = tos.next_token(next_token)? {
                print!("{t}");
                std::io::stdout().flush()?;
            }
            sampled += 1;
            if next_token == eos_token {
                break;
            };
        }
        if let Some(rest) = tos.decode_rest().map_err(fuel::Error::msg)? {
            print!("{rest}");
        }
        std::io::stdout().flush()?;
        let dt = start_post_prompt.elapsed();
        println!(
            "\n\n{:4} prompt tokens processed: {:.2} token/s",
            prompt_tokens.len(),
            prompt_tokens.len() as f64 / prompt_dt.as_secs_f64(),
        );
        println!(
            "{sampled:4} tokens generated: {:.2} token/s",
            sampled as f64 / dt.as_secs_f64(),
        );

        match prompt {
            Prompt::One(_) => break,
            Prompt::Interactive => {}
            Prompt::Chat => {
                pre_prompt_tokens = [prompt_tokens.as_slice(), all_tokens.as_slice()].concat()
            }
        }
    }

    Ok(())
}