Skip to main content

aurum_core/
window.rs

1//! Rolling-window helpers for host-driven “partial-like” decode loops.
2//!
3//! Aurum does not run a background streaming decoder. Hosts (e.g. dictation apps)
4//! can push PCM into a [`crate::pcm::PcmBuffer`], then use these policies to decide
5//! when to call [`crate::LocalWhisperProvider::transcribe_pcm`] on a slice.
6
7use crate::audio::WHISPER_SAMPLE_RATE;
8
9/// Policy aligned with common hold-to-talk UX.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub struct PartialWindowPolicy {
12    /// Minimum samples before a partial is worth running (~1 s default).
13    pub min_partial_samples: usize,
14    /// Max samples included in a partial slice (~15 s default).
15    pub window_samples: usize,
16    /// Suggested spacing between partial attempts (nanoseconds).
17    pub interval_nanos: u64,
18    /// RMS below this → skip partial (near-silence).
19    pub min_rms_bits: u32, // store as fixed-point * 1e6 for Eq; use helpers
20}
21
22impl PartialWindowPolicy {
23    /// Defaults suited to progressive hold-to-talk decode.
24    pub fn dictation() -> Self {
25        Self {
26            min_partial_samples: WHISPER_SAMPLE_RATE as usize, // 1 s
27            window_samples: WHISPER_SAMPLE_RATE as usize * 15, // 15 s
28            interval_nanos: 1_200_000_000,                     // 1.2 s
29            min_rms_bits: 500,                                 // 0.0005
30        }
31    }
32
33    pub fn min_rms(&self) -> f32 {
34        self.min_rms_bits as f32 / 1_000_000.0
35    }
36
37    pub fn with_min_rms(mut self, rms: f32) -> Self {
38        self.min_rms_bits = (rms.clamp(0.0, 1.0) * 1_000_000.0).round() as u32;
39        self
40    }
41
42    pub fn min_partial_secs(&self) -> f64 {
43        self.min_partial_samples as f64 / f64::from(WHISPER_SAMPLE_RATE)
44    }
45
46    pub fn window_secs(&self) -> f64 {
47        self.window_samples as f64 / f64::from(WHISPER_SAMPLE_RATE)
48    }
49
50    /// Enough audio accumulated for a partial attempt?
51    pub fn can_run_partial(&self, sample_count: usize) -> bool {
52        sample_count >= self.min_partial_samples
53    }
54
55    /// Most recent `window_samples` (or all if shorter).
56    pub fn slice_for_partial<'a>(&self, samples: &'a [f32]) -> &'a [f32] {
57        if samples.len() > self.window_samples {
58            &samples[samples.len() - self.window_samples..]
59        } else {
60            samples
61        }
62    }
63
64    /// Energy gate: skip near-silent buffers.
65    pub fn passes_energy_gate(&self, samples: &[f32]) -> bool {
66        if samples.is_empty() {
67            return false;
68        }
69        let sum: f32 = samples.iter().map(|s| s * s).sum();
70        let rms = (sum / samples.len() as f32).sqrt();
71        rms >= self.min_rms()
72    }
73
74    /// Combined: enough samples + energy.
75    pub fn should_run_partial(&self, samples: &[f32]) -> bool {
76        self.can_run_partial(samples.len()) && self.passes_energy_gate(samples)
77    }
78}
79
80/// Tracks wall-clock spacing between partial attempts.
81#[derive(Debug, Clone)]
82pub struct PartialClock {
83    policy: PartialWindowPolicy,
84    last_partial_at: Option<std::time::Instant>,
85}
86
87impl PartialClock {
88    pub fn new(policy: PartialWindowPolicy) -> Self {
89        Self {
90            policy,
91            last_partial_at: None,
92        }
93    }
94
95    pub fn policy(&self) -> &PartialWindowPolicy {
96        &self.policy
97    }
98
99    /// True if enough time has elapsed since the last partial (or never ran).
100    pub fn interval_elapsed(&self) -> bool {
101        match self.last_partial_at {
102            None => true,
103            Some(t) => t.elapsed().as_nanos() as u64 >= self.policy.interval_nanos,
104        }
105    }
106
107    /// Mark that a partial was started/completed now.
108    pub fn mark(&mut self) {
109        self.last_partial_at = Some(std::time::Instant::now());
110    }
111
112    /// Ready for another partial given full buffer samples.
113    /// Energy is evaluated on the same rolling window that would be decoded.
114    pub fn ready(&self, samples: &[f32]) -> bool {
115        if !self.interval_elapsed() || !self.policy.can_run_partial(samples.len()) {
116            return false;
117        }
118        let window = self.policy.slice_for_partial(samples);
119        self.policy.passes_energy_gate(window)
120    }
121
122    /// Slice + mark if ready; returns `None` if not ready.
123    pub fn take_partial_slice<'a>(&mut self, samples: &'a [f32]) -> Option<&'a [f32]> {
124        if !self.ready(samples) {
125            return None;
126        }
127        self.mark();
128        Some(self.policy.slice_for_partial(samples))
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn slice_takes_tail() {
138        let p = PartialWindowPolicy {
139            window_samples: 4,
140            min_partial_samples: 2,
141            interval_nanos: 0,
142            min_rms_bits: 0,
143        };
144        let s = [1., 2., 3., 4., 5., 6.];
145        assert_eq!(p.slice_for_partial(&s), &[3., 4., 5., 6.]);
146    }
147
148    #[test]
149    fn energy_gate() {
150        let p = PartialWindowPolicy::dictation().with_min_rms(0.01);
151        assert!(!p.passes_energy_gate(&[0.0; 1600]));
152        assert!(p.passes_energy_gate(&[0.5; 1600]));
153    }
154
155    #[test]
156    fn clock_interval() {
157        let mut c = PartialClock::new(PartialWindowPolicy {
158            min_partial_samples: 1,
159            window_samples: 100,
160            interval_nanos: 10_000_000_000, // 10s
161            min_rms_bits: 0,
162        });
163        let s = [0.1; 10];
164        assert!(c.ready(&s));
165        c.mark();
166        assert!(!c.ready(&s));
167    }
168}