onnx-asr/preprocessors/build.py at ffc1ed0b170e658eba6f5b1aa8960ef7eb38e4dc · istupakov/onnx-asr · GitHub

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
"""Build ONNX and NumPy preprocessors."""

from pathlib import Path

import numpy as np
import onnxscript

from preprocessors import gigaam, kaldi, nemo, resample, whisper


def save_onnx(
    function: onnxscript.OnnxFunction, filename: Path, version: str, input_size_limit: int = 1024 * 1024
) -> None:
    model = onnxscript.ir.from_proto(function.to_model_proto())

    model = onnxscript.optimizer.optimize(model, input_size_limit=input_size_limit)

    model.producer_name = "OnnxScript"
    model.producer_version = onnxscript.__version__
    model.metadata_props["model_author"] = "Ilya Stupakov"
    model.metadata_props["model_license"] = "MIT License"
    model.metadata_props["model_version"] = f"onnx-asr {version}"

    onnxscript.ir.save(model, filename)


def save_preprocessor_models(preprocessors_dir: Path, version: str) -> None:
    preprocessors = {
        "gigaam_v2.onnx": gigaam.GigaamPreprocessorV2,
        "gigaam_v3.onnx": gigaam.GigaamPreprocessorV3,
        "kaldi.onnx": kaldi.KaldiPreprocessorFast,
        "wespeaker.onnx": kaldi.WespeakerPreprocessor,
        "nemo80.onnx": nemo.NemoPreprocessor80,
        "nemo128.onnx": nemo.NemoPreprocessor128,
        "whisper80.onnx": whisper.WhisperPreprocessor80,
        "whisper128.onnx": whisper.WhisperPreprocessor128,
        "gigaam_v2_conv.onnx": gigaam.GigaamPreprocessorV2Conv,
        "gigaam_v3_conv.onnx": gigaam.GigaamPreprocessorV3Conv,
        "kaldi_conv.onnx": kaldi.KaldiPreprocessorFastConv,
        "nemo80_conv.onnx": nemo.NemoPreprocessor80Conv,
        "nemo128_conv.onnx": nemo.NemoPreprocessor128Conv,
        "whisper80_conv.onnx": whisper.WhisperPreprocessor80Conv,
        "whisper128_conv.onnx": whisper.WhisperPreprocessor128Conv,
    }
    for filename, model in preprocessors.items():
        save_onnx(model, preprocessors_dir.joinpath(filename), version)


def save_resampler_models(preprocessors_dir: Path, version: str) -> None:
    for orig_freq in [8_000, 11_025, 16_000, 22_050, 24_000, 32_000, 44_100, 48_000]:
        for new_freq in [8_000, 16_000]:
            if orig_freq != new_freq:
                save_onnx(
                    resample.create_resampler(orig_freq, new_freq),
                    preprocessors_dir.joinpath(f"resample_{orig_freq // 1000}_{new_freq // 1000}.onnx"),
                    version,
                    100,
                )


def save_fbanks(preprocessors_dir: Path) -> None:
    fbanks = {
        "gigaam_v2": gigaam.melscale_fbanks_v2,
        "gigaam_v3": gigaam.melscale_fbanks_v3,
        "gigaam_v3_window": gigaam.hann_window_v3,
        "kaldi": kaldi.kaldi_mel_banks,
        "wespeaker": kaldi.wespeaker_mel_banks,
        "nemo80": nemo.melscale_fbanks80,
        "nemo128": nemo.melscale_fbanks128,
        "whisper80": whisper.melscale_fbanks80,
        "whisper128": whisper.melscale_fbanks128,
    }
    np.savez_compressed(Path(preprocessors_dir, "fbanks"), allow_pickle=False, **fbanks)


def build(preprocessors_dir: Path, version: str) -> None:
    save_preprocessor_models(preprocessors_dir, version)
    save_resampler_models(preprocessors_dir, version)
    save_fbanks(preprocessors_dir)