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, Result};
use clap::{Parser, ValueEnum};
use std::io::Write;
use tokenizers::Tokenizer;

use fuel::lazy_glm4::{Glm4Activation, Glm4Config};
use fuel::lazy_quantized_glm4::QuantizedGlm4Model;
use fuel_transformers::generation::{LogitsProcessor, Sampling};

use fuel_examples::token_output_stream::TokenOutputStream;

const DEFAULT_PROMPT: &str = "Write a Rust function to calculate the factorial of a given number.";

#[derive(Clone, Debug, Copy, PartialEq, Eq, ValueEnum)]
enum Which {
    #[value(name = "q2k9b")]
    Q2k9b,
    #[value(name = "q2k32b")]
    Q2k32b,
    #[value(name = "q4k9b")]
    Q4k9b,
    #[value(name = "q4k32b")]
    Q4k32b,
}

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// GGUF file to load, typically a .gguf file generated by the quantize command from llama.cpp
    #[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 = "q4k9b")]
    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 = match self.which {
                    Which::Q2k9b => "THUDM/GLM-4-9B-0414",
                    Which::Q2k32b => "THUDM/GLM-4-32B-0414",
                    Which::Q4k9b => "THUDM/GLM-4-9B-0414",
                    Which::Q4k32b => "THUDM/GLM-4-32B-0414",
                };
                let api = api.model(repo.to_string());
                api.get("tokenizer.json")?
            }
        };
        Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg)
    }

    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, revision) = match self.which {
                    Which::Q2k9b => (
                        "unsloth/GLM-4-9B-0414-GGUF",
                        "GLM-4-9B-0414-Q2_K.gguf",
                        "main",
                    ),
                    Which::Q2k32b => (
                        "unsloth/GLM-4-32B-0414-GGUF",
                        "GLM-4-32B-0414-Q2_K.gguf",
                        "main",
                    ),
                    Which::Q4k9b => (
                        "unsloth/GLM-4-9B-0414-GGUF",
                        "GLM-4-9B-0414-Q4_K_M.gguf",
                        "main",
                    ),
                    Which::Q4k32b => (
                        "unsloth/GLM-4-32B-0414-GGUF",
                        "GLM-4-32B-0414-Q4_K_M.gguf",
                        "main",
                    ),
                };
                let api = hf_hub::api::sync::Api::new()?;
                api.repo(hf_hub::Repo::with_revision(
                    repo.to_string(),
                    hf_hub::RepoType::Model,
                    revision.to_string(),
                ))
                .get(filename)?
            }
        };
        Ok(model_path)
    }
}

fn format_size(size_in_bytes: usize) -> String {
    if size_in_bytes < 1_000 {
        format!("{}B", size_in_bytes)
    } 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)
    }
}

/// Build a `Glm4Config` by reading GGUF metadata for the `glm4`
/// architecture. Mirrors the keys used by the eager `quantized_glm4`
/// loader and infers `vocab_size` / `intermediate_size` from tensor
/// shapes when no explicit metadata key is present.
fn glm4_config_from_gguf(mc: &fuel::quantized::gguf_mmap::MmapedContent) -> Result<Glm4Config> {
    let meta = mc.metadata();
    let content = mc.content();

    let get_u32 = |k: &str| -> Result<u32> {
        meta.get(k)
            .ok_or_else(|| E::msg(format!("gguf metadata: missing {k:?}")))?
            .to_u32()
            .map_err(|e| E::msg(format!("gguf metadata {k:?}: {e:?}")))
    };
    let get_u32_opt = |k: &str| -> Option<u32> { meta.get(k).and_then(|v| v.to_u32().ok()) };
    let get_f32 = |k: &str| -> Result<f32> {
        meta.get(k)
            .ok_or_else(|| E::msg(format!("gguf metadata: missing {k:?}")))?
            .to_f32()
            .map_err(|e| E::msg(format!("gguf metadata {k:?}: {e:?}")))
    };
    let get_bool_opt = |k: &str| -> Option<bool> { meta.get(k).and_then(|v| v.to_bool().ok()) };

    let hidden_size = get_u32("glm4.embedding_length")? as usize;
    let num_hidden_layers = get_u32("glm4.block_count")? as usize;
    let num_attention_heads = get_u32("glm4.attention.head_count")? as usize;
    let num_key_value_heads = get_u32("glm4.attention.head_count_kv")? as usize;
    let head_dim = get_u32("glm4.attention.key_length")? as usize;
    let max_position_embeddings = get_u32("glm4.context_length")? as usize;
    let rms_norm_eps = get_f32("glm4.attention.layer_norm_rms_epsilon")? as f64;
    let rope_theta = get_f32("glm4.rope.freq_base")? as f64;

    // `glm4.feed_forward_length` carries `intermediate_size` on modern
    // converters; fall back to the `ffn_down` tensor's in-features dim
    // when absent (older converters or single-tensor split forms).
    let intermediate_size = match get_u32_opt("glm4.feed_forward_length") {
        Some(n) => n as usize,
        None => {
            let info = content
                .tensor_infos
                .get("blk.0.ffn_down.weight")
                .ok_or_else(|| {
                    E::msg("gguf: missing blk.0.ffn_down.weight; cannot derive intermediate_size")
                })?;
            // ffn_down has shape [hidden, inter] in GGUF; inter is dims()[1].
            let dims = info.shape.dims();
            if dims.len() != 2 {
                return Err(E::msg(format!(
                    "gguf blk.0.ffn_down.weight: expected 2D shape, got {dims:?}"
                )));
            }
            dims[1]
        }
    };

    // vocab_size from token_embd.weight shape [vocab, hidden].
    let vocab_size = content
        .tensor_infos
        .get("token_embd.weight")
        .ok_or_else(|| E::msg("gguf: missing token_embd.weight"))?
        .shape
        .dims()[0];

    // Optional metadata: partial rotary factor + attention bias.
    let partial_rotary_factor = meta
        .get("glm4.rope.scaling.factor")
        .and_then(|v| v.to_f32().ok())
        .map(|x| x as f64)
        .unwrap_or(0.5);
    let attention_bias = get_bool_opt("glm4.attention.bias").unwrap_or(true);

    // Tied embeddings: GGUF mirrors HF by omitting `output.weight`.
    let tie_word_embeddings = !content.tensor_infos.contains_key("output.weight");

    Ok(Glm4Config {
        vocab_size,
        hidden_size,
        intermediate_size,
        num_hidden_layers,
        num_attention_heads,
        num_key_value_heads,
        head_dim,
        partial_rotary_factor,
        attention_bias,
        max_position_embeddings,
        rope_theta,
        rms_norm_eps,
        hidden_activation: Glm4Activation::Silu,
        tie_word_embeddings,
    })
}

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 mmap once to derive the config + report header stats.
    let mc = fuel::quantized::gguf_mmap::MmapedContent::from_path(&model_path)
        .map_err(|e| e.with_path(&model_path))?;
    let header_done = start.elapsed();
    let mut total_size_in_bytes = 0;
    for tensor in mc.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",
        mc.content().tensor_infos.len(),
        format_size(total_size_in_bytes),
        header_done.as_secs_f32(),
    );

    let cfg = glm4_config_from_gguf(&mc)?;
    // Drop the first mmap before the loader re-opens (matches Phi-2
    // pattern: cheap double-mmap on modern OSes).
    drop(mc);

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

    let tokenizer = args.tokenizer()?;
    let mut tos = TokenOutputStream::new(tokenizer);
    let prompt_str = args
        .prompt
        .clone()
        .unwrap_or_else(|| DEFAULT_PROMPT.to_string());

    let prompt_str = format!("[gMASK]<sop><|user|>\n{}<|assistant|>", prompt_str);

    let tokens = tos
        .tokenizer()
        .encode(prompt_str, true)
        .map_err(anyhow::Error::msg)?;

    let tokens = tokens.get_ids().to_vec();

    let to_sample = args.sample_len.saturating_sub(1);

    let mut all_tokens: Vec<u32> = vec![];

    let mut logits_processor = {
        let temperature = args.temperature;
        let sampling = if temperature <= 0. {
            Sampling::ArgMax
        } else {
            match (args.top_k, args.top_p) {
                (None, None) => Sampling::All { temperature },
                (Some(k), None) => Sampling::TopK { k, temperature },
                (None, Some(p)) => Sampling::TopP { p, temperature },
                (Some(k), Some(p)) => Sampling::TopKThenTopP { k, p, temperature },
            }
        };
        LogitsProcessor::from_sampling(args.seed, sampling)
    };

    // Realize the Tensor logits and slice off the last token's row
    // as a plain `Vec<f32>` for the LogitsProcessor sampler.
    let vocab = cfg.vocab_size;
    let last_logits = |logits_flat: &[f32], seq: usize| -> Result<Vec<f32>> {
        let last_off = (seq - 1) * vocab;
        Ok(logits_flat[last_off..last_off + vocab].to_vec())
    };

    let start_prompt_processing = std::time::Instant::now();

    let mut next_token = if !args.split_prompt {
        let logits_lazy = model
            .forward(&tokens, 0)
            .map_err(|e| E::msg(format!("forward: {e}")))?;
        let logits_flat = logits_lazy.realize_f32();
        let logits = last_logits(&logits_flat, tokens.len())?;
        logits_processor.sample(&logits)?
    } else {
        let mut next_token = 0;
        for (pos, token) in tokens.iter().enumerate() {
            let single = [*token];
            let logits_lazy = model
                .forward(&single, pos)
                .map_err(|e| E::msg(format!("forward: {e}")))?;
            let logits_flat = logits_lazy.realize_f32();
            let logits = last_logits(&logits_flat, 1)?;
            next_token = logits_processor.sample(&logits)?
        }
        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()?;
    }

    let eos_token = *tos.tokenizer().get_vocab(true).get("<|user|>").unwrap();

    let start_post_prompt = std::time::Instant::now();

    let mut sampled = 0;
    for index in 0..to_sample {
        let single = [next_token];
        let logits_lazy = model
            .forward(&single, tokens.len() + index)
            .map_err(|e| E::msg(format!("forward: {e}")))?;
        let logits_flat = logits_lazy.realize_f32();
        let mut logits = last_logits(&logits_flat, 1)?;
        if args.repeat_penalty != 1. {
            let start_at = all_tokens.len().saturating_sub(args.repeat_last_n);
            fuel_transformers::utils::apply_repeat_penalty(
                &mut logits,
                args.repeat_penalty,
                &all_tokens[start_at..],
            );
        }
        next_token = logits_processor.sample(&logits)?;
        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",
        tokens.len(),
        tokens.len() as f64 / prompt_dt.as_secs_f64(),
    );
    println!(
        "{sampled:4} tokens generated: {:.2} token/s",
        sampled as f64 / dt.as_secs_f64(),
    );
    Ok(())
}