Lumotia/crates/transcription/tests/jfk_bench.rs

//! Benchmark: load the JFK WAV from disk, transcribe it via whisper-rs.
//! Reports cold-load time, transcribe time, RTF, peak RSS.
//!
//! Gated on env vars so it never runs in CI without setup:
//!   LUMOTIA_WHISPER_TEST_MODEL=/path/to/ggml-tiny.bin
//!   LUMOTIA_WHISPER_TEST_AUDIO=/path/to/jfk.wav

use std::env;
use std::time::Instant;

#[test]
fn jfk_transcription_benchmark() {
    let Ok(model_path) = env::var("LUMOTIA_WHISPER_TEST_MODEL") else {
        eprintln!("LUMOTIA_WHISPER_TEST_MODEL not set — skipping");
        return;
    };
    let Ok(audio_path) = env::var("LUMOTIA_WHISPER_TEST_AUDIO") else {
        eprintln!("LUMOTIA_WHISPER_TEST_AUDIO not set — skipping");
        return;
    };

    use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};

    eprintln!("[bench] loading WAV: {audio_path}");
    let bytes = std::fs::read(&audio_path).expect("read wav");
    // Minimal RIFF/WAV parse: skip the 44-byte canonical header for PCM-16-mono-16kHz.
    // Sanity-check magic bytes + format.
    assert_eq!(&bytes[0..4], b"RIFF", "expected RIFF");
    assert_eq!(&bytes[8..12], b"WAVE", "expected WAVE");
    let sample_rate = u32::from_le_bytes(bytes[24..28].try_into().unwrap());
    let channels = u16::from_le_bytes(bytes[22..24].try_into().unwrap());
    let bits = u16::from_le_bytes(bytes[34..36].try_into().unwrap());
    eprintln!(
        "[bench] wav spec: {} Hz, {} ch, {}-bit",
        sample_rate, channels, bits
    );
    assert_eq!(sample_rate, 16_000, "expected 16 kHz wav");
    assert_eq!(channels, 1, "expected mono");
    assert_eq!(bits, 16, "expected 16-bit PCM");

    let pcm = &bytes[44..];
    let samples: Vec<f32> = pcm
        .chunks_exact(2)
        .map(|c| i16::from_le_bytes([c[0], c[1]]) as f32 / 32768.0)
        .collect();
    let audio_secs = samples.len() as f64 / sample_rate as f64;
    eprintln!(
        "[bench] audio length: {} samples = {:.2}s",
        samples.len(),
        audio_secs
    );

    let rss_before_load_kb = read_rss_kb();
    eprintln!(
        "[bench] RSS before model load: {} MB",
        rss_before_load_kb / 1024
    );

    let load_start = Instant::now();
    let ctx = WhisperContext::new_with_params(&model_path, WhisperContextParameters::default())
        .expect("whisper model load");
    let load_dur = load_start.elapsed();
    eprintln!("[bench] model load: {:.2}s", load_dur.as_secs_f64());

    let rss_after_load_kb = read_rss_kb();
    eprintln!(
        "[bench] RSS after model load: {} MB (delta +{} MB)",
        rss_after_load_kb / 1024,
        (rss_after_load_kb.saturating_sub(rss_before_load_kb)) / 1024
    );

    let mut state = ctx.create_state().expect("whisper state");
    let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
    params.set_language(Some("en"));
    params.set_n_threads(6);
    params.set_print_special(false);
    params.set_print_progress(false);
    params.set_print_realtime(false);

    // Cold transcription (first run)
    let cold_start = Instant::now();
    state.full(params, &samples).expect("transcribe cold");
    let cold_dur = cold_start.elapsed();

    let n = state.full_n_segments();
    let mut full_text = String::new();
    for i in 0..n {
        let seg = state.get_segment(i).expect("get_segment");
        full_text.push_str(seg.to_str().unwrap_or(""));
    }
    eprintln!(
        "[bench] cold transcribe: {:.2}s ({} segments, RTF={:.3})",
        cold_dur.as_secs_f64(),
        n,
        cold_dur.as_secs_f64() / audio_secs
    );
    eprintln!("[bench] transcript: {}", full_text.trim());
    let rss_after_cold_kb = read_rss_kb();
    eprintln!("[bench] RSS after cold xc: {} MB", rss_after_cold_kb / 1024);

    // Warm transcription (second run, same state)
    let mut state2 = ctx.create_state().expect("whisper state 2");
    let mut params2 = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
    params2.set_language(Some("en"));
    params2.set_n_threads(6);
    params2.set_print_special(false);
    params2.set_print_progress(false);
    params2.set_print_realtime(false);
    let warm_start = Instant::now();
    state2.full(params2, &samples).expect("transcribe warm");
    let warm_dur = warm_start.elapsed();
    eprintln!(
        "[bench] warm transcribe: {:.2}s (RTF={:.3})",
        warm_dur.as_secs_f64(),
        warm_dur.as_secs_f64() / audio_secs
    );

    let rss_final_kb = read_rss_kb();
    eprintln!("[bench] RSS final: {} MB", rss_final_kb / 1024);

    eprintln!();
    eprintln!("=== SUMMARY ===");
    eprintln!("audio:       {:.2}s", audio_secs);
    eprintln!("model_load:  {:.2}s", load_dur.as_secs_f64());
    eprintln!(
        "cold xc:     {:.2}s  RTF={:.3}",
        cold_dur.as_secs_f64(),
        cold_dur.as_secs_f64() / audio_secs
    );
    eprintln!(
        "warm xc:     {:.2}s  RTF={:.3}",
        warm_dur.as_secs_f64(),
        warm_dur.as_secs_f64() / audio_secs
    );
    eprintln!("RSS peak:    {} MB", rss_final_kb / 1024);
}

fn read_rss_kb() -> u64 {
    let pid = std::process::id();
    let s = std::fs::read_to_string(format!("/proc/{pid}/status")).unwrap_or_default();
    for line in s.lines() {
        if let Some(rest) = line.strip_prefix("VmRSS:") {
            return rest
                .split_whitespace()
                .next()
                .and_then(|n| n.parse::<u64>().ok())
                .unwrap_or(0);
        }
    }
    0
}