parakeet-rs 0.4.2

Fast ASR & Speaker Diarization with NVIDIA Parakeet via ONNX
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

//! Parakeet Unified (offline + buffered streaming RNN-T) on burn.

use super::generated::{
    unified_decoder_joint, unified_decoder_joint_weights, unified_encoder, unified_encoder_weights,
};
use super::{array3, guard, ints, onnx, tensor3, vec_f32, vec_i64};
use crate::error::{Error, Result};
use crate::execution::ExecutionProvider;
use burn::tensor::Device;
use ndarray::{Array1, Array3, ArrayView3};
use std::path::Path;

pub(crate) struct UnifiedEncoder {
    model: unified_encoder::Model,
    device: Device,
}

impl UnifiedEncoder {
    pub(crate) fn load(path: &Path, provider: ExecutionProvider) -> Result<Self> {
        let device = super::device(provider)?;
        let file = onnx::OnnxFile::open(path)?;
        let mut model = guard("allocating the model", || {
            unified_encoder::Model::new(&device)
        })?;
        file.load(&mut model, unified_encoder_weights::WEIGHTS)?;
        Ok(Self { model, device })
    }

    /// `features` is `[1, n_mels, time]`; returns the encoder output `[1, dim, frames]` and the
    /// number of valid frames.
    pub(crate) fn run(&self, features: ArrayView3<f32>) -> Result<(Array3<f32>, i64)> {
        let time = features.dim().2;
        let (out, len) = guard("running the Unified encoder", || {
            let (out, len) = self.model.forward(
                tensor3(features, &self.device),
                ints(vec![time as i64], [1], &self.device),
            );
            (array3(out), vec_i64(len))
        })?;
        let len = *len?
            .first()
            .ok_or_else(|| Error::Model("empty encoder length".into()))?;
        Ok((out?, len))
    }
}

pub(crate) struct UnifiedDecoderJoint {
    model: unified_decoder_joint::Model,
    device: Device,
}

impl UnifiedDecoderJoint {
    pub(crate) fn load(path: &Path, provider: ExecutionProvider) -> Result<Self> {
        let device = super::device(provider)?;
        let file = onnx::OnnxFile::open(path)?;
        let mut model = guard("allocating the model", || {
            unified_decoder_joint::Model::new(&device)
        })?;
        file.load(&mut model, unified_decoder_joint_weights::WEIGHTS)?;
        Ok(Self { model, device })
    }

    /// One decoder/joint step: flat logits and the new LSTM states.
    pub(crate) fn step(
        &self,
        frame: ArrayView3<f32>,
        token: i32,
        state_1: ArrayView3<f32>,
        state_2: ArrayView3<f32>,
    ) -> Result<(Array1<f32>, Array3<f32>, Array3<f32>)> {
        guard("running the Unified decoder/joint", || {
            let d = &self.device;
            let (logits, _, s1, s2) = self.model.forward(
                tensor3(frame, d),
                ints(vec![i64::from(token)], [1, 1], d),
                ints(vec![1], [1], d),
                tensor3(state_1, d),
                tensor3(state_2, d),
            );
            Ok((Array1::from_vec(vec_f32(logits)?), array3(s1)?, array3(s2)?))
        })?
    }
}