aria_wake/livekit/
mod.rs

1//! livekit scores a whole two-second window at a time.
2
3use std::collections::VecDeque;
4
5pub use crate::{Config, ConfigError};
6
7use crate::frontend::{
8    AudioScale, Classifier, ClassifierSignature, Embedding, Error, FEATURE_LEN, FEATURE_WIDTH,
9    FEATURE_WINDOW_FRAMES, MEL_WINDOW_FRAMES, Mel, ModelFiles,
10};
11use aria_kit::wake_word::{WakeWordDetector, WakeWordMatch};
12
13pub const NAME: &str = "livekit";
14
15/// livekit exports with a dynamic batch dimension; openWakeWord pins it to one.
16const SIGNATURE: ClassifierSignature = ClassifierSignature {
17    input: "embeddings",
18    input_shape: &[-1, 16, 96],
19    output: "score",
20    output_shape: &[-1, 1],
21};
22
23const SCALE: AudioScale = AudioScale::Unit;
24
25/// Yields exactly the 16 embeddings the head reads.
26const WINDOW_SAMPLES: usize = 32_000;
27/// 320 ms between scored windows, as livekit's own listener uses.
28const HOP_SAMPLES: usize = 5_120;
29/// Mel frames between consecutive embedding windows.
30const EMBEDDING_STRIDE: usize = 8;
31
32pub struct LiveKitWakeWord {
33    mel: Mel,
34    embedding: Embedding,
35    classifier: Classifier,
36    threshold: f32,
37    window: VecDeque<i16>,
38    since_last_score: usize,
39}
40
41impl LiveKitWakeWord {
42    pub fn check_from(config: &Config) -> Result<ModelFiles, Error> {
43        let files = ModelFiles::beside(config.classifier())?;
44        Classifier::load(&files.classifier, SIGNATURE)?;
45        Ok(files)
46    }
47
48    pub fn load_from(config: &Config) -> Result<Self, Error> {
49        let files = Self::check_from(config)?;
50        Ok(Self {
51            mel: Mel::load(&files.melspectrogram, SCALE)?,
52            embedding: Embedding::load(&files.embedding)?,
53            classifier: Classifier::load(&files.classifier, SIGNATURE)?,
54            threshold: config.threshold(),
55            window: VecDeque::with_capacity(WINDOW_SAMPLES),
56            since_last_score: 0,
57        })
58    }
59
60    fn detect_inner(&mut self, canonical_samples: &[i16]) -> Result<WakeWordMatch, Error> {
61        let mut matched = WakeWordMatch::NoMatch;
62        for sample in canonical_samples {
63            self.window.push_back(*sample);
64            if self.window.len() > WINDOW_SAMPLES {
65                self.window.pop_front();
66            }
67            self.since_last_score += 1;
68            if self.since_last_score < HOP_SAMPLES || self.window.len() < WINDOW_SAMPLES {
69                continue;
70            }
71            self.since_last_score = 0;
72            if self.score()? >= self.threshold {
73                matched = WakeWordMatch::Detected;
74            }
75        }
76        Ok(matched)
77    }
78
79    /// Scores the buffered window the way livekit's own runtime does: mel over
80    /// the whole window, then the last 16 embeddings taken at stride 8.
81    fn score(&mut self) -> Result<f32, Error> {
82        let samples: Vec<i16> = self.window.iter().copied().collect();
83        let frames = self.mel.frames(&samples)?;
84        let last_start = frames
85            .len()
86            .checked_sub(MEL_WINDOW_FRAMES)
87            .map(|last| last / EMBEDDING_STRIDE * EMBEDDING_STRIDE)
88            .filter(|start| start / EMBEDDING_STRIDE + 1 >= FEATURE_WINDOW_FRAMES)
89            .ok_or_else(|| Error::Incompatible {
90                role: "melspectrogram",
91                detail: format!(
92                    "returned {} frames for a {WINDOW_SAMPLES}-sample window, too few for {FEATURE_WINDOW_FRAMES} embeddings",
93                    frames.len()
94                ),
95            })?;
96
97        let mut features: Vec<[f32; FEATURE_WIDTH]> = Vec::with_capacity(FEATURE_WINDOW_FRAMES);
98        for step in (0..FEATURE_WINDOW_FRAMES).rev() {
99            let start = last_start - step * EMBEDDING_STRIDE;
100            features.push(
101                self.embedding
102                    .feature(&frames[start..start + MEL_WINDOW_FRAMES])?,
103            );
104        }
105
106        let window: Vec<f32> = features.iter().flatten().copied().collect();
107        let window: [f32; FEATURE_LEN] = window.try_into().map_err(|_| Error::Incompatible {
108            role: "embedding",
109            detail: "built an incomplete classifier window".to_owned(),
110        })?;
111        self.classifier.score(&window)
112    }
113}
114
115impl WakeWordDetector for LiveKitWakeWord {
116    type Error = Error;
117
118    fn detect(&mut self, canonical_samples: &[i16]) -> Result<WakeWordMatch, Self::Error> {
119        match self.detect_inner(canonical_samples) {
120            Ok(matched) => Ok(matched),
121            Err(error) => {
122                self.reset();
123                Err(error)
124            }
125        }
126    }
127
128    fn reset(&mut self) {
129        self.window.clear();
130        self.since_last_score = 0;
131    }
132}
aria_wake/livekit/mod.rs

aria_wake/livekit/
mod.rs
