diff --git a/crates/ai-formatting/Cargo.toml b/crates/ai-formatting/Cargo.toml index 84601bf..754a14a 100644 --- a/crates/ai-formatting/Cargo.toml +++ b/crates/ai-formatting/Cargo.toml @@ -6,3 +6,4 @@ description = "Text post-processing pipeline: filler removal, British English co [dependencies] kon-core = { path = "../core" } +regex-lite = "0.1" diff --git a/crates/ai-formatting/src/lib.rs b/crates/ai-formatting/src/lib.rs index 11b3a8b..db91d63 100644 --- a/crates/ai-formatting/src/lib.rs +++ b/crates/ai-formatting/src/lib.rs @@ -1,2 +1,8 @@ -// kon-ai-formatting: Filler removal, British English conversion, -// text formatting, hallucination filtering, and paragraph breaks. +mod llm_client; +pub mod pipeline; +pub mod rule_based; + +pub use pipeline::{post_process_segments, FormatMode, PostProcessOptions}; +pub use rule_based::{ + format_text, is_hallucination, remove_fillers, to_british_english, +}; diff --git a/crates/ai-formatting/src/llm_client.rs b/crates/ai-formatting/src/llm_client.rs new file mode 100644 index 0000000..5c5ef9b --- /dev/null +++ b/crates/ai-formatting/src/llm_client.rs @@ -0,0 +1,2 @@ +// Stub for future LLM sidecar integration. +// Will implement TextProcessor trait when mistral.rs sidecar is added. diff --git a/crates/ai-formatting/src/pipeline.rs b/crates/ai-formatting/src/pipeline.rs new file mode 100644 index 0000000..eb423bb --- /dev/null +++ b/crates/ai-formatting/src/pipeline.rs @@ -0,0 +1,123 @@ +use kon_core::constants::SMART_PARAGRAPH_GAP_SECS; +use kon_core::types::Segment; + +use crate::rule_based; + +/// Post-processing options for a transcription pipeline run. +pub struct PostProcessOptions { + pub remove_fillers: bool, + pub british_english: bool, + pub anti_hallucination: bool, + pub format_mode: FormatMode, +} + +/// How aggressively to format the transcript text. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FormatMode { + Raw, + Clean, + Smart, +} + +impl FormatMode { + pub fn parse(s: &str) -> Self { + match s { + "Clean" => Self::Clean, + "Smart" => Self::Smart, + _ => Self::Raw, + } + } +} + +/// Apply all post-processing steps to a list of segments. +/// Modifies segments in place. Composed from individual pure functions. +pub fn post_process_segments( + segments: &mut Vec, + options: &PostProcessOptions, +) { + if options.anti_hallucination { + segments.retain(|seg| !rule_based::is_hallucination(&seg.text)); + } + + for seg in segments.iter_mut() { + if options.remove_fillers { + seg.text = rule_based::remove_fillers(&seg.text); + } + if options.british_english { + seg.text = rule_based::to_british_english(&seg.text); + } + if options.format_mode != FormatMode::Raw { + seg.text = rule_based::format_text(&seg.text); + } + } + + if options.format_mode == FormatMode::Smart && segments.len() > 1 { + for i in (1..segments.len()).rev() { + let gap = segments[i].start - segments[i - 1].end; + if gap > SMART_PARAGRAPH_GAP_SECS { + segments[i].text = + format!("\n\n{}", segments[i].text.trim_start()); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_segments() -> Vec { + vec![ + Segment { + start: 0.0, + end: 1.0, + text: "um hello world".into(), + }, + Segment { + start: 1.0, + end: 2.0, + text: "[blank_audio]".into(), + }, + Segment { + start: 5.0, + end: 6.0, + text: "organize the color scheme".into(), + }, + ] + } + + #[test] + fn post_process_applies_all_filters() { + let mut segments = make_segments(); + let options = PostProcessOptions { + remove_fillers: true, + british_english: true, + anti_hallucination: true, + format_mode: FormatMode::Clean, + }; + + post_process_segments(&mut segments, &options); + + assert_eq!(segments.len(), 2); + let lower0 = segments[0].text.to_lowercase(); + let lower1 = segments[1].text.to_lowercase(); + assert!(!lower0.contains("um")); + assert!(lower1.contains("organise")); + assert!(lower1.contains("colour")); + } + + #[test] + fn post_process_adds_paragraph_breaks_on_long_pauses() { + let mut segments = make_segments(); + let options = PostProcessOptions { + remove_fillers: false, + british_english: false, + anti_hallucination: false, + format_mode: FormatMode::Smart, + }; + + post_process_segments(&mut segments, &options); + + assert!(segments[2].text.starts_with("\n\n")); + } +} diff --git a/crates/ai-formatting/src/rule_based.rs b/crates/ai-formatting/src/rule_based.rs new file mode 100644 index 0000000..2c804e4 --- /dev/null +++ b/crates/ai-formatting/src/rule_based.rs @@ -0,0 +1,287 @@ +use std::sync::LazyLock; + +/// Compiled filler word regexes (built once, reused across calls). +/// Uses \b word boundaries instead of lookbehinds (regex-lite does not +/// support lookaround assertions). +static FILLER_REGEXES: LazyLock> = + LazyLock::new(|| { + let fillers = [ + "um", + "uh", + "er", + "ah", + "like", + "you know", + "sort of", + "kind of", + "I mean", + "basically", + "actually", + "literally", + ]; + fillers + .iter() + .filter_map(|filler| { + let escaped = regex_lite::escape(filler); + let pattern = + format!(r"(?i)\b{escaped}\b[,.]?\s*"); + regex_lite::Regex::new(&pattern).ok() + }) + .collect() + }); + +/// Remove common filler words from transcription text (case-insensitive). +pub fn remove_fillers(text: &str) -> String { + let mut result = text.to_string(); + + for re in FILLER_REGEXES.iter() { + result = re.replace_all(&result, " ").to_string(); + } + + while result.contains(" ") { + result = result.replace(" ", " "); + } + + result.trim().to_string() +} + +/// Replacement pairs for American → British English conversion. +static BRITISH_REPLACEMENTS: &[(&str, &str)] = &[ + // -ize → -ise (and inflected forms) + ("organize", "organise"), + ("recognize", "recognise"), + ("realize", "realise"), + ("analyze", "analyse"), + ("apologize", "apologise"), + ("authorize", "authorise"), + ("categorize", "categorise"), + ("characterize", "characterise"), + ("customize", "customise"), + ("digitize", "digitise"), + ("emphasize", "emphasise"), + ("finalize", "finalise"), + ("generalize", "generalise"), + ("harmonize", "harmonise"), + ("initialize", "initialise"), + ("maximize", "maximise"), + ("minimize", "minimise"), + ("modernize", "modernise"), + ("normalize", "normalise"), + ("optimize", "optimise"), + ("prioritize", "prioritise"), + ("revolutionize", "revolutionise"), + ("specialize", "specialise"), + ("standardize", "standardise"), + ("summarize", "summarise"), + ("utilize", "utilise"), + // -or → -our (word-boundary needed) + ("\\bcolor\\b", "colour"), + ("\\bfavor\\b", "favour"), + ("\\bhonor\\b", "honour"), + ("\\bhumor\\b", "humour"), + ("\\blabor\\b", "labour"), + ("\\bneighbor\\b", "neighbour"), + ("behavior", "behaviour"), + // -er → -re (word-boundary to avoid "parameter" etc.) + ("\\bcenter\\b", "centre"), + ("\\bfiber\\b", "fibre"), + ("\\bliter\\b", "litre"), + ("\\bmeter\\b", "metre"), + ("\\btheater\\b", "theatre"), + // -ense → -ence + ("defense", "defence"), + ("offense", "offence"), + // Other + ("\\bcatalog\\b", "catalogue"), + ("\\bdialog\\b", "dialogue"), +]; + +/// Convert American English spelling to British English (word-boundary aware). +pub fn to_british_english(text: &str) -> String { + let mut result = text.to_string(); + + for (us, uk) in BRITISH_REPLACEMENTS { + let pattern = if us.contains("\\b") { + format!("(?i){us}") + } else { + format!( + "(?i)\\b{}(?:d|s|r|rs)?\\b", + regex_lite::escape(us) + ) + }; + + if let Ok(re) = regex_lite::Regex::new(&pattern) { + result = re + .replace_all(&result, |caps: ®ex_lite::Captures| { + let matched = caps.get(0).unwrap().as_str(); + let us_base = us.replace("\\b", ""); + let base_len = us_base.len(); + let suffix = if matched.len() > base_len { + &matched[base_len..] + } else { + "" + }; + let first_char = matched.chars().next().unwrap(); + let uk_clean = uk.replace("\\b", ""); + if first_char.is_uppercase() { + let mut chars = uk_clean.chars(); + let upper_first: String = + chars.next().unwrap().to_uppercase().collect(); + format!( + "{}{}{}", + upper_first, + chars.collect::(), + suffix + ) + } else { + format!("{uk_clean}{suffix}") + } + }) + .to_string(); + } + } + + result +} + +/// Basic formatting: capitalise sentences, fix spacing, clean punctuation. +pub fn format_text(text: &str) -> String { + if text.is_empty() { + return String::new(); + } + + let mut result = String::with_capacity(text.len()); + let mut capitalise_next = true; + + let chars: Vec = text.chars().collect(); + let mut i = 0; + + while i < chars.len() { + let c = chars[i]; + + if c == ' ' && i + 1 < chars.len() && chars[i + 1] == ' ' { + i += 1; + continue; + } + + if capitalise_next && c.is_alphabetic() { + result.extend(c.to_uppercase()); + capitalise_next = false; + } else { + result.push(c); + } + + if c == '.' || c == '!' || c == '?' { + capitalise_next = true; + } + if c == '\n' { + capitalise_next = true; + } + + i += 1; + } + + result +} + +/// Known hallucination markers that should be filtered from transcriptions. +static HALLUCINATION_MARKERS: &[&str] = &[ + "[blank_audio]", + "[music]", + "[silence]", +]; + +static AUTO_THANKS_PHRASES: &[&str] = &[ + "thank you.", + "thanks.", + "you.", + "thank you for watching.", +]; + +/// Returns true if a segment's text looks like a hallucination. +pub fn is_hallucination(text: &str) -> bool { + let trimmed = text.trim().to_lowercase(); + if trimmed.is_empty() { + return true; + } + for marker in HALLUCINATION_MARKERS { + if trimmed.contains(marker) { + return true; + } + } + if trimmed.len() < 15 { + for phrase in AUTO_THANKS_PHRASES { + if trimmed == *phrase { + return true; + } + } + } + false +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn remove_fillers_strips_um_and_uh() { + let input = "So um I was thinking uh about this"; + let result = remove_fillers(input); + assert!(!result.contains("um")); + assert!(!result.contains("uh")); + } + + #[test] + fn remove_fillers_preserves_legitimate_words() { + let input = "The umbrella was actually useful"; + let result = remove_fillers(input); + assert!(result.contains("umbrella")); + assert!(result.contains("useful")); + } + + #[test] + fn to_british_english_converts_ize_to_ise() { + assert!(to_british_english("organize").contains("organise")); + assert!(to_british_english("realize").contains("realise")); + } + + #[test] + fn to_british_english_preserves_case() { + let result = to_british_english("Organize the files"); + assert!(result.starts_with("Organise")); + } + + #[test] + fn to_british_english_handles_colour() { + assert!(to_british_english("the color is red").contains("colour")); + } + + #[test] + fn format_text_capitalises_after_full_stops() { + let result = format_text("hello world. this is a test"); + assert!(result.starts_with('H')); + assert!(result.contains(". T")); + } + + #[test] + fn format_text_handles_empty_string() { + assert_eq!(format_text(""), ""); + } + + #[test] + fn is_hallucination_detects_blank_audio() { + assert!(is_hallucination("[blank_audio]")); + assert!(is_hallucination(" [music] ")); + } + + #[test] + fn is_hallucination_detects_auto_thanks() { + assert!(is_hallucination("Thank you.")); + assert!(is_hallucination("thanks.")); + } + + #[test] + fn is_hallucination_allows_real_text() { + assert!(!is_hallucination("The meeting is at three o'clock.")); + } +}