aria_kit/transcription/layer/
mod.rs

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

aria_kit/transcription/layer/
mod.rs
