feat(text): Add n-gram matching for multi-word custom word correction (#711)

* feat(text): Add n-gram matching for multi-word custom word correction

Improves custom word matching to handle speech artifacts where a single
word gets transcribed as multiple words (e.g., "Charge B" -> "ChargeBee",
"Chat G P T" -> "ChatGPT").

- Implements greedy n-gram matching (3 to 1 words) for better accuracy
- Adds length-based filtering (25% max difference) to prevent over-matching
- Extracts matching logic into reusable find_best_match function
- Preserves punctuation and case from original transcription
- Adds comprehensive test coverage for n-gram scenarios

* minor tweak

* format

---------

Co-authored-by: CJ Pais <cj@cjpais.com>

src-tauri/src/audio_toolkit/text.rs

@@ -3,12 +3,94 @@ use once_cell::sync::Lazy;
 use regex::Regex;
 use strsim::levenshtein;
 
+/// Builds an n-gram string by cleaning and concatenating words
+///
+/// Strips punctuation from each word, lowercases, and joins without spaces.
+/// This allows matching "Charge B" against "ChargeBee".
+fn build_ngram(words: &[&str]) -> String {
+    words
+        .iter()
+        .map(|w| {
+            w.trim_matches(|c: char| !c.is_alphanumeric())
+                .to_lowercase()
+        })
+        .collect::<Vec<_>>()
+        .concat()
+}
+
+/// Finds the best matching custom word for a candidate string
+///
+/// Uses Levenshtein distance and Soundex phonetic matching to find
+/// the best match above the given threshold.
+///
+/// # Arguments
+/// * `candidate` - The cleaned/lowercased candidate string to match
+/// * `custom_words` - Original custom words (for returning the replacement)
+/// * `custom_words_nospace` - Custom words with spaces removed, lowercased (for comparison)
+/// * `threshold` - Maximum similarity score to accept
+///
+/// # Returns
+/// The best matching custom word and its score, if any match was found
+fn find_best_match<'a>(
+    candidate: &str,
+    custom_words: &'a [String],
+    custom_words_nospace: &[String],
+    threshold: f64,
+) -> Option<(&'a String, f64)> {
+    if candidate.is_empty() || candidate.len() > 50 {
+        return None;
+    }
+
+    let mut best_match: Option<&String> = None;
+    let mut best_score = f64::MAX;
+
+    for (i, custom_word_nospace) in custom_words_nospace.iter().enumerate() {
+        // Skip if lengths are too different (optimization + prevents over-matching)
+        // Use percentage-based check: max 25% length difference (prevents n-grams from
+        // matching significantly shorter custom words, e.g., "openaigpt" vs "openai")
+        let len_diff = (candidate.len() as i32 - custom_word_nospace.len() as i32).abs() as f64;
+        let max_len = candidate.len().max(custom_word_nospace.len()) as f64;
+        let max_allowed_diff = (max_len * 0.25).max(2.0); // At least 2 chars difference allowed
+        if len_diff > max_allowed_diff {
+            continue;
+        }
+
+        // Calculate Levenshtein distance (normalized by length)
+        let levenshtein_dist = levenshtein(candidate, custom_word_nospace);
+        let max_len = candidate.len().max(custom_word_nospace.len()) as f64;
+        let levenshtein_score = if max_len > 0.0 {
+            levenshtein_dist as f64 / max_len
+        } else {
+            1.0
+        };
+
+        // Calculate phonetic similarity using Soundex
+        let phonetic_match = soundex(candidate, custom_word_nospace);
+
+        // Combine scores: favor phonetic matches, but also consider string similarity
+        let combined_score = if phonetic_match {
+            levenshtein_score * 0.3 // Give significant boost to phonetic matches
+        } else {
+            levenshtein_score
+        };
+
+        // Accept if the score is good enough (configurable threshold)
+        if combined_score < threshold && combined_score < best_score {
+            best_match = Some(&custom_words[i]);
+            best_score = combined_score;
+        }
+    }
+
+    best_match.map(|m| (m, best_score))
+}
+
 /// Applies custom word corrections to transcribed text using fuzzy matching
 ///
 /// This function corrects words in the input text by finding the best matches
 /// from a list of custom words using a combination of:
 /// - Levenshtein distance for string similarity
 /// - Soundex phonetic matching for pronunciation similarity
+/// - N-gram matching for multi-word speech artifacts (e.g., "Charge B" -> "ChargeBee")
 ///
 /// # Arguments
 /// * `text` - The input text to correct
@@ -25,74 +107,52 @@ pub fn apply_custom_words(text: &str, custom_words: &[String], threshold: f64) -
     // Pre-compute lowercase versions to avoid repeated allocations
     let custom_words_lower: Vec<String> = custom_words.iter().map(|w| w.to_lowercase()).collect();
 
+    // Pre-compute versions with spaces removed for n-gram comparison
+    let custom_words_nospace: Vec<String> = custom_words_lower
+        .iter()
+        .map(|w| w.replace(' ', ""))
+        .collect();
+
     let words: Vec<&str> = text.split_whitespace().collect();
-    let mut corrected_words = Vec::new();
+    let mut result = Vec::new();
+    let mut i = 0;
 
-    for word in words {
+    while i < words.len() {
-        let cleaned_word = word
+        let mut matched = false;
-            .trim_matches(|c: char| !c.is_alphabetic())
-            .to_lowercase();
 
-        if cleaned_word.is_empty() {
+        // Try n-grams from longest (3) to shortest (1) - greedy matching
-            corrected_words.push(word.to_string());
+        for n in (1..=3).rev() {
-            continue;
+            if i + n > words.len() {
-        }
-
-        // Skip extremely long words to avoid performance issues
-        if cleaned_word.len() > 50 {
-            corrected_words.push(word.to_string());
-            continue;
-        }
-
-        let mut best_match: Option<&String> = None;
-        let mut best_score = f64::MAX;
-
-        for (i, custom_word_lower) in custom_words_lower.iter().enumerate() {
-            // Skip if lengths are too different (optimization)
-            let len_diff = (cleaned_word.len() as i32 - custom_word_lower.len() as i32).abs();
-            if len_diff > 5 {
                 continue;
             }
 
-            // Calculate Levenshtein distance (normalized by length)
+            let ngram_words = &words[i..i + n];
-            let levenshtein_dist = levenshtein(&cleaned_word, custom_word_lower);
+            let ngram = build_ngram(ngram_words);
-            let max_len = cleaned_word.len().max(custom_word_lower.len()) as f64;
-            let levenshtein_score = if max_len > 0.0 {
-                levenshtein_dist as f64 / max_len
-            } else {
-                1.0
-            };
 
-            // Calculate phonetic similarity using Soundex
+            if let Some((replacement, _score)) =
-            let phonetic_match = soundex(&cleaned_word, custom_word_lower);
+                find_best_match(&ngram, custom_words, &custom_words_nospace, threshold)
+            {
+                // Extract punctuation from first and last words of the n-gram
+                let (prefix, _) = extract_punctuation(ngram_words[0]);
+                let (_, suffix) = extract_punctuation(ngram_words[n - 1]);
 
-            // Combine scores: favor phonetic matches, but also consider string similarity
+                // Preserve case from first word
-            let combined_score = if phonetic_match {
+                let corrected = preserve_case_pattern(ngram_words[0], replacement);
-                levenshtein_score * 0.3 // Give significant boost to phonetic matches
-            } else {
-                levenshtein_score
-            };
 
-            // Accept if the score is good enough (configurable threshold)
+                result.push(format!("{}{}{}", prefix, corrected, suffix));
-            if combined_score < threshold && combined_score < best_score {
+                i += n;
-                best_match = Some(&custom_words[i]);
+                matched = true;
-                best_score = combined_score;
+                break;
             }
         }
 
-        if let Some(replacement) = best_match {
+        if !matched {
-            // Preserve the original case pattern as much as possible
+            result.push(words[i].to_string());
-            let corrected = preserve_case_pattern(word, replacement);
+            i += 1;
-
-            // Preserve punctuation from original word
-            let (prefix, suffix) = extract_punctuation(word);
-            corrected_words.push(format!("{}{}{}", prefix, corrected, suffix));
-        } else {
-            corrected_words.push(word.to_string());
         }
     }
 
-    corrected_words.join(" ")
+    result.join(" ")
 }
 
 /// Preserves the case pattern of the original word when applying a replacement
@@ -112,11 +172,11 @@ fn preserve_case_pattern(original: &str, replacement: &str) -> String {
 
 /// Extracts punctuation prefix and suffix from a word
 fn extract_punctuation(word: &str) -> (&str, &str) {
-    let prefix_end = word.chars().take_while(|c| !c.is_alphabetic()).count();
+    let prefix_end = word.chars().take_while(|c| !c.is_alphanumeric()).count();
     let suffix_start = word
         .char_indices()
         .rev()
-        .take_while(|(_, c)| !c.is_alphabetic())
+        .take_while(|(_, c)| !c.is_alphanumeric())
         .count();
 
     let prefix = if prefix_end > 0 {
@@ -341,4 +401,61 @@ mod tests {
         let result = filter_transcription_output(text);
         assert_eq!(result, "no no is fine");
     }
+
+    #[test]
+    fn test_apply_custom_words_ngram_two_words() {
+        let text = "il cui nome è Charge B, che permette";
+        let custom_words = vec!["ChargeBee".to_string()];
+        let result = apply_custom_words(text, &custom_words, 0.5);
+        assert!(result.contains("ChargeBee,"));
+        assert!(!result.contains("Charge B"));
+    }
+
+    #[test]
+    fn test_apply_custom_words_ngram_three_words() {
+        let text = "use Chat G P T for this";
+        let custom_words = vec!["ChatGPT".to_string()];
+        let result = apply_custom_words(text, &custom_words, 0.5);
+        assert!(result.contains("ChatGPT"));
+    }
+
+    #[test]
+    fn test_apply_custom_words_prefers_longer_ngram() {
+        let text = "Open AI GPT model";
+        let custom_words = vec!["OpenAI".to_string(), "GPT".to_string()];
+        let result = apply_custom_words(text, &custom_words, 0.5);
+        assert_eq!(result, "OpenAI GPT model");
+    }
+
+    #[test]
+    fn test_apply_custom_words_ngram_preserves_case() {
+        let text = "CHARGE B is great";
+        let custom_words = vec!["ChargeBee".to_string()];
+        let result = apply_custom_words(text, &custom_words, 0.5);
+        assert!(result.contains("CHARGEBEE"));
+    }
+
+    #[test]
+    fn test_apply_custom_words_ngram_with_spaces_in_custom() {
+        // Custom word with space should also match against split words
+        let text = "using Mac Book Pro";
+        let custom_words = vec!["MacBook Pro".to_string()];
+        let result = apply_custom_words(text, &custom_words, 0.5);
+        assert!(result.contains("MacBook"));
+    }
+
+    #[test]
+    fn test_apply_custom_words_trailing_number_not_doubled() {
+        // Verify that trailing non-alpha chars (like numbers) aren't double-counted
+        // between build_ngram stripping them and extract_punctuation capturing them
+        let text = "use GPT4 for this";
+        let custom_words = vec!["GPT-4".to_string()];
+        let result = apply_custom_words(text, &custom_words, 0.5);
+        // Should NOT produce "GPT-44" (double-counting the trailing 4)
+        assert!(
+            !result.contains("GPT-44"),
+            "got double-counted result: {}",
+            result
+        );
+    }
 }