aria_kit/transcription/layer/
hesitations.rs

1use crate::cancel::CancelToken;
2
3use super::{Layer, TranscriptionError};
4
5const FILLERS: &[&str] = &[
6    "um", "umm", "uh", "uhh", "uhm", "er", "erm", "ah", "ahh", "hmm", "hm", "mm", "mhm",
7];
8
9pub struct RemoveHesitations;
10
11impl Layer for RemoveHesitations {
12    fn apply(&self, text: &mut String, _cancel: &CancelToken) -> Result<(), TranscriptionError> {
13        let mut kept: Vec<Word> = Vec::new();
14        let mut opening = String::new();
15        let mut capitalize_next = false;
16        let mut dropped = false;
17        for mut word in words(text) {
18            let core = word.core();
19            let filler = FILLERS.contains(&squeezed(&core).as_str());
20            if !filler {
21                if capitalize_next {
22                    word.capitalize();
23                    capitalize_next = false;
24                }
25                word.text.insert_str(0, &opening);
26                opening.clear();
27                kept.push(word);
28                continue;
29            }
30            dropped = true;
31            let spoken = word.text.trim_matches(|c: char| c.is_ascii_punctuation());
32            capitalize_next |= spoken.starts_with(char::is_uppercase);
33            let start = word.text.len()
34                - word
35                    .text
36                    .trim_start_matches(|c: char| c.is_ascii_punctuation())
37                    .len();
38            let end = word
39                .text
40                .trim_end_matches(|c: char| c.is_ascii_punctuation())
41                .len();
42            opening.extend(
43                word.text[..start]
44                    .chars()
45                    .filter(|c| matches!(c, '(' | '[' | '{' | '\"' | '\'')),
46            );
47            let mut punctuation = String::new();
48            for character in word.text[end..].chars() {
49                if matches!(
50                    (opening.chars().next_back(), character),
51                    (Some('('), ')')
52                        | (Some('['), ']')
53                        | (Some('{'), '}')
54                        | (Some('\"'), '\"')
55                        | (Some('\''), '\'')
56                ) {
57                    opening.pop();
58                } else {
59                    punctuation.push(character);
60                }
61            }
62            if let Some(last) = kept.last_mut() {
63                last.absorb(&punctuation, &word.trailing_space);
64            }
65        }
66        if !dropped {
67            return Ok(());
68        }
69        *text = kept.iter().map(Word::render).collect();
70        Ok(())
71    }
72}
73
74struct Word {
75    text: String,
76    trailing_space: String,
77}
78
79impl Word {
80    fn core(&self) -> String {
81        self.text
82            .trim_matches(|c: char| c.is_ascii_punctuation())
83            .to_lowercase()
84    }
85
86    fn capitalize(&mut self) {
87        let start = self.text.len()
88            - self
89                .text
90                .trim_start_matches(|c: char| c.is_ascii_punctuation())
91                .len();
92        let mut chars = self.text[start..].chars();
93        if let Some(first) = chars.next() {
94            self.text.replace_range(
95                start..start + first.len_utf8(),
96                &first.to_uppercase().to_string(),
97            );
98        }
99    }
100
101    fn absorb(&mut self, punctuation: &str, trailing_space: &str) {
102        if punctuation.contains(['.', '!', '?', ')', ']', '}', '"', '\'']) {
103            self.text
104                .truncate(self.text.trim_end_matches([',', ';', ':']).len());
105        }
106        if !self.text.ends_with(|c: char| c.is_ascii_punctuation()) {
107            self.text.push_str(punctuation);
108        } else {
109            self.text.extend(
110                punctuation
111                    .chars()
112                    .filter(|c| matches!(c, ')' | ']' | '}' | '"' | '\'')),
113            );
114        }
115        if !self.trailing_space.contains('\n') {
116            trailing_space.clone_into(&mut self.trailing_space);
117        }
118    }
119
120    fn render(&self) -> String {
121        format!("{}{}", self.text, self.trailing_space)
122    }
123}
124
125fn squeezed(core: &str) -> String {
126    let mut out = String::with_capacity(core.len());
127    let mut run = 0;
128    let mut last = None;
129    for c in core.chars() {
130        run = if last == Some(c) { run + 1 } else { 1 };
131        last = Some(c);
132        if run <= 2 {
133            out.push(c);
134        }
135    }
136    out
137}
138
139fn words(text: &str) -> Vec<Word> {
140    let mut words = Vec::new();
141    let mut rest = text.trim_start();
142    while !rest.is_empty() {
143        let end = rest.find(char::is_whitespace).unwrap_or(rest.len());
144        let (word, after) = rest.split_at(end);
145        let space_end = after
146            .find(|c: char| !c.is_whitespace())
147            .unwrap_or(after.len());
148        let (space, next) = after.split_at(space_end);
149        words.push(Word {
150            text: word.to_owned(),
151            trailing_space: space.to_owned(),
152        });
153        rest = next;
154    }
155    words
156}
aria_kit/transcription/layer/hesitations.rs

aria_kit/transcription/layer/
hesitations.rs
