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]>
116 lines
4.3 KiB
Rust
116 lines
4.3 KiB
Rust
//! Load the converted WavLM-SV safetensors and compute speaker embeddings
|
|
//! for a pair of input WAVs. Reports cosine similarity (the standard
|
|
//! speaker verification score: > ~0.5 = same speaker).
|
|
//!
|
|
//! Usage:
|
|
//! ```
|
|
//! cargo run -p rtx-csm --release --example wavlm_sv_demo -- \
|
|
//! --weights /tmp/wavlm_sv.safetensors \
|
|
//! --a /path/to/utt_a.wav \
|
|
//! --b /path/to/utt_b.wav
|
|
//! ```
|
|
|
|
use anyhow::{Context, Result};
|
|
use candle_core::{DType, Device, Tensor};
|
|
use clap::Parser;
|
|
use rtx_csm::{audio_io, wavlm_sv::WavLmSv};
|
|
use std::path::PathBuf;
|
|
|
|
const WAVLM_RATE: u32 = 16_000;
|
|
|
|
#[derive(Debug, Parser)]
|
|
#[command(name = "wavlm_sv_demo")]
|
|
struct Cli {
|
|
/// Converted safetensors (run `wavlm_sv_convert` first).
|
|
#[arg(long)]
|
|
weights: PathBuf,
|
|
/// First utterance.
|
|
#[arg(long)]
|
|
a: PathBuf,
|
|
/// Second utterance.
|
|
#[arg(long)]
|
|
b: PathBuf,
|
|
/// Force CPU device.
|
|
#[arg(long)]
|
|
cpu: bool,
|
|
/// Optional path to dump a JSON fingerprint of the embeddings (for
|
|
/// `scripts/wavlm_sv_parity.py` numerical comparison).
|
|
#[arg(long)]
|
|
parity_json: Option<PathBuf>,
|
|
}
|
|
|
|
fn main() -> Result<()> {
|
|
tracing_subscriber::fmt().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
|
|
};
|
|
|
|
let vb = unsafe {
|
|
candle_nn::VarBuilder::from_mmaped_safetensors(&[&cli.weights], DType::F32, &device)
|
|
}
|
|
.context("opening WavLM-SV safetensors")?;
|
|
let model = WavLmSv::new(vb).context("WavLmSv::new")?;
|
|
println!("loaded WavLM-SV from {}", cli.weights.display());
|
|
|
|
// Load utterances at 16 kHz mono.
|
|
let a_samples = audio_io::load_mono_at_rate(&cli.a, WAVLM_RATE)?;
|
|
let b_samples = audio_io::load_mono_at_rate(&cli.b, WAVLM_RATE)?;
|
|
println!(
|
|
"a: {} samples ({:.2}s), b: {} samples ({:.2}s)",
|
|
a_samples.len(),
|
|
a_samples.len() as f32 / WAVLM_RATE as f32,
|
|
b_samples.len(),
|
|
b_samples.len() as f32 / WAVLM_RATE as f32,
|
|
);
|
|
|
|
let a_norm = WavLmSv::normalize_waveform(&a_samples);
|
|
let b_norm = WavLmSv::normalize_waveform(&b_samples);
|
|
let xs_a = Tensor::from_slice(&a_norm, (1, 1, a_norm.len()), &device)?;
|
|
let xs_b = Tensor::from_slice(&b_norm, (1, 1, b_norm.len()), &device)?;
|
|
|
|
let emb_a = model.embed(&xs_a).context("embed a")?;
|
|
let emb_b = model.embed(&xs_b).context("embed b")?;
|
|
println!("embeddings: a={:?}, b={:?}", emb_a.dims(), emb_b.dims());
|
|
|
|
let sim = WavLmSv::cosine_similarity(&emb_a, &emb_b)?;
|
|
// cosine_similarity returns (B,) for (B, D) inputs; take the first element.
|
|
let sim_vec: Vec<f32> = sim.flatten_all()?.to_dtype(DType::F32)?.to_vec1()?;
|
|
let sim = sim_vec[0];
|
|
println!("cosine similarity = {sim:.4}");
|
|
if sim > 0.5 {
|
|
println!(" → likely same speaker");
|
|
} else {
|
|
println!(" → likely different speakers");
|
|
}
|
|
|
|
if let Some(parity_path) = cli.parity_json.as_ref() {
|
|
let emb_a_vec: Vec<f32> = emb_a.flatten_all()?.to_dtype(DType::F32)?.to_vec1()?;
|
|
let emb_b_vec: Vec<f32> = emb_b.flatten_all()?.to_dtype(DType::F32)?.to_vec1()?;
|
|
let norm = |v: &[f32]| (v.iter().map(|x| x * x).sum::<f32>()).sqrt();
|
|
let head = |v: &[f32], n: usize| v.iter().take(n).copied().collect::<Vec<_>>();
|
|
let tail = |v: &[f32], n: usize| v.iter().rev().take(n).copied().collect::<Vec<_>>();
|
|
let report = serde_json::json!({
|
|
"model": "rtx-csm port of microsoft/wavlm-base-plus-sv",
|
|
"wav_a": cli.a.display().to_string(),
|
|
"wav_b": cli.b.display().to_string(),
|
|
"embedding_dim": emb_a_vec.len(),
|
|
"cosine_similarity_rust": sim,
|
|
"embedding_a_norm": norm(&emb_a_vec),
|
|
"embedding_b_norm": norm(&emb_b_vec),
|
|
"embedding_a_head": head(&emb_a_vec, 8),
|
|
"embedding_a_tail": tail(&emb_a_vec, 8).into_iter().rev().collect::<Vec<_>>(),
|
|
"embedding_b_head": head(&emb_b_vec, 8),
|
|
"embedding_b_tail": tail(&emb_b_vec, 8).into_iter().rev().collect::<Vec<_>>(),
|
|
});
|
|
std::fs::write(parity_path, serde_json::to_string_pretty(&report)?)?;
|
|
println!("wrote parity fingerprint to {}", parity_path.display());
|
|
}
|
|
|
|
Ok(())
|
|
}
|