svod-model 0.1.0-alpha.3

Pretrained models inference abstraction.
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

//! Shared encoder-only JIT for [`GigaAm`]. Output is cast to fp32 so
//! the head-side path (CTC log-probs computed by `CtcHeadJit`, RN-T frames
//! consumed by the predictor/joint step JITs) sees a uniform dtype regardless
//! of whether the encoder ran in fp16, bf16, or fp32.
//!
//! The `jit_wrapper!` macro expands to `svod_model::jit::*` paths, so this
//! file needs the `extern crate self as svod_model;` binding in scope.

extern crate self as svod_model;

use snafu::ResultExt;
use svod_macros::jit_wrapper;

use super::model::GigaAm;
use crate::gigaam::error::TensorSnafu;

jit_wrapper! {
    GigaAmEncoderJit(GigaAm) {
        mel: Tensor,
        lengths: Tensor,

        vars {
            b: (1, model.config.max_batch_size),
            t: (1, model.config.max_mel_frames),
        }

        build(mel, lengths, b, t) {
            let out = model.encoder.forward_batch(mel, lengths, &b, &t)?;
            out.cast(svod_dtype::DType::Float32).context(TensorSnafu)
        }
    }
}