use crate::result::{AsrResult, InferenceMetrics, TokenPiece, WindowInference};
use crate::{Cancellation, Error, HaasEngine, InputLimits, Result, recording::SAMPLE_RATE};
use serde::Serialize;
use std::collections::VecDeque;
use std::{fmt, ops::Range};
use uuid::Uuid;
pub const ENCODER_STEP: usize = 1280;
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct TimedToken {
pub token_id: usize,
pub start_sample: u64,
pub confidence: f32,
pub duration_samples: u64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WindowConfig {
pub center_samples: usize,
pub left_samples: usize,
pub right_samples: usize,
}
impl WindowConfig {
fn legacy() -> Self {
Self {
center_samples: 163840,
left_samples: 40960,
right_samples: 40960,
}
}
pub fn interactive_for(limits: InputLimits) -> Result<Self> {
let legacy = Self::legacy();
if legacy.window_samples()? <= limits.maximum_samples() {
legacy.validate_for(limits)?;
return Ok(legacy);
}
let context = 2 * SAMPLE_RATE as usize;
let center_samples = limits.maximum_samples().checked_sub(2 * context)
.ok_or_else(|| Error::Invalid(format!("Model window too short for interactive context: limits={limits:?}, context={context}")))?
/ ENCODER_STEP * ENCODER_STEP;
let config = Self {
center_samples,
left_samples: context,
right_samples: context,
};
config.validate_for(limits)?;
Ok(config)
}
pub fn window_samples(&self) -> Result<usize> {
self.center_samples
.checked_add(self.left_samples)
.and_then(|samples| samples.checked_add(self.right_samples))
.ok_or_else(|| Error::Invalid(format!("ASR window length overflow: {self:?}")))
}
pub fn validate(&self) -> Result<()> {
self.window_samples()?;
let buffer_fits = self
.center_samples
.checked_mul(2)
.and_then(|samples| samples.checked_add(self.left_samples))
.and_then(|samples| samples.checked_add(self.right_samples))
.is_some();
if self.center_samples == 0
|| self.left_samples == 0
|| !buffer_fits
|| [self.center_samples, self.left_samples, self.right_samples]
.iter()
.any(|samples| !samples.is_multiple_of(ENCODER_STEP))
{
return Err(Error::Invalid(format!(
"Invalid ASR window structure or frame alignment: {self:?}"
)));
}
Ok(())
}
pub fn validate_for(&self, limits: InputLimits) -> Result<()> {
self.validate()?;
let span = self.window_samples()?;
if span > limits.maximum_samples()
|| self.center_samples < limits.minimum_samples()
|| self.left_samples < limits.minimum_samples()
{
return Err(Error::Invalid(format!(
"ASR windows exceed model limits: config={self:?}, span={span}, limits={limits:?}"
)));
}
Ok(())
}
pub fn maximum_buffered_samples(&self) -> usize {
2 * self.center_samples + self.left_samples + self.right_samples
}
}
#[derive(Clone, Debug)]
pub struct AsrConfig {
pub center_samples: usize,
pub left_samples: usize,
pub right_samples: usize,
pub partial_interval_samples: usize,
pub vad_min_silence_samples: usize,
}
impl AsrConfig {
pub fn interactive() -> Self {
Self::with_windows(WindowConfig::legacy())
}
fn with_windows(windows: WindowConfig) -> Self {
Self {
center_samples: windows.center_samples,
left_samples: windows.left_samples,
right_samples: windows.right_samples,
partial_interval_samples: 20480,
vad_min_silence_samples: 32768,
}
}
pub fn interactive_for(limits: InputLimits) -> Result<Self> {
let config = Self::with_windows(WindowConfig::interactive_for(limits)?);
config.validate_for(limits)?;
Ok(config)
}
pub fn window_config(&self) -> WindowConfig {
WindowConfig {
center_samples: self.center_samples,
left_samples: self.left_samples,
right_samples: self.right_samples,
}
}
pub fn window_samples(&self) -> Result<usize> {
self.window_config().window_samples()
}
pub fn validate(&self) -> Result<()> {
self.window_config().validate()?;
if self.vad_min_silence_samples < 8192
|| !self.vad_min_silence_samples.is_multiple_of(4096)
|| self.partial_interval_samples == 0
|| !self.partial_interval_samples.is_multiple_of(ENCODER_STEP)
{
return Err(Error::Invalid(format!(
"Invalid ASR endpoint or partial interval: {self:?}"
)));
}
Ok(())
}
pub fn validate_for(&self, limits: InputLimits) -> Result<()> {
self.validate()?;
self.window_config().validate_for(limits)?;
if self.vad_min_silence_samples > limits.maximum_samples()
|| self.partial_interval_samples < limits.minimum_samples()
{
return Err(Error::Invalid(format!(
"ASR endpoint or partial interval exceeds model limits: config={self:?}, limits={limits:?}"
)));
}
Ok(())
}
pub fn maximum_buffered_samples(&self) -> usize {
self.window_config().maximum_buffered_samples()
}
}
#[derive(Clone, Debug, Serialize)]
pub struct Transcript {
pub text: String,
pub tokens: Vec<TimedToken>,
pub audio_samples: u64,
}
impl Transcript {
pub fn audio_seconds(&self) -> f64 {
self.audio_samples as f64 / SAMPLE_RATE as f64
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum UpdateKind {
Confirmed,
Partial,
Final,
}
#[derive(Clone, Debug, Serialize)]
pub struct TokenRevision {
pub previous_token_count: usize,
pub retained_token_count: usize,
pub stable_token_count: usize,
pub appended_tokens: Vec<TimedToken>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize)]
#[serde(transparent)]
pub struct JobId(Uuid);
impl JobId {
pub(crate) fn new() -> Self {
Self(Uuid::new_v4())
}
}
impl fmt::Display for JobId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(formatter)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct RecognitionSpan {
pub window: Range<u64>,
pub center: Range<u64>,
}
#[derive(Clone, Debug, Serialize)]
pub struct JobUpdate {
pub job_id: JobId,
pub update: AsrUpdate,
}
#[derive(Clone, Debug, Serialize)]
pub struct AsrUpdate {
pub session_generation: u64,
pub revision: TokenRevision,
pub sequence: u64,
pub kind: UpdateKind,
pub confirmed_text: String,
pub provisional_text: String,
pub confirmed_samples: u64,
pub received_samples: u64,
pub recognition_span: Option<RecognitionSpan>,
}
pub(crate) trait WindowRecognizer {
fn input_limits(&self) -> InputLimits;
fn infer(&self, audio: &[f32], cancellation: &Cancellation) -> Result<Vec<TimedToken>>;
fn render(&self, tokens: &[TimedToken]) -> Result<String>;
fn stable_prefix(&self, _: &[TimedToken], proposed: usize) -> usize {
proposed
}
fn token_piece(&self, id: usize) -> Result<TokenPiece> {
Err(Error::Invalid(format!(
"Recognizer does not expose vocabulary pieces: id={id}"
)))
}
fn infer_with_metrics(
&self,
audio: &[f32],
original_samples: usize,
cancellation: &Cancellation,
) -> Result<WindowInference> {
let started = std::time::Instant::now();
let tokens = self.infer(audio, cancellation)?;
let mut metrics =
InferenceMetrics::measured(audio.len(), started.elapsed().as_secs_f64(), None);
metrics.audio_samples = original_samples as u64;
Ok(WindowInference { tokens, metrics })
}
}
#[derive(Clone)]
pub(crate) enum RecognizerRef<'a> {
Borrowed(&'a dyn WindowRecognizer),
Shared(std::rc::Rc<dyn WindowRecognizer + 'a>),
}
impl<'a> RecognizerRef<'a> {
fn get(&self) -> &(dyn WindowRecognizer + 'a) {
match self {
Self::Borrowed(value) => *value,
Self::Shared(value) => &**value,
}
}
}
impl<'a> std::ops::Deref for RecognizerRef<'a> {
type Target = dyn WindowRecognizer + 'a;
fn deref(&self) -> &Self::Target {
self.get()
}
}
impl WindowRecognizer for HaasEngine {
fn input_limits(&self) -> InputLimits {
HaasEngine::input_limits(self)
}
fn infer(&self, audio: &[f32], cancellation: &Cancellation) -> Result<Vec<TimedToken>> {
Ok(self
.transcribe_samples(audio, SAMPLE_RATE, cancellation)?
.token_timings)
}
fn token_piece(&self, id: usize) -> Result<TokenPiece> {
HaasEngine::token_piece(self, id)
}
fn infer_with_metrics(
&self,
audio: &[f32],
original_samples: usize,
cancellation: &Cancellation,
) -> Result<WindowInference> {
let started = std::time::Instant::now();
let result = self.transcribe_samples(audio, SAMPLE_RATE, cancellation)?;
Ok(WindowInference::from_recognition(
result,
audio.len(),
original_samples,
started.elapsed().as_secs_f64(),
))
}
fn render(&self, tokens: &[TimedToken]) -> Result<String> {
self.assets
.tokenizer
.decode(&tokens.iter().map(|t| t.token_id).collect::<Vec<_>>())
}
}
pub(crate) fn recognize_window(
recognizer: &dyn WindowRecognizer,
audio: &[f32],
cancellation: &Cancellation,
) -> Result<Vec<TimedToken>> {
Ok(recognize_window_with_metrics(recognizer, audio, cancellation)?.tokens)
}
pub(crate) fn recognize_window_with_metrics(
recognizer: &dyn WindowRecognizer,
audio: &[f32],
cancellation: &Cancellation,
) -> Result<WindowInference> {
cancellation.check()?;
if audio.is_empty() {
return Ok(WindowInference {
tokens: Vec::new(),
metrics: InferenceMetrics::empty(),
});
}
let limits = recognizer.input_limits();
if audio.len() > limits.maximum_samples() {
return Err(Error::Invalid(format!(
"ASR window exceeds model input: samples={}, limits={limits:?}",
audio.len()
)));
}
let padded;
let model_audio = if audio.len() < limits.minimum_samples() {
padded = {
let mut samples = audio.to_vec();
samples.resize(limits.minimum_samples(), 0.0);
samples
};
padded.as_slice()
} else {
audio
};
let result = recognizer.infer_with_metrics(model_audio, audio.len(), cancellation)?;
let mut tokens = result.tokens;
let mut metrics = result.metrics;
metrics.audio_samples = audio.len() as u64;
if tokens
.windows(2)
.any(|p| p[0].start_sample > p[1].start_sample)
|| tokens.iter().any(|t| {
t.start_sample >= model_audio.len() as u64
|| !t.confidence.is_finite()
|| !(0.0..=1.0).contains(&t.confidence)
|| t.duration_samples > (4 * ENCODER_STEP) as u64
})
{
return Err(Error::Invalid(
"Window recognizer returned invalid timing, confidence, or duration".into(),
));
}
tokens.retain(|t| t.start_sample < audio.len() as u64);
Ok(WindowInference { tokens, metrics })
}
pub(crate) struct WindowProcessor<'a> {
recognizer: RecognizerRef<'a>,
config: WindowConfig,
audio: VecDeque<f32>,
buffer_start: u64,
received: u64,
center_start: u64,
hypothesis: Vec<TimedToken>,
stable_count: usize,
stable_samples: u64,
published: Vec<TimedToken>,
published_stable: usize,
last_partial: u64,
sequence: u64,
peak_buffer: usize,
last_span: Option<RecognitionSpan>,
metrics: InferenceMetrics,
}
impl<'a> WindowProcessor<'a> {
#[cfg(test)]
pub(crate) fn new(recognizer: &'a dyn WindowRecognizer, config: WindowConfig) -> Result<Self> {
Self::from_recognizer(RecognizerRef::Borrowed(recognizer), config)
}
pub(crate) fn from_recognizer(
recognizer: RecognizerRef<'a>,
config: WindowConfig,
) -> Result<Self> {
config.validate_for(recognizer.input_limits())?;
Ok(Self {
recognizer,
config,
audio: VecDeque::new(),
buffer_start: 0,
received: 0,
center_start: 0,
hypothesis: Vec::new(),
stable_count: 0,
stable_samples: 0,
published: Vec::new(),
published_stable: 0,
last_partial: 0,
sequence: 0,
peak_buffer: 0,
last_span: None,
metrics: InferenceMetrics::empty(),
})
}
pub(crate) fn received(&self) -> u64 {
self.received
}
pub(crate) fn peak_buffer(&self) -> usize {
self.peak_buffer
}
pub(crate) fn available_capacity(&self) -> usize {
self.config.maximum_buffered_samples() - self.audio.len()
}
pub(crate) fn append(&mut self, audio: &[f32]) -> Result<()> {
if audio.len() > self.available_capacity() || audio.iter().any(|x| !x.is_finite()) {
return Err(Error::Invalid(format!(
"ASR input rejected: incoming={}, buffered={}, capacity={}, non_finite={}",
audio.len(),
self.audio.len(),
self.config.maximum_buffered_samples(),
audio.iter().any(|x| !x.is_finite())
)));
}
self.received = self
.received
.checked_add(audio.len() as u64)
.ok_or_else(|| Error::Invalid("ASR audio clock overflow".into()))?;
self.audio.extend(audio);
self.peak_buffer = self.peak_buffer.max(self.audio.len());
Ok(())
}
pub(crate) fn result(&self, transcript: &Transcript) -> Result<AsrResult> {
AsrResult::build(self.recognizer.get(), transcript, self.metrics.clone())
}
pub(crate) fn feed(
&mut self,
audio: &[f32],
cancellation: &Cancellation,
) -> Result<Vec<AsrUpdate>> {
let mut updates = Vec::new();
self.advance(audio, cancellation, |processor| {
let tokens = processor.hypothesis.clone();
updates.push(processor.update(UpdateKind::Confirmed, &tokens)?);
Ok(())
})?;
Ok(updates)
}
fn advance(
&mut self,
audio: &[f32],
cancellation: &Cancellation,
mut on_commit: impl FnMut(&mut Self) -> Result<()>,
) -> Result<()> {
if audio.iter().any(|v| !v.is_finite()) {
return Err(Error::Invalid("Non-finite ASR input".into()));
}
let mut offset = 0;
let ahead = (self.config.center_samples + self.config.right_samples) as u64;
while offset < audio.len() {
cancellation.check()?;
let count = self.available_capacity().min(audio.len() - offset);
if count == 0 {
return Err(Error::Invalid("ASR window buffer cannot advance".into()));
}
self.append(&audio[offset..offset + count])?;
offset += count;
while self.received - self.center_start >= ahead {
cancellation.check()?;
self.commit(self.center_start + ahead, cancellation)?;
on_commit(self)?;
}
}
Ok(())
}
fn window_tokens(
&mut self,
end: u64,
cancellation: &Cancellation,
) -> Result<(Vec<TimedToken>, u64)> {
let regular = self
.center_start
.saturating_sub(self.config.left_samples as u64);
let end_aligned =
end.saturating_sub((self.config.center_samples + self.config.left_samples) as u64);
let start = regular.min(end_aligned).max(self.buffer_start);
let first = usize::try_from(start - self.buffer_start)
.map_err(|e| Error::Invalid(e.to_string()))?;
let last =
usize::try_from(end - self.buffer_start).map_err(|e| Error::Invalid(e.to_string()))?;
if last > self.audio.len() || first > last {
return Err(Error::Invalid("ASR window outside buffered input".into()));
}
let audio: Vec<f32> = self.audio.range(first..last).copied().collect();
let inference = recognize_window_with_metrics(self.recognizer.get(), &audio, cancellation)?;
self.metrics = self.metrics.add(&inference.metrics)?;
let mut tokens = inference.tokens;
self.last_span = Some(RecognitionSpan {
window: start..end,
center: self.center_start
..(self.center_start + self.config.center_samples as u64).min(end),
});
for token in &mut tokens {
token.start_sample += start;
}
Ok((tokens, start))
}
fn revise(&self, tokens: &[TimedToken], window_start: u64, end: u64) -> Vec<TimedToken> {
let retained = self
.hypothesis
.partition_point(|t| t.start_sample < window_start)
.max(self.stable_count);
let mut revised = self.hypothesis[..retained].to_vec();
let tail = stitch_center(
&revised,
tokens,
window_start.max(self.stable_samples),
end,
window_start,
);
let mut previous = revised.last().map(|t| t.start_sample).unwrap_or(0);
for mut token in tail {
token.start_sample = token.start_sample.max(previous);
previous = token.start_sample;
revised.push(token);
}
revised
}
fn stabilize_through(&mut self, through: u64) -> Result<()> {
let full = self.recognizer.render(&self.hypothesis)?;
let mut count = self
.hypothesis
.partition_point(|t| t.start_sample < through)
.max(self.stable_count);
while count > self.stable_count {
let permitted = self
.recognizer
.stable_prefix(&self.hypothesis, count)
.max(self.stable_count);
if permitted < count {
count = permitted;
continue;
}
if full.starts_with(&self.recognizer.render(&self.hypothesis[..count])?) {
break;
}
count -= 1;
}
self.stable_count = count;
let safe = self
.hypothesis
.get(count)
.map(|t| t.start_sample.min(through))
.unwrap_or(through);
self.stable_samples = self.stable_samples.max(safe);
Ok(())
}
fn commit(&mut self, end: u64, cancellation: &Cancellation) -> Result<()> {
let center_end = (self.center_start + self.config.center_samples as u64).min(end);
let (tokens, start) = self.window_tokens(end, cancellation)?;
self.hypothesis = self.revise(&tokens, start, center_end);
self.center_start = center_end;
self.stabilize_through(center_end.saturating_sub(self.config.left_samples as u64))?;
self.last_partial = 0;
let keep = self
.center_start
.saturating_sub((self.config.center_samples + self.config.left_samples) as u64)
.max(self.buffer_start);
self.audio.drain(..(keep - self.buffer_start) as usize);
self.buffer_start = keep;
Ok(())
}
pub(crate) fn partial(
&mut self,
interval_samples: usize,
cancellation: &Cancellation,
) -> Result<Option<AsrUpdate>> {
if self.audio.len() < self.recognizer.input_limits().minimum_samples()
|| self.received == self.center_start
|| self.received - self.last_partial < interval_samples as u64
|| self.received < self.recognizer.input_limits().minimum_samples() as u64
{
return Ok(None);
}
let (tokens, start) = self.window_tokens(self.received, cancellation)?;
let preview = self.revise(&tokens, start, self.received);
self.last_partial = self.received;
Ok(Some(self.update(UpdateKind::Partial, &preview)?))
}
fn update(&mut self, kind: UpdateKind, tokens: &[TimedToken]) -> Result<AsrUpdate> {
let retained = self
.published
.iter()
.zip(tokens)
.take_while(|(a, b)| a == b)
.count();
if retained < self.published_stable
|| self.stable_count < self.published_stable
|| self.stable_count > tokens.len()
{
return Err(Error::Invalid(format!(
"Invalid streaming revision: retained={retained}, previous_stable={}, stable={}, tokens={}",
self.published_stable,
self.stable_count,
tokens.len()
)));
}
let confirmed_text = self.recognizer.render(&tokens[..self.stable_count])?;
let full = self.recognizer.render(tokens)?;
let provisional_text = full
.strip_prefix(&confirmed_text)
.ok_or_else(|| Error::Invalid("Tokenizer revised the stable transcript prefix".into()))?
.to_owned();
self.sequence = self
.sequence
.checked_add(1)
.ok_or_else(|| Error::Invalid("Streaming sequence overflow".into()))?;
let revision = TokenRevision {
previous_token_count: self.published.len(),
retained_token_count: retained,
stable_token_count: self.stable_count,
appended_tokens: tokens[retained..].to_vec(),
};
self.published = tokens.to_vec();
self.published_stable = self.stable_count;
Ok(AsrUpdate {
session_generation: 0,
revision,
sequence: self.sequence,
kind,
confirmed_text,
provisional_text,
confirmed_samples: self.stable_samples,
received_samples: self.received,
recognition_span: self.last_span.clone(),
})
}
pub(crate) fn finish_with_audio(
&mut self,
audio: &[f32],
cancellation: &Cancellation,
) -> Result<(Transcript, AsrUpdate)> {
self.advance(audio, cancellation, |_| Ok(()))?;
self.finish(cancellation)
}
pub(crate) fn finish(
&mut self,
cancellation: &Cancellation,
) -> Result<(Transcript, AsrUpdate)> {
self.flush_windows(cancellation)?;
self.stabilize_through(self.received)?;
let transcript = Transcript {
text: self.recognizer.render(&self.hypothesis)?,
tokens: self.hypothesis.clone(),
audio_samples: self.received,
};
let update = self.update(UpdateKind::Final, &transcript.tokens)?;
self.audio.clear();
Ok((transcript, update))
}
pub(crate) fn advance_quiet(&mut self, audio: &[f32]) -> Result<()> {
self.append(audio)?;
let discard = self.audio.len().saturating_sub(4096);
self.audio.drain(..discard);
self.buffer_start += discard as u64;
self.center_start = self.buffer_start;
Ok(())
}
pub(crate) fn finish_quiet(&mut self) {
self.audio.clear();
self.buffer_start = self.received;
self.center_start = self.received;
}
fn flush_windows(&mut self, cancellation: &Cancellation) -> Result<()> {
cancellation.check()?;
while self.center_start < self.received {
cancellation.check()?;
let end = (self.center_start
+ (self.config.center_samples + self.config.right_samples) as u64)
.min(self.received);
self.commit(end, cancellation)?;
}
Ok(())
}
pub(crate) fn confirm_boundary(&mut self, cancellation: &Cancellation) -> Result<AsrUpdate> {
self.flush_windows(cancellation)?;
let tokens = self.hypothesis.clone();
self.update(UpdateKind::Confirmed, &tokens)
}
pub(crate) fn utterance_boundary(&mut self, cancellation: &Cancellation) -> Result<AsrUpdate> {
self.flush_windows(cancellation)?;
self.stabilize_through(self.received)?;
let tokens = self.hypothesis.clone();
let update = self.update(UpdateKind::Confirmed, &tokens)?;
self.audio.clear();
self.buffer_start = self.received;
self.last_partial = 0;
Ok(update)
}
pub(crate) fn clear(&mut self) {
self.audio.clear();
self.hypothesis.clear();
self.published.clear();
self.stable_count = 0;
self.published_stable = 0;
self.last_span = None;
}
}
pub(crate) fn align_tokens(
a: &[TimedToken],
b: &[TimedToken],
tolerance: u64,
match_score: impl Fn(u64) -> (u64, i64),
) -> Vec<(usize, usize)> {
let mut previous = vec![(0u64, 0i64); b.len() + 1];
let mut steps = vec![vec![0u8; b.len() + 1]; a.len() + 1];
for i in 1..=a.len() {
let mut current = vec![(0u64, 0i64); b.len() + 1];
for j in 1..=b.len() {
let mut best = previous[j];
let mut step = 0;
if current[j - 1] > best {
best = current[j - 1];
step = 1;
}
let distance = a[i - 1].start_sample.abs_diff(b[j - 1].start_sample);
if a[i - 1].token_id == b[j - 1].token_id && distance <= tolerance {
let prev = previous[j - 1];
let score = match_score(distance);
let candidate = (prev.0 + score.0, prev.1 + score.1);
if candidate >= best {
best = candidate;
step = 2;
}
}
current[j] = best;
steps[i][j] = step;
}
previous = current;
}
let mut matches = Vec::new();
let (mut i, mut j) = (a.len(), b.len());
while i > 0 && j > 0 {
match steps[i][j] {
2 => {
matches.push((i - 1, j - 1));
i -= 1;
j -= 1;
}
1 => j -= 1,
_ => i -= 1,
}
}
matches.reverse();
matches
}
fn stitch_center(
old: &[TimedToken],
new: &[TimedToken],
start: u64,
end: u64,
window_start: u64,
) -> Vec<TimedToken> {
let first = old
.len()
.saturating_sub(64)
.max(old.partition_point(|t| t.start_sample < window_start));
let tail = &old[first..];
let prefix = new.partition_point(|t| t.start_sample <= start + SAMPLE_RATE as u64);
let matches = align_tokens(tail, &new[..prefix], SAMPLE_RATE as u64, |distance| {
(1, -(distance as i64))
});
let anchor = matches
.last()
.filter(|(i, _)| *i + 1 == tail.len())
.map(|(_, j)| j + 1);
let cut = if matches.len() >= 2 {
anchor.unwrap_or_else(|| new.partition_point(|t| t.start_sample < start))
} else {
new.partition_point(|t| t.start_sample < start)
};
new[cut..]
.iter()
.filter(|t| t.start_sample < end)
.cloned()
.collect()
}
#[cfg(test)]
pub(crate) mod test_support {
use super::*;
pub(crate) const FULL_WINDOW: InputLimits = InputLimits {
minimum_samples: 10240,
maximum_samples: 491520,
};
pub(crate) const SHORT_WINDOW: InputLimits = InputLimits {
minimum_samples: 10240,
maximum_samples: 240000,
};
use std::cell::Cell;
pub(crate) struct Markers {
pub(crate) calls: Cell<usize>,
pub(crate) largest_window: Cell<usize>,
limits: InputLimits,
}
impl Markers {
pub(crate) fn new() -> Self {
Self::with_limits(FULL_WINDOW)
}
pub(crate) fn with_limits(limits: InputLimits) -> Self {
Self {
calls: Cell::new(0),
largest_window: Cell::new(0),
limits,
}
}
}
impl WindowRecognizer for Markers {
fn input_limits(&self) -> InputLimits {
self.limits
}
fn infer(&self, audio: &[f32], c: &Cancellation) -> Result<Vec<TimedToken>> {
c.check()?;
self.limits.validate_samples(audio.len())?;
self.calls.set(self.calls.get() + 1);
self.largest_window
.set(self.largest_window.get().max(audio.len()));
Ok(audio
.iter()
.enumerate()
.filter(|(_, v)| **v >= 1.0)
.map(|(i, v)| TimedToken {
token_id: *v as usize,
start_sample: i as u64,
confidence: 0.95,
duration_samples: ENCODER_STEP as u64,
})
.collect())
}
fn token_piece(&self, id: usize) -> Result<TokenPiece> {
Ok(TokenPiece {
text: id.to_string(),
kind: crate::result::TokenKind::Normal,
})
}
fn render(&self, t: &[TimedToken]) -> Result<String> {
Ok(t.iter()
.map(|t| t.token_id.to_string())
.collect::<Vec<_>>()
.join(" "))
}
}
}
#[cfg(test)]
mod tests {
use super::test_support::{FULL_WINDOW, Markers, SHORT_WINDOW};
use super::*;
#[test]
fn window_geometry_is_independent_of_vad_and_preview_settings() {
for limits in [FULL_WINDOW, SHORT_WINDOW] {
let windows = WindowConfig::interactive_for(limits).unwrap();
let mut speech = AsrConfig::interactive_for(limits).unwrap();
assert_eq!(speech.window_config(), windows);
speech.vad_min_silence_samples = 0;
speech.partial_interval_samples = 0;
assert!(speech.validate_for(limits).is_err());
assert!(speech.window_config().validate_for(limits).is_ok());
let mut invalid = windows;
invalid.center_samples += 1;
assert!(invalid.validate_for(limits).is_err());
}
}
struct Revisable;
impl WindowRecognizer for Revisable {
fn input_limits(&self) -> InputLimits {
FULL_WINDOW
}
fn infer(&self, audio: &[f32], _: &Cancellation) -> Result<Vec<TimedToken>> {
let start = audio[0] as u64;
Ok([(1, 140000), (2, 150000), (3, 160000), (4, 360000)]
.into_iter()
.filter(|(id, time)| {
*time >= start
&& *time < start + audio.len() as u64
&& !(start == 0 && [2, 3].contains(id))
&& !(*id == 4 && audio.len() < 160000)
})
.map(|(token_id, time)| TimedToken {
token_id,
start_sample: time - start,
confidence: 0.95,
duration_samples: ENCODER_STEP as u64,
})
.collect())
}
fn render(&self, tokens: &[TimedToken]) -> Result<String> {
Ok(tokens
.iter()
.map(|t| t.token_id.to_string())
.collect::<Vec<_>>()
.join(" "))
}
}
#[test]
fn streaming_revises_the_editable_tail_and_backfills_final_audio() {
let recognizer = Revisable;
let mut processor = WindowProcessor::new(&recognizer, WindowConfig::legacy()).unwrap();
let audio: Vec<_> = (0..400000).map(|i| i as f32).collect();
let cancellation = Cancellation::new();
for part in audio.chunks(16000) {
processor.feed(part, &cancellation).unwrap();
}
let (result, _) = processor.finish(&cancellation).unwrap();
assert_eq!(result.text, "1 2 3 4");
}
#[test]
fn published_revisions_replay_exactly_and_never_retire_stable_tokens() {
let recognizer = Revisable;
let config = AsrConfig::interactive();
let mut processor = WindowProcessor::new(&recognizer, config.window_config()).unwrap();
let audio: Vec<_> = (0..400000).map(|i| i as f32).collect();
let cancellation = Cancellation::new();
let mut updates = Vec::new();
for part in audio.chunks(13007) {
updates.extend(processor.feed(part, &cancellation).unwrap());
if let Some(update) = processor
.partial(config.partial_interval_samples, &cancellation)
.unwrap()
{
updates.push(update);
}
}
let (transcript, final_update) = processor.finish(&cancellation).unwrap();
updates.push(final_update);
let mut held: Vec<TimedToken> = Vec::new();
let mut stable = 0;
let mut confirmed = String::new();
let mut confirmed_samples = 0;
let mut retirements = 0;
for (index, update) in updates.iter().enumerate() {
let edit = &update.revision;
assert_eq!(update.sequence as usize, index + 1);
assert_eq!(edit.previous_token_count, held.len());
assert!(edit.retained_token_count >= stable && edit.retained_token_count <= held.len());
retirements += usize::from(edit.retained_token_count < held.len());
held.truncate(edit.retained_token_count);
held.extend_from_slice(&edit.appended_tokens);
assert!(edit.stable_token_count >= stable && edit.stable_token_count <= held.len());
stable = edit.stable_token_count;
assert!(update.confirmed_text.starts_with(&confirmed));
confirmed = update.confirmed_text.clone();
assert!(
update.confirmed_samples >= confirmed_samples
&& update.confirmed_samples <= update.received_samples
);
confirmed_samples = update.confirmed_samples;
assert_eq!(
recognizer.render(&held).unwrap(),
format!("{}{}", update.confirmed_text, update.provisional_text)
);
}
assert!(retirements > 0);
assert_eq!(held, transcript.tokens);
assert_eq!(stable, held.len());
assert_eq!(transcript.text, "1 2 3 4");
assert!(processor.peak_buffer() <= config.maximum_buffered_samples());
}
struct Vanishing;
impl WindowRecognizer for Vanishing {
fn input_limits(&self) -> InputLimits {
FULL_WINDOW
}
fn infer(&self, audio: &[f32], cancellation: &Cancellation) -> Result<Vec<TimedToken>> {
cancellation.check()?;
Ok(if audio.len() < 48000 {
vec![TimedToken {
token_id: 9,
start_sample: 12800,
confidence: 0.9,
duration_samples: ENCODER_STEP as u64,
}]
} else {
Vec::new()
})
}
fn render(&self, tokens: &[TimedToken]) -> Result<String> {
Ok(if tokens.is_empty() {
String::new()
} else {
"9".into()
})
}
}
#[test]
fn an_empty_partial_revision_can_withdraw_the_entire_provisional_tail() {
let recognizer = Vanishing;
let mut processor = WindowProcessor::new(&recognizer, WindowConfig::legacy()).unwrap();
let cancellation = Cancellation::new();
processor.feed(&vec![0.0001; 32000], &cancellation).unwrap();
let first = processor.partial(20480, &cancellation).unwrap().unwrap();
assert_eq!(first.revision.appended_tokens.len(), 1);
processor.feed(&vec![0.0001; 32000], &cancellation).unwrap();
let second = processor.partial(20480, &cancellation).unwrap().unwrap();
assert_eq!(second.revision.previous_token_count, 1);
assert_eq!(second.revision.retained_token_count, 0);
assert_eq!(second.revision.stable_token_count, 0);
assert!(second.revision.appended_tokens.is_empty());
assert!(second.provisional_text.is_empty());
assert!(processor.finish(&cancellation).unwrap().0.tokens.is_empty());
}
#[test]
fn an_empty_final_revision_still_terminates_the_lineage() {
let recognizer = Markers::new();
let mut processor = WindowProcessor::new(&recognizer, WindowConfig::legacy()).unwrap();
let (_, update) = processor.finish(&Cancellation::new()).unwrap();
assert_eq!(update.kind, UpdateKind::Final);
assert_eq!(update.revision.previous_token_count, 0);
assert!(update.revision.appended_tokens.is_empty());
assert_eq!(update.revision.stable_token_count, 0);
assert!(update.recognition_span.is_none());
}
#[test]
fn revision_spans_describe_inference_input_not_total_received_audio() {
let recognizer = Markers::with_limits(SHORT_WINDOW);
let config = AsrConfig::interactive_for(SHORT_WINDOW).unwrap();
let mut processor = WindowProcessor::new(&recognizer, config.window_config()).unwrap();
let cancellation = Cancellation::new();
let count = config.center_samples + config.right_samples;
let events = processor.feed(&vec![0.0001; count], &cancellation).unwrap();
assert_eq!(events.len(), 1);
let first = events[0].recognition_span.as_ref().unwrap();
assert_eq!(first.window, 0..count as u64);
assert_eq!(first.center, 0..config.center_samples as u64);
processor.feed(&[0.0001; 4096], &cancellation).unwrap();
let next = processor
.partial(config.partial_interval_samples, &cancellation)
.unwrap()
.unwrap();
let span = next.recognition_span.as_ref().unwrap();
assert!(span.window.start > 0);
assert_eq!(span.window.end, (count + 4096) as u64);
assert_eq!(span.center.start, config.center_samples as u64);
assert_eq!(span.center.end, span.window.end);
assert!(span.window.end - span.window.start <= SHORT_WINDOW.maximum_samples() as u64);
}
#[test]
fn centers_keep_boundary_tokens_and_genuine_repetitions_once() {
let markers = Markers::new();
let config = AsrConfig::interactive();
let center = config.center_samples;
let len = center * 4 + 1;
let mut audio = vec![0.0; len];
let positions = [
0,
center - 1,
center,
center + 1,
2 * center,
3 * center,
4 * center,
];
for p in positions {
audio[p] = 7.0;
}
let mut processor = WindowProcessor::new(&markers, config.window_config()).unwrap();
let cancellation = Cancellation::new();
for part in audio.chunks(16000) {
processor.feed(part, &cancellation).unwrap();
}
let (result, last) = processor.finish(&cancellation).unwrap();
assert_eq!(
result
.tokens
.iter()
.map(|t| t.start_sample)
.collect::<Vec<_>>(),
positions.map(|p| p as u64)
);
assert_eq!(result.tokens.len(), 7);
assert_eq!(result.audio_samples, len as u64);
assert_eq!(last.kind, UpdateKind::Final);
assert!(last.provisional_text.is_empty());
assert!(processor.peak_buffer() <= config.maximum_buffered_samples());
}
#[test]
fn shifted_overlap_does_not_duplicate_a_word_or_drop_a_real_repetition() {
let token = |id, t| TimedToken {
token_id: id,
start_sample: t,
confidence: 0.95,
duration_samples: ENCODER_STEP as u64,
};
let old = vec![token(1, 9000), token(2, 10000)];
let new = vec![
token(1, 9500),
token(2, 11000),
token(2, 12000),
token(3, 13000),
];
let result = stitch_center(&old, &new, 10500, 20000, 8000);
assert_eq!(
result.iter().map(|t| t.token_id).collect::<Vec<_>>(),
[2, 3]
);
}
#[test]
fn model_profiles_preserve_legacy_windows_and_fit_fifteen_seconds() {
let legacy = AsrConfig::interactive();
let full = AsrConfig::interactive_for(FULL_WINDOW).unwrap();
assert_eq!(full.center_samples, legacy.center_samples);
assert_eq!(full.left_samples, legacy.left_samples);
assert_eq!(full.right_samples, legacy.right_samples);
assert_eq!(
full.partial_interval_samples,
legacy.partial_interval_samples
);
assert_eq!(full.vad_min_silence_samples, legacy.vad_min_silence_samples);
let short = AsrConfig::interactive_for(SHORT_WINDOW).unwrap();
assert_eq!(short.center_samples, 175360);
assert_eq!(short.left_samples, 32000);
assert_eq!(short.right_samples, 32000);
assert_eq!(short.window_samples().unwrap(), 239360);
short.validate_for(SHORT_WINDOW).unwrap();
assert!(legacy.validate_for(SHORT_WINDOW).is_err());
let original = legacy.window_samples().unwrap();
assert!(
WindowProcessor::new(&Markers::with_limits(SHORT_WINDOW), legacy.window_config())
.is_err()
);
assert_eq!(legacy.window_samples().unwrap(), original);
}
#[test]
fn model_limits_reject_oversized_windows_before_inference_and_pad_short_tails() {
let marker = Markers::with_limits(SHORT_WINDOW);
let cancellation = Cancellation::new();
assert!(recognize_window(&marker, &vec![0.0001; 240001], &cancellation).is_err());
assert_eq!(marker.calls.get(), 0);
let tokens = recognize_window(&marker, &[7.0, 0.0], &cancellation).unwrap();
assert_eq!(tokens.len(), 1);
assert_eq!(tokens[0].start_sample, 0);
assert_eq!(marker.largest_window.get(), 10240);
assert!(
recognize_window(&marker, &[], &cancellation)
.unwrap()
.is_empty()
);
assert_eq!(marker.calls.get(), 1);
}
#[test]
fn invalid_config_and_backpressure_are_explicit() {
let markers = Markers::new();
let mut config = AsrConfig::interactive();
config.center_samples = FULL_WINDOW.maximum_samples();
assert!(config.validate_for(FULL_WINDOW).is_err());
let mut p = WindowProcessor::new(&markers, WindowConfig::legacy()).unwrap();
assert!(p.append(&vec![0.0; p.available_capacity() + 1]).is_err());
assert_eq!(p.received(), 0);
}
}