aria_kit/transcription/layer/
mod.rs

1use std::collections::VecDeque;
2
3use crate::cancel::CancelToken;
4
5use super::{TranscriptUpdate, TranscriptionError, check_cancelled};
6
7mod backtrack;
8pub use backtrack::Backtrack;
9mod disfluencies;
10pub use disfluencies::{RemoveHesitations, RemoveRepeats};
11pub mod vocabulary;
12pub use vocabulary::Vocabulary;
13
14#[cfg(test)]
15mod tests;
16
17const MAX_TEXT_BYTES: usize = 64 * 1024;
18
19#[derive(Debug, thiserror::Error)]
20pub enum Error {
21    #[error("processed transcription exceeds 64 KiB")]
22    TooLarge,
23    #[error("a transcription layer failed")]
24    Failed(#[source] Box<dyn std::error::Error + Send + Sync>),
25}
26
27pub trait Layer: Send + Sync {
28    fn apply(&self, text: &mut String, cancel: &CancelToken) -> Result<(), TranscriptionError>;
29
30    fn enabled(&self) -> bool {
31        true
32    }
33}
34
35pub struct Layered<L, I> {
36    pub(super) inner: I,
37    pub(super) layer: L,
38}
39
40impl<L: Layer, I: Layer> Layered<L, I> {
41    pub fn new(inner: I, layer: L) -> Self {
42        Self { inner, layer }
43    }
44}
45
46impl<L: Layer, I: Layer> Layer for Layered<L, I> {
47    fn apply(&self, text: &mut String, cancel: &CancelToken) -> Result<(), TranscriptionError> {
48        apply(&self.inner, text, cancel)?;
49        apply(&self.layer, text, cancel)
50    }
51
52    fn enabled(&self) -> bool {
53        self.inner.enabled() || self.layer.enabled()
54    }
55}
56
57impl Layer for () {
58    fn apply(&self, _text: &mut String, _cancel: &CancelToken) -> Result<(), TranscriptionError> {
59        Ok(())
60    }
61
62    fn enabled(&self) -> bool {
63        false
64    }
65}
66
67impl<L: Layer> Layer for Option<L> {
68    fn apply(&self, text: &mut String, cancel: &CancelToken) -> Result<(), TranscriptionError> {
69        if let Some(layer) = self {
70            apply(layer, text, cancel)?;
71        }
72        Ok(())
73    }
74
75    fn enabled(&self) -> bool {
76        self.as_ref().is_some_and(Layer::enabled)
77    }
78}
79
80impl<L: Layer> Layer for &L {
81    fn apply(&self, text: &mut String, cancel: &CancelToken) -> Result<(), TranscriptionError> {
82        (**self).apply(text, cancel)
83    }
84
85    fn enabled(&self) -> bool {
86        (**self).enabled()
87    }
88}
89
90impl Layer for Vocabulary {
91    fn apply(&self, text: &mut String, cancel: &CancelToken) -> Result<(), TranscriptionError> {
92        if let Some(corrected) = self.correct_checked(text, &|bytes| {
93            check_cancelled(cancel)?;
94            check_size(bytes)
95        })? {
96            *text = corrected;
97        }
98        Ok(())
99    }
100}
101
102fn apply<L: Layer>(
103    layer: &L,
104    text: &mut String,
105    cancel: &CancelToken,
106) -> Result<(), TranscriptionError> {
107    check_cancelled(cancel)?;
108    if !layer.enabled() {
109        return Ok(());
110    }
111    check_size(text.len())?;
112    layer.apply(text, cancel)?;
113    check_cancelled(cancel)?;
114    check_size(text.len())
115}
116
117fn check_size(bytes: usize) -> Result<(), TranscriptionError> {
118    if bytes > MAX_TEXT_BYTES {
119        return Err(TranscriptionError::Processing(Error::TooLarge));
120    }
121    Ok(())
122}
123
124#[derive(Default)]
125pub(super) struct Output {
126    pub text: String,
127    pub pending: VecDeque<TranscriptUpdate>,
128    pub finished: bool,
129    pub retained_from: Option<u64>,
130}
131
132impl Output {
133    pub(super) fn push<L: Layer>(
134        &mut self,
135        mut update: TranscriptUpdate,
136        layers: &L,
137        cancel: &CancelToken,
138    ) -> Result<(), TranscriptionError> {
139        if let TranscriptUpdate::Finalized(segment) = &mut update {
140            // Matching stays within one segment. A phrase split across
141            // segments would need bounded holdback.
142            let mut text = segment.as_str().to_owned();
143            apply(layers, &mut text, cancel)?;
144            segment.rewrite(text.trim().to_owned());
145            if segment.as_str().is_empty() {
146                return Ok(());
147            }
148            if !self.text.is_empty() {
149                self.text.push(' ');
150            }
151            self.text.push_str(segment.as_str());
152        }
153        self.pending.push_back(update);
154        Ok(())
155    }
156}
aria_kit/transcription/layer/mod.rs

aria_kit/transcription/layer/
mod.rs
