mozu-audio 0.3.0

//! Greedy Haas TDT decoding with transactional KDA state updates.

use super::tensor::{TensorBuffer, TensorReader};
use super::window::{ENCODER_STEP, TimedToken};
use super::{
    Assets, Cancellation, Encoded, Error, Result, core_error, feature_provider, output_array,
};
use crate::vocabulary::{BiasState, CompiledVocabulary};
use objc2::{
    rc::{Retained, autoreleasepool},
    runtime::{AnyObject, ProtocolObject},
};
use objc2_core_ml::{
    MLDictionaryFeatureProvider, MLFeatureProvider, MLModel, MLMultiArray, MLMultiArrayDataType,
    MLPredictionOptions,
};
use objc2_foundation::{NSDictionary, NSString};

struct KdaState {
    memory: TensorBuffer<f32>,
    conv: TensorBuffer<f32>,
}

impl KdaState {
    fn zero() -> Result<Self> {
        Ok(Self {
            memory: TensorBuffer::zero(&[2, 1, 5, 64, 64])?,
            conv: TensorBuffer::zero(&[2, 3, 1, 320, 3])?,
        })
    }
}

/// The pending slot can be overwritten; only a nonblank decision commits it.
struct StatePair {
    slots: [KdaState; 2],
    active: usize,
    blank: usize,
}

impl StatePair {
    fn zero(blank: usize) -> Result<Self> {
        Ok(Self {
            slots: [KdaState::zero()?, KdaState::zero()?],
            active: 0,
            blank,
        })
    }

    fn accept(&mut self, token: usize) {
        if token != self.blank {
            self.active ^= 1;
        }
    }
}

struct Decision {
    label: usize,
    duration: usize,
    confidence: f32,
}

pub(super) enum DecoderModels {
    Fused(Retained<MLModel>),
    Split {
        prediction: Retained<MLModel>,
        joint: JointModel,
    },
}

#[derive(Clone, Copy)]
enum JointSchema {
    Logits,
    Decision,
}

pub(super) struct JointModel {
    model: Retained<MLModel>,
    schema: JointSchema,
}

impl DecoderModels {
    pub(super) fn require_logits(&self) -> Result<()> {
        if matches!(
            self,
            Self::Split {
                joint: JointModel {
                    schema: JointSchema::Decision,
                    ..
                },
                ..
            }
        ) {
            return Err(Error::Invalid("Vocabulary boosting requires a full-logit joint model; load joint.mlmodelc instead of joint-decision.mlmodelc".to_owned()));
        }
        Ok(())
    }
    pub(super) fn split(
        prediction: Retained<MLModel>,
        joint: Retained<MLModel>,
        path: &std::path::Path,
    ) -> Result<Self> {
        // SAFETY: Feature descriptions belong to the retained model.
        let description = unsafe { joint.modelDescription() };
        let inputs = unsafe { description.inputDescriptionsByName() };
        super::validate_array_features(
            &inputs,
            &[
                ("encoder", &[1, 1024], MLMultiArrayDataType::Float32),
                ("prediction", &[1, 640], MLMultiArrayDataType::Float32),
            ],
            &format!("joint inputs model={}", path.display()),
        )?;
        // SAFETY: Reading output metadata does not execute the model.
        let outputs = unsafe { description.outputDescriptionsByName() };
        let present = |name: &str| outputs.objectForKey(&NSString::from_str(name)).is_some();
        let schema = match (
            present("output"),
            present("token_id"),
            present("token_prob"),
            present("duration"),
        ) {
            (true, false, false, false) => {
                super::validate_array_features(
                    &outputs,
                    &[("output", &[1, 3078], MLMultiArrayDataType::Float32)],
                    &format!("joint logits model={}", path.display()),
                )?;
                JointSchema::Logits
            }
            (false, true, true, true) => {
                super::validate_array_features(
                    &outputs,
                    &[
                        ("token_id", &[1, 1, 1], MLMultiArrayDataType::Int32),
                        ("token_prob", &[1, 1, 1], MLMultiArrayDataType::Float32),
                        ("duration", &[1, 1, 1], MLMultiArrayDataType::Int32),
                    ],
                    &format!("joint decisions model={}", path.display()),
                )?;
                JointSchema::Decision
            }
            features => {
                return Err(Error::Invalid(format!(
                    "Unsupported joint outputs: model={}, present(output,token_id,token_prob,duration)={features:?}",
                    path.display()
                )));
            }
        };
        Ok(Self::Split {
            prediction,
            joint: JointModel {
                model: joint,
                schema,
            },
        })
    }
}

struct LogitOutputs {
    logits: TensorBuffer<f32>,
    reader: TensorReader<f32>,
}

struct DecisionOutputs {
    token: TensorBuffer<i32>,
    probability: TensorBuffer<f32>,
    duration: TensorBuffer<i32>,
    token_reader: TensorReader<i32>,
    probability_reader: TensorReader<f32>,
    duration_reader: TensorReader<i32>,
}

impl DecisionOutputs {
    fn new() -> Result<Self> {
        let token = TensorBuffer::zero(&[1, 1, 1])?;
        let probability = TensorBuffer::zero(&[1, 1, 1])?;
        let duration = TensorBuffer::zero(&[1, 1, 1])?;
        let token_reader = TensorReader::new(token.array())?;
        let probability_reader = TensorReader::new(probability.array())?;
        let duration_reader = TensorReader::new(duration.array())?;
        Ok(Self {
            token,
            probability,
            duration,
            token_reader,
            probability_reader,
            duration_reader,
        })
    }

    fn read(&mut self, blank: usize) -> Result<Decision> {
        self.token_reader.refresh(self.token.array())?;
        self.probability_reader.refresh(self.probability.array())?;
        self.duration_reader.refresh(self.duration.array())?;
        let label = self.token_reader.values()[0];
        let confidence = self.probability_reader.values()[0];
        let duration = self.duration_reader.values()[0];
        if !(0..3073).contains(&label)
            || !(0..5).contains(&duration)
            || !confidence.is_finite()
            || !(0.0..=1.0).contains(&confidence)
        {
            return Err(Error::Invalid(format!(
                "Invalid joint decision: token_id={label}, token_prob={confidence}, duration={duration}"
            )));
        }
        Ok(Decision {
            label: label as usize,
            duration: duration as usize,
            confidence: if label as usize == blank {
                0.0
            } else {
                confidence
            },
        })
    }
}

enum JointOutputs {
    Logits(Box<LogitOutputs>),
    Decision(Box<DecisionOutputs>),
}

impl JointOutputs {
    fn new(schema: JointSchema) -> Result<Self> {
        match schema {
            JointSchema::Logits => {
                let logits = TensorBuffer::zero(&[1, 3078])?;
                let reader = TensorReader::new(logits.array())?;
                Ok(Self::Logits(Box::new(LogitOutputs { logits, reader })))
            }
            JointSchema::Decision => Ok(Self::Decision(Box::new(DecisionOutputs::new()?))),
        }
    }

    fn logits(&self) -> Result<&MLMultiArray> {
        match self {
            Self::Logits(output) => Ok(output.logits.array()),
            Self::Decision(_) => Err(Error::Invalid(
                "Fused prediction requires joint logits".into(),
            )),
        }
    }

    fn with_backings<T>(&self, operation: impl FnOnce(&[(&str, &MLMultiArray)]) -> T) -> T {
        match self {
            Self::Logits(output) => operation(&[("output", output.logits.array())]),
            Self::Decision(output) => operation(&[
                ("token_id", output.token.array()),
                ("token_prob", output.probability.array()),
                ("duration", output.duration.array()),
            ]),
        }
    }

    fn read_selected(&mut self, blank: usize, selector: &impl LabelSelector) -> Result<Decision> {
        match self {
            Self::Logits(output) => {
                output.reader.refresh(output.logits.array())?;
                let logits = output.reader.values();
                let label = selector.select(&logits[..3073], blank)?;
                let duration = first_argmax(&logits[3073..])?;
                let confidence = if label == blank {
                    0.0
                } else {
                    label_probability(&logits[..3073], label)?
                };
                Ok(Decision {
                    label,
                    duration,
                    confidence,
                })
            }
            Self::Decision(output) => output.read(blank),
        }
    }
}

fn output_options(backings: &[(&str, &MLMultiArray)]) -> Retained<MLPredictionOptions> {
    let names: Vec<_> = backings
        .iter()
        .map(|(name, _)| NSString::from_str(name))
        .collect();
    let keys: Vec<&NSString> = names.iter().map(|name| &**name).collect();
    let arrays: Vec<&AnyObject> = backings.iter().map(|(_, array)| array.as_ref()).collect();
    let dictionary = NSDictionary::from_slices(&keys, &arrays);
    // SAFETY: Output backings are retained MLMultiArray objects matching the FP32 model contract.
    unsafe {
        let options = MLPredictionOptions::new();
        options.setOutputBackings(&dictionary);
        options
    }
}

fn run_model(
    model: &MLModel,
    provider: &MLDictionaryFeatureProvider,
    options: &MLPredictionOptions,
    backings: &[(&str, &MLMultiArray)],
    operation: &str,
) -> Result<Retained<ProtocolObject<dyn MLFeatureProvider>>> {
    // SAFETY: Providers and options retain distinct accepted and pending state backings.
    let output = unsafe {
        model.predictionFromFeatures_options_error(ProtocolObject::from_ref(provider), options)
    }
    .map_err(|error| core_error(operation, error))?;
    for (name, expected) in backings {
        let actual = output_array(&output, name)?;
        if !std::ptr::eq(&*actual, *expected) {
            return Err(Error::Invalid(format!(
                "Core ML did not use the supplied output backing: operation={operation}, output={name}"
            )));
        }
    }
    Ok(output)
}

/// Blank decisions retain both the prediction vector and its uncommitted state.
enum PredictionCache {
    Empty,
    Ready,
}

impl PredictionCache {
    fn needs_prediction(&self) -> bool {
        matches!(self, Self::Empty)
    }

    fn compute(&mut self, prediction: impl FnOnce() -> Result<()>) -> Result<()> {
        if self.needs_prediction() {
            prediction()?;
            *self = Self::Ready;
        }
        Ok(())
    }

    fn accept(&mut self, token: usize, blank: usize) {
        if token != blank {
            *self = Self::Empty;
        }
    }
}

enum Inference<'a> {
    Fused {
        model: &'a MLModel,
        providers: [Retained<MLDictionaryFeatureProvider>; 2],
        options: [Retained<MLPredictionOptions>; 2],
    },
    Split {
        prediction_model: &'a MLModel,
        joint_model: &'a MLModel,
        prediction: TensorBuffer<f32>,
        cache: PredictionCache,
        providers: [Retained<MLDictionaryFeatureProvider>; 2],
        options: [Retained<MLPredictionOptions>; 2],
        joint_provider: Retained<MLDictionaryFeatureProvider>,
        joint_options: Retained<MLPredictionOptions>,
    },
}

impl<'a> Inference<'a> {
    fn new(
        models: &'a DecoderModels,
        embedding: &MLMultiArray,
        frame: &MLMultiArray,
        outputs: &JointOutputs,
        states: &StatePair,
    ) -> Result<Self> {
        match models {
            DecoderModels::Fused(model) => {
                let logits = outputs.logits()?;
                let provider = |index: usize| {
                    feature_provider(&[
                        ("x", embedding),
                        ("encoder", frame),
                        ("memory", states.slots[index].memory.array()),
                        ("conv", states.slots[index].conv.array()),
                    ])
                };
                let options = |index: usize| {
                    output_options(&[
                        ("output", logits),
                        ("next_state", states.slots[index ^ 1].memory.array()),
                        ("next_conv", states.slots[index ^ 1].conv.array()),
                    ])
                };
                Ok(Self::Fused {
                    model,
                    providers: [provider(0)?, provider(1)?],
                    options: [options(0), options(1)],
                })
            }
            DecoderModels::Split {
                prediction: prediction_model,
                joint: joint_model,
            } => {
                let prediction = TensorBuffer::zero(&[1, 640])?;
                let provider = |index: usize| {
                    feature_provider(&[
                        ("x", embedding),
                        ("memory", states.slots[index].memory.array()),
                        ("conv", states.slots[index].conv.array()),
                    ])
                };
                let options = |index: usize| {
                    output_options(&[
                        ("prediction", prediction.array()),
                        ("next_state", states.slots[index ^ 1].memory.array()),
                        ("next_conv", states.slots[index ^ 1].conv.array()),
                    ])
                };
                let providers = [provider(0)?, provider(1)?];
                let options = [options(0), options(1)];
                let joint_provider =
                    feature_provider(&[("encoder", frame), ("prediction", prediction.array())])?;
                let joint_options = outputs.with_backings(output_options);
                Ok(Self::Split {
                    prediction_model,
                    joint_model: &joint_model.model,
                    prediction,
                    cache: PredictionCache::Empty,
                    providers,
                    options,
                    joint_provider,
                    joint_options,
                })
            }
        }
    }

    fn needs_embedding(&self) -> bool {
        match self {
            Self::Fused { .. } => true,
            Self::Split { cache, .. } => cache.needs_prediction(),
        }
    }

    fn accept(&mut self, token: usize, blank: usize) {
        if let Self::Split { cache, .. } = self {
            cache.accept(token, blank);
        }
    }

    fn run(&mut self, states: &StatePair, outputs: &JointOutputs) -> Result<()> {
        let active = states.active;
        let pending = &states.slots[active ^ 1];
        match self {
            Self::Fused {
                model,
                providers,
                options,
            } => {
                let logits = outputs.logits()?;
                run_model(
                    model,
                    &providers[active],
                    &options[active],
                    &[
                        ("output", logits),
                        ("next_state", pending.memory.array()),
                        ("next_conv", pending.conv.array()),
                    ],
                    "Haas fused prediction",
                )?;
            }
            Self::Split {
                prediction_model,
                joint_model,
                prediction,
                cache,
                providers,
                options,
                joint_provider,
                joint_options,
            } => {
                cache.compute(|| {
                    run_model(
                        prediction_model,
                        &providers[active],
                        &options[active],
                        &[
                            ("prediction", prediction.array()),
                            ("next_state", pending.memory.array()),
                            ("next_conv", pending.conv.array()),
                        ],
                        "Haas prediction",
                    )
                    .map(drop)
                })?;
                outputs.with_backings(|backings| {
                    run_model(
                        joint_model,
                        joint_provider,
                        joint_options,
                        backings,
                        "Haas joint",
                    )
                })?;
            }
        }
        Ok(())
    }
}

struct DecodeWorkspace<'a> {
    inference: Inference<'a>,
    encoded: &'a Encoded,
    assets: &'a Assets,
    embedding: TensorBuffer<f32>,
    frame: TensorBuffer<f32>,
    outputs: JointOutputs,
    states: StatePair,
}

impl<'a> DecodeWorkspace<'a> {
    fn new(models: &'a DecoderModels, encoded: &'a Encoded, assets: &'a Assets) -> Result<Self> {
        let embedding = TensorBuffer::zero(&[1, assets.metadata.embedding_width])?;
        let frame = TensorBuffer::zero(&[1, 1024])?;
        let schema = match models {
            DecoderModels::Fused(_) => JointSchema::Logits,
            DecoderModels::Split { joint, .. } => joint.schema,
        };
        let outputs = JointOutputs::new(schema)?;
        let states = StatePair::zero(assets.metadata.blank_id)?;
        let inference =
            Inference::new(models, embedding.array(), frame.array(), &outputs, &states)?;
        Ok(Self {
            inference,
            encoded,
            assets,
            embedding,
            frame,
            outputs,
            states,
        })
    }

    fn step(
        &mut self,
        label: usize,
        frame: usize,
        selector: &impl LabelSelector,
    ) -> Result<Decision> {
        if label >= self.assets.metadata.embedding_rows || frame >= self.encoded.valid_frames {
            return Err(Error::Invalid(format!(
                "Invalid Haas step: label={label}, frame={frame}, valid_frames={}",
                self.encoded.valid_frames
            )));
        }
        let width = self.assets.metadata.embedding_width;
        if self.inference.needs_embedding() {
            self.embedding.replace(
                self.assets.embeddings[label * width..(label + 1) * width]
                    .iter()
                    .copied(),
            )?;
        }
        self.frame
            .replace((0..1024).map(|channel| {
                self.encoded.values[channel * self.encoded.storage_frames + frame]
            }))?;
        self.inference.run(&self.states, &self.outputs)?;
        self.outputs
            .read_selected(self.assets.metadata.blank_id, selector)
    }
}

fn first_argmax(values: &[f32]) -> Result<usize> {
    if values.is_empty() || values.iter().any(|v| !v.is_finite()) {
        return Err(Error::Invalid("Non-finite or empty joint logits".into()));
    }
    Ok((1..values.len()).fold(0, |best, i| if values[i] > values[best] { i } else { best }))
}
/// Posterior over label classes only; TDT duration logits have a separate distribution.
pub(super) fn label_probability(logits: &[f32], label: usize) -> Result<f32> {
    let maximum = logits[first_argmax(logits)?] as f64;
    let selected = logits.get(label).ok_or_else(|| {
        Error::Invalid(format!(
            "Label probability index out of bounds: label={label}, classes={}",
            logits.len()
        ))
    })?;
    let total: f64 = logits.iter().map(|v| (*v as f64 - maximum).exp()).sum();
    Ok(((*selected as f64 - maximum).exp() / total) as f32)
}

pub(super) fn decode(
    models: &DecoderModels,
    encoded: &Encoded,
    assets: &Assets,
    cancellation: &Cancellation,
) -> Result<Vec<TimedToken>> {
    decode_selected(models, encoded, assets, cancellation, &mut Greedy)
}

trait LabelSelector {
    fn select(&self, logits: &[f32], blank: usize) -> Result<usize>;
    fn accept(&mut self, token: usize);
}
struct Greedy;
impl LabelSelector for Greedy {
    fn select(&self, logits: &[f32], _: usize) -> Result<usize> {
        first_argmax(logits)
    }
    fn accept(&mut self, _: usize) {}
}
impl LabelSelector for BiasState<'_> {
    fn select(&self, logits: &[f32], blank: usize) -> Result<usize> {
        BiasState::select(self, logits, blank)
    }
    fn accept(&mut self, token: usize) {
        BiasState::accept(self, token);
    }
}
pub(super) fn decode_biased(
    models: &DecoderModels,
    encoded: &Encoded,
    assets: &Assets,
    vocabulary: &CompiledVocabulary,
    cancellation: &Cancellation,
) -> Result<Vec<TimedToken>> {
    models.require_logits()?;
    decode_selected(
        models,
        encoded,
        assets,
        cancellation,
        &mut BiasState::new(vocabulary),
    )
}
fn decode_selected(
    models: &DecoderModels,
    encoded: &Encoded,
    assets: &Assets,
    cancellation: &Cancellation,
    selector: &mut impl LabelSelector,
) -> Result<Vec<TimedToken>> {
    let mut workspace = DecodeWorkspace::new(models, encoded, assets)?;
    let mut label = assets.metadata.blank_id;
    let (mut frame, mut symbols, mut iterations) = (0usize, 0usize, 0usize);
    let mut tokens = Vec::new();
    while frame < encoded.valid_frames {
        cancellation.check()?;
        iterations += 1;
        if iterations > encoded.valid_frames * 11 {
            return Err(Error::Invalid(format!(
                "TDT failed to advance: frame={frame}, iterations={iterations}"
            )));
        }
        let decision = autoreleasepool(|_| workspace.step(label, frame, selector))?;
        workspace.states.accept(decision.label);
        workspace
            .inference
            .accept(decision.label, assets.metadata.blank_id);
        if decision.label == assets.metadata.blank_id {
            frame += decision.duration.max(1);
            symbols = 0;
            continue;
        }
        tokens.push(TimedToken {
            token_id: decision.label,
            start_sample: (frame * ENCODER_STEP) as u64,
            confidence: decision.confidence,
            duration_samples: (decision.duration * ENCODER_STEP) as u64,
        });
        selector.accept(decision.label);
        label = decision.label;
        if decision.duration > 0 {
            frame += decision.duration;
            symbols = 0;
        } else {
            symbols += 1;
            if symbols >= 10 {
                frame += 1;
                symbols = 0;
            }
        }
    }
    Ok(tokens)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tensor::read_array;

    #[test]
    fn scalar_decisions_refresh_all_fields_and_reject_invalid_outputs() {
        autoreleasepool(|_| {
            let mut output = DecisionOutputs::new().unwrap();
            output.token.replace([37].into_iter()).unwrap();
            output.probability.replace([0.625].into_iter()).unwrap();
            output.duration.replace([4].into_iter()).unwrap();
            let decision = output.read(3072).unwrap();
            assert_eq!(
                (decision.label, decision.duration, decision.confidence),
                (37, 4, 0.625)
            );
            output.token.replace([3072].into_iter()).unwrap();
            assert_eq!(output.read(3072).unwrap().confidence, 0.0);
            for (token, probability, duration) in [
                (3073, 0.5, 0),
                (-1, 0.5, 0),
                (1, f32::NAN, 0),
                (1, 1.1, 0),
                (1, 0.5, -1),
                (1, 0.5, 5),
            ] {
                output.token.replace([token].into_iter()).unwrap();
                output
                    .probability
                    .replace([probability].into_iter())
                    .unwrap();
                output.duration.replace([duration].into_iter()).unwrap();
                assert!(output.read(3072).is_err());
            }
            output.token.replace([12].into_iter()).unwrap();
            output.probability.replace([0.75].into_iter()).unwrap();
            output.duration.replace([2].into_iter()).unwrap();
            let decision = output.read(3072).unwrap();
            assert_eq!(
                (decision.label, decision.duration, decision.confidence),
                (12, 2, 0.75)
            );
        });
    }

    #[test]
    fn prediction_runs_once_across_blanks_and_again_for_each_emission() {
        let mut cache = PredictionCache::Empty;
        let mut calls = 0;
        let decisions = [3072, 3072, 12, 3072, 3072, 12, 3072];
        let expected_calls = [1, 1, 1, 2, 2, 2, 3];
        for (token, expected) in decisions.into_iter().zip(expected_calls) {
            cache
                .compute(|| {
                    calls += 1;
                    Ok(())
                })
                .unwrap();
            assert_eq!(calls, expected);
            cache.accept(token, 3072);
        }
        assert_eq!(calls, 3);
    }

    #[test]
    fn failed_prediction_and_new_utterance_cannot_reuse_cached_state() {
        let mut failed = PredictionCache::Empty;
        assert!(
            failed
                .compute(|| Err(Error::Invalid("Inference failed".into())))
                .is_err()
        );
        assert!(failed.needs_prediction());
        failed.compute(|| Ok(())).unwrap();
        assert!(!failed.needs_prediction());
        let mut next_utterance = PredictionCache::Empty;
        let mut calls = 0;
        next_utterance
            .compute(|| {
                calls += 1;
                Ok(())
            })
            .unwrap();
        assert_eq!(calls, 1);
        assert!(!next_utterance.needs_prediction());
    }

    #[test]
    fn blank_keeps_accepted_state_and_nonblank_commits_pending_state() {
        autoreleasepool(|_| {
            let mut states = StatePair::zero(3072).unwrap();
            states.slots[1]
                .memory
                .replace(std::iter::repeat_n(7.0, 2 * 5 * 64 * 64))
                .unwrap();
            states.slots[1]
                .conv
                .replace(std::iter::repeat_n(3.0, 2 * 3 * 320 * 3))
                .unwrap();
            states.accept(3072);
            assert_eq!(states.active, 0);
            assert!(
                read_array::<f32>(states.slots[states.active].memory.array())
                    .unwrap()
                    .iter()
                    .all(|&v| v == 0.0)
            );
            states.accept(10);
            assert_eq!(states.active, 1);
            assert!(
                read_array::<f32>(states.slots[states.active].memory.array())
                    .unwrap()
                    .iter()
                    .all(|&v| v == 7.0)
            );
            assert!(
                read_array::<f32>(states.slots[states.active].conv.array())
                    .unwrap()
                    .iter()
                    .all(|&v| v == 3.0)
            );
            states.slots[0]
                .memory
                .replace(std::iter::repeat_n(9.0, 2 * 5 * 64 * 64))
                .unwrap();
            states.accept(3072);
            assert!(
                read_array::<f32>(states.slots[states.active].memory.array())
                    .unwrap()
                    .iter()
                    .all(|&v| v == 7.0)
            );
            states.accept(11);
            assert!(
                read_array::<f32>(states.slots[states.active].memory.array())
                    .unwrap()
                    .iter()
                    .all(|&v| v == 9.0)
            );
        });
    }
}