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:
parent
e25c6043fb
commit
7703d10f5f
1 changed files with 172 additions and 55 deletions
|
|
@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue