osobhandClaude Opus 4.7
a5cedfb46a
8-gen bench (4 emotions × 2 corpora) at seed=42 against firdhokk Whisper-LV3: target RAVDESS CREMA-D happy happy (0.999) ✓ happy (0.999) ✓ angry neutral (0.92) sad (0.99) fearful happy (0.998) fearful (0.984) ✓ sad angry (0.99) fearful (0.99) CREMA-D 2/4 vs RAVDESS 1/4. Larger / more naturalistic corpus produces more class-pure fearful direction. Neither corpus solves angry or sad — recipe shifts into 'vague expressivity' rather than class-specific corners. Practical: prefer CREMA-D when available; A/B both per emotion if class precision matters. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
424 lines
13 KiB
Rust
424 lines
13 KiB
Rust
//! 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::{GenerateOptions, Generator, PostProcess, Segment, WerResult, audio_io, compute_wer};
|
||
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.clone())?;
|
||
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(())
|
||
}
|