use std::collections::VecDeque;
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
use super::*;
use crate::audio::Recording;
use crate::audio::testing::recorded;
use crate::transcription::{
Cadence, ChunkStream, Emitted, Hypothesis, Language, Session, Stream, TimedWord, Transcriber,
Transcription, TranscriptionAudio,
};
struct Scripted(Mutex<VecDeque<Hypothesis>>);
impl Scripted {
fn new(answers: impl IntoIterator<Item = impl Into<String>>) -> Self {
Self(Mutex::new(
answers.into_iter().map(Hypothesis::new).collect(),
))
}
}
impl Transcriber for Scripted {
fn validate_language(&self, _: &Language, _: &CancelToken) -> Result<(), TranscriptionError> {
Ok(())
}
fn transcribe_with(
&self,
_: TranscriptionAudio<'_>,
_: &Language,
_: &CancelToken,
) -> Result<Hypothesis, TranscriptionError> {
Ok(self
.0
.lock()
.expect("script lock")
.pop_front()
.expect("script covers decodes"))
}
}
struct Replace(&'static str, &'static str);
impl Layer for Replace {
fn apply(&self, text: &mut String, _: &CancelToken) -> Result<(), TranscriptionError> {
*text = text.replace(self.0, self.1);
Ok(())
}
}
#[test]
fn layers_apply_in_declaration_order_and_options_preserve_the_raw_path() {
let (_directory, mut recording) = recorded(&[1, 2]);
let transcription = Transcription::complete_samples(Scripted::new(["one"]))
.with(Replace("one", "two"))
.with(Replace("two", "three"));
assert_eq!(
transcription
.transcribe(&mut recording, &Language::Automatic)
.unwrap()
.as_str(),
"three"
);
for enabled in [false, true] {
let (_directory, mut recording) = recorded(&[1, 2]);
let transcription = Transcription::complete_samples(Scripted::new(["Um, hello."]))
.with(enabled.then_some(RemoveHesitations));
assert_eq!(
transcription
.transcribe(&mut recording, &Language::Automatic)
.unwrap()
.as_str(),
if enabled { "Hello." } else { "Um, hello." }
);
}
}
#[test]
fn complete_and_streaming_sessions_apply_real_layers() {
let vocabulary = Vocabulary::new(["Cargo.toml"]).unwrap();
let (_directory, mut recording) = recorded(&[1, 2]);
let transcription = Transcription::complete_samples(Scripted::new([
"Um, open cargo dot toml. They had had enough.",
]))
.with(RemoveHesitations)
.with(&vocabulary);
assert_eq!(
transcription
.transcribe(&mut recording, &Language::Automatic)
.unwrap()
.as_str(),
"Open Cargo.toml. They had had enough."
);
let (_directory, recording) = recorded(&[1; 1600]);
let transcription = Transcription::agreement(
Cadence::from_millis(100, 20_000),
Scripted::new(["Um, open cargo dot toml."]),
)
.with(RemoveHesitations)
.with(vocabulary);
let mut session = Session::over(&transcription, &Language::Automatic, recording).unwrap();
let segment = session.next_finalized().unwrap().unwrap();
assert_eq!(segment.as_str(), "Open Cargo.toml.");
assert!(session.next_finalized().unwrap().is_none());
assert_eq!(session.stop().unwrap().as_str(), segment.as_str());
}
struct WavScripted(&'static str);
impl Transcriber for WavScripted {
fn validate_language(&self, _: &Language, _: &CancelToken) -> Result<(), TranscriptionError> {
Ok(())
}
fn transcribe_with(
&self,
audio: TranscriptionAudio<'_>,
_: &Language,
_: &CancelToken,
) -> Result<Hypothesis, TranscriptionError> {
let wav = audio.require_wav("scripted")?;
assert!(wav.path().unwrap().is_file());
Ok(Hypothesis::new(self.0))
}
}
#[test]
fn whole_wav_layers_process_text_and_can_remove_all_speech() {
for (input, expected) in [
("Um, open cargo dot toml.", Some("Open Cargo.toml.")),
("Um, uh.", None),
] {
let (_directory, recording) = recorded(&[1, 2]);
let transcription = Transcription::complete_wav(WavScripted(input))
.with(RemoveHesitations)
.with(Vocabulary::new(["Cargo.toml"]).unwrap());
let session = Session::over(&transcription, &Language::Automatic, recording).unwrap();
let result = session.stop();
match expected {
Some(text) => assert_eq!(result.unwrap().as_str(), text),
None => assert!(matches!(result, Err(TranscriptionError::NoSpeech))),
}
}
}
struct NativeScripted;
impl Transcriber for NativeScripted {
fn validate_language(&self, _: &Language, _: &CancelToken) -> Result<(), TranscriptionError> {
Ok(())
}
fn transcribe_with(
&self,
_: TranscriptionAudio<'_>,
_: &Language,
_: &CancelToken,
) -> Result<Hypothesis, TranscriptionError> {
panic!("native transcription must use the chunk stream")
}
fn open_stream(
&self,
_: &Language,
_: &CancelToken,
) -> Result<Option<Box<dyn ChunkStream>>, TranscriptionError> {
Ok(Some(Box::new(ScriptedStream(VecDeque::from([
("um car", false),
("go dot toml done", true),
("uh", false),
])))))
}
}
struct ScriptedStream(VecDeque<(&'static str, bool)>);
impl ChunkStream for ScriptedStream {
fn chunk_samples(&self) -> usize {
1600
}
fn feed(&mut self, chunk: &[f32]) -> Result<Emitted, TranscriptionError> {
assert_eq!(chunk.len(), self.chunk_samples());
let (text, joins_previous) = self.0.pop_front().expect("script covers chunks and flush");
Ok(Emitted {
words: Hypothesis::new(text),
joins_previous,
})
}
}
#[test]
fn native_layers_process_joined_words_and_flush_with_independent_opt_outs() {
for (hesitations, vocabulary, expected) in [
(false, false, vec!["um", "cargo dot toml", "done uh"]),
(true, false, vec!["cargo dot toml", "done"]),
(false, true, vec!["um", "Cargo.toml", "done uh"]),
(true, true, vec!["Cargo.toml", "done"]),
] {
for drain_tail in [false, true] {
let (_directory, mut recording) = recorded(&[1; 3200]);
let transcription = Transcription::stream(Stream::Native, NativeScripted)
.with(hesitations.then_some(RemoveHesitations))
.with(vocabulary.then(|| Vocabulary::new(["Cargo.toml"]).unwrap()));
let mut operation = transcription.start(&Language::Automatic).unwrap();
for text in expected
.iter()
.take(if drain_tail { expected.len() } else { 1 })
{
let segment = operation.next_finalized(&mut recording).unwrap().unwrap();
assert_eq!(segment.as_str(), *text);
let span = segment.span.unwrap();
assert_eq!(
recording.span(span.clone()).len() as u64,
span.end - span.start
);
recording.release_before(operation.retained_from());
}
if drain_tail {
assert!(operation.next_finalized(&mut recording).unwrap().is_none());
assert_eq!(operation.completed().unwrap().as_str(), expected.join(" "));
}
assert_eq!(
operation.finish(&mut recording).unwrap().as_str(),
expected.join(" ")
);
}
}
}
#[test]
fn vocabulary_does_not_rewrite_across_finalized_segments() {
let (mut recording, mut capture) = Recording::open();
let transcription = Transcription::agreement(
Cadence::from_millis(100, 20_000),
Scripted::new(["cargo", "cargo dot toml"]),
)
.with(Vocabulary::new(["Cargo.toml"]).unwrap());
let mut operation = transcription.start(&Language::Automatic).unwrap();
capture.push(&[0.5; 1600]);
operation.next_update(&mut recording).unwrap();
capture.push(&[0.5; 1600]);
capture.finish();
recording.fill().unwrap();
let first = operation.next_finalized(&mut recording).unwrap().unwrap();
let tail = operation.next_finalized(&mut recording).unwrap().unwrap();
assert_eq!(first.as_str(), "cargo");
assert_eq!(tail.as_str(), "dot toml");
assert_eq!(
operation.finish(&mut recording).unwrap().as_str(),
"cargo dot toml"
);
}
#[test]
fn removing_every_word_finishes_without_redecoding() {
let (_directory, mut recording) = recorded(&[1, 2]);
let transcription =
Transcription::complete_samples(Scripted::new(["Um, uh, hmm."])).with(RemoveHesitations);
let mut operation = transcription.start(&Language::Automatic).unwrap();
assert!(operation.next_finalized(&mut recording).unwrap().is_none());
assert!(operation.completed().is_none());
assert!(matches!(
operation.finish(&mut recording),
Err(TranscriptionError::NoSpeech)
));
}
#[test]
fn a_suppressed_streaming_segment_does_not_end_the_stream() {
let (mut recording, mut capture) = Recording::open();
let transcription = Transcription::agreement(
Cadence::from_millis(100, 20_000),
Scripted::new(["um", "um hello"]),
)
.with(RemoveHesitations);
let mut operation = transcription.start(&Language::Automatic).unwrap();
capture.push(&[0.5; 1600]);
assert!(matches!(
operation.next_update(&mut recording).unwrap(),
Some(TranscriptUpdate::Provisional(text)) if text == "um"
));
capture.push(&[0.5; 1600]);
capture.finish();
recording.fill().unwrap();
assert_eq!(
operation
.next_finalized(&mut recording)
.unwrap()
.unwrap()
.as_str(),
"hello"
);
assert_eq!(operation.completed().unwrap().as_str(), "hello");
assert!(operation.next_finalized(&mut recording).unwrap().is_none());
}
#[test]
fn completion_without_a_final_segment_keeps_processed_text() {
let (mut recording, mut capture) = Recording::open();
let transcription = Transcription::agreement(
Cadence::from_millis(100, 20_000),
Scripted::new(["um hello", "um hello", "um hello"]),
)
.with(RemoveHesitations);
let mut operation = transcription.start(&Language::Automatic).unwrap();
capture.push(&[0.5; 1600]);
operation.next_update(&mut recording).unwrap();
capture.push(&[0.5; 1600]);
assert_eq!(
operation
.next_finalized(&mut recording)
.unwrap()
.unwrap()
.as_str(),
"hello"
);
capture.finish();
assert_eq!(operation.finish(&mut recording).unwrap().as_str(), "hello");
}
struct CountApplications(AtomicUsize);
impl Layer for CountApplications {
fn apply(&self, text: &mut String, _: &CancelToken) -> Result<(), TranscriptionError> {
self.0.fetch_add(1, Ordering::Relaxed);
text.push('!');
Ok(())
}
}
#[test]
fn completion_includes_the_processed_tail_without_reapplying_layers() {
for drain_tail in [false, true] {
let calls = CountApplications(AtomicUsize::new(0));
let (mut recording, mut capture) = Recording::open();
let transcription = Transcription::agreement(
Cadence::from_millis(100, 20_000),
Scripted::new(["one two", "one two three"]),
)
.with(&calls);
let mut operation = transcription.start(&Language::Automatic).unwrap();
capture.push(&[0.5; 1600]);
assert!(matches!(
operation.next_update(&mut recording).unwrap(),
Some(TranscriptUpdate::Provisional(text)) if text == "one two"
));
capture.push(&[0.5; 1600]);
capture.finish();
recording.fill().unwrap();
let first = operation.next_finalized(&mut recording).unwrap().unwrap();
assert_eq!(first.as_str(), "one two!");
assert_eq!(operation.completed().unwrap().as_str(), "one two! three!");
if drain_tail {
let tail = operation.next_finalized(&mut recording).unwrap().unwrap();
assert_eq!(tail.as_str(), "three!");
assert_eq!(
[first.as_str(), tail.as_str()].join(" "),
operation.completed().unwrap().as_str()
);
}
assert_eq!(
operation.finish(&mut recording).unwrap().as_str(),
"one two! three!"
);
assert_eq!(calls.0.load(Ordering::Relaxed), 2);
}
}
#[test]
fn processed_segments_keep_their_audio_spans_when_recording_is_released() {
let timings = vec![
TimedWord {
start: 0.0,
end: 0.025,
},
TimedWord {
start: 0.025,
end: 0.05,
},
TimedWord {
start: 0.05,
end: 0.075,
},
];
let scripted = Scripted(Mutex::new(VecDeque::from([
Hypothesis::with_words("um one", timings[..2].to_vec()),
Hypothesis::with_words("um one two", timings),
])));
let (mut recording, mut capture) = Recording::open();
let transcription = Transcription::agreement(Cadence::from_millis(100, 20_000), scripted)
.with(RemoveHesitations);
let mut operation = transcription.start(&Language::Automatic).unwrap();
capture.push(&[0.5; 1600]);
operation.next_update(&mut recording).unwrap();
capture.push(&[0.5; 1600]);
capture.finish();
recording.fill().unwrap();
let first = operation.next_finalized(&mut recording).unwrap().unwrap();
assert_eq!(first.as_str(), "one");
assert_eq!(first.span, Some(0..800));
let retained = operation.retained_from();
recording.release_before(retained);
let tail = operation.next_finalized(&mut recording).unwrap().unwrap();
assert_eq!(tail.as_str(), "two");
assert_eq!(tail.span, Some(800..1200));
let span = tail.span.unwrap();
assert!(retained <= span.start);
assert_eq!(recording.span(span).len(), 400);
}
struct StopOnTail(bool);
impl Layer for StopOnTail {
fn apply(&self, text: &mut String, cancel: &CancelToken) -> Result<(), TranscriptionError> {
if text == "three" {
if self.0 {
cancel.cancel();
} else {
return Err(TranscriptionError::Processing(Error::Failed(Box::new(
std::io::Error::other("rejected"),
))));
}
}
Ok(())
}
}
#[test]
fn cancellation_or_failure_in_the_tail_never_exposes_raw_completion() {
for cancel_on_tail in [false, true] {
let (mut recording, mut capture) = Recording::open();
let transcription = Transcription::agreement(
Cadence::from_millis(100, 20_000),
Scripted::new(["one two", "one two three"]),
)
.with(StopOnTail(cancel_on_tail));
let cancel = CancelToken::new();
let mut operation = transcription
.start_with(&Language::Automatic, &cancel)
.unwrap();
capture.push(&[0.5; 1600]);
operation.next_update(&mut recording).unwrap();
capture.push(&[0.5; 1600]);
capture.finish();
recording.fill().unwrap();
let error = operation.next_finalized(&mut recording).unwrap_err();
if cancel_on_tail {
assert!(matches!(error, TranscriptionError::Cancelled));
} else {
assert!(matches!(
error,
TranscriptionError::Processing(Error::Failed(_))
));
}
assert!(operation.completed().is_none());
assert!(operation.finish(&mut recording).is_err());
assert!(operation.completed().is_none());
}
}
#[test]
fn hesitations_preserve_sentence_punctuation() {
for (input, expected) in [
("Well, um. Next.", "Well. Next."),
("Well, uh? Next.", "Well? Next."),
("Go. Um, so then", "Go. So then"),
] {
let mut text = input.to_owned();
RemoveHesitations
.apply(&mut text, &CancelToken::never())
.unwrap();
assert_eq!(text, expected);
}
}
#[test]
fn hesitations_preserve_delimiters_around_surviving_words() {
for (input, expected) in [
("\"Um, hello.\"", "\"Hello.\""),
("(um hello)", "(hello)"),
("go (um uh now)", "go (now)"),
("go (um uh) now", "go now"),
("(Um) hello", "Hello"),
("(hello, um)", "(hello)"),
("[well; uh] next", "[well] next"),
("(um hello, uh)", "(hello)"),
("\"Hello. um\"", "\"Hello.\""),
] {
let mut text = input.to_owned();
RemoveHesitations
.apply(&mut text, &CancelToken::never())
.unwrap();
assert_eq!(text, expected);
}
}
#[test]
fn vocabulary_layer_rejects_personal_expansion_before_replacing_text() {
let rules = vocabulary::PersonalRules::new(
Vec::new(),
vec![("x".to_owned(), "y".repeat(MAX_TEXT_BYTES + 1))],
);
let vocabulary = Vocabulary::compile(&[], rules).unwrap();
let mut text = "x".to_owned();
assert!(matches!(
vocabulary.apply(&mut text, &CancelToken::never()),
Err(TranscriptionError::Processing(Error::TooLarge))
));
assert_eq!(text, "x");
}
struct Expand(usize);
impl Layer for Expand {
fn apply(&self, text: &mut String, _: &CancelToken) -> Result<(), TranscriptionError> {
*text = "x".repeat(self.0);
Ok(())
}
}
#[test]
fn text_bounds_cover_each_segment_before_and_after_a_layer() {
for (input_bytes, output_bytes) in [
(MAX_TEXT_BYTES, MAX_TEXT_BYTES),
(MAX_TEXT_BYTES + 1, 1),
(1, MAX_TEXT_BYTES + 1),
] {
let (_directory, mut recording) = recorded(&[1, 2]);
let transcription =
Transcription::complete_samples(Scripted::new(["x".repeat(input_bytes)]))
.with(Expand(output_bytes));
let result = transcription.transcribe(&mut recording, &Language::Automatic);
if input_bytes <= MAX_TEXT_BYTES && output_bytes <= MAX_TEXT_BYTES {
assert_eq!(result.unwrap().as_str().len(), output_bytes);
} else {
assert!(matches!(
result,
Err(TranscriptionError::Processing(Error::TooLarge))
));
}
}
}
#[test]
fn a_session_accumulates_past_the_segment_bound() {
let (mut recording, mut capture) = Recording::open();
let transcription = Transcription::agreement(
Cadence::from_millis(100, 20_000),
Scripted::new(["one", "one two"]),
)
.with(Expand(MAX_TEXT_BYTES / 2));
let mut operation = transcription.start(&Language::Automatic).unwrap();
capture.push(&[0.5; 1600]);
operation.next_update(&mut recording).unwrap();
capture.push(&[0.5; 1600]);
capture.finish();
recording.fill().unwrap();
while operation.next_finalized(&mut recording).unwrap().is_some() {}
assert_eq!(
operation.completed().unwrap().as_str().len(),
MAX_TEXT_BYTES + 1
);
}