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

//! Text normalization for CSM input.
//!
//! Why this exists: Sesame CSM hangs (or produces gibberish) on certain text
//! shapes — `(parenthetical)`, `10:30`-style times, mismatched unicode forms,
//! literal `[N]` strings inside user content (which collide with our speaker
//! prefix format). See SesameAILabs/csm issue #141.
//!
//! Defaults are conservative. Each rule can be disabled if a downstream
//! component (e.g. an SSML-style normalizer) handles it instead.

use crate::error::{CsmError, Result};
use std::sync::OnceLock;
use unicode_normalization::UnicodeNormalization;

#[derive(Debug, Clone, Copy)]
pub struct TextNormalize {
    /// Apply NFC unicode normalization (combining marks → composed forms).
    pub unicode_nfc: bool,
    /// Replace user-content '[' and ']' with '(' and ')' to avoid colliding with
    /// the `[<speaker>]` speaker prefix format used internally.
    pub escape_brackets: bool,
    /// Convert `HH:MM` patterns to spelled-out time (e.g. "10:30" → "ten thirty").
    pub spell_times: bool,
    /// Strip zero-width characters (ZWSP, ZWJ, BOM, etc.) — they confuse the BPE.
    pub strip_zero_width: bool,
    /// Hard cap on character count after normalization. Inputs above this fail
    /// fast rather than producing pathological generation. The long-form
    /// chunker (Item 8) operates below this limit.
    pub max_chars: usize,
}

impl Default for TextNormalize {
    fn default() -> Self {
        Self {
            unicode_nfc: true,
            escape_brackets: true,
            spell_times: true,
            strip_zero_width: true,
            max_chars: 600,
        }
    }
}

impl TextNormalize {
    pub fn passthrough() -> Self {
        Self {
            unicode_nfc: false,
            escape_brackets: false,
            spell_times: false,
            strip_zero_width: false,
            max_chars: usize::MAX,
        }
    }

    pub fn apply(&self, text: &str) -> Result<String> {
        let mut s = text.to_string();
        if self.unicode_nfc {
            s = s.nfc().collect();
        }
        if self.strip_zero_width {
            s = strip_zero_width(&s);
        }
        if self.escape_brackets {
            s = s.replace('[', "(").replace(']', ")");
        }
        if self.spell_times {
            s = spell_times(&s);
        }
        let trimmed = s.trim();
        if trimmed.is_empty() {
            return Err(CsmError::Config(
                "text is empty or whitespace-only after normalization".into(),
            ));
        }
        if trimmed.chars().count() > self.max_chars {
            return Err(CsmError::Config(format!(
                "text length {} exceeds max_chars {} (use long-form chunker)",
                trimmed.chars().count(),
                self.max_chars
            )));
        }
        Ok(trimmed.to_string())
    }
}

fn strip_zero_width(s: &str) -> String {
    s.chars()
        .filter(|c| {
            !matches!(
                *c,
                '\u{200B}' // ZWSP
                | '\u{200C}' // ZWNJ
                | '\u{200D}' // ZWJ
                | '\u{FEFF}' // BOM
                | '\u{2060}' // word joiner
            )
        })
        .collect()
}

fn time_regex() -> &'static regex::Regex {
    static R: OnceLock<regex::Regex> = OnceLock::new();
    // Match HH:MM with optional am/pm; HH 0-23, MM 00-59. Bounded to word edges.
    R.get_or_init(|| regex::Regex::new(r"\b([01]?\d|2[0-3]):([0-5]\d)(\s*(?i)(am|pm))?\b").unwrap())
}

fn spell_times(s: &str) -> String {
    let re = time_regex();
    re.replace_all(s, |caps: &regex::Captures<'_>| {
        let h: u32 = caps[1].parse().unwrap_or(0);
        let m: u32 = caps[2].parse().unwrap_or(0);
        let ampm = caps.get(4).map(|m| m.as_str().to_ascii_lowercase());
        let mut out = number_to_words(h);
        if m == 0 {
            if h == 12 && ampm.as_deref() == Some("pm") {
                out = "noon".into();
            } else if h == 0 || (h == 12 && ampm.as_deref() == Some("am")) {
                out = "midnight".into();
            } else {
                out.push_str(" o'clock");
            }
        } else if m < 10 {
            out.push_str(" oh ");
            out.push_str(&number_to_words(m));
        } else {
            out.push(' ');
            out.push_str(&number_to_words(m));
        }
        if let Some(s) = ampm {
            out.push(' ');
            out.push_str(if s == "am" { "ay em" } else { "pee em" });
        }
        out
    })
    .into_owned()
}

/// English spelling for integers 0..100. Wider ranges get the literal digit
/// fall-through. Sufficient for time-of-day normalization.
fn number_to_words(n: u32) -> String {
    const SMALL: [&str; 20] = [
        "zero",
        "one",
        "two",
        "three",
        "four",
        "five",
        "six",
        "seven",
        "eight",
        "nine",
        "ten",
        "eleven",
        "twelve",
        "thirteen",
        "fourteen",
        "fifteen",
        "sixteen",
        "seventeen",
        "eighteen",
        "nineteen",
    ];
    const TENS: [&str; 10] = [
        "", "", "twenty", "thirty", "forty", "fifty", "sixty", "seventy", "eighty", "ninety",
    ];
    if n < 20 {
        SMALL[n as usize].into()
    } else if n < 100 {
        let t = (n / 10) as usize;
        let u = (n % 10) as usize;
        if u == 0 {
            TENS[t].into()
        } else {
            format!("{}-{}", TENS[t], SMALL[u])
        }
    } else {
        n.to_string()
    }
}

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

    #[test]
    fn empty_input_fails() {
        let r = TextNormalize::default().apply("");
        assert!(r.is_err());
        let r = TextNormalize::default().apply("   \t\n  ");
        assert!(r.is_err());
    }

    #[test]
    fn passthrough_keeps_brackets() {
        let s = TextNormalize::passthrough()
            .apply("hello [0] world")
            .unwrap();
        assert_eq!(s, "hello [0] world");
    }

    #[test]
    fn brackets_get_escaped_to_parens() {
        let s = TextNormalize::default().apply("hello [0] world").unwrap();
        assert_eq!(s, "hello (0) world");
    }

    #[test]
    fn spell_times_o_clock_and_minutes() {
        let s = TextNormalize::default().apply("be there at 10:30").unwrap();
        assert!(s.contains("ten thirty"), "got: {s}");
        let s = TextNormalize::default()
            .apply("breakfast at 8:00 am")
            .unwrap();
        assert!(s.contains("eight o'clock"), "got: {s}");
    }

    #[test]
    fn spell_times_oh_minutes() {
        let s = TextNormalize::default().apply("call me at 9:05").unwrap();
        assert!(s.contains("nine oh five"), "got: {s}");
    }

    #[test]
    fn noon_and_midnight() {
        let s = TextNormalize::default()
            .apply("see you at 12:00 pm")
            .unwrap();
        assert!(s.contains("noon"), "got: {s}");
        let s = TextNormalize::default().apply("at 12:00 am").unwrap();
        assert!(s.contains("midnight"), "got: {s}");
    }

    #[test]
    fn zero_width_stripped() {
        let s = TextNormalize::default()
            .apply("hello\u{200B}world")
            .unwrap();
        assert_eq!(s, "helloworld");
    }

    #[test]
    fn over_length_fails_loudly() {
        let too_long = "a ".repeat(600);
        let r = TextNormalize::default().apply(&too_long);
        assert!(r.is_err());
    }
}