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>
This commit is contained in:
Stefano Verna 2026-02-05 11:29:25 +01:00 committed by GitHub
parent e25c6043fb
commit 7703d10f5f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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 {
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
);
}
}