From 7703d10f5fa84047b08be497a682bb9fb5bc31fb Mon Sep 17 00:00:00 2001 From: Stefano Verna Date: Thu, 5 Feb 2026 11:29:25 +0100 Subject: [PATCH] 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 --- src-tauri/src/audio_toolkit/text.rs | 227 +++++++++++++++++++++------- 1 file changed, 172 insertions(+), 55 deletions(-) diff --git a/src-tauri/src/audio_toolkit/text.rs b/src-tauri/src/audio_toolkit/text.rs index 3c04776..0f780bc 100644 --- a/src-tauri/src/audio_toolkit/text.rs +++ b/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::>() + .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 = custom_words.iter().map(|w| w.to_lowercase()).collect(); + // Pre-compute versions with spaces removed for n-gram comparison + let custom_words_nospace: Vec = 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 { - let cleaned_word = word - .trim_matches(|c: char| !c.is_alphabetic()) - .to_lowercase(); + while i < words.len() { + let mut matched = false; - if cleaned_word.is_empty() { - corrected_words.push(word.to_string()); - continue; - } - - // 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 { + // Try n-grams from longest (3) to shortest (1) - greedy matching + for n in (1..=3).rev() { + if i + n > words.len() { continue; } - // Calculate Levenshtein distance (normalized by length) - let levenshtein_dist = levenshtein(&cleaned_word, custom_word_lower); - 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 - }; + let ngram_words = &words[i..i + n]; + let ngram = build_ngram(ngram_words); - // Calculate phonetic similarity using Soundex - let phonetic_match = soundex(&cleaned_word, custom_word_lower); + if let Some((replacement, _score)) = + 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 - let combined_score = if phonetic_match { - levenshtein_score * 0.3 // Give significant boost to phonetic matches - } else { - levenshtein_score - }; + // Preserve case from first word + let corrected = preserve_case_pattern(ngram_words[0], replacement); - // 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; + result.push(format!("{}{}{}", prefix, corrected, suffix)); + i += n; + matched = true; + break; } } - if let Some(replacement) = best_match { - // Preserve the original case pattern as much as possible - let corrected = preserve_case_pattern(word, replacement); - - // 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()); + if !matched { + result.push(words[i].to_string()); + i += 1; } } - 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 + ); + } }