rustytorch/crates/models/rtx-csm/src/speaker_sim.rs

//! Speaker similarity scoring.
//!
//! Three implementations:
//!
//!   - [`WavLmSimilarity`] — production. Wraps the in-process
//!     `wavlm_sv::WavLmSv` model (microsoft/wavlm-base-plus-sv X-vector
//!     head, 100M params). Run `wavlm_sv_convert` once to produce the
//!     safetensors, then `WavLmSimilarity::load(&path, &device)`.
//!   - [`CosineSimilarityFromEmbeddings`] — caller-provides embeddings,
//!     we just cosine. Useful when embeddings are pre-computed elsewhere
//!     (e.g. ECAPA-TDNN Python sidecar) or you want to cache them.
//!   - [`SpectralCentroidSimilarity`] — pure-Rust weak baseline from a
//!     handful of cheap acoustic features. Useful for quick consistency
//!     checks but NOT a real speaker fingerprint.

use crate::error::{CsmError, Result};

pub trait SpeakerSimilarity {
    /// Returns cosine similarity in `[-1, 1]` between two utterances.
    /// Implementations may take either raw PCM or pre-computed embeddings.
    fn score(&self, a: &[f32], b: &[f32]) -> Result<f32>;
}

/// Cheap pure-Rust baseline: pools spectral centroid, RMS, ZCR over the
/// utterance and compares as a 3-vector. Useful for catching gross drift
/// (utterance suddenly sounds totally different) — NOT a real speaker
/// fingerprint. Documented as a weak baseline only.
pub struct SpectralCentroidSimilarity;

impl SpeakerSimilarity for SpectralCentroidSimilarity {
    fn score(&self, a: &[f32], b: &[f32]) -> Result<f32> {
        let fa = acoustic_summary(a);
        let fb = acoustic_summary(b);
        Ok(cosine(&fa, &fb))
    }
}

fn acoustic_summary(samples: &[f32]) -> [f32; 3] {
    if samples.is_empty() {
        return [0.0, 0.0, 0.0];
    }
    let n = samples.len() as f32;
    let rms = (samples.iter().map(|x| x * x).sum::<f32>() / n).sqrt();
    let zcr = samples
        .windows(2)
        .filter(|w| (w[0] >= 0.0) != (w[1] >= 0.0))
        .count() as f32
        / (n - 1.0).max(1.0);
    // Crude spectral centroid via |y| weighted by index. Good enough for the
    // weak-baseline use case; not a real STFT.
    let mut weighted = 0.0f32;
    let mut total = 0.0f32;
    for (i, s) in samples.iter().enumerate() {
        let m = s.abs();
        weighted += i as f32 * m;
        total += m;
    }
    let centroid = if total > 1e-9 { weighted / total } else { 0.0 };
    [rms, zcr, centroid / n]
}

fn cosine(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if na < 1e-9 || nb < 1e-9 {
        return 0.0;
    }
    dot / (na * nb)
}

/// User supplies pre-computed embeddings (e.g. from a Python `speechbrain`
/// sidecar that ran ECAPA-TDNN or WavLM-SV). We just compute cosine.
pub struct CosineSimilarityFromEmbeddings;

impl CosineSimilarityFromEmbeddings {
    pub fn cosine_of_embeddings(a: &[f32], b: &[f32]) -> Result<f32> {
        if a.len() != b.len() {
            return Err(CsmError::Shape(format!(
                "embedding dim mismatch: {} vs {}",
                a.len(),
                b.len()
            )));
        }
        Ok(cosine(a, b))
    }
}

/// Real WavLM-Base+ SV scorer. Wraps the `wavlm_sv::WavLmSv` model loaded
/// from a converted safetensors file (run `wavlm_sv_convert` once to
/// produce that). Computes 512-d embeddings for each input via the full
/// 12-layer transformer + x-vector head, then cosines them.
///
/// Inputs MUST be 16 kHz mono. We don't resample inside `score` — caller
/// is responsible for getting the rate right (see `audio_io::resample`).
pub struct WavLmSimilarity {
    model: crate::wavlm_sv::WavLmSv,
    device: candle_core::Device,
}

impl WavLmSimilarity {
    /// Load a converted WavLM-SV safetensors file. Pass the Metal/CUDA/CPU
    /// device you want inference to run on.
    pub fn load(
        safetensors: impl AsRef<std::path::Path>,
        device: &candle_core::Device,
    ) -> Result<Self> {
        let model = crate::wavlm_sv::load_from_safetensors(safetensors, device)?;
        Ok(Self {
            model,
            device: device.clone(),
        })
    }

    /// Compute the 512-d speaker embedding for a single utterance. Useful
    /// when you want to cache embeddings for repeated comparisons.
    pub fn embed(&self, samples: &[f32]) -> Result<Vec<f32>> {
        self.model.embed_samples(samples, &self.device)
    }
}

impl SpeakerSimilarity for WavLmSimilarity {
    fn score(&self, a: &[f32], b: &[f32]) -> Result<f32> {
        let ea = self.model.embed_samples(a, &self.device)?;
        let eb = self.model.embed_samples(b, &self.device)?;
        CosineSimilarityFromEmbeddings::cosine_of_embeddings(&ea, &eb)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cosine_identity_is_one() {
        let v = vec![1.0, 2.0, 3.0];
        assert!((cosine(&v, &v) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn cosine_orthogonal_is_zero() {
        let a = vec![1.0, 0.0];
        let b = vec![0.0, 1.0];
        assert!((cosine(&a, &b)).abs() < 1e-6);
    }

    #[test]
    fn cosine_of_embeddings_dim_check() {
        let r = CosineSimilarityFromEmbeddings::cosine_of_embeddings(&[1.0], &[1.0, 2.0]);
        assert!(r.is_err());
    }

    #[test]
    fn spectral_baseline_self_similar() {
        let a: Vec<f32> = (0..1000).map(|i| (i as f32 * 0.01).sin()).collect();
        let s = SpectralCentroidSimilarity.score(&a, &a).unwrap();
        // Same signal → cosine ~ 1
        assert!(s > 0.99, "self similarity {s} too low");
    }

    // WavLmSimilarity now requires actual converted weights; tested via
    // examples/wavlm_sv_demo.rs (cosine on real CSM speech pairs).
}