parakeet-rs 0.4.3

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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

//! Multitalker (speaker-conditioned cache-aware streaming ASR) on burn.

use super::generated::{
    multitalker_decoder_joint, multitalker_decoder_joint_weights, multitalker_encoder,
    multitalker_encoder_weights,
};
use super::{array3, array4, guard, ints, onnx, tensor2, tensor3, tensor4, vec_f32, vec_i64};
use crate::error::{Error, Result};
use crate::execution::ExecutionProvider;
use burn::tensor::Device;
use ndarray::{Array1, Array3, Array4, ArrayView2, ArrayView3, ArrayView4};
use std::path::Path;

pub(crate) struct MultitalkerEncoder {
    model: multitalker_encoder::Model,
    device: Device,
}

/// One streaming encoder step: output `[1, hidden, frames]`, its length, and the new caches.
pub(crate) type EncoderStep = (Array3<f32>, i64, Array4<f32>, Array4<f32>, Array1<i64>);

impl MultitalkerEncoder {
    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", || {
            multitalker_encoder::Model::new(&device)
        })?;
        file.load(&mut model, multitalker_encoder_weights::WEIGHTS)?;
        Ok(Self { model, device })
    }

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn run(
        &self,
        features: ArrayView3<f32>,
        length: i64,
        cache_last_channel: ArrayView4<f32>,
        cache_last_time: ArrayView4<f32>,
        cache_last_channel_len: &Array1<i64>,
        spk_targets: ArrayView2<f32>,
        bg_spk_targets: ArrayView2<f32>,
    ) -> Result<EncoderStep> {
        guard("running the Multitalker encoder", || {
            let d = &self.device;
            let (encoded, len, channel, time, channel_len) = self.model.forward(
                tensor3(features, d),
                ints(vec![length], [1], d),
                tensor4(cache_last_channel, d),
                tensor4(cache_last_time, d),
                ints(
                    cache_last_channel_len.to_vec(),
                    [cache_last_channel_len.len()],
                    d,
                ),
                tensor2(spk_targets, d),
                tensor2(bg_spk_targets, d),
            );
            let len = *vec_i64(len)?
                .first()
                .ok_or_else(|| Error::Model("empty encoded_len".into()))?;
            Ok((
                array3(encoded)?,
                len,
                array4(channel)?,
                array4(time)?,
                Array1::from_vec(vec_i64(channel_len)?),
            ))
        })?
    }
}

pub(crate) struct MultitalkerDecoderJoint {
    model: multitalker_decoder_joint::Model,
    device: Device,
}

impl MultitalkerDecoderJoint {
    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", || {
            multitalker_decoder_joint::Model::new(&device)
        })?;
        file.load(&mut model, multitalker_decoder_joint_weights::WEIGHTS)?;
        Ok(Self { model, device })
    }

    /// One decoder/joint step for an encoder frame `[1, 1, hidden]`: logits and new 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 Multitalker decoder/joint", || {
            let d = &self.device;
            let (logits, _, s1, s2) = self.model.forward(
                tensor3(frame, d),
                ints(vec![i64::from(token)], [1, 1], d),
                tensor3(state_1, d),
                tensor3(state_2, d),
            );
            Ok((Array1::from_vec(vec_f32(logits)?), array3(s1)?, array3(s2)?))
        })?
    }
}