aria_wake/
frontend.rs

1//! Shared mel, embedding, and classifier stages for the wake engines.
2
3use std::path::{Path, PathBuf};
4
5use ort::{
6    session::{Session, SessionOutputs},
7    value::{Outlet, Tensor, TensorElementType, ValueType},
8};
9
10use aria_kit::backend::{ModelError, validate_model_file};
11
12pub(crate) const FEATURE_LEN: usize = FEATURE_WINDOW_FRAMES * FEATURE_WIDTH;
13
14pub(crate) const MEL_WINDOW_FRAMES: usize = 76;
15pub(crate) const MEL_BINS: usize = 32;
16pub(crate) const FEATURE_WINDOW_FRAMES: usize = 16;
17pub(crate) const FEATURE_WIDTH: usize = 96;
18const MEL_SCALE: f32 = 10.0;
19const MEL_OFFSET: f32 = 2.0;
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub(crate) enum AudioScale {
23    /// Raw `i16` magnitudes, as openWakeWord trains.
24    #[cfg(any(feature = "openwakeword", test))]
25    Int16,
26    /// Normalized to `[-1, 1]`, as livekit-wakeword trains.
27    #[cfg(any(feature = "livekit", test))]
28    Unit,
29}
30
31impl AudioScale {
32    fn apply(self, sample: i16) -> f32 {
33        match self {
34            #[cfg(any(feature = "openwakeword", test))]
35            Self::Int16 => f32::from(sample),
36            #[cfg(any(feature = "livekit", test))]
37            Self::Unit => aria_kit::audio::f32_from_pcm16(sample),
38        }
39    }
40}
41
42/// The two shared models sit beside the classifier, so naming the classifier
43/// names all three.
44pub struct ModelFiles {
45    pub melspectrogram: PathBuf,
46    pub embedding: PathBuf,
47    pub classifier: PathBuf,
48}
49
50impl ModelFiles {
51    pub fn beside(classifier: &Path) -> Result<Self, Error> {
52        let classifier = trusted_model(classifier, "classifier")?;
53        let directory = classifier.parent().ok_or_else(|| Error::Incompatible {
54            role: "classifier",
55            detail: "its canonical path has no parent directory".to_owned(),
56        })?;
57        Ok(Self {
58            melspectrogram: trusted_model(
59                &directory.join("melspectrogram.onnx"),
60                "melspectrogram",
61            )?,
62            embedding: trusted_model(&directory.join("embedding_model.onnx"), "embedding")?,
63            classifier,
64        })
65    }
66}
67
68/// Applies openWakeWord's `x / 10 + 2` transform.
69pub(crate) struct Mel {
70    session: Session,
71    scale: AudioScale,
72}
73
74impl Mel {
75    pub(crate) fn load(path: &Path, scale: AudioScale) -> Result<Self, Error> {
76        let session = load(path, "melspectrogram")?;
77        validate_model(
78            &session,
79            "melspectrogram",
80            TensorSignature::new("input", &[-1, -1]),
81            TensorSignature::new("output", &[-1, 1, -1, 32]),
82        )?;
83        Ok(Self { session, scale })
84    }
85
86    pub(crate) fn frames(&mut self, samples: &[i16]) -> Result<Vec<[f32; MEL_BINS]>, Error> {
87        let audio: Vec<f32> = samples.iter().map(|s| self.scale.apply(*s)).collect();
88        let sample_count = i64::try_from(audio.len()).map_err(|_| Error::Incompatible {
89            role: "melspectrogram",
90            detail: "audio window is too large".to_owned(),
91        })?;
92        let input = tensor(&[1, sample_count], audio)?;
93        let mut outputs = self
94            .session
95            .run(ort::inputs!["input" => input])
96            .map_err(|error| Error::Inference(error.into()))?;
97        let values = output(&mut outputs, "melspectrogram", "output", None)?;
98        if values.is_empty() || !values.len().is_multiple_of(MEL_BINS) {
99            return Err(Error::Incompatible {
100                role: "melspectrogram",
101                detail: format!("returned {} values, not complete mel frames", values.len()),
102            });
103        }
104        values
105            .as_chunks::<MEL_BINS>()
106            .0
107            .iter()
108            .map(|frame| {
109                let transformed: Vec<f32> = frame
110                    .iter()
111                    .map(|value| value / MEL_SCALE + MEL_OFFSET)
112                    .collect();
113                transformed.try_into().map_err(|_| Error::Incompatible {
114                    role: "melspectrogram",
115                    detail: "returned an incomplete mel frame".to_owned(),
116                })
117            })
118            .collect()
119    }
120}
121
122pub(crate) struct Embedding {
123    session: Session,
124}
125
126impl Embedding {
127    pub(crate) fn load(path: &Path) -> Result<Self, Error> {
128        let session = load(path, "embedding")?;
129        validate_model(
130            &session,
131            "embedding",
132            TensorSignature::new("input_1", &[-1, 76, 32, 1]),
133            TensorSignature::new("conv2d_19", &[-1, 1, 1, 96]),
134        )?;
135        Ok(Self { session })
136    }
137
138    pub(crate) fn feature(
139        &mut self,
140        window: &[[f32; MEL_BINS]],
141    ) -> Result<[f32; FEATURE_WIDTH], Error> {
142        if window.len() != MEL_WINDOW_FRAMES {
143            return Err(Error::Incompatible {
144                role: "embedding",
145                detail: format!(
146                    "given {} mel frames, expected {MEL_WINDOW_FRAMES}",
147                    window.len()
148                ),
149            });
150        }
151        let spectrogram: Vec<f32> = window.iter().flatten().copied().collect();
152        let input = tensor(
153            &[1, MEL_WINDOW_FRAMES as i64, MEL_BINS as i64, 1],
154            spectrogram,
155        )?;
156        let mut outputs = self
157            .session
158            .run(ort::inputs!["input_1" => input])
159            .map_err(|error| Error::Inference(error.into()))?;
160        output(
161            &mut outputs,
162            "embedding",
163            "conv2d_19",
164            Some(&[1, 1, 1, FEATURE_WIDTH as i64]),
165        )?
166        .try_into()
167        .map_err(|values: Vec<f32>| Error::Incompatible {
168            role: "embedding",
169            detail: format!(
170                "returned {} features, expected {FEATURE_WIDTH}",
171                values.len()
172            ),
173        })
174    }
175}
176
177pub(crate) struct Classifier {
178    session: Session,
179    signature: ClassifierSignature,
180}
181
182impl Classifier {
183    pub(crate) fn load(path: &Path, signature: ClassifierSignature) -> Result<Self, Error> {
184        let session = load(path, "classifier")?;
185        validate_model(
186            &session,
187            "classifier",
188            TensorSignature::new(signature.input, signature.input_shape),
189            TensorSignature::new(signature.output, signature.output_shape),
190        )?;
191        let mut classifier = Self { session, signature };
192        classifier.score(&[0.0; FEATURE_LEN])?;
193        Ok(classifier)
194    }
195
196    /// The probability is bounds-checked because a mismatched head can return
197    /// numbers a threshold would silently accept.
198    pub(crate) fn score(&mut self, window: &[f32; FEATURE_LEN]) -> Result<f32, Error> {
199        let signature = self.signature;
200        let input = tensor(
201            &[1, FEATURE_WINDOW_FRAMES as i64, FEATURE_WIDTH as i64],
202            window.to_vec(),
203        )?;
204        let mut outputs = self
205            .session
206            .run(ort::inputs![signature.input => input])
207            .map_err(|error| Error::Inference(error.into()))?;
208        let values = output(&mut outputs, "classifier", signature.output, Some(&[1, 1]))?;
209        let [probability] = values.as_slice() else {
210            return Err(Error::Incompatible {
211                role: "classifier",
212                detail: format!("returned {} probabilities, expected one", values.len()),
213            });
214        };
215        if !probability.is_finite() || !(0.0..=1.0).contains(probability) {
216            return Err(Error::InvalidProbability(*probability));
217        }
218        Ok(*probability)
219    }
220}
221
222#[derive(Debug, Clone, Copy)]
223pub(crate) struct ClassifierSignature {
224    pub(crate) input: &'static str,
225    pub(crate) input_shape: &'static [i64],
226    pub(crate) output: &'static str,
227    pub(crate) output_shape: &'static [i64],
228}
229
230// ONNX Runtime errors are boxed so the pinned `ort` version stays out of the
231// public API.
232#[derive(Debug, thiserror::Error)]
233pub enum Error {
234    #[error("the {role} wake-word model cannot be used")]
235    Model {
236        role: &'static str,
237        #[source]
238        source: ModelError,
239    },
240
241    #[error("failed to load the {role} wake-word model")]
242    Load {
243        role: &'static str,
244        #[source]
245        source: Box<dyn std::error::Error + Send + Sync>,
246    },
247
248    #[error("wake-word inference failed")]
249    Inference(#[source] Box<dyn std::error::Error + Send + Sync>),
250
251    #[error("the {role} wake-word model is incompatible: {detail}")]
252    Incompatible { role: &'static str, detail: String },
253
254    #[error("the wake-word classifier returned invalid probability {0}")]
255    InvalidProbability(f32),
256}
257
258fn trusted_model(path: &Path, role: &'static str) -> Result<PathBuf, Error> {
259    validate_model_file(path.to_path_buf()).map_err(|source| Error::Model { role, source })
260}
261
262fn load(path: &Path, role: &'static str) -> Result<Session, Error> {
263    Session::builder()
264        .and_then(|mut builder| builder.commit_from_file(path))
265        .map_err(|source| Error::Load {
266            role,
267            source: source.into(),
268        })
269}
270
271struct TensorSignature<'a> {
272    name: &'a str,
273    shape: &'a [i64],
274}
275
276impl<'a> TensorSignature<'a> {
277    fn new(name: &'a str, shape: &'a [i64]) -> Self {
278        Self { name, shape }
279    }
280}
281
282fn validate_model(
283    session: &Session,
284    role: &'static str,
285    input: TensorSignature<'_>,
286    output: TensorSignature<'_>,
287) -> Result<(), Error> {
288    validate_outlets(session.inputs(), role, "input", input)?;
289    validate_outlets(session.outputs(), role, "output", output)
290}
291
292fn validate_outlets(
293    outlets: &[Outlet],
294    role: &'static str,
295    direction: &'static str,
296    expected: TensorSignature<'_>,
297) -> Result<(), Error> {
298    let [outlet] = outlets else {
299        return Err(Error::Incompatible {
300            role,
301            detail: format!("expected one {direction}, found {}", outlets.len()),
302        });
303    };
304    if outlet.name() != expected.name {
305        return Err(Error::Incompatible {
306            role,
307            detail: format!(
308                "expected {direction} {}, found {}",
309                expected.name,
310                outlet.name()
311            ),
312        });
313    }
314    match outlet.dtype() {
315        ValueType::Tensor {
316            ty: TensorElementType::Float32,
317            shape,
318            ..
319        } if &shape[..] == expected.shape => Ok(()),
320        actual => Err(Error::Incompatible {
321            role,
322            detail: format!(
323                "expected {direction} {} to be a float tensor shaped {:?}, found {actual}",
324                expected.name, expected.shape
325            ),
326        }),
327    }
328}
329
330fn output(
331    outputs: &mut SessionOutputs<'_>,
332    role: &'static str,
333    name: &str,
334    expected_shape: Option<&[i64]>,
335) -> Result<Vec<f32>, Error> {
336    let value = outputs.remove(name).ok_or_else(|| Error::Incompatible {
337        role,
338        detail: format!("returned no {name} output"),
339    })?;
340    let (shape, values) = value
341        .try_extract_tensor::<f32>()
342        .map_err(|error| Error::Inference(error.into()))?;
343    if expected_shape.is_some_and(|expected| &shape[..] != expected) {
344        return Err(Error::Incompatible {
345            role,
346            detail: format!("returned tensor shaped {shape}, expected {expected_shape:?}"),
347        });
348    }
349    Ok(values.to_vec())
350}
351
352fn tensor(shape: &[i64], values: Vec<f32>) -> Result<Tensor<f32>, Error> {
353    Tensor::from_array((shape, values)).map_err(|error| Error::Inference(error.into()))
354}
355
356#[cfg(test)]
357mod tests {
358    use super::*;
359
360    #[test]
361    fn int16_audio_passes_through_and_unit_audio_is_normalized() {
362        assert_eq!(AudioScale::Int16.apply(16_384), 16_384.0);
363        assert_eq!(AudioScale::Unit.apply(16_384), 0.5);
364    }
365}
366
367pub struct Config {
368    classifier: PathBuf,
369    threshold: f32,
370}
371
372impl Config {
373    pub fn new(classifier: impl Into<PathBuf>, threshold: f32) -> Result<Self, ConfigError> {
374        if !threshold.is_finite() || threshold <= 0.0 || threshold > 1.0 {
375            return Err(ConfigError::InvalidThreshold(threshold));
376        }
377        Ok(Self {
378            classifier: classifier.into(),
379            threshold,
380        })
381    }
382
383    pub fn classifier(&self) -> &Path {
384        &self.classifier
385    }
386
387    pub fn threshold(&self) -> f32 {
388        self.threshold
389    }
390}
391
392#[derive(Debug, thiserror::Error)]
393pub enum ConfigError {
394    #[error("wake threshold must be finite, greater than zero, and at most one; got {0}")]
395    InvalidThreshold(f32),
396}
aria_wake/frontend.rs

aria_wake/
frontend.rs
