414 lines
15 KiB
Rust
414 lines
15 KiB
Rust
//! Per-pair TTS quality eval: speaker similarity + WER + amplitude.
|
|
//!
|
|
//! Reads a JSONL where each row has at minimum `ref_wav` and `gen_wav`,
|
|
//! optionally `ref_text` (used as the reference transcript for WER). The
|
|
//! reference text comes from the manifest because the reference audio
|
|
//! might not have a clean public transcript (e.g. context-conditioning
|
|
//! reference clips).
|
|
//!
|
|
//! Designed as the metric foundation for Sprint 1 of the Phase 8 roadmap.
|
|
//! Originally TTSDS2 (arXiv 2506.19441) was on the menu — its install on
|
|
//! Python 3.12 + torchaudio 2.x is broken (`torchaudio.sox_effects` was
|
|
//! removed; `pyannote.audio` 3.1 calls `set_audio_backend`; openai-whisper
|
|
//! pinned dep needs `pkg_resources` which requires `setuptools<81`). For
|
|
//! now we score with what we already own end-to-end in candle:
|
|
//!
|
|
//! - **speaker_cosine**: WavLM-SV (microsoft/wavlm-base-plus-sv)
|
|
//! - **wer**: Moonshine v2 transcribe vs `ref_text`
|
|
//! - **peak_db / rms_db**: full-band amplitude of `gen_wav`
|
|
//!
|
|
//! Output: one JSONL line per input, with the input's fields plus a
|
|
//! `metrics` object. Stable column order = appendable to a CSV/jq pipeline.
|
|
//!
|
|
//! ```bash
|
|
//! target/release/examples/quality_eval \
|
|
//! --in /tmp/i2d/eval_pairs.jsonl \
|
|
//! --wavlm-sv /tmp/wavlm_sv.safetensors \
|
|
//! --out /tmp/i2d/scores.jsonl
|
|
//! ```
|
|
|
|
use anyhow::{Context, Result};
|
|
use candle_core::{Device, Tensor};
|
|
use clap::Parser;
|
|
use hf_hub::api::sync::Api;
|
|
use rtx_csm::speaker_sim::{SpeakerSimilarity, WavLmSimilarity};
|
|
use rtx_csm::{audio_io, moonshine};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::io::{BufRead, BufWriter, Write};
|
|
use std::path::PathBuf;
|
|
|
|
#[derive(Debug, Parser)]
|
|
struct Cli {
|
|
/// Input JSONL with rows: `{"ref_wav": "...", "gen_wav": "...", "ref_text": "..."}`.
|
|
/// `ref_text` is optional; without it WER is `null`.
|
|
#[arg(long = "in")]
|
|
input: PathBuf,
|
|
/// Output JSONL. One line per input row, with the input's fields plus
|
|
/// a `metrics` object.
|
|
#[arg(long)]
|
|
out: PathBuf,
|
|
/// Converted WavLM-SV safetensors. See `examples/wavlm_sv_convert`.
|
|
#[arg(long)]
|
|
wavlm_sv: PathBuf,
|
|
/// Max decoded tokens for Moonshine (capped at 194 by config).
|
|
#[arg(long, default_value_t = 120)]
|
|
max_tokens: usize,
|
|
/// Optional path to the emotion2vec_plus_base `.pt` checkpoint. When
|
|
/// set, each row also gets emotion-classifier metrics. Pair with
|
|
/// `--target-emotion` to score the probability of a specific class.
|
|
///
|
|
/// **Caveat (2026-04-30):** the emotion2vec_plus_base checkpoint as
|
|
/// loaded through our pickle path produces saturated outputs —
|
|
/// every audio sample classifies as "Surprised" with probability
|
|
/// ≥ 0.99, INCLUDING ground-truth RAVDESS clips with explicit
|
|
/// neutral / happy / angry / etc. labels. The metric is wired
|
|
/// correctly (verified against the trait `EmotionDetector::classify`
|
|
/// path) but the underlying classifier collapses to one dominant
|
|
/// class on most input. Useful for verifying the emotion2vec
|
|
/// pipeline works at all; not useful for verifying generation
|
|
/// quality. Treat the output as diagnostic, not as ground truth.
|
|
#[arg(long)]
|
|
emotion2vec: Option<PathBuf>,
|
|
/// Target emotion to score the probability of. One of: angry,
|
|
/// disgusted, fearful, happy, neutral, excited, sad, surprised.
|
|
/// Ignored unless `--emotion2vec` is also set. See caveat above.
|
|
#[arg(long)]
|
|
target_emotion: Option<String>,
|
|
/// Force CPU device.
|
|
#[arg(long)]
|
|
cpu: bool,
|
|
}
|
|
|
|
fn emotion_index(name: &str) -> Option<usize> {
|
|
// Maps a CLI emotion name to emotion2vec class index. Order from
|
|
// src/emotion2vec.rs::Classifier::tag_for_class:
|
|
match name.trim().to_lowercase().as_str() {
|
|
"angry" => Some(0),
|
|
"disgust" | "disgusted" => Some(1),
|
|
"fearful" => Some(2),
|
|
"happy" => Some(3),
|
|
"neutral" => Some(4),
|
|
"excited" | "other" => Some(5),
|
|
"sad" => Some(6),
|
|
"surprised" | "calm" => Some(7), // emotion2vec lumps calm in here
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
fn softmax(logits: &[f32]) -> Vec<f32> {
|
|
let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
|
let exp: Vec<f32> = logits.iter().map(|&v| (v - max).exp()).collect();
|
|
let sum: f32 = exp.iter().sum();
|
|
exp.into_iter().map(|v| v / sum.max(1e-12)).collect()
|
|
}
|
|
|
|
#[derive(Debug, Deserialize, Serialize)]
|
|
struct PairRow {
|
|
ref_wav: PathBuf,
|
|
gen_wav: PathBuf,
|
|
#[serde(skip_serializing_if = "Option::is_none", default)]
|
|
ref_text: Option<String>,
|
|
/// Free-form per-row metadata (iteration index, model id, etc).
|
|
/// Round-tripped unchanged into the output row.
|
|
#[serde(flatten, default)]
|
|
extra: serde_json::Map<String, serde_json::Value>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
struct ScoredRow {
|
|
ref_wav: PathBuf,
|
|
gen_wav: PathBuf,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
ref_text: Option<String>,
|
|
#[serde(flatten)]
|
|
extra: serde_json::Map<String, serde_json::Value>,
|
|
metrics: Metrics,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
struct Metrics {
|
|
/// Cosine in [-1, 1]. Higher = more similar speaker identity.
|
|
speaker_cosine: f32,
|
|
/// Word error rate in [0, 1+]. None if no `ref_text`.
|
|
wer: Option<f32>,
|
|
/// Moonshine's transcript of `gen_wav` (always emitted; useful for
|
|
/// debugging).
|
|
transcript: String,
|
|
/// Generated audio characteristics.
|
|
gen_peak_db: f32,
|
|
gen_rms_db: f32,
|
|
/// Probability of the target emotion (None unless --emotion2vec +
|
|
/// --target-emotion are passed).
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
target_emotion_prob: Option<f32>,
|
|
/// Argmax emotion class label per emotion2vec.
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
top_emotion: Option<String>,
|
|
/// Probability of the argmax emotion class.
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
top_emotion_prob: Option<f32>,
|
|
}
|
|
|
|
fn main() -> Result<()> {
|
|
tracing_subscriber::fmt().with_target(false).init();
|
|
let cli = Cli::parse();
|
|
|
|
let device = if cli.cpu {
|
|
Device::Cpu
|
|
} else if candle_core::utils::metal_is_available() {
|
|
Device::new_metal(0)?
|
|
} else {
|
|
Device::Cpu
|
|
};
|
|
tracing::info!("device: {device:?}");
|
|
|
|
// Speaker similarity model.
|
|
let sv = WavLmSimilarity::load(&cli.wavlm_sv, &device)?;
|
|
tracing::info!("loaded WavLM-SV from {}", cli.wavlm_sv.display());
|
|
|
|
// Moonshine encoder + decoder + tokenizer (all from HF cache).
|
|
let api = Api::new()?;
|
|
let repo = api.model("UsefulSensors/moonshine-tiny".to_string());
|
|
let m_weights = repo.get("model.safetensors")?;
|
|
let m_tok = repo.get("tokenizer.json")?;
|
|
let m_cfg = moonshine::MoonshineConfig::tiny();
|
|
let (m_enc, m_dec) = moonshine::load_full(&m_weights, &device, &m_cfg)?;
|
|
let m_tokenizer =
|
|
moonshine::load_tokenizer(&m_tok).map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
|
|
tracing::info!("loaded Moonshine-tiny");
|
|
|
|
// Optional emotion2vec for emotion-classification metrics.
|
|
let target_idx = cli.target_emotion.as_deref().and_then(emotion_index);
|
|
if cli.emotion2vec.is_some() && cli.target_emotion.is_some() && target_idx.is_none() {
|
|
anyhow::bail!(
|
|
"--target-emotion '{}' is unrecognized — must be angry, disgusted, fearful, happy, neutral, excited, sad, or surprised",
|
|
cli.target_emotion.as_deref().unwrap_or("")
|
|
);
|
|
}
|
|
let emo_model = if let Some(path) = cli.emotion2vec.as_ref() {
|
|
let m = rtx_csm::emotion2vec::Emotion2Vec::load_from_pickle(path, &device)?;
|
|
tracing::info!(
|
|
"loaded emotion2vec from {} (target='{}', class_idx={:?})",
|
|
path.display(),
|
|
cli.target_emotion.as_deref().unwrap_or("(none)"),
|
|
target_idx,
|
|
);
|
|
Some(m)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Stream the input JSONL line by line so big batches don't blow memory.
|
|
let f = std::fs::File::open(&cli.input)
|
|
.with_context(|| format!("open input {}", cli.input.display()))?;
|
|
let reader = std::io::BufReader::new(f);
|
|
|
|
if let Some(parent) = cli.out.parent() {
|
|
std::fs::create_dir_all(parent).ok();
|
|
}
|
|
let out_file = std::fs::File::create(&cli.out)
|
|
.with_context(|| format!("create out {}", cli.out.display()))?;
|
|
let mut out_w = BufWriter::new(out_file);
|
|
|
|
let mut n_rows = 0usize;
|
|
let mut sum_cos = 0.0f64;
|
|
let mut sum_wer = 0.0f64;
|
|
let mut n_wer = 0usize;
|
|
for (idx, line) in reader.lines().enumerate() {
|
|
let line = line?;
|
|
if line.trim().is_empty() {
|
|
continue;
|
|
}
|
|
let row: PairRow =
|
|
serde_json::from_str(&line).with_context(|| format!("parse row {idx}: {line}"))?;
|
|
|
|
// Speaker similarity at 16 kHz. WavLM-SV's TDNN front-end has
|
|
// kernel size 5 with stride 2 chains, so it requires at least a
|
|
// few hundred samples of audio. Below ~0.25 s the model errors
|
|
// out — score those as -1.0 (sentinel) so the eval still emits
|
|
// a row instead of bailing on the whole batch.
|
|
let ref_16k = audio_io::load_mono_at_rate(&row.ref_wav, 16_000)?;
|
|
let gen_16k = audio_io::load_mono_at_rate(&row.gen_wav, 16_000)?;
|
|
const MIN_SV_SAMPLES: usize = 4000; // 0.25 s at 16 kHz
|
|
let speaker_cosine = if gen_16k.len() < MIN_SV_SAMPLES || ref_16k.len() < MIN_SV_SAMPLES {
|
|
tracing::warn!(
|
|
"row {idx}: gen={}s ref={}s — too short for WavLM-SV; cosine=-1",
|
|
gen_16k.len() as f32 / 16_000.0,
|
|
ref_16k.len() as f32 / 16_000.0,
|
|
);
|
|
-1.0f32
|
|
} else {
|
|
sv.score(&ref_16k, &gen_16k)? as f32
|
|
};
|
|
|
|
// Transcribe gen_wav via Moonshine (16 kHz).
|
|
let pcm_t = Tensor::from_vec(gen_16k.clone(), (1, 1, gen_16k.len()), &device)?;
|
|
let enc_out = m_enc.forward(&pcm_t)?;
|
|
let token_ids = m_dec.generate_cached(&enc_out, &m_cfg, cli.max_tokens)?;
|
|
let transcript = m_tokenizer
|
|
.decode(&token_ids, /* skip_special */ true)
|
|
.map_err(|e| anyhow::anyhow!("detok: {e}"))?
|
|
.trim()
|
|
.to_string();
|
|
|
|
let wer = row
|
|
.ref_text
|
|
.as_ref()
|
|
.map(|t| word_error_rate(t, &transcript));
|
|
|
|
// Amplitude of gen_wav at native rate.
|
|
let gen_native = audio_io::load_mono_24k(&row.gen_wav)?;
|
|
let (peak_db, rms_db) = peak_rms_db(&gen_native);
|
|
|
|
// Optional emotion classification on gen_wav at 16 kHz.
|
|
let (target_emotion_prob, top_emotion, top_emotion_prob) = if let Some(emo) =
|
|
emo_model.as_ref()
|
|
{
|
|
if gen_16k.len() < MIN_SV_SAMPLES {
|
|
(None, None, None)
|
|
} else {
|
|
use rtx_csm::ser::EmotionDetector;
|
|
// Run forward to get logits (the EmotionDetector trait
|
|
// gives us only argmax via classify; we need the full
|
|
// probability vector).
|
|
let n = gen_16k.len();
|
|
let mean = gen_16k.iter().sum::<f32>() / n as f32;
|
|
let var = gen_16k.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / n as f32;
|
|
let std = var.sqrt().max(1e-7);
|
|
let normed: Vec<f32> = gen_16k.iter().map(|x| (x - mean) / std).collect();
|
|
let audio = Tensor::from_vec(normed, (1, 1, n), &device)?;
|
|
let logits_t = emo.forward(&audio)?;
|
|
let logits = logits_t.flatten_all()?.to_vec1::<f32>()?;
|
|
let probs = softmax(&logits);
|
|
let target_p = target_idx.and_then(|i| probs.get(i).copied());
|
|
let (top_i, top_p) = probs
|
|
.iter()
|
|
.enumerate()
|
|
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
|
|
.map(|(i, &p)| (i, p))
|
|
.unwrap_or((0, 0.0));
|
|
// Use the trait's tag for human-readable label.
|
|
let top_label = emo
|
|
.classify(&gen_16k)
|
|
.map(|l| format!("{l:?}"))
|
|
.unwrap_or_else(|_| format!("class_{top_i}"));
|
|
let _ = top_i;
|
|
(target_p, Some(top_label), Some(top_p))
|
|
}
|
|
} else {
|
|
(None, None, None)
|
|
};
|
|
|
|
let scored = ScoredRow {
|
|
ref_wav: row.ref_wav,
|
|
gen_wav: row.gen_wav,
|
|
ref_text: row.ref_text,
|
|
extra: row.extra,
|
|
metrics: Metrics {
|
|
speaker_cosine,
|
|
wer,
|
|
transcript: transcript.clone(),
|
|
gen_peak_db: peak_db,
|
|
gen_rms_db: rms_db,
|
|
target_emotion_prob,
|
|
top_emotion,
|
|
top_emotion_prob,
|
|
},
|
|
};
|
|
writeln!(out_w, "{}", serde_json::to_string(&scored)?)?;
|
|
out_w.flush()?;
|
|
|
|
n_rows += 1;
|
|
sum_cos += speaker_cosine as f64;
|
|
if let Some(w) = wer {
|
|
sum_wer += w as f64;
|
|
n_wer += 1;
|
|
}
|
|
tracing::info!(
|
|
"row {idx}: cos={speaker_cosine:.3} wer={} peak={peak_db:.2} rms={rms_db:.2}",
|
|
wer.map(|w| format!("{w:.3}")).unwrap_or_else(|| "-".into())
|
|
);
|
|
}
|
|
|
|
if n_rows == 0 {
|
|
anyhow::bail!("no rows processed (empty input?)");
|
|
}
|
|
eprintln!("--- summary ---");
|
|
eprintln!("rows {n_rows}");
|
|
eprintln!("speaker_cosine {:.3} (mean)", sum_cos / n_rows as f64);
|
|
if n_wer > 0 {
|
|
eprintln!(
|
|
"wer {:.3} (mean over {n_wer})",
|
|
sum_wer / n_wer as f64
|
|
);
|
|
} else {
|
|
eprintln!("wer (no ref_text in any row)");
|
|
}
|
|
eprintln!("wrote {}", cli.out.display());
|
|
Ok(())
|
|
}
|
|
|
|
/// Word error rate: Levenshtein on whitespace-split, lowercased,
|
|
/// punctuation-stripped tokens. Returns errors / max(ref_len, 1).
|
|
fn word_error_rate(reference: &str, hypothesis: &str) -> f32 {
|
|
let r = tokenize(reference);
|
|
let h = tokenize(hypothesis);
|
|
if r.is_empty() {
|
|
return if h.is_empty() { 0.0 } else { 1.0 };
|
|
}
|
|
let n = r.len();
|
|
let m = h.len();
|
|
let mut prev: Vec<usize> = (0..=m).collect();
|
|
let mut curr = vec![0usize; m + 1];
|
|
for i in 1..=n {
|
|
curr[0] = i;
|
|
for j in 1..=m {
|
|
let cost = if r[i - 1] == h[j - 1] { 0 } else { 1 };
|
|
curr[j] = (prev[j] + 1).min(curr[j - 1] + 1).min(prev[j - 1] + cost);
|
|
}
|
|
std::mem::swap(&mut prev, &mut curr);
|
|
}
|
|
prev[m] as f32 / n as f32
|
|
}
|
|
|
|
fn tokenize(s: &str) -> Vec<String> {
|
|
let mut buf = String::with_capacity(s.len());
|
|
for c in s.chars() {
|
|
if c.is_alphanumeric() || c.is_whitespace() {
|
|
for lc in c.to_lowercase() {
|
|
buf.push(lc);
|
|
}
|
|
} else {
|
|
buf.push(' ');
|
|
}
|
|
}
|
|
buf.split_whitespace().map(String::from).collect()
|
|
}
|
|
|
|
fn peak_rms_db(pcm: &[f32]) -> (f32, f32) {
|
|
if pcm.is_empty() {
|
|
return (-100.0, -100.0);
|
|
}
|
|
let mut peak = 0.0f32;
|
|
let mut sum_sq = 0.0f64;
|
|
for &s in pcm {
|
|
let a = s.abs();
|
|
if a > peak {
|
|
peak = a;
|
|
}
|
|
sum_sq += (s as f64) * (s as f64);
|
|
}
|
|
let rms = (sum_sq / pcm.len() as f64).sqrt() as f32;
|
|
let peak_db = if peak < 1e-6 {
|
|
-100.0
|
|
} else {
|
|
20.0 * peak.log10()
|
|
};
|
|
let rms_db = if rms < 1e-6 {
|
|
-100.0
|
|
} else {
|
|
20.0 * rms.log10()
|
|
};
|
|
(peak_db, rms_db)
|
|
}
|