Fangjun Kuang
Committed by GitHub

test exported sense voice models (#1147)

@@ -158,6 +158,7 @@ def main(): @@ -158,6 +158,7 @@ def main():
158 meta_data = { 158 meta_data = {
159 "lfr_window_size": lfr_window_size, 159 "lfr_window_size": lfr_window_size,
160 "lfr_window_shift": lfr_window_shift, 160 "lfr_window_shift": lfr_window_shift,
  161 + "normalize_samples": 0, # input should be in the range [-32768, 32767]
161 "neg_mean": neg_mean, 162 "neg_mean": neg_mean,
162 "inv_stddev": inv_stddev, 163 "inv_stddev": inv_stddev,
163 "model_type": "sense_voice_ctc", 164 "model_type": "sense_voice_ctc",
@@ -35,3 +35,28 @@ echo "pwd: $PWD" @@ -35,3 +35,28 @@ echo "pwd: $PWD"
35 ./show-info.py 35 ./show-info.py
36 36
37 ls -lh 37 ls -lh
  38 +
  39 +# Download test wavs
  40 +curl -SL -O https://huggingface.co/csukuangfj/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17/resolve/main/test_wavs/en.wav
  41 +curl -SL -O https://huggingface.co/csukuangfj/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17/resolve/main/test_wavs/zh.wav
  42 +curl -SL -O https://huggingface.co/csukuangfj/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17/resolve/main/test_wavs/ja.wav
  43 +curl -SL -O https://huggingface.co/csukuangfj/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17/resolve/main/test_wavs/ko.wav
  44 +curl -SL -O https://huggingface.co/csukuangfj/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17/resolve/main/test_wavs/yue.wav
  45 +
  46 +for m in model.onnx model.int8.onnx; do
  47 + for w in en zh ja ko yue; do
  48 + echo "----------test $m $w.wav----------"
  49 +
  50 + echo "without inverse text normalization, lang auto"
  51 + ./test.py --model $m --tokens ./tokens.txt --wav $w.wav --use-itn 0
  52 +
  53 + echo "with inverse text normalization, lang auto"
  54 + ./test.py --model $m --tokens ./tokens.txt --wav $w.wav --use-itn 1
  55 +
  56 + echo "without inverse text normalization, lang $w"
  57 + ./test.py --model $m --tokens ./tokens.txt --wav $w.wav --use-itn 0 --lang $w
  58 +
  59 + echo "with inverse text normalization, lang $w"
  60 + ./test.py --model $m --tokens ./tokens.txt --wav $w.wav --use-itn 1 --lang $w
  61 + done
  62 +done
  1 +#!/usr/bin/env python3
  2 +# Copyright 2024 Xiaomi Corp. (authors: Fangjun Kuang)
  3 +
  4 +import argparse
  5 +from typing import Tuple
  6 +
  7 +import kaldi_native_fbank as knf
  8 +import numpy as np
  9 +import onnxruntime
  10 +import onnxruntime as ort
  11 +import soundfile as sf
  12 +import torch
  13 +
  14 +
  15 +def get_args():
  16 + parser = argparse.ArgumentParser(
  17 + formatter_class=argparse.ArgumentDefaultsHelpFormatter
  18 + )
  19 +
  20 + parser.add_argument(
  21 + "--model",
  22 + type=str,
  23 + required=True,
  24 + help="Path to model.onnx",
  25 + )
  26 +
  27 + parser.add_argument(
  28 + "--tokens",
  29 + type=str,
  30 + required=True,
  31 + help="Path to tokens.txt",
  32 + )
  33 +
  34 + parser.add_argument(
  35 + "--wave",
  36 + type=str,
  37 + required=True,
  38 + help="The input wave to be recognized",
  39 + )
  40 +
  41 + parser.add_argument(
  42 + "--language",
  43 + type=str,
  44 + default="auto",
  45 + help="the language of the input wav file. Supported values: zh, en, ja, ko, yue, auto",
  46 + )
  47 +
  48 + parser.add_argument(
  49 + "--use-itn",
  50 + type=int,
  51 + default=0,
  52 + help="1 to use inverse text normalization. 0 to not use inverse text normalization",
  53 + )
  54 +
  55 + return parser.parse_args()
  56 +
  57 +
  58 +class OnnxModel:
  59 + def __init__(self, filename):
  60 + session_opts = ort.SessionOptions()
  61 + session_opts.inter_op_num_threads = 1
  62 + session_opts.intra_op_num_threads = 1
  63 +
  64 + self.session_opts = session_opts
  65 +
  66 + self.model = ort.InferenceSession(
  67 + filename,
  68 + sess_options=self.session_opts,
  69 + providers=["CPUExecutionProvider"],
  70 + )
  71 +
  72 + meta = self.model.get_modelmeta().custom_metadata_map
  73 +
  74 + self.window_size = int(meta["lfr_window_size"]) # lfr_m
  75 + self.window_shift = int(meta["lfr_window_shift"]) # lfr_n
  76 +
  77 + lang_zh = int(meta["lang_zh"])
  78 + lang_en = int(meta["lang_en"])
  79 + lang_ja = int(meta["lang_ja"])
  80 + lang_ko = int(meta["lang_ko"])
  81 + lang_auto = int(meta["lang_auto"])
  82 +
  83 + self.lang_id = {
  84 + "zh": lang_zh,
  85 + "en": lang_en,
  86 + "ja": lang_ja,
  87 + "ko": lang_ko,
  88 + "auto": lang_auto,
  89 + }
  90 + self.with_itn = int(meta["with_itn"])
  91 + self.without_itn = int(meta["without_itn"])
  92 +
  93 + neg_mean = meta["neg_mean"].split(",")
  94 + neg_mean = list(map(lambda x: float(x), neg_mean))
  95 +
  96 + inv_stddev = meta["inv_stddev"].split(",")
  97 + inv_stddev = list(map(lambda x: float(x), inv_stddev))
  98 +
  99 + self.neg_mean = np.array(neg_mean, dtype=np.float32)
  100 + self.inv_stddev = np.array(inv_stddev, dtype=np.float32)
  101 +
  102 + def __call__(self, x, x_length, language, text_norm):
  103 + logits = self.model.run(
  104 + [
  105 + self.model.get_outputs()[0].name,
  106 + ],
  107 + {
  108 + self.model.get_inputs()[0].name: x.numpy(),
  109 + self.model.get_inputs()[1].name: x_length.numpy(),
  110 + self.model.get_inputs()[2].name: language.numpy(),
  111 + self.model.get_inputs()[3].name: text_norm.numpy(),
  112 + },
  113 + )[0]
  114 +
  115 + return torch.from_numpy(logits)
  116 +
  117 +
  118 +def load_audio(filename: str) -> Tuple[np.ndarray, int]:
  119 + data, sample_rate = sf.read(
  120 + filename,
  121 + always_2d=True,
  122 + dtype="float32",
  123 + )
  124 + data = data[:, 0] # use only the first channel
  125 + samples = np.ascontiguousarray(data)
  126 + return samples, sample_rate
  127 +
  128 +
  129 +def load_tokens(filename):
  130 + ans = dict()
  131 + i = 0
  132 + with open(filename, encoding="utf-8") as f:
  133 + for line in f:
  134 + ans[i] = line.strip().split()[0]
  135 + i += 1
  136 + return ans
  137 +
  138 +
  139 +def compute_feat(
  140 + samples,
  141 + sample_rate,
  142 + neg_mean: np.ndarray,
  143 + inv_stddev: np.ndarray,
  144 + window_size: int = 7, # lfr_m
  145 + window_shift: int = 6, # lfr_n
  146 +):
  147 + opts = knf.FbankOptions()
  148 + opts.frame_opts.dither = 0
  149 + opts.frame_opts.snip_edges = False
  150 + opts.frame_opts.window_type = "hamming"
  151 + opts.frame_opts.samp_freq = sample_rate
  152 + opts.mel_opts.num_bins = 80
  153 +
  154 + online_fbank = knf.OnlineFbank(opts)
  155 + online_fbank.accept_waveform(sample_rate, (samples * 32768).tolist())
  156 + online_fbank.input_finished()
  157 +
  158 + features = np.stack(
  159 + [online_fbank.get_frame(i) for i in range(online_fbank.num_frames_ready)]
  160 + )
  161 + assert features.data.contiguous is True
  162 + assert features.dtype == np.float32, features.dtype
  163 +
  164 + T = (features.shape[0] - window_size) // window_shift + 1
  165 + features = np.lib.stride_tricks.as_strided(
  166 + features,
  167 + shape=(T, features.shape[1] * window_size),
  168 + strides=((window_shift * features.shape[1]) * 4, 4),
  169 + )
  170 +
  171 + features = (features + neg_mean) * inv_stddev
  172 +
  173 + return features
  174 +
  175 +
  176 +def main():
  177 + args = get_args()
  178 + print(vars(args))
  179 + samples, sample_rate = load_audio(args.wave)
  180 + if sample_rate != 16000:
  181 + import librosa
  182 +
  183 + samples = librosa.resample(samples, orig_sr=sample_rate, target_sr=16000)
  184 + sample_rate = 16000
  185 +
  186 + model = OnnxModel(filename=args.model)
  187 +
  188 + features = compute_feat(
  189 + samples=samples,
  190 + sample_rate=sample_rate,
  191 + neg_mean=model.neg_mean,
  192 + inv_stddev=model.inv_stddev,
  193 + window_size=model.window_size,
  194 + window_shift=model.window_shift,
  195 + )
  196 +
  197 + features = torch.from_numpy(features).unsqueeze(0)
  198 + features_length = torch.tensor([features.size(1)], dtype=torch.int32)
  199 +
  200 + language = model.lang_id["auto"]
  201 + if args.language in model.lang_id:
  202 + language = model.lang_id[args.language]
  203 + else:
  204 + print(f"Invalid language: '{args.language}'")
  205 + print("Use auto")
  206 +
  207 + if args.use_itn:
  208 + text_norm = model.with_itn
  209 + else:
  210 + text_norm = model.without_itn
  211 +
  212 + language = torch.tensor([language], dtype=torch.int32)
  213 + text_norm = torch.tensor([text_norm], dtype=torch.int32)
  214 +
  215 + logits = model(
  216 + x=features,
  217 + x_length=features_length,
  218 + language=language,
  219 + text_norm=text_norm,
  220 + )
  221 +
  222 + idx = logits.squeeze(0).argmax(dim=-1)
  223 + # idx is of shape (T,)
  224 + idx = torch.unique_consecutive(idx)
  225 +
  226 + blank_id = 0
  227 + idx = idx[idx != blank_id].tolist()
  228 +
  229 + tokens = load_tokens(args.tokens)
  230 + text = "".join([tokens[i] for i in idx])
  231 +
  232 + text = text.replace("▁", " ")
  233 + print(text)
  234 +
  235 +
  236 +if __name__ == "__main__":
  237 + main()