scriptrs 0.1.0

Rust transcription with native CoreML Parakeet v2 inference
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use std::path::Path;

use crate::constants::{VAD_CONTEXT_SAMPLES, VAD_STATE_SIZE, VAD_WINDOW_SAMPLES};
use crate::coreml::{CoreMlInput, CoreMlModel, CoreMlTensor};
use crate::error::TranscriptionError;
use crate::long_form::planner::VadConfig;

#[derive(Debug, Clone)]
pub(crate) struct SileroVad {
    model: VadModelInner,
}

impl SileroVad {
    pub(crate) fn new(model_path: &Path) -> Result<Self, TranscriptionError> {
        Ok(Self {
            model: VadModelInner::new(model_path)?,
        })
    }

    pub(crate) fn process(
        &self,
        audio: &[f32],
        config: &VadConfig,
    ) -> Result<Vec<f32>, TranscriptionError> {
        self.model.process(audio, config)
    }
}

#[derive(Debug, Clone)]
enum VadModelInner {
    #[cfg(target_os = "macos")]
    CoreMl(CoreMlModel),
    #[cfg(not(target_os = "macos"))]
    Unsupported,
}

impl VadModelInner {
    fn new(model_path: &Path) -> Result<Self, TranscriptionError> {
        #[cfg(target_os = "macos")]
        {
            Ok(Self::CoreMl(CoreMlModel::new(model_path)?))
        }
        #[cfg(not(target_os = "macos"))]
        {
            let _ = model_path;
            Err(TranscriptionError::UnsupportedPlatform)
        }
    }

    fn process(&self, audio: &[f32], _config: &VadConfig) -> Result<Vec<f32>, TranscriptionError> {
        match self {
            #[cfg(target_os = "macos")]
            Self::CoreMl(model) => process_coreml_vad(model, audio),
            #[cfg(not(target_os = "macos"))]
            Self::Unsupported => Err(TranscriptionError::UnsupportedPlatform),
        }
    }
}

#[cfg(target_os = "macos")]
fn process_coreml_vad(model: &CoreMlModel, audio: &[f32]) -> Result<Vec<f32>, TranscriptionError> {
    if audio.is_empty() {
        return Ok(Vec::new());
    }

    let mut probabilities = Vec::with_capacity(audio.len().div_ceil(VAD_WINDOW_SAMPLES));
    let mut hidden_state = vec![0.0f32; VAD_STATE_SIZE];
    let mut cell_state = vec![0.0f32; VAD_STATE_SIZE];
    let mut context = vec![0.0f32; VAD_CONTEXT_SAMPLES];

    for chunk_start in (0..audio.len()).step_by(VAD_WINDOW_SAMPLES) {
        let chunk_end = (chunk_start + VAD_WINDOW_SAMPLES).min(audio.len());
        let mut chunk = audio[chunk_start..chunk_end].to_vec();
        if chunk.len() < VAD_WINDOW_SAMPLES {
            let last = chunk.last().copied().unwrap_or(0.0);
            chunk.resize(VAD_WINDOW_SAMPLES, last);
        }

        let next_context = chunk[chunk.len() - VAD_CONTEXT_SAMPLES..].to_vec();
        let mut model_input = Vec::with_capacity(VAD_WINDOW_SAMPLES + VAD_CONTEXT_SAMPLES);
        model_input.extend_from_slice(&context);
        model_input.extend_from_slice(&chunk);

        let outputs = model.predict(
            &[
                CoreMlInput::F32 {
                    name: "audio",
                    values: &model_input,
                    shape: &[1, 1, VAD_WINDOW_SAMPLES + VAD_CONTEXT_SAMPLES],
                },
                CoreMlInput::F32 {
                    name: "h",
                    values: &hidden_state,
                    shape: &[1, 1, VAD_STATE_SIZE],
                },
                CoreMlInput::F32 {
                    name: "c",
                    values: &cell_state,
                    shape: &[1, 1, VAD_STATE_SIZE],
                },
            ],
            &["probability", "h_out", "c_out"],
        )?;

        let probability = outputs
            .get("probability")
            .and_then(|tensor| tensor.data.first())
            .copied()
            .ok_or_else(|| {
                TranscriptionError::CoreMl("VAD output `probability` was empty".to_owned())
            })?;
        hidden_state = take_state(outputs.get("h_out"), VAD_STATE_SIZE)?;
        cell_state = take_state(outputs.get("c_out"), VAD_STATE_SIZE)?;
        probabilities.push(probability);
        context = next_context;
    }

    Ok(probabilities)
}

#[cfg(target_os = "macos")]
fn take_state(
    tensor: Option<&CoreMlTensor>,
    state_size: usize,
) -> Result<Vec<f32>, TranscriptionError> {
    let tensor = tensor.ok_or_else(|| {
        TranscriptionError::CoreMl("missing recurrent VAD state output".to_owned())
    })?;
    if tensor.data.len() < state_size {
        return Err(TranscriptionError::CoreMl(format!(
            "invalid recurrent VAD state length: {}",
            tensor.data.len()
        )));
    }
    Ok(tensor.data[..state_size].to_vec())
}

#[cfg(test)]
mod tests {
    use crate::constants::SAMPLE_RATE;

    #[test]
    fn sample_rate_constant_matches_expected_value() {
        assert_eq!(SAMPLE_RATE, 16_000);
    }
}