inferencelayer 0.2.4

//! `lfm2-generate` — minimal greedy generation on a full (unsharded) model: the real-checkpoint
//! smoke for any supported architecture. Prints the completion and the rates.
//!
//! ```text
//! lfm2-generate --model <dir> --prompt "..." [--max-tokens 32] [--chat] [--solo]
//! ```
//!
//! The prompt goes through the SCHEDULER's batched prefill ([`generate_once`]), which packs prompt
//! tokens as columns per dispatch. `--solo` selects [`Lfm2Gpu::decode`] instead — the M=1 loop that
//! ingests the prompt one token at a time. That path stays because it is the byte-exact oracle the
//! tests compare against, but it is NOT what you want for a long prompt: it is what made a
//! 2,883-token prompt on the 4B claim-extractor take 147 s (~20 tok/s), never touching prefill.
use anyhow::{Context, Result};
use inferencelayer::server::{RequestParams, generate_once};

fn main() -> Result<()> {
    let mut model = None;
    let mut prompt = None;
    let mut max_tokens = 32usize;
    let mut chat = false;
    let mut solo = false;
    let mut args = std::env::args().skip(1);
    while let Some(a) = args.next() {
        match a.as_str() {
            "--model" => model = args.next(),
            "--prompt" => prompt = args.next(),
            "--max-tokens" => max_tokens = args.next().context("--max-tokens value")?.parse()?,
            "--chat" => chat = true,
            "--solo" => solo = true,
            other => anyhow::bail!("unknown argument {other}"),
        }
    }
    let model = model.context("--model <dir> required")?;
    let prompt = prompt.context("--prompt required")?;
    let dir = std::path::Path::new(&model);
    let ctx = inferencelayer::GpuCtx::new()?;
    eprintln!("backend: {}", ctx.backend);
    let t0 = std::time::Instant::now();
    let w = inferencelayer::Weights::load(&ctx, dir)?;
    eprintln!(
        "loaded {:?} ({} layers) in {:.1}s",
        w.cfg.arch,
        w.cfg.n_layers,
        t0.elapsed().as_secs_f64()
    );
    let gpu = inferencelayer::Lfm2Gpu::new(&ctx, w);
    let tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
        .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
    let text = if chat {
        // Qwen3 no-think: pre-fill an empty <think></think> so a small model answers directly
        // instead of rambling in reasoning it can't finish within the token budget.
        format!(
            "<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
        )
    } else {
        prompt.clone()
    };
    let ids = tok
        .encode(text.as_str(), true)
        .map_err(|e| anyhow::anyhow!("encode: {e}"))?
        .get_ids()
        .to_vec();
    eprintln!("prompt: {} tokens", ids.len());
    let out = if solo {
        let (out, rate, prefill_ms) = gpu.decode(&ctx, &ids, max_tokens)?;
        eprintln!(
            "[--solo: M=1 loop, prompt ingested one token at a time] \
             decode {} tokens at {rate:.1} tok/s (prompt ingest {prefill_ms:.0} ms)",
            out.len()
        );
        out
    } else {
        let eos = eos_ids(dir);
        let (out, st) = generate_once(
            &gpu,
            &ctx,
            ids.clone(),
            max_tokens,
            RequestParams::default(),
            eos,
        )?;
        eprintln!(
            "prefill: {} tokens in {:.2} s ({:.0} tok/s) | decode: {} tokens at {:.1} tok/s",
            ids.len(),
            st.prefill_s,
            st.prefill_tok_s,
            out.len(),
            st.decode_tok_s
        );
        out
    };
    println!("{}", tok.decode(&out, true).unwrap_or_default());
    Ok(())
}

/// EOS ids from the checkpoint's `generation_config.json` / `config.json`. Without them a
/// generation runs to `max_tokens` even after the model says it is done — harmless for a smoke
/// test, wrong for anything real.
fn eos_ids(dir: &std::path::Path) -> Vec<u32> {
    let mut out = Vec::new();
    for name in ["generation_config.json", "config.json"] {
        let Ok(bytes) = std::fs::read(dir.join(name)) else {
            continue;
        };
        let Ok(v) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
            continue;
        };
        match v.get("eos_token_id") {
            Some(serde_json::Value::Number(n)) => out.extend(n.as_u64().map(|x| x as u32)),
            Some(serde_json::Value::Array(a)) => {
                out.extend(a.iter().filter_map(|x| x.as_u64()).map(|x| x as u32))
            }
            _ => {}
        }
        if !out.is_empty() {
            break;
        }
    }
    out
}