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

//! Word Error Rate computation.
//!
//! Pure Rust, no model dependencies. Used by the bench harness to score
//! generated audio against the ground-truth prompt after running an external
//! ASR (Whisper, etc.).
//!
//! WER = (substitutions + deletions + insertions) / reference_word_count
//!
//! Computed via Levenshtein distance over word-level tokenization. Light
//! normalization (lowercase, strip punctuation) so trivial casing/punctuation
//! mismatches don't inflate the score.

/// Standard NIST-style normalization:
///   - lowercase
///   - keep only `[a-z0-9' ]`, replace everything else with space
///   - collapse runs of whitespace
pub fn normalize_for_wer(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut last_was_space = true;
    for c in s.chars() {
        let c = c.to_ascii_lowercase();
        let keep = c.is_ascii_alphanumeric() || c == '\'';
        if keep {
            out.push(c);
            last_was_space = false;
        } else if !last_was_space {
            out.push(' ');
            last_was_space = true;
        }
    }
    out.trim().to_string()
}

pub fn tokenize(s: &str) -> Vec<&str> {
    s.split_whitespace().collect()
}

/// Compute the Levenshtein-distance-based word error rate.
/// Returns a struct with the raw counts so the caller can aggregate / format.
pub fn wer(reference: &str, hypothesis: &str) -> WerResult {
    let r = normalize_for_wer(reference);
    let h = normalize_for_wer(hypothesis);
    let r_tokens = tokenize(&r);
    let h_tokens = tokenize(&h);

    let (subs, dels, ins) = lev_align(&r_tokens, &h_tokens);
    let n = r_tokens.len();
    WerResult {
        substitutions: subs,
        deletions: dels,
        insertions: ins,
        reference_words: n,
        hypothesis_words: h_tokens.len(),
    }
}

/// Levenshtein alignment producing per-edit counts.
fn lev_align(r: &[&str], h: &[&str]) -> (usize, usize, usize) {
    let n = r.len();
    let m = h.len();
    if n == 0 {
        return (0, 0, m);
    }
    if m == 0 {
        return (0, n, 0);
    }
    // dp[i][j] = (cost, subs, dels, ins)
    #[derive(Clone, Copy)]
    struct Cell {
        cost: usize,
        s: usize,
        d: usize,
        i: usize,
    }
    impl Cell {
        const fn new() -> Self {
            Self {
                cost: 0,
                s: 0,
                d: 0,
                i: 0,
            }
        }
    }
    let mut prev = vec![Cell::new(); m + 1];
    let mut curr = vec![Cell::new(); m + 1];
    #[allow(clippy::needless_range_loop)]
    for j in 0..=m {
        prev[j] = Cell {
            cost: j,
            s: 0,
            d: 0,
            i: j,
        };
    }
    for i in 1..=n {
        curr[0] = Cell {
            cost: i,
            s: 0,
            d: i,
            i: 0,
        };
        for j in 1..=m {
            if r[i - 1] == h[j - 1] {
                curr[j] = prev[j - 1];
            } else {
                let sub = Cell {
                    cost: prev[j - 1].cost + 1,
                    s: prev[j - 1].s + 1,
                    d: prev[j - 1].d,
                    i: prev[j - 1].i,
                };
                let del = Cell {
                    cost: prev[j].cost + 1,
                    s: prev[j].s,
                    d: prev[j].d + 1,
                    i: prev[j].i,
                };
                let ins = Cell {
                    cost: curr[j - 1].cost + 1,
                    s: curr[j - 1].s,
                    d: curr[j - 1].d,
                    i: curr[j - 1].i + 1,
                };
                curr[j] = [sub, del, ins].into_iter().min_by_key(|c| c.cost).unwrap();
            }
        }
        std::mem::swap(&mut prev, &mut curr);
    }
    let final_cell = prev[m];
    (final_cell.s, final_cell.d, final_cell.i)
}

#[derive(Debug, Clone, Copy, serde::Serialize)]
pub struct WerResult {
    pub substitutions: usize,
    pub deletions: usize,
    pub insertions: usize,
    pub reference_words: usize,
    pub hypothesis_words: usize,
}

impl WerResult {
    pub fn errors(&self) -> usize {
        self.substitutions + self.deletions + self.insertions
    }
    pub fn rate(&self) -> f32 {
        if self.reference_words == 0 {
            return if self.hypothesis_words == 0 { 0.0 } else { 1.0 };
        }
        self.errors() as f32 / self.reference_words as f32
    }
}

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

    #[test]
    fn perfect_match_zero_wer() {
        let r = wer("the cat sat on the mat", "the cat sat on the mat");
        assert_eq!(r.rate(), 0.0);
    }

    #[test]
    fn substitution_counted() {
        let r = wer("the cat sat", "the dog sat");
        assert_eq!(r.substitutions, 1);
        assert_eq!(r.deletions, 0);
        assert_eq!(r.insertions, 0);
        assert!((r.rate() - 1.0 / 3.0).abs() < 1e-5);
    }

    #[test]
    fn insertion_counted() {
        let r = wer("the cat sat", "the cat sat down");
        assert_eq!(r.insertions, 1);
        assert!((r.rate() - 1.0 / 3.0).abs() < 1e-5);
    }

    #[test]
    fn deletion_counted() {
        let r = wer("the big cat sat", "the cat sat");
        assert_eq!(r.deletions, 1);
        assert!((r.rate() - 1.0 / 4.0).abs() < 1e-5);
    }

    #[test]
    fn case_and_punctuation_normalized() {
        let r = wer("Hello, world!", "hello world");
        assert_eq!(r.rate(), 0.0);
    }

    #[test]
    fn empty_reference_with_hypothesis_is_full_error() {
        let r = wer("", "anything");
        assert_eq!(r.rate(), 1.0);
    }
}