use core::{
num::{NonZeroU32, NonZeroUsize},
sync::atomic::AtomicBool,
time::Duration,
};
use smol_str::{SmolStr, format_smolstr};
use crate::{
core::{OovDetection, OovResolution, UnalignedCause, UnitAlignment, UnitJob, UnitOutcome},
runner::aligner::{
algorithm::{
compose::DEFAULT_MAX_INTRA_SILENT_RUN,
encode::{validate_stride_extent, validate_vocab_dim},
errors::{EmissionsError, EmissionsFailure},
},
core::{
AlignerCore, AlignerCoreLoadError, PreparedChunk, WAV2VEC2_RECEPTIVE_FIELD_SAMPLES,
WAV2VEC2_WORD_DELIMITER, capture_vocab_size, declared_unk_token_id, detect_blank_token_id,
load_tokenizer_bytes_with_compat, validate_word_delimiter_present,
},
emissions_api::{Emissions, EncoderOutput, OutputClock, SpeechCoverage, SpeechSpans},
normalizer::DynTextNormalizer,
normalizers::default_normalizer_for,
},
types::{AlignmentError, Lang, WorkFailure},
};
const DEFAULT_HOP_SAMPLES: NonZeroU32 = match NonZeroU32::new(320) {
Some(v) => v,
None => unreachable!(),
};
fn to_emissions_error(err: WorkFailure, stage: Stage) -> EmissionsError {
let neutral = |f: &crate::types::AlignmentFailure| EmissionsFailure::new(f.message().clone());
match err {
WorkFailure::Alignment(inner) => match inner {
AlignmentError::ModelInference(ref f) => match stage {
Stage::Prepare => EmissionsError::NonFiniteAudio(neutral(f)),
Stage::Finish => EmissionsError::Config(neutral(f)),
},
AlignmentError::Normalization(ref f) => EmissionsError::Normalization(neutral(f)),
AlignmentError::Tokenization(ref f) | AlignmentError::EmptyText(ref f) => {
EmissionsError::Tokenization(neutral(f))
}
AlignmentError::SemanticOutOfVocab(ref f) => EmissionsError::SemanticOutOfVocab(neutral(f)),
AlignmentError::NoAlignmentPath(ref f) => EmissionsError::NoAlignmentPath(neutral(f)),
AlignmentError::Aborted(ref f) | AlignmentError::Abandoned(ref f) => {
EmissionsError::Aborted(neutral(f))
}
AlignmentError::Geometry(ref f) => EmissionsError::Geometry(neutral(f)),
},
WorkFailure::WorkerHang(_) => EmissionsError::Aborted(EmissionsFailure::new(format_smolstr!(
"aborted via abort_flag before completing"
))),
other @ (WorkFailure::Asr(_) | WorkFailure::LanguageUnsupported(_)) => {
EmissionsError::Config(EmissionsFailure::new(format_smolstr!(
"internal call chain produced an unexpected WorkFailure variant ({other:?}); this \
is a bug in the seam, not in the caller's input"
)))
}
}
}
#[derive(Clone, Copy)]
enum Stage {
Prepare,
Finish,
}
fn load_error(err: AlignerCoreLoadError) -> EmissionsError {
EmissionsError::Config(EmissionsFailure::new(err.message().clone()))
}
pub struct EmissionsAligner {
core: AlignerCore,
}
impl EmissionsAligner {
#[must_use]
pub fn builder(language: Lang, tokenizer_json: &[u8]) -> EmissionsAlignerBuilder {
EmissionsAlignerBuilder {
language,
tokenizer_json: tokenizer_json.to_vec(),
normalizer: None,
hop_samples: DEFAULT_HOP_SAMPLES,
word_delimiter: SmolStr::new_static(WAV2VEC2_WORD_DELIMITER),
letter_case: LetterCase::Upper,
receptive_field_samples: WAV2VEC2_RECEPTIVE_FIELD_SAMPLES,
min_speech_coverage: SpeechCoverage::DEFAULT,
max_intra_silent_run: DEFAULT_MAX_INTRA_SILENT_RUN,
blank_token_id: None,
}
}
#[must_use]
pub const fn vocab_size(&self) -> NonZeroUsize {
self.core.vocab_size()
}
#[must_use]
pub const fn blank_token_id(&self) -> u32 {
self.core.blank_token_id()
}
#[must_use]
pub const fn hop_samples(&self) -> NonZeroU32 {
self.core.hop_samples()
}
#[must_use]
pub fn word_delimiter(&self) -> &str {
self.core.word_delimiter()
}
#[must_use]
pub const fn letter_case(&self) -> LetterCase {
if self.core.vocab_uppercase_only() {
LetterCase::Upper
} else {
LetterCase::AsWritten
}
}
#[must_use]
pub const fn receptive_field_samples(&self) -> NonZeroU32 {
self.core.receptive_field_samples()
}
#[must_use]
pub const fn language(&self) -> &Lang {
self.core.language()
}
#[must_use]
pub const fn min_speech_coverage(&self) -> SpeechCoverage {
self.core.min_speech_coverage()
}
#[must_use]
pub const fn max_intra_silent_run(&self) -> Duration {
self.core.max_intra_silent_run()
}
pub fn detect_oov(&self, text: &str) -> Result<OovDetection, EmissionsError> {
let events = self
.core
.detect_oov(text)
.map_err(|e| to_emissions_error(e, Stage::Prepare))?;
Ok(OovDetection::of_text(
text,
self.core.language().clone(),
events,
self.core.id().get(),
))
}
pub fn detect_oov_unit(&self, job: &UnitJob) -> Result<OovDetection, EmissionsError> {
self
.core
.detect_job(job, false)
.map_err(|e| to_emissions_error(e, Stage::Prepare))
}
pub fn detect_oov_unit_as_fallback(&self, job: &UnitJob) -> Result<OovDetection, EmissionsError> {
self
.core
.detect_job(job, true)
.map_err(|e| to_emissions_error(e, Stage::Prepare))
}
pub fn prepare<'a>(
&self,
samples: &[f32],
speech: &SpeechSpans,
text: &'a str,
resolution: OovResolution,
clock: OutputClock,
abort_flag: &AtomicBool,
) -> Result<PreparedChunk<'a>, EmissionsError> {
let oov_decisions = self
.core
.accept(&resolution, text)
.map_err(|e| to_emissions_error(e, Stage::Prepare))?;
let expected = self.core.language().clone();
let preparation = self
.core
.prepare(samples, speech, text, oov_decisions, &expected, abort_flag)
.map_err(|e| to_emissions_error(e, Stage::Prepare))?;
Ok(PreparedChunk::new(preparation, clock))
}
pub fn finish(
&self,
prepared: PreparedChunk<'_>,
emissions: Emissions,
abort_flag: &AtomicBool,
) -> Result<UnitAlignment, EmissionsError> {
let (prepared, clock) = prepared.into_parts();
if !self.core.owns(&prepared) {
return Err(EmissionsError::AlignerMismatch(EmissionsFailure::new(
format_smolstr!(
"this PreparedChunk was produced by a DIFFERENT EmissionsAligner. It carries \
token ids, a word map, and OOV decisions resolved against that aligner's tokenizer, \
blank id, and language — none of which need match this one's, even when the vocab \
sizes and hops are identical. Call `finish` on the same aligner that called `prepare`."
),
)));
}
if emissions.preparation() != prepared.id() {
return Err(EmissionsError::PreparationMismatch(EmissionsFailure::new(
SmolStr::new_static(
"these Emissions were made through another PreparedChunk. Emissions answer the one \
chunk they were made through, whatever their shape: make them with \
`prepared.encode_with` on the chunk you \
finish.",
),
)));
}
if prepared.is_trivial() {
return Ok(UnitAlignment::Unaligned(UnalignedCause::NoAlignableText));
}
let language = self.core.language();
validate_stride_extent(
emissions.frames(),
self.core.hop_samples().get(),
self.core.receptive_field_samples().get(),
prepared.encoder_input().len(),
language,
)
.map_err(|e| EmissionsError::StrideMismatch(work_failure_message(e)))?;
validate_vocab_dim(
emissions.vocab().get(),
self.core.vocab_size().get(),
language,
)
.map_err(|e| EmissionsError::VocabMismatch(work_failure_message(e)))?;
self
.core
.finish(
prepared,
emissions.inner(),
clock.chunk_first_sample_in_stream(),
|start, end| clock.range(start, end),
abort_flag,
)
.map_err(|e| to_emissions_error(e, Stage::Finish))
}
}
impl EmissionsAligner {
pub fn align_unit<E: From<EmissionsError>>(
&self,
job: UnitJob,
resolution: OovResolution,
encoder: impl FnOnce(&[f32]) -> Result<EncoderOutput, E>,
abort_flag: &AtomicBool,
) -> Result<UnitOutcome, E> {
let place = job.place();
let invalid = |error: crate::runner::aligner::emissions_api::SpanError| {
EmissionsError::Config(EmissionsFailure::new(format_smolstr!(
"the unit's place in the stream is invalid: {error}"
)))
};
let speech = SpeechSpans::from_time_ranges(&place.sub_segments).map_err(invalid)?;
let clock = OutputClock::new(
place.first_sample_in_stream,
place.output_tb,
place.base_pts_out_anchor,
)
.map_err(invalid)?;
let decisions = self
.core
.accept_job(&resolution, &job)
.map_err(|e| to_emissions_error(e, Stage::Prepare))?;
let preparation = match self
.core
.prepare(
job.samples(),
&speech,
job.text(),
decisions,
job.language(),
abort_flag,
)
.map_err(|e| to_emissions_error(e, Stage::Prepare))
{
Ok(preparation) => preparation,
Err(error) => return unit_failure(job, error),
};
let prepared = PreparedChunk::new(preparation, clock);
let emissions = prepared.encode_with(encoder)?;
match self.finish(prepared, emissions, abort_flag) {
Ok(alignment) => Ok(job.answer(alignment)),
Err(error) => unit_failure(job, error),
}
}
}
fn unit_failure<E: From<EmissionsError>>(
job: UnitJob,
error: EmissionsError,
) -> Result<UnitOutcome, E> {
let language = job.language().clone();
let failure =
|f: &EmissionsFailure| crate::types::AlignmentFailure::new(f.message().clone(), language);
let cause = match &error {
EmissionsError::NoAlignmentPath(f) => AlignmentError::NoAlignmentPath(failure(f)),
EmissionsError::SemanticOutOfVocab(f) => AlignmentError::SemanticOutOfVocab(failure(f)),
_ => return Err(error.into()),
};
Ok(job.answer(UnitAlignment::Unaligned(UnalignedCause::Failed(cause))))
}
fn work_failure_message(err: WorkFailure) -> EmissionsFailure {
match err {
WorkFailure::Alignment(
AlignmentError::ModelInference(f)
| AlignmentError::Tokenization(f)
| AlignmentError::Normalization(f)
| AlignmentError::NoAlignmentPath(f)
| AlignmentError::EmptyText(f)
| AlignmentError::SemanticOutOfVocab(f)
| AlignmentError::Aborted(f)
| AlignmentError::Abandoned(f)
| AlignmentError::Geometry(f),
) => EmissionsFailure::new(f.message().clone()),
other => EmissionsFailure::new(format_smolstr!("{other:?}")),
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum LetterCase {
#[default]
Upper,
AsWritten,
}
pub struct EmissionsAlignerBuilder {
language: Lang,
tokenizer_json: Vec<u8>,
normalizer: Option<DynTextNormalizer>,
hop_samples: NonZeroU32,
word_delimiter: SmolStr,
letter_case: LetterCase,
receptive_field_samples: NonZeroU32,
min_speech_coverage: SpeechCoverage,
max_intra_silent_run: Duration,
blank_token_id: Option<u32>,
}
impl EmissionsAlignerBuilder {
#[must_use]
pub fn normalizer(mut self, normalizer: DynTextNormalizer) -> Self {
self.normalizer = Some(normalizer);
self
}
#[must_use]
pub const fn hop_samples(mut self, hop: NonZeroU32) -> Self {
self.hop_samples = hop;
self
}
#[must_use]
pub fn word_delimiter(mut self, token: &str) -> Self {
self.word_delimiter = SmolStr::new(token);
self
}
#[must_use]
pub const fn letter_case(mut self, case: LetterCase) -> Self {
self.letter_case = case;
self
}
#[must_use]
pub const fn receptive_field_samples(mut self, samples: NonZeroU32) -> Self {
self.receptive_field_samples = samples;
self
}
#[must_use]
pub const fn min_speech_coverage(mut self, value: SpeechCoverage) -> Self {
self.min_speech_coverage = value;
self
}
#[must_use]
pub const fn max_intra_silent_run(mut self, value: Duration) -> Self {
self.max_intra_silent_run = value;
self
}
#[must_use]
pub const fn blank_token_id(mut self, id: u32) -> Self {
self.blank_token_id = Some(id);
self
}
pub fn build(self) -> Result<EmissionsAligner, EmissionsError> {
let tokenizer = load_tokenizer_bytes_with_compat(&self.tokenizer_json, "tokenizer.json")
.map_err(load_error)?;
let blank_token_id = match self.blank_token_id {
Some(id) => id,
None => detect_blank_token_id(&tokenizer).ok_or_else(|| {
EmissionsError::Config(EmissionsFailure::new(format_smolstr!(
"tokenizer has no <pad> / [PAD] / <blank> entry; cannot determine the CTC blank \
token. Supply it explicitly with `.blank_token_id(id)`."
)))
})?,
};
let normalizer = match self.normalizer {
Some(n) => n,
None => default_normalizer_for(&self.language).ok_or_else(|| {
EmissionsError::Config(EmissionsFailure::new(format_smolstr!(
"no default text normalizer for {:?}; supply one with `.normalizer(..)`",
self.language
)))
})?,
};
let unk_token_id = declared_unk_token_id(&tokenizer);
validate_word_delimiter_present(
&tokenizer,
normalizer.use_word_delimiter(),
&self.word_delimiter,
)
.map_err(load_error)?;
let tokenizer_vocab_size = capture_vocab_size(&tokenizer).ok_or_else(|| {
EmissionsError::Config(EmissionsFailure::new(format_smolstr!(
"tokenizer reports a zero-size vocab; a CTC vocabulary must contain at least the \
blank token"
)))
})?;
Ok(EmissionsAligner {
core: AlignerCore::from_parts(
tokenizer,
self.language,
normalizer,
self.hop_samples,
self.word_delimiter,
self.receptive_field_samples,
blank_token_id,
unk_token_id,
matches!(self.letter_case, LetterCase::Upper),
tokenizer_vocab_size,
self.min_speech_coverage,
self.max_intra_silent_run,
),
})
}
}
#[cfg(test)]
mod tests;