open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/falcon_decode.rs

//! Wave 58 bring-up: Falcon-H1 PARALLEL hybrid end-to-end
//! on the real checkpoint — third hybrid topology through
//! the wave 47 SSD primitives, first RoPE hybrid.
//!
//! ```bash
//! cargo run --release --features rayon --example falcon_decode -- \
//!     /path/to/snapshot-dir [n_tokens] [prompt text]
//! ```

use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer;
use jouleclaw_loader_gguf::falcon_h1::{FalconH1Config, FalconH1Model};
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
use std::path::PathBuf;
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: falcon_decode <snapshot-dir> [n_tokens] [prompt]");
        std::process::exit(1);
    }
    let dir = PathBuf::from(&args[1]);
    let n_tokens: usize = args.get(2).map(|s| s.parse()).transpose()?.unwrap_or(8);
    let prompt_text = args.get(3).cloned()
        .unwrap_or_else(|| "The capital of France is".to_string());

    let cfg_bytes = std::fs::read(dir.join("config.json"))?;
    let config: FalconH1Config = serde_json::from_slice(&cfg_bytes)?;
    println!(
        "config: {} layers (attention ∥ mamba in EVERY layer), hidden {}, vocab {}, rope_theta {:.0e}",
        config.num_hidden_layers, config.hidden_size, config.vocab_size, config.rope_theta,
    );

    let t0 = Instant::now();
    let st = read_safetensors_file(dir.join("model.safetensors"))?;
    println!("checkpoint mapped in {:.2}s ({} tensors)", t0.elapsed().as_secs_f64(), st.tensors.len());
    let model = FalconH1Model::bind(config, st)?;
    let tokenizer = BpeTokenizer::from_file(dir.join("tokenizer.json"))?;
    println!("tokenizer: {} vocab", tokenizer.vocab_size());

    let prompt = tokenizer.encode(&prompt_text);
    println!("prompt {prompt_text:?} → {prompt:?}");

    let mut state = model.new_state();
    let t_pre = Instant::now();
    let mut logits_next = model.forward_prompt(&prompt, &mut state);
    let dt_pre = t_pre.elapsed();
    println!(
        "prefill: {} tokens in {:.0} ms ({:.1} tok/s)",
        prompt.len(),
        dt_pre.as_secs_f64() * 1e3,
        prompt.len() as f64 / dt_pre.as_secs_f64(),
    );
    let mut generated: Vec<u32> = Vec::new();
    println!("greedy decode:");
    for step in 0..n_tokens {
        let t = Instant::now();
        let logits = if step == 0 {
            std::mem::take(&mut logits_next)
        } else {
            model.forward_token(*generated.last().unwrap(), &mut state)
        };
        let dt = t.elapsed();
        let finite = logits.iter().all(|v| v.is_finite());
        let (argmax, max) = logits.iter().enumerate()
            .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| {
                if v > bv { (i, v) } else { (bi, bv) }
            });
        println!(
            "  step {step:>2}: out={argmax:>6}  max={max:>8.3}  finite={finite}  {:.0} ms",
            dt.as_secs_f64() * 1e3,
        );
        if !finite {
            eprintln!("FAILED: non-finite logits");
            std::process::exit(1);
        }
        generated.push(argmax as u32);
    }
    println!("\ngenerated text: {:?}", tokenizer.decode(&generated));
    #[cfg(feature = "fabric")]
    {
        let entries = jouleclaw_loader_gguf::fabric_dispatch::ledger_snapshot();
        if !entries.is_empty() {
            let total_gf: f64 = entries.iter().map(|e| e.gflops).sum();
            println!("\nfabric attribution (whole run):");
            for e in &entries {
                let uj = e
                    .projected_uj
                    .map(|u| format!("{:.1} J projected", u / 1e6))
                    .unwrap_or_else(|| "no calibration".into());
                println!(
                    "  {:>18}: {:>6} ops, {:>9.1} GFLOP ({:>4.1}%), {}",
                    e.member_id, e.ops, e.gflops,
                    100.0 * e.gflops / total_gf, uj,
                );
            }
        }
    }
    println!("PASS");
    Ok(())
}