add itn (#41)

Author: cjpais

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "transcribe-rs"
-version = "0.2.5"
+version = "0.2.6"
 edition = "2021"
 description = "A simple library to help you transcribe audio"
 license = "MIT"
@@ -19,8 +19,11 @@ whisperfile = ["dep:ureq"]
 # Remote engines
 openai = ["dep:async-openai", "dep:tokio", "dep:async-trait"]
 
+# Post-processing
+itn = ["dep:nemo-text-processing"]
+
 # Convenience
-all = ["whisper", "parakeet", "moonshine", "sense_voice", "whisperfile", "openai"]
+all = ["whisper", "parakeet", "moonshine", "sense_voice", "whisperfile", "openai", "itn"]
 
 [dependencies]
 # Always required
@@ -47,6 +50,9 @@ base64 = { version = "0.22", optional = true }
 # Whisperfile
 ureq = { version = "3", optional = true }
 
+# ITN (Inverse Text Normalization)
+nemo-text-processing = { git = "https://github.com/FluidInference/text-processing-rs", optional = true }
+
 # OpenAI
 tokio = { version = "1.47.1", features = ["rt-multi-thread"], optional = true }
 async-openai = { version = "0.29.3", optional = true }

Cargo.toml

@@ -0,0 +1,74 @@
+//! Inverse Text Normalization (ITN) post-processing.
+//!
+//! Converts spoken-form text (e.g. "twenty three dollars") into written-form
+//! (e.g. "$23") using rules from the `nemo-text-processing` crate.
+
+use crate::TranscriptionResult;
+use nemo_text_processing::normalize_sentence;
+
+/// Apply inverse text normalization to a transcription result.
+///
+/// Normalizes `result.text` and, if present, each segment's text.
+///
+/// # Example
+///
+/// ```ignore
+/// use transcribe_rs::itn::apply_itn;
+///
+/// let mut result = engine.transcribe_file(&path, None)?;
+/// apply_itn(&mut result);
+/// println!("{}", result.text); // written-form output
+/// ```
+pub fn apply_itn(result: &mut TranscriptionResult) {
+    result.text = normalize_sentence(&result.text);
+    if let Some(segments) = &mut result.segments {
+        for segment in segments.iter_mut() {
+            segment.text = normalize_sentence(&segment.text);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::TranscriptionSegment;
+
+    #[test]
+    fn test_apply_itn_normalizes_text_and_segments() {
+        let mut result = TranscriptionResult {
+            text: "twenty three dollars".to_string(),
+            segments: Some(vec![
+                TranscriptionSegment {
+                    start: 0.0,
+                    end: 1.0,
+                    text: "twenty three dollars".to_string(),
+                },
+                TranscriptionSegment {
+                    start: 1.0,
+                    end: 2.0,
+                    text: "one hundred fifty two".to_string(),
+                },
+            ]),
+        };
+
+        apply_itn(&mut result);
+
+        assert_eq!(result.text, "$23");
+        let segments = result.segments.unwrap();
+        assert_eq!(segments[0].text, "$23");
+        assert_eq!(segments[1].text, "152");
+    }
+
+    #[test]
+    fn test_apply_itn_no_segments() {
+        let mut result = TranscriptionResult {
+            text: "twenty three dollars".to_string(),
+            segments: None,
+        };
+
+        apply_itn(&mut result);
+
+        assert_eq!(result.text, "$23");
+        assert!(result.segments.is_none());
+    }
+}

src/itn.rs

@@ -60,6 +60,11 @@ pub mod remote;
 #[cfg(feature = "openai")]
 pub use remote::RemoteTranscriptionEngine;
 
+#[cfg(feature = "itn")]
+pub mod itn;
+#[cfg(feature = "itn")]
+pub use itn::apply_itn;
+
 use std::path::Path;
 
 /// The result of a transcription operation.