fuel-examples 0.19.1

// 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_llama_full::LlamaFullConfig;
use fuel::lazy_phi3::Phi3Config;
use fuel::lazy_quantized_llama::QuantizedLlama3Model;
use fuel::lazy_quantized_phi3::QuantizedPhi3Model;
use fuel_transformers::generation::{LogitsProcessor, Sampling};

use fuel_examples::token_output_stream::TokenOutputStream;

const DEFAULT_PROMPT: &str = "Write a function to count prime numbers up to N. ";

#[derive(Clone, Debug, Copy, PartialEq, Eq, ValueEnum)]
enum Which {
    #[value(name = "phi-2")]
    Phi2,
    #[value(name = "phi-3")]
    Phi3,
    /// Alternative implementation of phi-3, based on llama.
    #[value(name = "phi-3b")]
    Phi3b,
    #[value(name = "phi-4")]
    Phi4,
}

#[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 = "phi-3b")]
    which: Which,

    /// Unused in the lazy port: the optimizer selects the attention
    /// kernel from the graph, not from a CLI flag. Retained for CLI
    /// compatibility with the eager binary.
    #[arg(long)]
    use_flash_attn: bool,
}

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::Phi2 => "microsoft/phi-2",
                    Which::Phi3 | Which::Phi3b => "microsoft/Phi-3-mini-4k-instruct",
                    Which::Phi4 => "microsoft/phi-4",
                };
                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::Phi2 => ("TheBloke/phi-2-GGUF", "phi-2.Q4_K_M.gguf", "main"),
                    Which::Phi3 => (
                        "microsoft/Phi-3-mini-4k-instruct-gguf",
                        "Phi-3-mini-4k-instruct-q4.gguf",
                        "main",
                    ),
                    Which::Phi3b => (
                        "microsoft/Phi-3-mini-4k-instruct-gguf",
                        "Phi-3-mini-4k-instruct-q4.gguf",
                        "5eef2ce24766d31909c0b269fe90c817a8f263fb",
                    ),
                    Which::Phi4 => ("microsoft/phi-4-gguf", "phi-4-q4.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!("{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)
    }
}

/// Lazy variant of the original eager `Model` enum. Each variant
/// holds a `QuantizedXModel` whose `forward(tokens, start_pos)` returns
/// a `Tensor` of shape `(1, seq, vocab_size)`.
enum Model {
    Phi3(QuantizedPhi3Model, usize), // (model, vocab_size)
    Phi3b(QuantizedLlama3Model, usize),
}

/// Intermediate config bundle — built from GGUF metadata, then handed
/// to the matching `from_gguf` constructor. Decoupling config
/// extraction from the weight load lets us drop the metadata-only mmap
/// before re-opening the file for the (much larger) weight pass.
enum ModelCfg {
    Phi3(Phi3Config),
    Phi3b(LlamaFullConfig),
}

impl Model {
    /// Forward the model over the given token slice at `start_pos` and
    /// return the per-token logits of the LAST position as a plain
    /// `Vec<f32>` suitable for `LogitsProcessor::sample`.
    fn forward_last_logits(&self, tokens: &[u32], start_pos: usize) -> anyhow::Result<Vec<f32>> {
        let (logits_flat, vocab) = match self {
            Self::Phi3(m, v) => (
                m.forward(tokens, start_pos).map_err(E::msg)?.realize_f32(),
                *v,
            ),
            Self::Phi3b(m, v) => (
                m.forward(tokens, start_pos).map_err(E::msg)?.realize_f32(),
                *v,
            ),
        };
        let seq = tokens.len();
        if seq == 0 {
            anyhow::bail!("forward_last_logits: empty token slice");
        }
        let last_off = (seq - 1) * vocab;
        Ok(logits_flat[last_off..last_off + vocab].to_vec())
    }
}

/// Read a metadata field as a `usize` (auto-upcasting smaller integers
/// via `Value::to_u64`).
fn md_usize(
    md: &std::collections::HashMap<String, fuel::quantized::gguf_file::Value>,
    key: &str,
) -> anyhow::Result<usize> {
    let v = md
        .get(key)
        .ok_or_else(|| E::msg(format!("gguf metadata: missing {key:?}")))?;
    Ok(v.to_u64().map_err(E::msg)? as usize)
}

/// Read a metadata field as an `f64` (auto-upcasting from `f32`).
fn md_f64(
    md: &std::collections::HashMap<String, fuel::quantized::gguf_file::Value>,
    key: &str,
) -> anyhow::Result<f64> {
    let v = md
        .get(key)
        .ok_or_else(|| E::msg(format!("gguf metadata: missing {key:?}")))?;
    if let Ok(x) = v.to_f64() {
        Ok(x)
    } else if let Ok(x) = v.to_f32() {
        Ok(x as f64)
    } else {
        anyhow::bail!("gguf metadata: {key:?} is not a float")
    }
}

/// Build a `Phi3Config` by reading GGUF metadata. Mirrors the eager
/// `quantized_phi3::ModelWeights::from_gguf` extraction.
fn phi3_config_from_gguf(
    md: &std::collections::HashMap<String, fuel::quantized::gguf_file::Value>,
) -> anyhow::Result<Phi3Config> {
    let num_attention_heads = md_usize(md, "phi3.attention.head_count")?;
    let num_key_value_heads = md_usize(md, "phi3.attention.head_count_kv")?;
    let num_hidden_layers = md_usize(md, "phi3.block_count")?;
    let hidden_size = md_usize(md, "phi3.embedding_length")?;
    let max_position_embeddings = md_usize(md, "phi3.context_length")?;
    let intermediate_size = md_usize(md, "phi3.feed_forward_length")?;
    let rms_norm_eps = md_f64(md, "phi3.attention.layer_norm_rms_epsilon")?;
    let rope_theta = md_f64(md, "phi3.rope.freq_base").unwrap_or(10_000.0);
    let vocab_size = md_usize(md, "phi3.vocab_size")
        .ok()
        .or_else(|| {
            // Fallback: use the token-list metadata length if present.
            md.get("tokenizer.ggml.tokens")
                .and_then(|v| v.to_vec().ok().map(|arr| arr.len()))
        })
        .ok_or_else(|| E::msg("gguf metadata: cannot determine vocab_size"))?;

    Ok(Phi3Config {
        vocab_size,
        hidden_size,
        intermediate_size,
        num_hidden_layers,
        num_attention_heads,
        num_key_value_heads,
        max_position_embeddings,
        rope_theta,
        rms_norm_eps,
        tie_word_embeddings: false,
    })
}

/// Build a `LlamaFullConfig` for the Phi3b path (Phi-3 via the
/// llama-family loader). Reads the same `phi3.*` metadata keys as
/// `phi3_config_from_gguf` and projects them into the LLaMA shape.
fn llama_config_from_phi3_gguf(
    md: &std::collections::HashMap<String, fuel::quantized::gguf_file::Value>,
) -> anyhow::Result<LlamaFullConfig> {
    let num_attention_heads = md_usize(md, "phi3.attention.head_count")?;
    let num_key_value_heads = md_usize(md, "phi3.attention.head_count_kv")?;
    let num_hidden_layers = md_usize(md, "phi3.block_count")?;
    let hidden_size = md_usize(md, "phi3.embedding_length")?;
    let max_position_embeddings = md_usize(md, "phi3.context_length")?;
    let intermediate_size = md_usize(md, "phi3.feed_forward_length")?;
    let rms_norm_eps = md_f64(md, "phi3.attention.layer_norm_rms_epsilon")?;
    let rope_theta = md_f64(md, "phi3.rope.freq_base").unwrap_or(10_000.0);
    let vocab_size = md_usize(md, "phi3.vocab_size")
        .ok()
        .or_else(|| {
            md.get("tokenizer.ggml.tokens")
                .and_then(|v| v.to_vec().ok().map(|arr| arr.len()))
        })
        .ok_or_else(|| E::msg("gguf metadata: cannot determine vocab_size"))?;
    let head_dim = hidden_size / num_attention_heads;

    Ok(LlamaFullConfig {
        hidden_size,
        intermediate_size,
        vocab_size,
        num_hidden_layers,
        num_attention_heads,
        num_key_value_heads,
        head_dim,
        rms_norm_eps,
        rope_theta,
        max_position_embeddings,
        bos_token_id: None,
        eos_token_id: None,
        rope_scaling: None,
        tie_word_embeddings: false,
    })
}

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
    };

    // --use_flash_attn is accepted for CLI compatibility with the eager
    // binary, but the lazy port's attention path is fixed. Touch the
    // flag so clippy doesn't flag it as unused.
    let _ = args.use_flash_attn;

    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
    );

    // Phi-2 is not yet ported to a lazy module; fail early with a
    // descriptive message so users know to wait for / request a
    // `lazy_quantized_phi` port.
    if args.which == Which::Phi2 {
        anyhow::bail!(
            "lazy quantized-phi binary: Phi-2 is not yet supported. \
             No `lazy_quantized_phi` module exists. Use `--which phi-3`, \
             `--which phi-3b`, or `--which phi-4`, or wait for a \
             `lazy_quantized_phi` port."
        );
    }

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

    // Open the GGUF once to compute the displayed size + extract the
    // model config from the metadata. The actual weight load is done
    // by `QuantizedXModel::from_gguf(path, &cfg)`, which re-opens the
    // file internally (the second mmap is cheap — the OS page cache
    // is hot from this first pass).
    let model_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 model_content = mmaped.content();
        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(),
        );
        match args.which {
            Which::Phi2 => unreachable!("phi-2 was rejected above"),
            Which::Phi3 | Which::Phi4 => ModelCfg::Phi3(
                phi3_config_from_gguf(&model_content.metadata)
                    .map_err(|e| E::msg(format!("phi3 config from gguf: {e}")))?,
            ),
            Which::Phi3b => ModelCfg::Phi3b(
                llama_config_from_phi3_gguf(&model_content.metadata)
                    .map_err(|e| E::msg(format!("llama config from phi3 gguf: {e}")))?,
            ),
        }
    };

    let model = match model_cfg {
        ModelCfg::Phi3(cfg) => {
            let vocab = cfg.vocab_size;
            let m = QuantizedPhi3Model::from_gguf(&model_path, &cfg)
                .map_err(|e| E::msg(format!("QuantizedPhi3Model::from_gguf: {e}")))?;
            Model::Phi3(m, vocab)
        }
        ModelCfg::Phi3b(cfg) => {
            let vocab = cfg.vocab_size;
            let m = QuantizedLlama3Model::from_gguf(&model_path, &cfg)
                .map_err(|e| E::msg(format!("QuantizedLlama3Model::from_gguf: {e}")))?;
            Model::Phi3b(m, vocab)
        }
    };
    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.unwrap_or_else(|| DEFAULT_PROMPT.to_string());
    print!("{}", 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![];
    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)
    };

    let start_prompt_processing = std::time::Instant::now();
    let mut next_token = if !args.split_prompt {
        let logits = model.forward_last_logits(&tokens, 0)?;
        logits_processor.sample(&logits)?
    } else {
        let mut next_token = 0;
        for (pos, token) in tokens.iter().enumerate() {
            let logits = model.forward_last_logits(&[*token], pos)?;
            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("<|endoftext|>")
        .unwrap();
    let start_post_prompt = std::time::Instant::now();
    let mut sampled = 0;
    for index in 0..to_sample {
        let mut logits = model.forward_last_logits(&[next_token], tokens.len() + index)?;
        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(())
}