rustytorch/crates/models/rtx-csm/examples/bench.rs

//! Benchmark + eval harness for rtx-csm.
//!
//! Reads a prompts file (one prompt per line), generates a WAV per prompt,
//! and emits a JSON manifest you can feed to any external scorer (Whisper for
//! WER, WavLM for speaker similarity, TTSDS2, NISQA, etc.).
//!
//! Bundled prompts from the Harvard Sentences set (phonetically balanced) plus
//! a handful of CSM-specific stress tests (brackets, times, repeated phrases).
//!
//! Example:
//! ```
//! cargo run -p rtx-csm --release --features metal --example bench -- \
//!     --out-dir /tmp/csm-bench --prompts harvard
//! ```
//!
//! The manifest is `out-dir/manifest.json`. Run your scorer over its `samples`
//! list; `prompt_text` is the ground truth for WER.

use anyhow::Result;
use clap::{Parser, ValueEnum};
use rtx_csm::{audio_io, compute_wer, GenerateOptions, Generator, PostProcess, Segment, WerResult};
use serde::Serialize;
use std::fs;
use std::path::{Path, PathBuf};
use std::time::Instant;

#[derive(Debug, Clone, Copy, ValueEnum)]
enum PromptSet {
    /// 10 phonetically balanced Harvard sentences (set 1, lines 1–10).
    Harvard,
    /// Brackets, times, abbreviations, repeated phrases — stress tests.
    Stress,
    /// Both sets concatenated.
    All,
}

#[derive(Debug, Parser)]
#[command(name = "csm-bench", about = "Generate eval set + manifest for rtx-csm")]
struct Cli {
    #[arg(long, default_value = "/tmp/csm-bench")]
    out_dir: PathBuf,

    #[arg(long, default_value = "harvard")]
    prompts: PromptSet,

    /// Optional path to a custom prompts file (one prompt per line). Overrides --prompts.
    #[arg(long)]
    prompts_file: Option<PathBuf>,

    #[arg(long, default_value_t = 0)]
    speaker: u32,

    #[arg(long, default_value_t = 8000)]
    max_audio_ms: u32,

    #[arg(long, default_value_t = 0.9)]
    temperature: f64,

    #[arg(long, default_value_t = 50)]
    top_k: usize,

    #[arg(long, default_value_t = 0.9)]
    top_p: f64,

    #[arg(long, default_value_t = 42)]
    seed: u64,

    #[arg(long)]
    cpu: bool,

    #[arg(long)]
    raw: bool,

    /// Score each generated WAV against the prompt with WER. Requires `whisper`
    /// (or `whisper-cpp`) on PATH unless `--ground-truth-dir` is provided.
    #[arg(long)]
    score: bool,

    /// Whisper CLI to invoke (only used when --score is set).
    #[arg(long, default_value = "whisper")]
    whisper_bin: String,

    /// Whisper model size — "tiny.en" / "base.en" / "small.en" — passed via --model.
    #[arg(long, default_value = "tiny.en")]
    whisper_model: String,

    /// Optional dir containing pre-computed transcripts (named `sample_NNN.txt`),
    /// e.g. from a separate ASR run. Skips the whisper invocation entirely.
    #[arg(long)]
    ground_truth_dir: Option<PathBuf>,

    /// Use in-process whisper-rs instead of shelling out. Requires the
    /// `asr` (or `asr-metal`) feature flag at build time.
    #[arg(long)]
    asr_inproc: bool,
}

const HARVARD_PROMPTS: &[&str] = &[
    "The birch canoe slid on the smooth planks.",
    "Glue the sheet to the dark blue background.",
    "It's easy to tell the depth of a well.",
    "These days a chicken leg is a rare dish.",
    "Rice is often served in round bowls.",
    "The juice of lemons makes fine punch.",
    "The box was thrown beside the parked truck.",
    "The hogs were fed chopped corn and garbage.",
    "Four hours of steady work faced us.",
    "A large size in stockings is hard to sell.",
];

const STRESS_PROMPTS: &[&str] = &[
    "Meet me at 10:30 in the morning.",
    "It is now 3:05 pm exactly.",
    "Go to 12:00 am tonight.",
    "The list contains apples bananas and oranges.",
    "She said hello and waved goodbye.",
    "I am thinking deeply about this question.",
    "Yes yes yes yes yes yes.",
    "He laughed loudly at the joke.",
];

#[derive(Serialize)]
struct ManifestEntry {
    index: usize,
    prompt: String,
    wav: String,
    samples: usize,
    duration_s: f32,
    generation_ms: u128,
    realtime_factor: f32,
    #[serde(skip_serializing_if = "Option::is_none")]
    asr_text: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    wer: Option<WerEntry>,
}

#[derive(Serialize)]
struct WerEntry {
    rate: f32,
    substitutions: usize,
    deletions: usize,
    insertions: usize,
    reference_words: usize,
    hypothesis_words: usize,
}

impl From<WerResult> for WerEntry {
    fn from(r: WerResult) -> Self {
        Self {
            rate: r.rate(),
            substitutions: r.substitutions,
            deletions: r.deletions,
            insertions: r.insertions,
            reference_words: r.reference_words,
            hypothesis_words: r.hypothesis_words,
        }
    }
}

#[derive(Serialize)]
struct Manifest<'a> {
    model: &'a str,
    sample_rate: u32,
    speaker: u32,
    seed: u64,
    temperature: f64,
    top_k: usize,
    top_p: f64,
    post_processing: bool,
    samples: Vec<ManifestEntry>,
    total_audio_s: f32,
    total_generation_s: f32,
    avg_realtime_factor: f32,
    #[serde(skip_serializing_if = "Option::is_none")]
    aggregate_wer: Option<f32>,
}

/// Run external Whisper CLI on a WAV file, return the transcript text or an
/// error if the binary isn't available / failed.
fn whisper_transcribe(
    whisper_bin: &str,
    model: &str,
    wav: &Path,
    out_dir: &Path,
) -> Result<String> {
    let status = std::process::Command::new(whisper_bin)
        .arg(wav)
        .args(["--model", model, "--output_format", "txt", "--output_dir"])
        .arg(out_dir)
        .arg("--language")
        .arg("en")
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status();
    let status = match status {
        Ok(s) => s,
        Err(e) => anyhow::bail!("failed to invoke {whisper_bin}: {e}"),
    };
    if !status.success() {
        anyhow::bail!("{whisper_bin} exited with {status}");
    }
    let stem = wav.file_stem().and_then(|s| s.to_str()).unwrap_or("");
    let txt_path = out_dir.join(format!("{stem}.txt"));
    Ok(fs::read_to_string(txt_path)?.trim().to_string())
}

/// Read pre-computed transcript at `<dir>/<stem>.txt`.
fn read_ground_truth_transcript(dir: &Path, wav: &Path) -> Result<String> {
    let stem = wav.file_stem().and_then(|s| s.to_str()).unwrap_or("");
    let txt = dir.join(format!("{stem}.txt"));
    Ok(fs::read_to_string(&txt)?.trim().to_string())
}

#[cfg(feature = "asr")]
fn inproc_whisper(pcm_24k: &[f32]) -> Result<String> {
    use std::sync::OnceLock;
    static ASR: OnceLock<rtx_csm::asr::WhisperAsr> = OnceLock::new();
    let asr = ASR.get_or_init(|| {
        rtx_csm::asr::WhisperAsr::load_default().expect("failed to load default whisper model")
    });
    Ok(asr.transcribe_24k(pcm_24k)?)
}

#[cfg(not(feature = "asr"))]
fn inproc_whisper(_pcm_24k: &[f32]) -> Result<String> {
    anyhow::bail!("--asr-inproc requires building with --features asr (or asr-metal)")
}

fn load_prompts(cli: &Cli) -> Result<Vec<String>> {
    if let Some(p) = &cli.prompts_file {
        let body = fs::read_to_string(p)?;
        return Ok(body
            .lines()
            .map(|l| l.trim().to_string())
            .filter(|l| !l.is_empty() && !l.starts_with('#'))
            .collect());
    }
    let static_set: Vec<&str> = match cli.prompts {
        PromptSet::Harvard => HARVARD_PROMPTS.to_vec(),
        PromptSet::Stress => STRESS_PROMPTS.to_vec(),
        PromptSet::All => HARVARD_PROMPTS
            .iter()
            .chain(STRESS_PROMPTS.iter())
            .copied()
            .collect(),
    };
    Ok(static_set.into_iter().map(String::from).collect())
}

fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();
    fs::create_dir_all(&cli.out_dir)?;

    let prompts = load_prompts(&cli)?;
    if prompts.is_empty() {
        anyhow::bail!("no prompts to run");
    }
    eprintln!("running {} prompts → {}", prompts.len(), cli.out_dir.display());

    let device = if cli.cpu {
        candle_core::Device::Cpu
    } else {
        Generator::default_device()?
    };
    let mut generator = Generator::load_csm_1b(&device)?;
    let post = if cli.raw { PostProcess::disabled() } else { PostProcess::default() };
    let opts = GenerateOptions {
        max_audio_ms: cli.max_audio_ms,
        temperature: cli.temperature,
        top_k: cli.top_k,
        top_p: cli.top_p,
        seed: cli.seed,
        ..GenerateOptions::default()
    };

    let mut entries: Vec<ManifestEntry> = Vec::with_capacity(prompts.len());
    let mut total_audio_s = 0.0f32;
    let mut total_gen_s = 0.0f32;

    let no_context: Vec<Segment> = Vec::new();
    for (i, prompt) in prompts.iter().enumerate() {
        let wav_name = format!("sample_{:03}.wav", i);
        let wav_path = cli.out_dir.join(&wav_name);

        let t0 = Instant::now();
        let mut pcm = generator.generate(prompt, cli.speaker, &no_context, opts)?;
        let gen_ms = t0.elapsed().as_millis();
        post.apply(&mut pcm, generator.config.sample_rate)?;
        audio_io::write_wav_24k_mono(&wav_path, &pcm)?;

        let samples = pcm.len();
        let duration_s = samples as f32 / generator.config.sample_rate as f32;
        let realtime_factor = if duration_s > 0.0 {
            (gen_ms as f32 / 1000.0) / duration_s
        } else {
            0.0
        };
        // Optional WER scoring.
        let (asr_text, wer_entry) = if cli.score {
            let asr_result = if let Some(gt_dir) = cli.ground_truth_dir.as_ref() {
                read_ground_truth_transcript(gt_dir, &wav_path)
            } else if cli.asr_inproc {
                inproc_whisper(&pcm)
            } else {
                whisper_transcribe(&cli.whisper_bin, &cli.whisper_model, &wav_path, &cli.out_dir)
            };
            match asr_result {
                Ok(text) => {
                    let w = compute_wer(prompt, &text);
                    (Some(text), Some(WerEntry::from(w)))
                }
                Err(e) => {
                    eprintln!("    asr failed: {e}");
                    (None, None)
                }
            }
        } else {
            (None, None)
        };
        let wer_str = wer_entry
            .as_ref()
            .map(|w| format!(" wer={:.2}", w.rate))
            .unwrap_or_default();
        eprintln!(
            "  [{:>2}/{}] {:.2}s audio in {} ms ({:.2}× realtime){} — {}",
            i + 1,
            prompts.len(),
            duration_s,
            gen_ms,
            realtime_factor,
            wer_str,
            prompt.chars().take(48).collect::<String>(),
        );
        total_audio_s += duration_s;
        total_gen_s += gen_ms as f32 / 1000.0;
        entries.push(ManifestEntry {
            index: i,
            prompt: prompt.clone(),
            wav: wav_name,
            samples,
            duration_s,
            generation_ms: gen_ms,
            realtime_factor,
            asr_text,
            wer: wer_entry,
        });
    }

    let avg_rtf = if total_audio_s > 0.0 {
        total_gen_s / total_audio_s
    } else {
        0.0
    };
    // Aggregate WER: micro-average over all scored samples.
    let aggregate_wer = if cli.score {
        let mut errs = 0usize;
        let mut refs = 0usize;
        for e in &entries {
            if let Some(w) = e.wer.as_ref() {
                errs += w.substitutions + w.deletions + w.insertions;
                refs += w.reference_words;
            }
        }
        if refs > 0 { Some(errs as f32 / refs as f32) } else { None }
    } else {
        None
    };
    let manifest = Manifest {
        model: "sesame/csm-1b",
        sample_rate: generator.config.sample_rate,
        speaker: cli.speaker,
        seed: cli.seed,
        temperature: cli.temperature,
        top_k: cli.top_k,
        top_p: cli.top_p,
        post_processing: !cli.raw,
        samples: entries,
        total_audio_s,
        total_generation_s: total_gen_s,
        avg_realtime_factor: avg_rtf,
        aggregate_wer,
    };
    let manifest_path = cli.out_dir.join("manifest.json");
    fs::write(&manifest_path, serde_json::to_string_pretty(&manifest)?)?;
    eprintln!(
        "\nwrote {} samples + manifest.json\ntotal audio: {:.1}s, total gen: {:.1}s, avg {:.2}× realtime",
        manifest.samples.len(),
        total_audio_s,
        total_gen_s,
        avg_rtf,
    );
    if let Some(w) = aggregate_wer {
        eprintln!("aggregate WER: {:.3}", w);
    }
    eprintln!("\nNext: feed manifest.json to your scorer of choice. Examples:");
    eprintln!(
        "  whisper {}/sample_*.wav --model base.en --output_format json",
        cli.out_dir.display()
    );
    eprintln!("  python -m ttsds.benchmark --manifest {}", manifest_path.display());
    Ok(())
}