main.rs

//! `ember` — a from-scratch LLM inference engine.
//!
//! Loads a Qwen2.5 model and generates text on the CPU with a hand-written
//! transformer forward pass. This file wires the CLI and tokenizer to the
//! decode loop; the numeric work lives in the sibling modules.

#![allow(dead_code)]

mod attention;
mod chat;
mod config;
mod model;
mod ops;
mod quant;
mod sample;
mod tensor;

use std::io::Write;
use std::path::PathBuf;
use std::time::{Instant, SystemTime, UNIX_EPOCH};

use anyhow::{anyhow, Context, Result};
use clap::Parser;
use tokenizers::Tokenizer;

use config::Config;
use model::Model;
use quant::Quant;
use sample::{Rng, Sampler};

/// Token ids that end generation (Qwen2.5: <|endoftext|> and <|im_end|>).
const EOS_IDS: [u32; 2] = [151643, 151645];

/// Run a local LLM on your CPU.
#[derive(Parser, Debug)]
#[command(name = "ember", version, about)]
struct Args {
    /// Prompt to complete.
    #[arg(short, long, default_value = "The capital of France is")]
    prompt: String,

    /// Directory with `config.json`, `model.safetensors`, and `tokenizer.json`.
    #[arg(short, long, default_value = "weights")]
    model: PathBuf,

    /// Maximum number of tokens to generate.
    #[arg(short = 'n', long, default_value_t = 128)]
    max_tokens: usize,

    /// Sampling temperature (0 ⇒ greedy).
    #[arg(short, long, default_value_t = 0.0)]
    temperature: f32,

    /// Nucleus (top-p) cutoff, used when `temperature > 0`.
    #[arg(long, default_value_t = 0.95)]
    top_p: f32,

    /// Weight quantization: `none` (f32), `int8`, or `int4`.
    #[arg(short, long, default_value = "none")]
    quant: String,

    /// Treat the prompt as a chat turn (Qwen ChatML template).
    #[arg(long)]
    chat: bool,

    /// System prompt used in `--chat` mode.
    #[arg(long, default_value = "You are a helpful assistant.")]
    system: String,

    /// Benchmark throughput on a randomly-weighted model (no weights needed).
    #[arg(long)]
    bench: bool,

    /// Measure perplexity (output-quality metric) over a fixed passage.
    #[arg(long)]
    perplexity: bool,
}

/// A fixed English passage used for the perplexity metric.
const PPL_TEXT: &str = "Machine learning is a branch of artificial intelligence that \
focuses on building systems that learn from data. Instead of following explicitly \
programmed rules, these systems identify patterns and make decisions with minimal human \
intervention. As the amount of available data grows, machine learning models continue to \
improve in accuracy and capability.";

fn main() -> Result<()> {
    let args = Args::parse();

    let scheme = Quant::parse(&args.quant)
        .ok_or_else(|| anyhow!("unknown --quant '{}' (expected none, int8, or int4)", args.quant))?;

    if args.bench {
        return run_bench(scheme, &args);
    }

    eprintln!("loading model from {} ({}) ...", args.model.display(), args.quant);
    let load_start = Instant::now();
    let model = Model::load(&args.model, scheme)
        .with_context(|| format!("loading model from {}", args.model.display()))?;
    let tokenizer = Tokenizer::from_file(args.model.join("tokenizer.json"))
        .map_err(|e| anyhow!("loading tokenizer: {e}"))?;
    eprintln!(
        "loaded {} layers, {:.0} MB weights, in {:.1}s",
        model.config.num_hidden_layers,
        model.weight_bytes() as f64 / 1e6,
        load_start.elapsed().as_secs_f32()
    );

    if args.perplexity {
        return perplexity(&model, &tokenizer, &args.quant);
    }

    let sampler = if args.temperature <= 0.0 {
        Sampler::Greedy
    } else {
        Sampler::TopP { temperature: args.temperature, top_p: args.top_p }
    };
    let mut rng = Rng::new(seed());

    // In chat mode, wrap the prompt in the ChatML template; otherwise complete it raw.
    let full_prompt = if args.chat {
        chat::chatml(&args.system, &args.prompt)
    } else {
        args.prompt.clone()
    };
    let encoding = tokenizer
        .encode(full_prompt.as_str(), false)
        .map_err(|e| anyhow!("tokenizing prompt: {e}"))?;
    let prompt_ids = encoding.get_ids();
    if prompt_ids.is_empty() {
        return Err(anyhow!("empty prompt"));
    }

    let mut cache = model.new_cache();
    let stdout = std::io::stdout();
    let mut out = stdout.lock();

    // Prefill: feed every prompt token, keeping the logits after the last one.
    // In completion mode we echo the prompt; in chat mode we stream only the reply.
    if !args.chat {
        print!("{}", args.prompt);
        out.flush().ok();
    }
    let prefill_start = Instant::now();
    let mut logits = Vec::new();
    for (pos, &id) in prompt_ids.iter().enumerate() {
        logits = model.forward(id, pos, &mut cache);
        cache.advance();
    }
    let prefill_secs = prefill_start.elapsed().as_secs_f32();

    // Decode: sample, emit, feed back.
    let decode_start = Instant::now();
    let mut pos = prompt_ids.len();
    let mut generated = 0usize;
    for _ in 0..args.max_tokens {
        let next = sampler.sample(&logits, &mut rng);
        if EOS_IDS.contains(&next) {
            break;
        }
        let piece = tokenizer
            .decode(&[next], false)
            .map_err(|e| anyhow!("decoding token: {e}"))?;
        print!("{piece}");
        out.flush().ok();

        logits = model.forward(next, pos, &mut cache);
        cache.advance();
        pos += 1;
        generated += 1;
    }
    println!();

    let decode_secs = decode_start.elapsed().as_secs_f32();
    let prefill_toks = prompt_ids.len();
    eprintln!(
        "\nprefill: {prefill_toks} tok in {:.0} ms ({:.1} tok/s)  |  decode: {generated} tok, {:.1} tok/s ({:.0} ms/tok)",
        prefill_secs * 1000.0,
        prefill_toks as f32 / prefill_secs.max(1e-6),
        generated as f32 / decode_secs.max(1e-6),
        decode_secs * 1000.0 / generated.max(1) as f32,
    );
    Ok(())
}

/// Measure perplexity — `exp(mean negative log-likelihood)` — over a fixed
/// passage. Lower is better; it quantifies how much (if anything) quantization
/// degrades the model's predictions.
fn perplexity(model: &Model, tokenizer: &Tokenizer, label: &str) -> Result<()> {
    let ids = tokenizer
        .encode(PPL_TEXT, false)
        .map_err(|e| anyhow!("tokenizing: {e}"))?
        .get_ids()
        .to_vec();

    let mut cache = model.new_cache();
    let mut nll = 0.0f64;
    let mut scored = 0usize;

    let start = Instant::now();
    let mut logits = model.forward(ids[0], 0, &mut cache);
    cache.advance();
    for pos in 1..ids.len() {
        // Log-probability the model assigned to the token that actually came next.
        let target = ids[pos] as usize;
        let mut max = f32::NEG_INFINITY;
        for &l in &logits {
            if l > max {
                max = l;
            }
        }
        let mut sum_exp = 0.0f32;
        for &l in &logits {
            sum_exp += (l - max).exp();
        }
        let log_prob = (logits[target] - max) - sum_exp.ln();
        nll += -(log_prob as f64);
        scored += 1;

        logits = model.forward(ids[pos], pos, &mut cache);
        cache.advance();
    }

    let ppl = (nll / scored as f64).exp();
    eprintln!("scored {scored} tokens in {:.1}s", start.elapsed().as_secs_f32());
    println!("{label:>5} | perplexity {ppl:7.3} | {scored} tokens");
    Ok(())
}

/// Benchmark throughput on a randomly-weighted model of the Qwen2.5-0.5B shape.
fn run_bench(scheme: Quant, args: &Args) -> Result<()> {
    let config = Config::preset_qwen2_5_0_5b();
    eprintln!("building random {} model (Qwen2.5-0.5B shape) ...", args.quant);
    let built = Instant::now();
    let model = Model::random(config, scheme);
    let mb = model.weight_bytes() as f64 / 1e6;
    eprintln!("built {mb:.0} MB in {:.1}s", built.elapsed().as_secs_f32());

    let mut cache = model.new_cache();
    for pos in 0..3 {
        let _ = model.forward(0, pos, &mut cache); // warm up
        cache.advance();
    }
    cache.clear();

    let n = args.max_tokens.max(8);
    let start = Instant::now();
    for pos in 0..n {
        let _ = model.forward((pos % 100) as u32, pos, &mut cache);
        cache.advance();
    }
    let secs = start.elapsed().as_secs_f32();
    println!(
        "{:>5} | {:6.0} MB | {:6.1} tok/s",
        args.quant,
        mb,
        n as f32 / secs.max(1e-6)
    );
    Ok(())
}

/// A time-based seed for the sampler's RNG.
fn seed() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_nanos() as u64)
        .unwrap_or(0x1234_5678)
}