尝试解决fp16溢出问题

README.md

@@ -28,12 +28,10 @@ SenseVoice是具有音频理解能力的音频基础模型， 包括语音识别
 2. 安装依赖
 
 ```bash
-pip install kaldi_native_fbank onnxruntime sentencepiece soundfile pyyaml numpy<2
+pip install kaldi_native_fbank onnxruntime sentencepiece soundfile pyyaml "numpy<2" rknn-toolkit-lite2
 ```
 
-你还需要手动安装rknn-toolkit2-lite2. 
-
-3. 运行
+1. 运行
 
 ```bash
 python ./sensevoice_rknn.py --audio_file output.wav
@@ -47,7 +45,7 @@ ffmpeg -i input.mp3 -f wav -acodec pcm_s16le -ac 1 -ar 16000 output.wav
 
 ## RKNN模型转换
 
-你需要提前安装rknn-toolkit2 v2.1.0或更高版本. 
+你需要提前安装rknn-toolkit2, 测试可用的版本为2.3.3a25，可从https://console.zbox.filez.com/l/I00fc3 下载(密码为"rknn")
 
 1. 下载或转换onnx模型
 
@@ -58,13 +56,13 @@ ffmpeg -i input.mp3 -f wav -acodec pcm_s16le -ac 1 -ar 16000 output.wav
 
 2. 转换为rknn模型
 ```bash
-python convert_rknn.py 
+python convert_rknn.py ./sense-voice-encoder.onnx
 ```
 
 ## 已知问题
 
-- RKNN2使用fp16推理时可能会出现溢出，导致结果为inf，可以尝试修改输入数据的缩放比例来解决.  
-  在`sensevoice_rknn.py`中将`SPEECH_SCALE`设置为更小的值.
+- ~~RKNN2使用fp16推理时可能会出现溢出，导致结果为inf，可以尝试修改输入数据的缩放比例来解决.  
+  在`sensevoice_rknn.py`中将`SPEECH_SCALE`设置为更小的值.~~ (现在应该已经通过模型内部插入缩放算子解决了)
 
 ## 参考
 - [FunAudioLLM/SenseVoiceSmall](https://huggingface.co/FunAudioLLM/SenseVoiceSmall)
@@ -88,11 +86,9 @@ Currently, SenseVoice-small supports multilingual speech recognition, emotion re
 2. Install dependencies
 
 ```bash
-pip install kaldi_native_fbank onnxruntime sentencepiece soundfile pyyaml numpy<2
+pip install kaldi_native_fbank onnxruntime sentencepiece soundfile pyyaml "numpy<2" rknn-toolkit-lite2
 ```
 
-You also need to manually install rknn-toolkit2-lite2.
-
 3. Run
 
 ```bash
@@ -107,7 +103,7 @@ ffmpeg -i input.mp3 -f wav -acodec pcm_s16le -ac 1 -ar 16000 output.wav
 
 ## RKNN Model Conversion
 
-You need to install rknn-toolkit2 v2.1.0 or higher in advance.
+You need to install rknn-toolkit2 in advance. The tested working version is 2.3.3a25, which can be downloaded from https://console.zbox.filez.com/l/I00fc3 (password: "rknn").
 
 1. Download or convert the ONNX model
 
@@ -118,13 +114,13 @@ The model file should be named 'sense-voice-encoder.onnx' and placed in the same
 
 2. Convert to RKNN model
 ```bash
-python convert_rknn.py 
+python convert_rknn.py ./sense-voice-encoder.onnx
 ```
 
 ## Known Issues
 
-- When using fp16 inference with RKNN2, overflow may occur, resulting in inf values. You can try modifying the scaling ratio of the input data to resolve this.  
-  Set `SPEECH_SCALE` to a smaller value in `sensevoice_rknn.py`.
+- ~~When using fp16 inference with RKNN2, overflow may occur, resulting in inf values. You can try modifying the scaling ratio of the input data to resolve this.  
+  Set `SPEECH_SCALE` to a smaller value in `sensevoice_rknn.py`.~~ (This issue should now be resolved by inserting scaling operators inside the model.)
 
 ## References
 - [FunAudioLLM/SenseVoiceSmall](https://huggingface.co/FunAudioLLM/SenseVoiceSmall)

convert_rknn.py

@@ -2,6 +2,8 @@
 # coding: utf-8
 
 import os
+import re
+from typing import Optional, Set
 from rknn.api import RKNN
 from math import exp
 from sys import exit
@@ -18,28 +20,158 @@ os.chdir(os.path.dirname(os.path.abspath(__file__)))
 
 speech_length = 171
 
-def convert_encoder():
+def _remove_file(path: str, *, keep: Optional[Set[str]] = None) -> None:
+    if not path:
+        return
+    keep_paths: Set[str] = {os.path.abspath(item) for item in keep} if keep else set()
+    normalized = os.path.abspath(path)
+    if keep_paths and normalized in keep_paths:
+        return
+    if not os.path.exists(normalized):
+        return
+    try:
+        os.remove(normalized)
+        print(f'cleaned temp model: {normalized}')
+    except OSError as err:
+        print(f'warning: failed to remove {normalized}: {err}')
+
+def _with_suffix(path: str, suffix: str) -> str:
+    stem, ext = os.path.splitext(path)
+    return f"{stem}{suffix}{ext}"
+
+def _sanitize_name(name: str) -> str:
+    return re.sub(r'[^0-9A-Za-z_]', '_', name)
+
+def _insert_div_node(model: onnx.ModelProto, tensor_name: str, divisor: float = 16.0) -> bool:
+    graph = model.graph
+
+    for node in graph.node:
+        if node.op_type == 'Div' and tensor_name in node.output:
+            return False
+
+    producer_index = None
+    output_index = None
+    for idx, node in enumerate(graph.node):
+        for out_idx, output in enumerate(node.output):
+            if output == tensor_name:
+                producer_index = idx
+                output_index = out_idx
+                producer_node = node
+                break
+        if producer_index is not None:
+            break
+
+    if producer_index is None:
+        raise RuntimeError(f"Producer node for tensor {tensor_name} not found.")
+
+    pre_div_output = f"{tensor_name}_pre_div"
+    producer_node.output[output_index] = pre_div_output
+
+    sanitized = _sanitize_name(tensor_name)
+    const_output = f"{sanitized}_div_const"
+    const_node_name = f"{sanitized}_DivConst"
+    div_node_name = f"{sanitized}_Div"
+
+    const_tensor = onnx.helper.make_tensor(
+        name=f"{const_node_name}_value",
+        data_type=onnx.TensorProto.FLOAT,
+        dims=[],
+        vals=[divisor],
+    )
+
+    const_node = onnx.helper.make_node(
+        'Constant',
+        inputs=[],
+        outputs=[const_output],
+        value=const_tensor,
+        name=const_node_name,
+    )
+
+    div_node = onnx.helper.make_node(
+        'Div',
+        inputs=[pre_div_output, const_output],
+        outputs=[tensor_name],
+        name=div_node_name,
+    )
+
+    graph.node.insert(producer_index + 1, const_node)
+    graph.node.insert(producer_index + 2, div_node)
+    return True
+
+def _scale_initializer(model: onnx.ModelProto, initializer_name: str, divisor: float = 16.0) -> bool:
+    for idx, initializer in enumerate(model.graph.initializer):
+        if initializer.name == initializer_name:
+            data = onh.to_array(initializer).astype(np.float32, copy=False)
+            scaled = data / divisor
+            model.graph.initializer[idx].CopyFrom(onh.from_array(scaled, name=initializer_name))
+            return True
+    return False
+
+def convert_encoder(model_path: str):
     rknn = RKNN(verbose=True)
 
-    ONNX_MODEL=f"sense-voice-encoder.onnx"
-    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
-    DATASET="dataset.txt"
-    QUANTIZE=False
+    ONNX_MODEL = os.path.abspath(model_path)
+    if not os.path.isfile(ONNX_MODEL):
+        print(f'Model file not found: {model_path}')
+        exit(1)
+    if not ONNX_MODEL.lower().endswith('.onnx'):
+        print(f'Model file must be an ONNX file: {model_path}')
+        exit(1)
+
+    RKNN_MODEL = os.path.splitext(ONNX_MODEL)[0] + ".rknn"
+    DATASET = "dataset.txt"
+    QUANTIZE = False
+    original_model = ONNX_MODEL
+    preserve_files: Set[str] = {original_model}
+
 
+    print('--> Patching model to avoid overflow issue')
+    base_model = onnx.load(ONNX_MODEL)
+    modified = False
+    for layer_idx in range(48, 49): 
+        for target in [
+            f'/encoders.{layer_idx}/feed_forward/activation/Relu_output_0',
+            f'/encoders.{layer_idx}/norm2/Cast_output_0',
+        ]:
+            modified |= _insert_div_node(base_model, target, divisor=2.0)
+    bias_scaled = False
+    if modified:
+        for layer_idx in range(48, 49): 
+            bias_scaled |= _scale_initializer(base_model, f'model.encoders.{layer_idx}.feed_forward.w_2.bias', divisor=2.0)
+    div_model_path = _with_suffix(ONNX_MODEL, "_div")
+    onnx.save(base_model, div_model_path)
+    if os.path.exists(div_model_path):
+        previous_model = ONNX_MODEL
+        ONNX_MODEL = div_model_path
+        _remove_file(previous_model, keep=preserve_files)
+    if modified:
+        if bias_scaled:
+            print('done (created div-adjusted model and scaled bias)')
+        else:
+            print('done (created div-adjusted model; bias initializer not found)')
+    else:
+        print('done (div nodes already present)')
+        
     #开局先给���来个大惊喜，rknn做第一步常量折叠的时候就会在这个子图里报错，所以要单独拿出来先跑一遍
     #然后把这个子图的输出结果保存下来喂给rknn
-    onnx.utils.extract_model(ONNX_MODEL, "extract_model.onnx", ['speech_lengths'], ['/make_pad_mask/Cast_2_output_0'])
-    sess = ort.InferenceSession("extract_model.onnx", providers=['CPUExecutionProvider'])
+    extract_model_path = os.path.join(os.getcwd(), "extract_model.onnx")
+    onnx.utils.extract_model(ONNX_MODEL, extract_model_path, ['speech_lengths'], ['/make_pad_mask/Cast_2_output_0'])
+    sess = ort.InferenceSession(extract_model_path, providers=['CPUExecutionProvider'])
     extract_result = sess.run(None, {"speech_lengths": np.array([speech_length], dtype=np.int64)})[0]
+    _remove_file(extract_model_path)
 
-    # 删掉模型最后的多余transpose, 速度从365ms提升到350ms
+    # 删掉模型最后的多余transpose, 速度从365ms提升到259ms
+    edited_model_path = _with_suffix(ONNX_MODEL, "_edited")
     ret = onnx_edit(model = ONNX_MODEL,
-        export_path = ONNX_MODEL.replace(".onnx", "_edited.onnx"),
+        export_path = edited_model_path,
         # # 1, len, 25055 -> 1, 25055, 1, len   # 这个是坏的, 我真服了，
-        # outputs_transform = {'encoder_out': 'a,b,c->a,c,1,b'},
-        outputs_transform = {'encoder_out': 'a,b,c->a,c,b'},
+        outputs_transform = {'encoder_out': 'a,b,c->a,c,1,b'},
+        # outputs_transform = {'encoder_out': 'a,b,c->a,c,b'},
     )
-    ONNX_MODEL = ONNX_MODEL.replace(".onnx", "_edited.onnx")
+    if os.path.exists(edited_model_path):
+        previous_model = ONNX_MODEL
+        ONNX_MODEL = edited_model_path
+        _remove_file(previous_model, keep=preserve_files)
 
     # pre-process config
     print('--> Config model')
@@ -48,8 +180,9 @@ def convert_encoder():
 
     # Load ONNX model
     print("--> Loading model")
+    current_model_path = ONNX_MODEL
     ret = rknn.load_onnx(
-        model=ONNX_MODEL,
+        model=current_model_path,
         inputs=["speech", "/make_pad_mask/Cast_2_output_0"],
         input_size_list=[[1, speech_length, 560], [extract_result.shape[0], extract_result.shape[1]]],
         input_initial_val=[None, extract_result],
@@ -60,6 +193,7 @@ def convert_encoder():
         print('Load model failed!')
         exit(ret)
     print('done')
+    _remove_file(current_model_path, keep=preserve_files)
 
     # Build model
     print('--> Building model')
@@ -76,20 +210,14 @@ def convert_encoder():
         print('Export RKNN model failed!')
         exit(ret)
     print('done')
+    # 精度分析(可选)
+    # rknn.accuracy_analysis(inputs=["input_content.npy"], target="rk3588", device_id=None)
 
-# usage: python convert_rknn.py encoder|all
+# usage: python convert_rknn.py path/to/model.onnx
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("model", type=str, help="model to convert", choices=["encoder", "all"], nargs='?')
+    parser.add_argument("model_path", type=str, help="path to source ONNX model")
     args = parser.parse_args()
-    if args.model is None:
-        args.model = "all"
 
-    if args.model == "encoder":
-        convert_encoder()
-    elif args.model == "all":
-        convert_encoder()
-    else:
-        print(f"Unknown model: {args.model}")
-        exit(1)
+    convert_encoder(args.model_path)

sense-voice-encoder.rknn

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8db70c1a8d4887e35dff55ab0f5d8da283d32359bd1599ece51eb81f99a6f468
-size 485687354
+oid sha256:198e9c843ed0ffd8500f3b045e47ef5cbc9a4fc36401f6d51c853dacd8129d98
+size 484090175

sensevoice_rknn.py

@@ -28,7 +28,7 @@ from rknnlite.api.rknn_lite import RKNNLite
 
 RKNN_INPUT_LEN = 171
 
-SPEECH_SCALE = 1/2  # 因为是fp16推理，如果中间结果太大可能会溢出变inf，所以需要缩放一下
+SPEECH_SCALE = 1 
 
 class VadOrtInferRuntimeSession:
     def __init__(self, config, root_dir: Path):
@@ -72,7 +72,7 @@ class VadOrtInferRuntimeSession:
             )
 
     def __call__(
-        self, input_content: List[Union[np.ndarray, np.ndarray]]
+        self, input_content
     ) -> np.ndarray:
         if isinstance(input_content, list):
             input_dict = {
@@ -97,10 +97,10 @@ class VadOrtInferRuntimeSession:
     ):
         return [v.name for v in self.session.get_outputs()]
 
-    def get_character_list(self, key: str = "character"):
+    def get_character_list(self, key = "character"):
         return self.meta_dict[key].splitlines()
 
-    def have_key(self, key: str = "character") -> bool:
+    def have_key(self, key = "character") -> bool:
         self.meta_dict = self.session.get_modelmeta().custom_metadata_map
         if key in self.meta_dict.keys():
             return True
@@ -196,10 +196,10 @@ class OrtInferRuntimeSession:
     ):
         return [v.name for v in self.session.get_outputs()]
 
-    def get_character_list(self, key: str = "character"):
+    def get_character_list(self, key = "character"):
         return self.meta_dict[key].splitlines()
 
-    def have_key(self, key: str = "character") -> bool:
+    def have_key(self, key = "character") -> bool:
         self.meta_dict = self.session.get_modelmeta().custom_metadata_map
         if key in self.meta_dict.keys():
             return True
@@ -250,7 +250,7 @@ class SenseVoiceInferenceSession:
         self.sp = spm.SentencePieceProcessor()
         self.sp.load(bpe_model_file)
 
-    def __call__(self, speech, language: int, use_itn: bool) -> np.ndarray:
+    def __call__(self, speech, language, use_itn: bool) -> np.ndarray:
         language_query = self.embedding[[[language]]]
 
         # 14 means with itn, 15 means without itn
@@ -274,6 +274,7 @@ class SenseVoiceInferenceSession:
         input_content = np.pad(input_content, ((0, 0), (0, RKNN_INPUT_LEN - input_content.shape[1]), (0, 0)))
         print("padded shape:", input_content.shape)
         start_time = time.time()
+        np.save("input_content.npy",input_content)
         encoder_out = self.encoder.inference(inputs=[input_content])[0]
         end_time = time.time()
         print(f"encoder inference time: {end_time - start_time:.2f} seconds")
@@ -308,14 +309,14 @@ class WavFrontend:
 
     def __init__(
         self,
-        cmvn_file: str = None,
-        fs: int = 16000,
-        window: str = "hamming",
-        n_mels: int = 80,
-        frame_length: int = 25,
-        frame_shift: int = 10,
-        lfr_m: int = 7,
-        lfr_n: int = 6,
+        cmvn_file = None,
+        fs = 16000,
+        window = "hamming",
+        n_mels = 80,
+        frame_length = 25,
+        frame_shift = 10,
+        lfr_m = 7,
+        lfr_n = 6,
         dither: float = 0,
         **kwargs,
     ) -> None:
@@ -367,7 +368,7 @@ class WavFrontend:
         feat_len = np.array(feat.shape[0]).astype(np.int32)
         return feat, feat_len
 
-    def load_audio(self, filename: str) -> Tuple[np.ndarray, int]:
+    def load_audio(self, filename) -> Tuple[np.ndarray, int]:
         data, sample_rate = sf.read(
             filename,
             always_2d=True,
@@ -383,7 +384,7 @@ class WavFrontend:
         return samples, sample_rate
 
     @staticmethod
-    def apply_lfr(inputs: np.ndarray, lfr_m: int, lfr_n: int) -> np.ndarray:
+    def apply_lfr(inputs: np.ndarray, lfr_m, lfr_n) -> np.ndarray:
         LFR_inputs = []
 
         T = inputs.shape[0]
@@ -417,7 +418,7 @@ class WavFrontend:
         inputs = (inputs + means) * vars
         return inputs
 
-    def get_features(self, inputs: Union[str, np.ndarray]) -> Tuple[np.ndarray, int]:
+    def get_features(self, inputs: Union[str, np.ndarray]):
         if isinstance(inputs, str):
             inputs, _ = self.load_audio(inputs)
 
@@ -504,35 +505,35 @@ class VadDetectMode(Enum):
 class VADXOptions:
     def __init__(
         self,
-        sample_rate: int = 16000,
-        detect_mode: int = VadDetectMode.kVadMutipleUtteranceDetectMode.value,
-        snr_mode: int = 0,
-        max_end_silence_time: int = 800,
-        max_start_silence_time: int = 3000,
+        sample_rate = 16000,
+        detect_mode = VadDetectMode.kVadMutipleUtteranceDetectMode.value,
+        snr_mode = 0,
+        max_end_silence_time = 800,
+        max_start_silence_time = 3000,
         do_start_point_detection: bool = True,
         do_end_point_detection: bool = True,
-        window_size_ms: int = 200,
-        sil_to_speech_time_thres: int = 150,
-        speech_to_sil_time_thres: int = 150,
+        window_size_ms = 200,
+        sil_to_speech_time_thres = 150,
+        speech_to_sil_time_thres = 150,
         speech_2_noise_ratio: float = 1.0,
-        do_extend: int = 1,
-        lookback_time_start_point: int = 200,
-        lookahead_time_end_point: int = 100,
-        max_single_segment_time: int = 60000,
-        nn_eval_block_size: int = 8,
-        dcd_block_size: int = 4,
-        snr_thres: int = -100.0,
-        noise_frame_num_used_for_snr: int = 100,
-        decibel_thres: int = -100.0,
+        do_extend = 1,
+        lookback_time_start_point = 200,
+        lookahead_time_end_point = 100,
+        max_single_segment_time = 60000,
+        nn_eval_block_size = 8,
+        dcd_block_size = 4,
+        snr_thres = -100.0,
+        noise_frame_num_used_for_snr = 100,
+        decibel_thres = -100.0,
         speech_noise_thres: float = 0.6,
         fe_prior_thres: float = 1e-4,
-        silence_pdf_num: int = 1,
+        silence_pdf_num = 1,
         sil_pdf_ids: List[int] = [0],
         speech_noise_thresh_low: float = -0.1,
         speech_noise_thresh_high: float = 0.3,
         output_frame_probs: bool = False,
-        frame_in_ms: int = 10,
-        frame_length_ms: int = 25,
+        frame_in_ms = 10,
+        frame_length_ms = 25,
     ):
         self.sample_rate = sample_rate
         self.detect_mode = detect_mode
@@ -595,10 +596,10 @@ class E2EVadFrameProb(object):
 class WindowDetector(object):
     def __init__(
         self,
-        window_size_ms: int,
-        sil_to_speech_time: int,
-        speech_to_sil_time: int,
-        frame_size_ms: int,
+        window_size_ms,
+        sil_to_speech_time,
+        speech_to_sil_time,
+        frame_size_ms,
     ):
         self.window_size_ms = window_size_ms
         self.sil_to_speech_time = sil_to_speech_time
@@ -633,7 +634,7 @@ class WindowDetector(object):
         return int(self.win_size_frame)
 
     def detect_one_frame(
-        self, frameState: FrameState, frame_count: int
+        self, frameState: FrameState, frame_count
     ) -> AudioChangeState:
         cur_frame_state = FrameState.kFrameStateSil
         if frameState == FrameState.kFrameStateSpeech:
@@ -773,7 +774,7 @@ class E2EVadModel:
 
         return scores[1:]
 
-    def pop_data_buf_till_frame(self, frame_idx: int) -> None:  # need check again
+    def pop_data_buf_till_frame(self, frame_idx) -> None:  # need check again
         while self.data_buf_start_frame < frame_idx:
             if self.data_buf_size >= int(
                 self.vad_opts.frame_in_ms * self.vad_opts.sample_rate / 1000
@@ -787,8 +788,8 @@ class E2EVadModel:
 
     def pop_data_to_output_buf(
         self,
-        start_frm: int,
-        frm_cnt: int,
+        start_frm,
+        frm_cnt,
         first_frm_is_start_point: bool,
         last_frm_is_end_point: bool,
         end_point_is_sent_end: bool,
@@ -849,18 +850,18 @@ class E2EVadModel:
         if last_frm_is_end_point:
             cur_seg.contain_seg_end_point = True
 
-    def on_silence_detected(self, valid_frame: int):
+    def on_silence_detected(self, valid_frame):
         self.lastest_confirmed_silence_frame = valid_frame
         if self.vad_state_machine == VadStateMachine.kVadInStateStartPointNotDetected:
             self.pop_data_buf_till_frame(valid_frame)
         # silence_detected_callback_
         # pass
 
-    def on_voice_detected(self, valid_frame: int) -> None:
+    def on_voice_detected(self, valid_frame) -> None:
         self.latest_confirmed_speech_frame = valid_frame
         self.pop_data_to_output_buf(valid_frame, 1, False, False, False)
 
-    def on_voice_start(self, start_frame: int, fake_result: bool = False) -> None:
+    def on_voice_start(self, start_frame, fake_result: bool = False) -> None:
         if self.vad_opts.do_start_point_detection:
             pass
         if self.confirmed_start_frame != -1:
@@ -878,7 +879,7 @@ class E2EVadModel:
             )
 
     def on_voice_end(
-        self, end_frame: int, fake_result: bool, is_last_frame: bool
+        self, end_frame, fake_result: bool, is_last_frame: bool
     ) -> None:
         for t in range(self.latest_confirmed_speech_frame + 1, end_frame):
             self.on_voice_detected(t)
@@ -896,7 +897,7 @@ class E2EVadModel:
         self.number_end_time_detected += 1
 
     def maybe_on_voice_end_last_frame(
-        self, is_final_frame: bool, cur_frm_idx: int
+        self, is_final_frame: bool, cur_frm_idx
     ) -> None:
         if is_final_frame:
             self.on_voice_end(cur_frm_idx, False, True)
@@ -913,7 +914,7 @@ class E2EVadModel:
             )
         return vad_latency
 
-    def get_frame_state(self, t: int) -> FrameState:
+    def get_frame_state(self, t) -> FrameState:
         frame_state = FrameState.kFrameStateInvalid
         cur_decibel = self.decibel[t - self.decibel_offset]
         cur_snr = cur_decibel - self.noise_average_decibel
@@ -1010,7 +1011,7 @@ class E2EVadModel:
         waveform: np.ndarray,
         in_cache: list = None,
         is_final: bool = False,
-        max_end_sil: int = 800,
+        max_end_sil = 800,
     ) -> Tuple[List[List[List[int]]], Dict[str, np.ndarray]]:
         feats = [feats]
         if in_cache is None:
@@ -1059,7 +1060,7 @@ class E2EVadModel:
         waveform: np.ndarray,
         in_cache: list = None,
         is_final: bool = False,
-        max_end_sil: int = 800,
+        max_end_sil = 800,
     ):
         feats = [feats]
         states = []
@@ -1116,7 +1117,7 @@ class E2EVadModel:
         return 0
 
     def detect_one_frame(
-        self, cur_frm_state: FrameState, cur_frm_idx: int, is_final_frame: bool
+        self, cur_frm_state: FrameState, cur_frm_idx, is_final_frame: bool
     ) -> None:
         tmp_cur_frm_state = FrameState.kFrameStateInvalid
         if cur_frm_state == FrameState.kFrameStateSpeech:
@@ -1267,7 +1268,7 @@ class E2EVadModel:
 
 
 class FSMNVad(object):
-    def __init__(self, config_dir: str):
+    def __init__(self, config_dir):
         config_dir = Path(config_dir)
         self.config = read_yaml(config_dir / "fsmn-config.yaml")
         self.frontend = WavFrontend(
@@ -1332,16 +1333,16 @@ languages = {"auto": 0, "zh": 3, "en": 4, "yue": 7, "ja": 11, "ko": 12, "nospeec
 formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
 logging.basicConfig(format=formatter, level=logging.INFO)
 
-def main():
+def parse_args():
     arg_parser = argparse.ArgumentParser(description="Sense Voice")
     arg_parser.add_argument("-a", "--audio_file", required=True, type=str, help="Model")
     download_model_path = os.path.dirname(__file__)
     arg_parser.add_argument(
         "-dp",
-        "--download_path",
+        "--download_path", 
         default=download_model_path,
         type=str,
-        help="dir path of resource downloaded",
+        help="dir path of resource downloaded"
     )
     arg_parser.add_argument("-d", "--device", default=-1, type=int, help="Device")
     arg_parser.add_argument(
@@ -1351,35 +1352,36 @@ def main():
         "-l",
         "--language",
         choices=languages.keys(),
-        default="auto",
+        default="auto", 
         type=str,
-        help="Language",
+        help="Language"
     )
     arg_parser.add_argument("--use_itn", action="store_true", help="Use ITN")
-    args = arg_parser.parse_args()
+    return arg_parser.parse_args()
 
-    front = WavFrontend(os.path.join(download_model_path, "am.mvn"))
+def main(audio_file, download_path, device, num_threads, language, use_itn):
+    front = WavFrontend(os.path.join(download_path, "am.mvn"))
 
     model = SenseVoiceInferenceSession(
-        os.path.join(download_model_path, "embedding.npy"),
+        os.path.join(download_path, "embedding.npy"),
         os.path.join(
-            download_model_path,
+            download_path,
             "sense-voice-encoder.rknn",
         ),
-        os.path.join(download_model_path, "chn_jpn_yue_eng_ko_spectok.bpe.model"),
-        args.device,
-        args.num_threads,
+        os.path.join(download_path, "chn_jpn_yue_eng_ko_spectok.bpe.model"),
+        device,
+        num_threads,
     )
     waveform, _sample_rate = sf.read(
-        args.audio_file,
+        audio_file,
         dtype="float32",
         always_2d=True
     )
 
-    logging.info(f"Audio {args.audio_file} is {len(waveform) / _sample_rate} seconds, {waveform.shape[1]} channel")
-    # load vad model
+    logging.info(f"Audio {audio_file} is {len(waveform) / _sample_rate} seconds, {waveform.shape[1]} channel")
+    # load vad model 
     start = time.time()
-    vad = FSMNVad(download_model_path)
+    vad = FSMNVad(download_path)
     for channel_id, channel_data in enumerate(waveform.T):
         segments = vad.segments_offline(channel_data)
         results = ""
@@ -1387,16 +1389,20 @@ def main():
             audio_feats = front.get_features(channel_data[part[0] * 16 : part[1] * 16])
             asr_result = model(
                 audio_feats[None, ...],
-                language=languages[args.language],
-                use_itn=args.use_itn,
+                language=languages[language],
+                use_itn=use_itn,
             )
             logging.info(f"[Channel {channel_id}] [{part[0] / 1000}s - {part[1] / 1000}s] {asr_result}")
+            results += asr_result
+        logging.info(f"Results: {results}")
         vad.vad.all_reset_detection()
     decoding_time = time.time() - start
     logging.info(f"Decoder audio takes {decoding_time} seconds")
     logging.info(f"The RTF is {decoding_time/(waveform.shape[1] * len(waveform) / _sample_rate)}.")
+    return results
 
 
 if __name__ == "__main__":
-    main()
+    args = parse_args()
+    main(args.audio_file, args.download_path, args.device, args.num_threads, args.language, args.use_itn)