First commit

.gitattributes

@@ -29,6 +29,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 祈Inory
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,151 @@
----
-license: mit
----
+# wenet.axera
+
+WeNet on Axera.
+
+预转换好的模型在[Release](https://github.com/ml-inory/wenet.axera/releases/download/v1.0/axmodel.tar.gz)  基于AIShell数据集的100条语音量化
+
+## 安装依赖
+
+x86：
+
+```bash
+pip install -r requirements_x86.txt
+```
+
+开发板：
+
+```bash
+pip install -r requirements_ax.txt
+```
+
+## 导出 ONNX
+
+```bash
+python export_onnx.py \
+  --pretrained_model_dir pretrained/aishell_u2pp_conformer_exp \
+  --output_onnx_dir onnx_model
+```
+
+如果不指定参数，默认会下载并使用 `pretrained/aishell_u2pp_conformer_exp`。
+
+输出：
+
+```text
+onnx_model/encoder_offline.onnx
+onnx_model/encoder_online.onnx
+onnx_model/decoder.onnx
+```
+
+## 生成量化数据集
+
+```bash
+wget https://github.com/ml-inory/whisper.axera/releases/download/v1.0/datasets.zip
+unzip datasets.zip -d datasets
+```
+
+```bash
+python generate_data.py \
+  -i datasets \
+  --config pretrained/aishell_u2pp_conformer_exp/train.yaml \
+  --vocab pretrained/aishell_u2pp_conformer_exp/units.txt \
+  --onnx_dir onnx_model \
+  --calib_data_path calibration_dataset \
+  --max_num 100
+```
+
+## 转 axmodel
+
+转换：
+
+```bash
+pulsar2 build \
+  --input onnx_model/encoder_offline.onnx \
+  --config config_encoder_offline.json \
+  --output_dir axmodel/encoder_offline \
+  --output_name encoder_offline.axmodel \
+  --target_hardware AX650 \
+  --npu_mode NPU3
+
+pulsar2 build \
+  --input onnx_model/encoder_online.onnx \
+  --config config_encoder_online.json \
+  --output_dir axmodel/encoder_online \
+  --output_name encoder_online.axmodel \
+  --target_hardware AX650 \
+  --npu_mode NPU3
+
+pulsar2 build \
+  --input onnx_model/decoder.onnx \
+  --config config_decoder.json \
+  --output_dir axmodel/decoder \
+  --output_name decoder.axmodel \
+  --target_hardware AX650 \
+  --npu_mode NPU3
+```
+
+## 跑 ONNX
+
+Offline CTC：
+
+```bash
+python run_ort.py \
+  -i demo.wav \
+  --config pretrained/aishell_u2pp_conformer_exp/train.yaml \
+  --vocab pretrained/aishell_u2pp_conformer_exp/units.txt \
+  --onnx_dir onnx_model \
+  --mode ctc_prefix_beam_search \
+  --calib_data_path ""
+```
+
+Online CTC：
+
+```bash
+python run_ort.py \
+  -i demo.wav \
+  --online \
+  --config pretrained/aishell_u2pp_conformer_exp/train.yaml \
+  --vocab pretrained/aishell_u2pp_conformer_exp/units.txt \
+  --onnx_dir onnx_model \
+  --mode ctc_prefix_beam_search \
+  --calib_data_path ""
+```
+
+Attention rescoring：
+
+```bash
+python run_ort.py \
+  -i demo.wav \
+  --config pretrained/aishell_u2pp_conformer_exp/train.yaml \
+  --vocab pretrained/aishell_u2pp_conformer_exp/units.txt \
+  --onnx_dir onnx_model \
+  --mode attention_rescoring \
+  --calib_data_path ""
+```
+
+## 跑 axmodel
+
+在开发板仓库目录运行。
+
+Offline CTC：
+
+```bash
+python3 run_ax.py -i demo.wav --mode ctc_prefix_beam_search
+```
+
+Online CTC：
+
+```bash
+python3 run_ax.py -i demo.wav --online --mode ctc_prefix_beam_search
+```
+
+Offline attention rescoring：
+
+```bash
+python3 run_ax.py -i demo.wav --mode attention_rescoring
+```
+
+Online attention rescoring：
+
+```bash
+python3 run_ax.py -i demo.wav --online --mode attention_rescoring
+```

ax_common.py

@@ -0,0 +1,625 @@
+import multiprocessing
+import wave
+
+import numpy as np
+import yaml
+
+
+def _log_add(*values):
+    values = [value for value in values if value != -float("inf")]
+    if not values:
+        return -float("inf")
+    max_value = max(values)
+    return max_value + np.log(sum(np.exp(value - max_value)
+                                  for value in values))
+
+
+def _map_sentence(sent, vocabulary, greedy=False, blank_id=0):
+    mapped = []
+    prev = None
+    for token in sent:
+        token = int(token)
+        if greedy and token == prev:
+            prev = token
+            continue
+        prev = token
+        if token == blank_id or token < 0 or token >= len(vocabulary):
+            continue
+        piece = vocabulary[token]
+        if piece.startswith("<") and piece.endswith(">"):
+            continue
+        mapped.append(piece)
+    return "".join(mapped)
+
+
+def map_batch(batch_sents, vocabulary, num_processes, greedy=False,
+              blank_id=0):
+    del num_processes
+    return [_map_sentence(sent, vocabulary, greedy, blank_id)
+            for sent in batch_sents]
+
+
+def _ctc_prefix_beam_search(log_probs_seq, log_probs_idx, beam_size,
+                            blank_id):
+    beam = {(): (0.0, -float("inf"))}
+    for frame_probs, frame_ids in zip(log_probs_seq, log_probs_idx):
+        frame_probs = np.asarray(frame_probs, dtype=np.float32)
+        frame_probs = frame_probs - _log_add(*frame_probs.tolist())
+        next_beam = {}
+        for prefix, (prob_blank, prob_non_blank) in beam.items():
+            for prob, token in zip(frame_probs, frame_ids):
+                token = int(token)
+                prob = float(prob)
+                next_prob_blank, next_prob_non_blank = next_beam.get(
+                    prefix, (-float("inf"), -float("inf")))
+                if token == blank_id:
+                    next_beam[prefix] = (
+                        _log_add(next_prob_blank, prob_blank + prob,
+                                 prob_non_blank + prob),
+                        next_prob_non_blank,
+                    )
+                    continue
+
+                last = prefix[-1] if prefix else None
+                if token == last:
+                    next_beam[prefix] = (
+                        next_prob_blank,
+                        _log_add(next_prob_non_blank, prob_non_blank + prob),
+                    )
+                    new_prefix = prefix + (token, )
+                    nb_blank, nb_non_blank = next_beam.get(
+                        new_prefix, (-float("inf"), -float("inf")))
+                    next_beam[new_prefix] = (
+                        nb_blank,
+                        _log_add(nb_non_blank, prob_blank + prob),
+                    )
+                else:
+                    new_prefix = prefix + (token, )
+                    nb_blank, nb_non_blank = next_beam.get(
+                        new_prefix, (-float("inf"), -float("inf")))
+                    next_beam[new_prefix] = (
+                        nb_blank,
+                        _log_add(nb_non_blank, prob_blank + prob,
+                                 prob_non_blank + prob),
+                    )
+        beam = dict(sorted(
+            next_beam.items(),
+            key=lambda item: _log_add(item[1][0], item[1][1]),
+            reverse=True)[:beam_size])
+    return [(_log_add(prob_blank, prob_non_blank), prefix)
+            for prefix, (prob_blank, prob_non_blank) in sorted(
+                beam.items(),
+                key=lambda item: _log_add(item[1][0], item[1][1]),
+                reverse=True)]
+
+
+def ctc_beam_search_decoder_batch(batch_log_probs_seq,
+                                  batch_log_probs_idx,
+                                  batch_root_trie,
+                                  batch_start,
+                                  beam_size,
+                                  num_processes,
+                                  blank_id=0,
+                                  space_id=-1,
+                                  cutoff_prob=0.999,
+                                  ext_scorer=None):
+    del batch_root_trie, batch_start, num_processes, space_id
+    del cutoff_prob, ext_scorer
+    return [
+        _ctc_prefix_beam_search(log_probs_seq, log_probs_idx, beam_size,
+                                blank_id)
+        for log_probs_seq, log_probs_idx in zip(batch_log_probs_seq,
+                                                batch_log_probs_idx)
+    ]
+
+
+def load_config(config_path):
+    with open(config_path, "r") as fin:
+        return yaml.load(fin, Loader=yaml.FullLoader)
+
+
+def load_vocab(vocab_path):
+    vocabulary = []
+    char_dict = {}
+    with open(vocab_path, "r") as fin:
+        for line in fin:
+            arr = line.strip().split()
+            assert len(arr) == 2
+            char_dict[int(arr[1])] = arr[0]
+            vocabulary.append(arr[0])
+    return vocabulary, char_dict
+
+
+def load_wav(audio_file):
+    with wave.open(audio_file, "rb") as wav_file:
+        sample_rate = wav_file.getframerate()
+        num_channels = wav_file.getnchannels()
+        sample_width = wav_file.getsampwidth()
+        frames = wav_file.readframes(wav_file.getnframes())
+
+    if sample_width == 1:
+        waveform = np.frombuffer(frames, dtype=np.uint8).astype(np.float32)
+        waveform -= 128.0
+    elif sample_width == 2:
+        waveform = np.frombuffer(frames, dtype="<i2").astype(np.float32)
+    elif sample_width == 4:
+        waveform = np.frombuffer(frames, dtype="<i4").astype(np.float32)
+    else:
+        raise ValueError(f"Unsupported wav sample width: {sample_width}")
+
+    if num_channels > 1:
+        waveform = waveform.reshape(-1, num_channels).mean(axis=1)
+    return waveform, sample_rate
+
+
+def resample_linear(waveform, orig_sr, target_sr):
+    if orig_sr == target_sr:
+        return waveform
+    duration = waveform.shape[0] / float(orig_sr)
+    target_len = int(round(duration * target_sr))
+    if target_len <= 1:
+        return waveform
+    src_pos = np.linspace(0, waveform.shape[0] - 1, target_len)
+    return np.interp(src_pos, np.arange(waveform.shape[0]),
+                     waveform).astype(np.float32)
+
+
+def hz_to_mel(freq):
+    return 1127.0 * np.log1p(freq / 700.0)
+
+
+def mel_to_hz(mel):
+    return 700.0 * np.expm1(mel / 1127.0)
+
+
+def mel_filterbank(num_mel_bins, n_fft, sample_rate):
+    low_mel = hz_to_mel(20.0)
+    high_mel = hz_to_mel(sample_rate / 2.0)
+    mel_points = np.linspace(low_mel, high_mel, num_mel_bins + 2)
+    hz_points = mel_to_hz(mel_points)
+    bins = np.floor((n_fft + 1) * hz_points / sample_rate).astype(np.int32)
+    fbanks = np.zeros((num_mel_bins, n_fft // 2 + 1), dtype=np.float32)
+    for i in range(num_mel_bins):
+        left, center, right = bins[i], bins[i + 1], bins[i + 2]
+        if center > left:
+            fbanks[i, left:center] = (
+                np.arange(left, center) - left) / float(center - left)
+        if right > center:
+            fbanks[i, center:right] = (
+                right - np.arange(center, right)) / float(right - center)
+    return fbanks
+
+
+def numpy_fbank(waveform,
+                sample_rate=16000,
+                num_mel_bins=80,
+                frame_length=25,
+                frame_shift=10):
+    frame_size = int(round(sample_rate * frame_length / 1000.0))
+    frame_step = int(round(sample_rate * frame_shift / 1000.0))
+    if waveform.shape[0] < frame_size:
+        waveform = np.pad(waveform, (0, frame_size - waveform.shape[0]))
+    num_frames = 1 + (waveform.shape[0] - frame_size) // frame_step
+    frames = np.lib.stride_tricks.as_strided(
+        waveform,
+        shape=(num_frames, frame_size),
+        strides=(waveform.strides[0] * frame_step, waveform.strides[0]),
+    ).copy()
+    frames *= np.hamming(frame_size).astype(np.float32)
+    n_fft = 1
+    while n_fft < frame_size:
+        n_fft <<= 1
+    power = np.abs(np.fft.rfft(frames, n=n_fft))**2
+    fbanks = mel_filterbank(num_mel_bins, n_fft, sample_rate)
+    mel_energies = np.maximum(np.dot(power, fbanks.T), np.finfo(np.float32).eps)
+    return np.log(mel_energies).astype(np.float32)
+
+
+def compute_feats(audio_file, sr=16000):
+    waveform, sample_rate = load_wav(audio_file)
+    waveform = resample_linear(waveform.astype(np.float32), sample_rate, sr)
+    return numpy_fbank(waveform, sample_rate=sr).reshape(1, -1, 80)
+
+
+def pad_array_along_axis(array, pad_width, axis, mode="constant", **kwargs):
+    if array.shape[axis] >= pad_width:
+        return array
+    full_pad_width = [(0, 0)] * array.ndim
+    full_pad_width[axis] = (0, pad_width - array.shape[axis])
+    return np.pad(array, pad_width=full_pad_width, mode=mode, **kwargs)
+
+
+def numpy_topk(array, k, axis=-1, largest=True):
+    if largest:
+        partitioned_indices = np.argpartition(array, -k, axis=axis)
+        topk_indices = np.take(partitioned_indices, range(-k, 0), axis=axis)
+    else:
+        partitioned_indices = np.argpartition(array, k, axis=axis)
+        topk_indices = np.take(partitioned_indices, range(0, k), axis=axis)
+
+    topk_values = np.take_along_axis(array, topk_indices, axis=axis)
+    sorted_indices_in_topk = np.argsort(topk_values, axis=axis)
+    if largest:
+        sorted_indices_in_topk = np.flip(sorted_indices_in_topk, axis=axis)
+    sorted_topk_values = np.take_along_axis(topk_values,
+                                            sorted_indices_in_topk,
+                                            axis=axis)
+    sorted_topk_indices = np.take_along_axis(topk_indices,
+                                             sorted_indices_in_topk,
+                                             axis=axis)
+    return sorted_topk_values, sorted_topk_indices
+
+
+def ctc_decoding(beam_log_probs,
+                 beam_log_probs_idx,
+                 encoder_out_lens,
+                 vocabulary,
+                 mode="ctc_prefix_beam_search"):
+    beam_size = beam_log_probs.shape[-1]
+    batch_size = beam_log_probs.shape[0]
+    num_processes = min(multiprocessing.cpu_count(), batch_size)
+    hyps = []
+    score_hyps = []
+
+    if mode == "ctc_greedy_search":
+        if beam_size == 1:
+            log_probs_idx = beam_log_probs_idx.squeeze(-1)
+        else:
+            log_probs_idx = beam_log_probs_idx[:, :, 0]
+        batch_sents = []
+        for idx, seq in enumerate(log_probs_idx):
+            batch_sents.append(seq[0:encoder_out_lens[idx]].tolist())
+        hyps = map_batch(batch_sents, vocabulary, num_processes, True, 0)
+    elif mode in ("ctc_prefix_beam_search", "attention_rescoring"):
+        batch_log_probs_seq_list = beam_log_probs.tolist()
+        batch_log_probs_idx_list = beam_log_probs_idx.tolist()
+        batch_len_list = encoder_out_lens.tolist()
+        batch_log_probs_seq = []
+        batch_log_probs_ids = []
+        batch_start = []
+        batch_root = []
+        for i in range(len(batch_len_list)):
+            num_sent = batch_len_list[i]
+            batch_log_probs_seq.append(batch_log_probs_seq_list[i][0:num_sent])
+            batch_log_probs_ids.append(batch_log_probs_idx_list[i][0:num_sent])
+            batch_root.append(None)
+            batch_start.append(True)
+        score_hyps = ctc_beam_search_decoder_batch(batch_log_probs_seq,
+                                                   batch_log_probs_ids,
+                                                   batch_root,
+                                                   batch_start,
+                                                   beam_size,
+                                                   num_processes,
+                                                   0, -2, 0.99999)
+        if mode == "ctc_prefix_beam_search":
+            for cand_hyps in score_hyps:
+                hyps.append(cand_hyps[0][1])
+            hyps = map_batch(hyps, vocabulary, num_processes, False, 0)
+    return hyps, score_hyps
+
+
+def has_higher_scored_collapsed_repeat(hyp, kept_hyps):
+    for i in range(1, len(hyp)):
+        if hyp[i] != hyp[i - 1]:
+            continue
+        collapsed = hyp[:i] + hyp[i + 1:]
+        if collapsed in kept_hyps:
+            return True
+    return False
+
+
+def make_decoder_inputs(encoder_out,
+                        encoder_out_lens,
+                        beam_log_probs,
+                        beam_log_probs_idx,
+                        vocabulary,
+                        sos,
+                        eos,
+                        decoder_len):
+    _, score_hyps = ctc_decoding(beam_log_probs, beam_log_probs_idx,
+                                 encoder_out_lens, vocabulary,
+                                 "attention_rescoring")
+    ignore_id = -1
+    beam_size = beam_log_probs.shape[-1]
+    batch_size = beam_log_probs.shape[0]
+    ctc_score, all_hyps = [], []
+    for hyps in score_hyps:
+        filtered_hyps = []
+        kept_hyps = set()
+        for score, hyp in hyps:
+            hyp = tuple(hyp)
+            if has_higher_scored_collapsed_repeat(hyp, kept_hyps):
+                continue
+            filtered_hyps.append((score, hyp))
+            kept_hyps.add(hyp)
+            if len(filtered_hyps) == beam_size:
+                break
+        hyps = filtered_hyps
+        cur_len = len(hyps)
+        if len(hyps) < beam_size:
+            hyps += (beam_size - cur_len) * [(-float("inf"), (0,))]
+        cur_ctc_score = []
+        for hyp in hyps:
+            cur_ctc_score.append(hyp[0])
+            all_hyps.append(list(hyp[1]))
+        ctc_score.append(cur_ctc_score)
+    ctc_score = np.array(ctc_score, dtype=np.float32)
+
+    max_len = decoder_len - 2
+    hyps_pad_sos_eos = np.ones((batch_size, beam_size, max_len + 2),
+                               dtype=np.int64) * ignore_id
+    r_hyps_pad_sos_eos = np.ones((batch_size, beam_size, max_len + 2),
+                                 dtype=np.int64) * ignore_id
+    hyps_lens_sos = np.ones((batch_size, beam_size), dtype=np.int32)
+    k = 0
+    for i in range(batch_size):
+        for j in range(beam_size):
+            cand = all_hyps[k][:max_len]
+            length = len(cand) + 2
+            hyps_pad_sos_eos[i][j][0:length] = [sos] + cand + [eos]
+            r_hyps_pad_sos_eos[i][j][0:length] = [sos] + cand[::-1] + [eos]
+            hyps_lens_sos[i][j] = len(cand) + 1
+            k += 1
+
+    if decoder_len > encoder_out.shape[1]:
+        encoder_out = np.pad(encoder_out,
+                             [(0, 0),
+                              (0, decoder_len - encoder_out.shape[1]),
+                              (0, 0)],
+                             mode="constant",
+                             constant_values=0)
+    elif decoder_len < encoder_out.shape[1]:
+        encoder_out = encoder_out[:, :decoder_len, :]
+
+    return {
+        "encoder_out": encoder_out,
+        "encoder_out_lens": np.full(batch_size,
+                                    fill_value=decoder_len,
+                                    dtype=np.int32),
+        "hyps_pad_sos_eos": hyps_pad_sos_eos.astype(np.int32),
+        "hyps_lens_sos": hyps_lens_sos,
+        "r_hyps_pad_sos_eos": r_hyps_pad_sos_eos.astype(np.int32),
+        "ctc_score": ctc_score,
+    }, all_hyps
+
+
+def make_offline_inputs(feats, seq_len):
+    feats = feats[:, :seq_len, :]
+    speech_lengths = np.array([feats.shape[1]], dtype=np.int32)
+    if feats.shape[1] < seq_len:
+        feats = pad_array_along_axis(feats, pad_width=seq_len, axis=1)
+    return {"speech": feats, "speech_lengths": speech_lengths}
+
+
+def make_online_initial_state(configs,
+                              batch_size=1,
+                              decoding_chunk_size=16,
+                              num_decoding_left_chunks=5):
+    subsampling = 4
+    context = 7
+    stride = subsampling * decoding_chunk_size
+    decoding_window = (decoding_chunk_size - 1) * subsampling + context
+    required_cache_size = decoding_chunk_size * num_decoding_left_chunks
+
+    output_size = configs["encoder_conf"]["output_size"]
+    num_layers = configs["encoder_conf"]["num_blocks"]
+    cnn_module_kernel = configs["encoder_conf"].get("cnn_module_kernel", 1) - 1
+    head = configs["encoder_conf"]["attention_heads"]
+    d_k = configs["encoder_conf"]["output_size"] // head
+
+    state = {
+        "att_cache": np.zeros((batch_size, num_layers, head,
+                               required_cache_size, d_k * 2),
+                              dtype=np.float32),
+        "cnn_cache": np.zeros((batch_size, num_layers, output_size,
+                               cnn_module_kernel),
+                              dtype=np.float32),
+        "cache_mask": np.zeros((batch_size, 1, required_cache_size),
+                               dtype=np.float32),
+        "offset": np.zeros((batch_size, 1), dtype=np.int32),
+    }
+    params = {
+        "batch_size": batch_size,
+        "context": context,
+        "stride": stride,
+        "decoding_window": decoding_window,
+    }
+    return state, params
+
+
+def make_online_encoder_input(feats, cur, params, state):
+    batch_size = params["batch_size"]
+    decoding_window = params["decoding_window"]
+    end = min(cur + decoding_window, feats.shape[1])
+    chunk_xs = feats[:, cur:end, :]
+    if chunk_xs.shape[1] < decoding_window:
+        chunk_xs = pad_array_along_axis(chunk_xs,
+                                        pad_width=decoding_window,
+                                        axis=1)
+    chunk_xs = chunk_xs.astype(np.float32)
+    chunk_lens = np.full(batch_size,
+                         fill_value=chunk_xs.shape[1],
+                         dtype=np.int32)
+    return {
+        "chunk_xs": chunk_xs,
+        "chunk_lens": chunk_lens,
+        "offset": state["offset"],
+        "att_cache": state["att_cache"],
+        "cnn_cache": state["cnn_cache"],
+        "cache_mask": state["cache_mask"],
+    }
+
+
+def output_value(outputs, name):
+    if name in outputs:
+        return outputs[name]
+    r_name = "r_" + name
+    if r_name in outputs:
+        return outputs[r_name]
+    raise KeyError(name)
+
+
+def update_online_state(state, outputs):
+    state["offset"] = output_value(outputs, "offset")
+    state["att_cache"] = output_value(outputs, "att_cache")
+    state["cnn_cache"] = output_value(outputs, "cnn_cache")
+    state["cache_mask"] = output_value(outputs, "cache_mask")
+
+
+class AxModel:
+
+    def __init__(self, path):
+        from axengine import InferenceSession
+
+        self.session = InferenceSession(path)
+        self.output_names = [item.name for item in self.session.get_outputs()]
+
+    def run(self, input_feed):
+        output_values = self.session.run(self.output_names, input_feed)
+        return dict(zip(self.output_names, output_values))
+
+
+class WenetAXRunner:
+
+    def __init__(self,
+                 config_path,
+                 vocab_path,
+                 encoder_offline_path="axmodel/encoder_offline/encoder_offline.axmodel",
+                 encoder_online_path="axmodel/encoder_online/encoder_online.axmodel",
+                 decoder_path="axmodel/decoder/decoder.axmodel",
+                 offline_seq_len=1024,
+                 decoder_len=32,
+                 decoding_chunk_size=16,
+                 num_decoding_left_chunks=5,
+                 batch_size=1):
+        self.config_path = config_path
+        self.vocab_path = vocab_path
+        self.encoder_offline_path = encoder_offline_path
+        self.encoder_online_path = encoder_online_path
+        self.decoder_path = decoder_path
+        self.offline_seq_len = offline_seq_len
+        self.decoder_len = decoder_len
+        self.decoding_chunk_size = decoding_chunk_size
+        self.num_decoding_left_chunks = num_decoding_left_chunks
+        self.batch_size = batch_size
+
+        self.configs = load_config(config_path)
+        self.vocabulary, self.char_dict = load_vocab(vocab_path)
+        self.eos = self.sos = len(self.char_dict) - 1
+
+        self._offline_encoder = None
+        self._online_encoder = None
+        self._decoder = None
+
+    @property
+    def offline_encoder(self):
+        if self._offline_encoder is None:
+            self._offline_encoder = AxModel(self.encoder_offline_path)
+        return self._offline_encoder
+
+    @property
+    def online_encoder(self):
+        if self._online_encoder is None:
+            self._online_encoder = AxModel(self.encoder_online_path)
+        return self._online_encoder
+
+    @property
+    def decoder(self):
+        if self._decoder is None:
+            self._decoder = AxModel(self.decoder_path)
+        return self._decoder
+
+    def compute_feats(self, audio_file):
+        return compute_feats(audio_file)
+
+    def run_offline_encoder(self, feats):
+        encoder_input = make_offline_inputs(feats, self.offline_seq_len)
+        speech_lengths = encoder_input["speech_lengths"]
+        outputs = self.offline_encoder.run(encoder_input)
+        encoder_out_lens = outputs["encoder_out_lens"].astype(np.int32)
+        encoder_out_lens[0] = np.ones([speech_lengths[0]],
+                                      dtype=np.int32)[2::2][2::2].sum()
+        beam_log_probs, beam_log_probs_idx = numpy_topk(
+            outputs["ctc_log_probs"], k=10)
+        return {
+            "encoder_out": outputs["encoder_out"],
+            "encoder_out_lens": encoder_out_lens,
+            "ctc_log_probs": outputs["ctc_log_probs"],
+            "beam_log_probs": beam_log_probs,
+            "beam_log_probs_idx": beam_log_probs_idx,
+        }
+
+    def run_online_encoder(self, feats):
+        state, online_params = make_online_initial_state(
+            self.configs, self.batch_size, self.decoding_chunk_size,
+            self.num_decoding_left_chunks)
+        encoder_out = []
+        beam_log_probs = []
+        beam_log_probs_idx = []
+        num_frames = feats.shape[1]
+
+        for cur in range(0, num_frames - online_params["context"] + 1,
+                         online_params["stride"]):
+            encoder_input = make_online_encoder_input(feats, cur,
+                                                      online_params, state)
+            outputs = self.online_encoder.run(encoder_input)
+            update_online_state(state, outputs)
+            encoder_out.append(outputs["chunk_out"])
+            beam_log_probs.append(outputs["log_probs"])
+            beam_log_probs_idx.append(outputs["log_probs_idx"].astype(np.int32))
+
+        return {
+            "encoder_out": np.concatenate(encoder_out, axis=1),
+            "encoder_out_lens": np.full(self.batch_size,
+                                        fill_value=sum(
+                                            out.shape[1]
+                                            for out in encoder_out),
+                                        dtype=np.int32),
+            "beam_log_probs": np.concatenate(beam_log_probs, axis=1),
+            "beam_log_probs_idx": np.concatenate(beam_log_probs_idx, axis=1),
+        }
+
+    def ctc_decode(self, encoder_outputs, mode):
+        return ctc_decoding(encoder_outputs["beam_log_probs"],
+                            encoder_outputs["beam_log_probs_idx"],
+                            encoder_outputs["encoder_out_lens"],
+                            self.vocabulary, mode)
+
+    def run_decoder(self, encoder_outputs):
+        decoder_input, all_hyps = make_decoder_inputs(
+            encoder_outputs["encoder_out"],
+            encoder_outputs["encoder_out_lens"],
+            encoder_outputs["beam_log_probs"],
+            encoder_outputs["beam_log_probs_idx"],
+            self.vocabulary,
+            self.sos,
+            self.eos,
+            self.decoder_len,
+        )
+        best_index = self.decoder.run(decoder_input)["best_index"].astype(
+            np.int32)
+        beam_size = encoder_outputs["beam_log_probs"].shape[-1]
+        num_processes = min(multiprocessing.cpu_count(), best_index.shape[0])
+        best_sents = []
+        k = 0
+        for idx in best_index:
+            best_sents.append(all_hyps[k:k + beam_size][idx])
+            k += beam_size
+        hyps = map_batch(best_sents, self.vocabulary, num_processes)
+        return "".join(hyps)
+
+    def transcribe(self,
+                   audio_file,
+                   online=False,
+                   mode="ctc_prefix_beam_search"):
+        feats = self.compute_feats(audio_file)
+        if online:
+            encoder_outputs = self.run_online_encoder(feats)
+        else:
+            encoder_outputs = self.run_offline_encoder(feats)
+
+        if mode == "attention_rescoring":
+            return self.run_decoder(encoder_outputs)
+
+        hyps, _ = self.ctc_decode(encoder_outputs, mode)
+        return "".join(hyps) if hyps else ""

config_decoder.json

@@ -0,0 +1,88 @@
+{
+    "model_type": "ONNX",
+    "npu_mode": "NPU1",
+    "quant": {
+      "input_configs": [
+        {
+            "tensor_name": "encoder_out",
+            "calibration_dataset": "./calibration_dataset/encoder_out.tar.gz",
+            "calibration_size": -1,
+            "calibration_format": "Numpy"
+        },
+        {
+            "tensor_name": "encoder_out_lens",
+            "calibration_dataset": "./calibration_dataset/encoder_out_lens.tar.gz",
+            "calibration_size": -1,
+            "calibration_format": "Numpy"
+        },
+        {
+            "tensor_name": "hyps_pad_sos_eos",
+            "calibration_dataset": "./calibration_dataset/hyps_pad_sos_eos.tar.gz",
+            "calibration_size": -1,
+            "calibration_format": "Numpy"
+        },
+        {
+          "tensor_name": "hyps_lens_sos",
+          "calibration_dataset": "./calibration_dataset/hyps_lens_sos.tar.gz",
+          "calibration_size": -1,
+          "calibration_format": "Numpy"
+        },
+        {
+          "tensor_name": "r_hyps_pad_sos_eos",
+          "calibration_dataset": "./calibration_dataset/r_hyps_pad_sos_eos.tar.gz",
+          "calibration_size": -1,
+          "calibration_format": "Numpy"
+        },
+        {
+          "tensor_name": "ctc_score",
+          "calibration_dataset": "./calibration_dataset/ctc_score.tar.gz",
+          "calibration_size": -1,
+          "calibration_format": "Numpy"
+        }
+      ],
+      "layer_configs": [
+        {
+            "start_tensor_names": ["DEFAULT"], 
+            "end_tensor_names": ["DEFAULT"], 
+            "data_type": "U8"
+        } 
+      ],
+      "calibration_method": "MinMax",
+      "conv_bias_data_type": "FP32",
+      "disable_auto_refine_scale": true,
+      "enable_smooth_quant": true,
+      "precision_analysis": true,
+      "precision_analysis_method": "EndToEnd",
+      "transformer_opt_level": 1
+    },
+    "input_processors": [
+      {
+        "tensor_name": "encoder_out",
+        "src_dtype": "FP32"
+      },
+      {
+        "tensor_name": "encoder_out_lens",
+        "src_dtype": "S32"
+      },
+      {
+        "tensor_name": "hyps_pad_sos_eos",
+        "src_dtype": "S32"
+      },
+      {
+        "tensor_name": "hyps_lens_sos",
+        "src_dtype": "S32"
+      },
+      {
+        "tensor_name": "r_hyps_pad_sos_eos",
+        "src_dtype": "S32"
+      },
+      {
+        "tensor_name": "ctc_score",
+        "src_dtype": "FP32"
+      }
+    ],
+    "compiler": {
+      "check": 0
+    }
+  }
+  

config_encoder_offline.json

@@ -0,0 +1,48 @@
+{
+    "model_type": "ONNX",
+    "npu_mode": "NPU3",
+    "quant": {
+      "input_configs": [
+        {
+            "tensor_name": "speech",
+            "calibration_dataset": "./calibration_dataset/speech.tar.gz",
+            "calibration_size": -1,
+            "calibration_format": "Numpy"
+        },
+        {
+          "tensor_name": "speech_lengths",
+          "calibration_dataset": "./calibration_dataset/speech_lengths.tar.gz",
+          "calibration_size": -1,
+          "calibration_format": "Numpy"
+      }
+      ],
+      "layer_configs": [
+        {
+            "start_tensor_names": ["DEFAULT"], 
+            "end_tensor_names": ["DEFAULT"], 
+            "data_type": "U16"
+        } 
+      ],
+      "calibration_method": "MinMax",
+      "conv_bias_data_type": "FP32",
+      "disable_auto_refine_scale": true,
+      "enable_smooth_quant": true,
+      "precision_analysis": true,
+      "precision_analysis_method": "EndToEnd",
+      "transformer_opt_level": 1
+    },
+    "input_processors": [
+      {
+        "tensor_name": "speech",
+        "src_dtype": "FP32"
+      },
+      {
+        "tensor_name": "speech_lengths",
+        "src_dtype": "S32"
+      }
+    ],
+    "compiler": {
+      "check": 0
+    }
+  }
+  

config_encoder_online.json

@@ -0,0 +1,88 @@
+{
+    "model_type": "ONNX",
+    "npu_mode": "NPU3",
+    "quant": {
+      "input_configs": [
+        {
+            "tensor_name": "chunk_xs",
+            "calibration_dataset": "./calibration_dataset/chunk_xs.tar.gz",
+            "calibration_size": -1,
+            "calibration_format": "Numpy"
+        },
+        {
+          "tensor_name": "chunk_lens",
+          "calibration_dataset": "./calibration_dataset/chunk_lens.tar.gz",
+          "calibration_size": -1,
+          "calibration_format": "Numpy"
+        },
+        {
+          "tensor_name": "offset",
+          "calibration_dataset": "./calibration_dataset/offset.tar.gz",
+          "calibration_size": -1,
+          "calibration_format": "Numpy"
+        },
+        {
+          "tensor_name": "att_cache",
+          "calibration_dataset": "./calibration_dataset/att_cache.tar.gz",
+          "calibration_size": -1,
+          "calibration_format": "Numpy"
+        },
+        {
+          "tensor_name": "cnn_cache",
+          "calibration_dataset": "./calibration_dataset/cnn_cache.tar.gz",
+          "calibration_size": -1,
+          "calibration_format": "Numpy"
+        },
+        {
+          "tensor_name": "cache_mask",
+          "calibration_dataset": "./calibration_dataset/cache_mask.tar.gz",
+          "calibration_size": -1,
+          "calibration_format": "Numpy"
+        }
+      ],
+      "layer_configs": [
+        {
+            "start_tensor_names": ["DEFAULT"], 
+            "end_tensor_names": ["DEFAULT"], 
+            "data_type": "U16"
+        } 
+      ],
+      "calibration_method": "MinMax",
+      "conv_bias_data_type": "FP32",
+      "disable_auto_refine_scale": true,
+      "enable_smooth_quant": true,
+      "precision_analysis": true,
+      "precision_analysis_method": "EndToEnd",
+      "transformer_opt_level": 1
+    },
+    "input_processors": [
+      {
+        "tensor_name": "chunk_xs",
+        "src_dtype": "FP32"
+      },
+      {
+        "tensor_name": "chunk_lens",
+        "src_dtype": "S32"
+      },
+      {
+        "tensor_name": "offset",
+        "src_dtype": "S32"
+      },
+      {
+        "tensor_name": "att_cache",
+        "src_dtype": "FP32"
+      },
+      {
+        "tensor_name": "cnn_cache",
+        "src_dtype": "FP32"
+      },
+      {
+        "tensor_name": "cache_mask",
+        "src_dtype": "FP32"
+      }
+    ],
+    "compiler": {
+      "check": 0
+    }
+  }
+  

demo.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46dbc998c9d1d48111267c40741dd3200f2e5bcf4075f8c4c97f4451160dce50
+size 134570

export_onnx.py

@@ -0,0 +1,1723 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import argparse
+import copy
+import logging
+import os
+import sys
+import tarfile
+import urllib.request
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import yaml
+from wenet.transformer.ctc import CTC
+from wenet.transformer.decoder import TransformerDecoder
+from wenet.transformer.encoder import BaseEncoder
+from wenet.utils.init_model import init_model
+from wenet.utils.mask import make_pad_mask
+from typing import List, Tuple
+
+try:
+    import onnx
+    import onnxruntime
+    from onnx import helper, numpy_helper
+    from onnxsim import simplify
+except ImportError:
+    print("Please install onnxruntime!")
+    sys.exit(1)
+
+logger = logging.getLogger(__file__)
+logger.setLevel(logging.INFO)
+
+DEFAULT_PRETRAINED_MODEL_URL = (
+    "https://huggingface.co/openspeech/wenet-models/resolve/main/"
+    "aishell_u2pp_conformer_exp.tar.gz")
+DEFAULT_PRETRAINED_MODEL_DIR = "pretrained/aishell_u2pp_conformer_exp"
+
+
+def safe_extract_tar(tar, output_dir):
+    output_dir = os.path.abspath(output_dir)
+    for member in tar.getmembers():
+        member_path = os.path.abspath(os.path.join(output_dir, member.name))
+        if not member_path.startswith(output_dir + os.sep):
+            raise RuntimeError(f"Unsafe tar member path: {member.name}")
+    tar.extractall(output_dir)
+
+
+def download_file(url, output_path):
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    print(f"Downloading pretrained model from {url}")
+    print(f"Saving to {output_path}")
+    urllib.request.urlretrieve(url, output_path)
+
+
+def prepare_pretrained_model(args):
+    model_dir = args.pretrained_model_dir
+    archive_dir = os.path.dirname(model_dir.rstrip(os.sep)) or "."
+    archive_path = os.path.join(
+        archive_dir, os.path.basename(model_dir.rstrip(os.sep)) + ".tar.gz")
+
+    if not os.path.exists(model_dir):
+        if not os.path.exists(archive_path):
+            download_file(args.pretrained_model_url, archive_path)
+        print(f"Extracting pretrained model to {archive_dir}")
+        with tarfile.open(archive_path, "r:gz") as tar:
+            safe_extract_tar(tar, archive_dir)
+
+    args.config = os.path.join(model_dir, "train.yaml")
+    args.checkpoint = os.path.join(model_dir, "final.pt")
+    args.cmvn_file = os.path.join(model_dir, "global_cmvn")
+
+    missing = [path for path in (args.config, args.checkpoint)
+               if not os.path.exists(path)]
+    if missing:
+        raise FileNotFoundError(
+            "Missing pretrained model files: " + ", ".join(missing))
+
+    print(f"Using config: {args.config}")
+    print(f"Using checkpoint: {args.checkpoint}")
+    if os.path.exists(args.cmvn_file):
+        print(f"Using CMVN: {args.cmvn_file}")
+
+
+def _constant_node_value(node):
+    if node is None or node.op_type != "Constant":
+        return None
+    for attr in node.attribute:
+        if attr.name == "value":
+            return numpy_helper.to_array(attr.t)
+    return None
+
+
+def _attribute_value(attr):
+    return helper.get_attribute_value(attr)
+
+
+def _get_attr(node, name, default=None):
+    for attr in node.attribute:
+        if attr.name == name:
+            return _attribute_value(attr)
+    return default
+
+
+def _cast_to_onnx_dtype(value, to_dtype):
+    tensor_type = onnx.TensorProto.DataType.Name(to_dtype).lower()
+    dtype_map = {
+        "float": np.float32,
+        "double": np.float64,
+        "float16": np.float16,
+        "int64": np.int64,
+        "int32": np.int32,
+        "int16": np.int16,
+        "int8": np.int8,
+        "uint64": np.uint64,
+        "uint32": np.uint32,
+        "uint16": np.uint16,
+        "uint8": np.uint8,
+        "bool": np.bool_,
+    }
+    if tensor_type not in dtype_map:
+        return None
+    return value.astype(dtype_map[tensor_type])
+
+
+def _shape_from_value_info(value_info):
+    if not value_info.type.HasField("tensor_type"):
+        return None
+    if not value_info.type.tensor_type.shape.dim:
+        return None
+    shape = []
+    for dim in value_info.type.tensor_type.shape.dim:
+        if dim.HasField("dim_value") and dim.dim_value > 0:
+            shape.append(dim.dim_value)
+        else:
+            return None
+    return tuple(shape)
+
+
+def _collect_static_shapes(model):
+    inferred = onnx.shape_inference.infer_shapes(model)
+    shapes = {}
+    for value_info in list(inferred.graph.input) + list(
+            inferred.graph.value_info) + list(inferred.graph.output):
+        shape = _shape_from_value_info(value_info)
+        if shape is not None:
+            shapes[value_info.name] = shape
+    for initializer in inferred.graph.initializer:
+        shapes[initializer.name] = tuple(initializer.dims)
+    return shapes
+
+
+def _eval_static_node(node, inputs, static_shapes):
+    if node.op_type == "Constant":
+        return _constant_node_value(node)
+    if node.op_type != "Shape" and any(value is None for value in inputs):
+        return None
+
+    try:
+        if node.op_type == "Add":
+            return np.add(inputs[0], inputs[1])
+        if node.op_type == "Sub":
+            return np.subtract(inputs[0], inputs[1])
+        if node.op_type == "Mul":
+            return np.multiply(inputs[0], inputs[1])
+        if node.op_type == "Div":
+            return np.divide(inputs[0], inputs[1])
+        if node.op_type == "Equal":
+            return np.equal(inputs[0], inputs[1])
+        if node.op_type == "Greater":
+            return np.greater(inputs[0], inputs[1])
+        if node.op_type == "GreaterOrEqual":
+            return np.greater_equal(inputs[0], inputs[1])
+        if node.op_type == "Less":
+            return np.less(inputs[0], inputs[1])
+        if node.op_type == "LessOrEqual":
+            return np.less_equal(inputs[0], inputs[1])
+        if node.op_type == "Where":
+            return np.where(inputs[0], inputs[1], inputs[2])
+        if node.op_type == "Concat":
+            axis = _get_attr(node, "axis", 0)
+            return np.concatenate(inputs, axis=axis)
+        if node.op_type == "Unsqueeze":
+            axes = _get_attr(node, "axes")
+            if axes is None and len(inputs) > 1:
+                axes = inputs[1]
+            axes = tuple(int(axis) for axis in np.asarray(axes).reshape(-1))
+            return np.expand_dims(inputs[0], axis=axes)
+        if node.op_type == "Squeeze":
+            axes = _get_attr(node, "axes")
+            if axes is None and len(inputs) > 1:
+                axes = inputs[1]
+            if axes is None:
+                return np.squeeze(inputs[0])
+            axes = tuple(int(axis) for axis in np.asarray(axes).reshape(-1))
+            return np.squeeze(inputs[0], axis=axes)
+        if node.op_type == "Cast":
+            return _cast_to_onnx_dtype(inputs[0], _get_attr(node, "to"))
+        if node.op_type == "Reshape":
+            return np.reshape(inputs[0], tuple(int(i) for i in inputs[1]))
+        if node.op_type == "Shape":
+            if inputs[0] is not None:
+                shape = inputs[0].shape
+            else:
+                shape = static_shapes.get(node.input[0])
+            if shape is None:
+                return None
+            return np.asarray(shape, dtype=np.int64)
+        if node.op_type == "Slice":
+            data = inputs[0]
+            starts = np.asarray(inputs[1]).reshape(-1)
+            ends = np.asarray(inputs[2]).reshape(-1)
+            axes = (np.asarray(inputs[3]).reshape(-1)
+                    if len(inputs) > 3 and inputs[3] is not None else
+                    np.arange(len(starts)))
+            steps = (np.asarray(inputs[4]).reshape(-1)
+                     if len(inputs) > 4 and inputs[4] is not None else
+                     np.ones(len(starts), dtype=np.int64))
+            slices = [slice(None)] * data.ndim
+            for start, end, axis, step in zip(starts, ends, axes, steps):
+                axis = int(axis)
+                start = int(start)
+                end = int(end)
+                step = int(step)
+                if end >= np.iinfo(np.int32).max:
+                    end = None
+                if end <= np.iinfo(np.int32).min:
+                    end = None
+                slices[axis] = slice(start, end, step)
+            return data[tuple(slices)]
+        if node.op_type == "Gather":
+            axis = _get_attr(node, "axis", 0)
+            return np.take(inputs[0], inputs[1], axis=axis)
+    except Exception:
+        return None
+    return None
+
+
+def _constant_node(output_name, value, name):
+    const_tensor = numpy_helper.from_array(np.asarray(value),
+                                           name=output_name + "_value")
+    return helper.make_node("Constant",
+                            inputs=[],
+                            outputs=[output_name],
+                            name=name,
+                            value=const_tensor)
+
+
+def _node_attributes(node):
+    return {attr.name: helper.get_attribute_value(attr) for attr in node.attribute}
+
+
+def _copy_node(node, inputs=None, outputs=None, name=None):
+    copied = copy.deepcopy(node)
+    if inputs is not None:
+        del copied.input[:]
+        copied.input.extend(inputs)
+    if outputs is not None:
+        del copied.output[:]
+        copied.output.extend(outputs)
+    if name is not None:
+        copied.name = name
+    return copied
+
+
+def _producer_map(model):
+    return {output: node for node in model.graph.node for output in node.output}
+
+
+def _unsqueeze_greater_equal_pattern(producer, value_name):
+    unsqueeze = producer.get(value_name)
+    if unsqueeze is None or unsqueeze.op_type != "Unsqueeze":
+        return None, None
+    compare = producer.get(unsqueeze.input[0])
+    if compare is None or compare.op_type != "GreaterOrEqual":
+        return None, None
+    return unsqueeze, compare
+
+
+def rewrite_pulsar2_bool_not(onnx_path):
+    """Remove simple Not nodes that Pulsar2 quantization can cast to float.
+
+    The encoder mask contains Not(Unsqueeze(GreaterOrEqual(...))) and another
+    Not over a sliced version of that mask. Pulsar2 can quantize the Not input
+    to FP32 and then fail because bitwise Not only accepts bool/integer tensors.
+    Rewriting those patterns keeps the graph boolean-equivalent without Not.
+    """
+    model = onnx.load(onnx_path)
+    producer = _producer_map(model)
+    rewritten = 0
+    new_nodes = []
+
+    for node in model.graph.node:
+        if node.op_type != "Not":
+            new_nodes.append(node)
+            continue
+
+        compare = producer.get(node.input[0])
+        if compare is not None and compare.op_type == "GreaterOrEqual":
+            less = helper.make_node("Less",
+                                    inputs=list(compare.input),
+                                    outputs=list(node.output),
+                                    name=node.name + "_less",
+                                    **_node_attributes(compare))
+            new_nodes.append(less)
+            rewritten += 1
+            continue
+
+        unsqueeze, compare = _unsqueeze_greater_equal_pattern(
+            producer, node.input[0])
+        if unsqueeze is not None:
+            less_output = node.output[0] + "_less"
+            less = helper.make_node("Less",
+                                    inputs=list(compare.input),
+                                    outputs=[less_output],
+                                    name=node.name + "_less",
+                                    **_node_attributes(compare))
+            rewritten_unsqueeze = _copy_node(
+                unsqueeze,
+                inputs=[less_output] + list(unsqueeze.input[1:]),
+                outputs=list(node.output),
+                name=node.name + "_unsqueeze")
+            new_nodes.extend([less, rewritten_unsqueeze])
+            rewritten += 1
+            continue
+
+        slice_1 = producer.get(node.input[0])
+        slice_0 = producer.get(slice_1.input[0]) if slice_1 else None
+        inner_not = producer.get(slice_0.input[0]) if slice_0 else None
+        if (slice_1 is not None and slice_1.op_type == "Slice"
+                and slice_0 is not None and slice_0.op_type == "Slice"
+                and inner_not is not None and inner_not.op_type == "Not"):
+            unsqueeze, _ = _unsqueeze_greater_equal_pattern(
+                producer, inner_not.input[0])
+            if unsqueeze is not None:
+                slice_0_output = node.output[0] + "_slice0"
+                rewritten_slice_0 = _copy_node(
+                    slice_0,
+                    inputs=[unsqueeze.output[0]] + list(slice_0.input[1:]),
+                    outputs=[slice_0_output],
+                    name=node.name + "_slice0")
+                rewritten_slice_1 = _copy_node(
+                    slice_1,
+                    inputs=[slice_0_output] + list(slice_1.input[1:]),
+                    outputs=list(node.output),
+                    name=node.name + "_slice1")
+                new_nodes.extend([rewritten_slice_0, rewritten_slice_1])
+                rewritten += 1
+                continue
+
+        new_nodes.append(node)
+
+    if rewritten:
+        del model.graph.node[:]
+        model.graph.node.extend(new_nodes)
+        onnx.checker.check_model(model)
+        onnx.save(model, onnx_path)
+        print(f"Rewrote {rewritten} bool Not node(s) in {onnx_path}")
+
+
+def rewrite_pulsar2_bool_and(onnx_path):
+    """Replace boolean And with arithmetic comparison for Pulsar2 quantization."""
+    model = onnx.load(onnx_path)
+    rewritten = 0
+    new_nodes = []
+
+    for node in model.graph.node:
+        if node.op_type != "And" or len(node.input) != 2 or len(
+                node.output) != 1:
+            new_nodes.append(node)
+            continue
+
+        left = node.output[0] + "_left_i32"
+        right = node.output[0] + "_right_i32"
+        added = node.output[0] + "_sum"
+        threshold = node.output[0] + "_threshold"
+        new_nodes.append(
+            helper.make_node("Cast",
+                             inputs=[node.input[0]],
+                             outputs=[left],
+                             name=node.name + "_cast_left",
+                             to=onnx.TensorProto.INT32))
+        new_nodes.append(
+            helper.make_node("Cast",
+                             inputs=[node.input[1]],
+                             outputs=[right],
+                             name=node.name + "_cast_right",
+                             to=onnx.TensorProto.INT32))
+        new_nodes.append(
+            helper.make_node("Add",
+                             inputs=[left, right],
+                             outputs=[added],
+                             name=node.name + "_add"))
+        new_nodes.append(
+            _constant_node(threshold, np.asarray(1, dtype=np.int32),
+                           node.name + "_threshold"))
+        new_nodes.append(
+            helper.make_node("Greater",
+                             inputs=[added, threshold],
+                             outputs=list(node.output),
+                             name=node.name + "_greater"))
+        rewritten += 1
+
+    if rewritten:
+        del model.graph.node[:]
+        model.graph.node.extend(new_nodes)
+        onnx.checker.check_model(model)
+        onnx.save(model, onnx_path)
+        print(f"Rewrote {rewritten} bool And node(s) in {onnx_path}")
+
+
+def simplify_pulsar2_onnx(onnx_path):
+    model = onnx.load(onnx_path)
+    sim_model, ok = simplify(model)
+    if not ok:
+        raise RuntimeError(f"onnxsim failed to validate {onnx_path}")
+    onnx.checker.check_model(sim_model)
+    onnx.save(sim_model, onnx_path)
+    print(f"Simplified {onnx_path} for Pulsar2")
+
+
+def fold_static_pulsar2_subgraphs(onnx_path):
+    """Fold static ONNX patterns that Pulsar2 5.0 cannot infer reliably.
+
+    Pulsar2 5.0 can fail shape inference on ConstantOfShape when its input is
+    a constant tensor value instead of an initializer. The legacy exporter emits
+    this pattern for masks/padding in the encoder graphs. It can also fail when
+    an Expand shape is produced by a constant-only subgraph such as
+    Mul/Equal/Where. Fold those static pieces before handing the model to
+    Pulsar2.
+    """
+    model = onnx.load(onnx_path)
+    static_shapes = _collect_static_shapes(model)
+    constants = {
+        initializer.name: numpy_helper.to_array(initializer)
+        for initializer in model.graph.initializer
+    }
+    folded = 0
+    new_nodes = []
+    for node in model.graph.node:
+        inputs = [constants.get(name) for name in node.input]
+        if node.op_type == "ConstantOfShape" and node.input:
+            shape_value = inputs[0]
+            if shape_value is not None:
+                fill_value = np.array(0, dtype=np.float32)
+                for attr in node.attribute:
+                    if attr.name == "value":
+                        fill_value = numpy_helper.to_array(attr.t)
+                        break
+                shape = tuple(int(dim)
+                              for dim in np.asarray(shape_value).reshape(-1))
+                value = np.full(shape,
+                                fill_value.reshape(-1)[0],
+                                dtype=fill_value.dtype)
+            else:
+                value = None
+        else:
+            value = _eval_static_node(node, inputs, static_shapes)
+
+        if value is None or len(node.output) != 1:
+            new_nodes.append(node)
+            continue
+
+        constants[node.output[0]] = value
+        new_nodes.append(_constant_node(node.output[0], value, node.name))
+        folded += 1
+
+    if folded:
+        del model.graph.node[:]
+        model.graph.node.extend(new_nodes)
+        onnx.checker.check_model(model)
+        onnx.save(model, onnx_path)
+        print(f"Folded {folded} static node(s) in {onnx_path}")
+
+
+class Encoder(torch.nn.Module):
+
+    def __init__(self, encoder: BaseEncoder, ctc: CTC, beam_size: int = 10):
+        super().__init__()
+        self.encoder = encoder
+        self.ctc = ctc
+        self.beam_size = beam_size
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+    ):
+        """Encoder
+        Args:
+            speech: (Batch, Length, ...)
+            speech_lengths: (Batch, )
+        Returns:
+            encoder_out: B x T x F
+            encoder_out_lens: B
+            ctc_log_probs: B x T x V
+            beam_log_probs: B x T x beam_size
+            beam_log_probs_idx: B x T x beam_size
+        """
+        encoder_out, encoder_mask = self.encoder(speech, speech_lengths, -1,
+                                                 -1)
+        encoder_out_lens = encoder_mask.squeeze(1).sum(1)
+        # ctc_log_probs = self.ctc.log_softmax(encoder_out)
+        ctc_log_probs = self.ctc.linear(encoder_out)
+        encoder_out_lens = encoder_out_lens.int()
+        beam_log_probs, beam_log_probs_idx = torch.topk(ctc_log_probs,
+                                                        self.beam_size,
+                                                        dim=2)
+        return (
+            encoder_out,
+            encoder_out_lens,
+            ctc_log_probs,
+            beam_log_probs,
+            beam_log_probs_idx,
+        )
+
+
+class StreamingEncoder(torch.nn.Module):
+
+    def __init__(
+        self,
+        model,
+        required_cache_size,
+        beam_size,
+        transformer=False,
+        return_ctc_logprobs=False,
+    ):
+        super().__init__()
+        self.ctc = model.ctc
+        self.subsampling_rate = model.encoder.embed.subsampling_rate
+        self.embed = model.encoder.embed
+        self.global_cmvn = model.encoder.global_cmvn
+        self.required_cache_size = required_cache_size
+        self.beam_size = beam_size
+        self.encoder = model.encoder
+        self.transformer = transformer
+        self.return_ctc_logprobs = return_ctc_logprobs
+
+    def forward(self, chunk_xs, chunk_lens, offset, att_cache, cnn_cache,
+                cache_mask):
+        """Streaming Encoder
+        Args:
+            xs (torch.Tensor): chunk input, with shape (b, time, mel-dim),
+                where `time == (chunk_size - 1) * subsample_rate + \
+                        subsample.right_context + 1`
+            offset (torch.Tensor): offset with shape (b, 1)
+                        1 is retained for triton deployment
+            required_cache_size (int): cache size required for next chunk
+                compuation
+                > 0: actual cache size
+                <= 0: not allowed in streaming gpu encoder                   `
+            att_cache (torch.Tensor): cache tensor for KEY & VALUE in
+                transformer/conformer attention, with shape
+                (b, elayers, head, cache_t1, d_k * 2), where
+                `head * d_k == hidden-dim` and
+                `cache_t1 == chunk_size * num_decoding_left_chunks`.
+            cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer,
+                (b, elayers, b, hidden-dim, cache_t2), where
+                `cache_t2 == cnn.lorder - 1`
+            cache_mask: (torch.Tensor): cache mask with shape (b, required_cache_size)
+                 in a batch of request, each request may have different
+                 history cache. Cache mask is used to indidate the effective
+                 cache for each request
+        Returns:
+            torch.Tensor: log probabilities of ctc output and cutoff by beam size
+                with shape (b, chunk_size, beam)
+            torch.Tensor: index of top beam size probabilities for each timestep
+                with shape (b, chunk_size, beam)
+            torch.Tensor: output of current input xs,
+                with shape (b, chunk_size, hidden-dim).
+            torch.Tensor: new attention cache required for next chunk, with
+                same shape (b, elayers, head, cache_t1, d_k * 2)
+                as the original att_cache
+            torch.Tensor: new conformer cnn cache required for next chunk, with
+                same shape as the original cnn_cache.
+            torch.Tensor: new cache mask, with same shape as the original
+                cache mask
+        """
+        offset = offset.squeeze(1)
+        T = chunk_xs.size(1)
+        chunk_mask = ~make_pad_mask(chunk_lens, T).unsqueeze(1)
+        # B X 1 X T
+        chunk_mask = chunk_mask.to(chunk_xs.dtype)
+        # transpose batch & num_layers dim
+        att_cache = torch.transpose(att_cache, 0, 1)
+        cnn_cache = torch.transpose(cnn_cache, 0, 1)
+
+        # rewrite encoder.forward_chunk
+        # <---------forward_chunk START--------->
+        xs = self.global_cmvn(chunk_xs)
+        # chunk mask is important for batch inferencing since
+        # different sequence in a batch has different length
+        xs, pos_emb, chunk_mask = self.embed(xs, chunk_mask, offset)
+        cache_size = att_cache.size(3)  # required cache size
+        masks = torch.cat((cache_mask, chunk_mask), dim=2)
+        index = offset - cache_size
+
+        pos_emb = self.embed.position_encoding(index, cache_size + xs.size(1))
+        pos_emb = pos_emb.to(dtype=xs.dtype)
+
+        next_cache_start = -self.required_cache_size
+        r_cache_mask = masks[:, :, next_cache_start:]
+
+        r_att_cache = []
+        r_cnn_cache = []
+        for i, layer in enumerate(self.encoder.encoders):
+            i_kv_cache = att_cache[i]
+            size = att_cache.size(-1) // 2
+            kv_cache = (i_kv_cache[:, :, :, :size], i_kv_cache[:, :, :, size:])
+            xs, _, new_kv_cache, new_cnn_cache = layer(
+                xs,
+                masks,
+                pos_emb,
+                att_cache=kv_cache,
+                cnn_cache=cnn_cache[i],
+            )
+            #   shape(new_att_cache) is (B, head, attention_key_size, d_k * 2),
+            #   shape(new_cnn_cache) is (B, hidden-dim, cache_t2)
+            new_att_cache = torch.cat(new_kv_cache, dim=-1)
+            r_att_cache.append(
+                new_att_cache[:, :, next_cache_start:, :].unsqueeze(1))
+            if not self.transformer:
+                r_cnn_cache.append(new_cnn_cache.unsqueeze(1))
+        if self.encoder.normalize_before:
+            chunk_out = self.encoder.after_norm(xs)
+        else:
+            chunk_out = xs
+
+        r_att_cache = torch.cat(r_att_cache, dim=1)  # concat on layers idx
+        if not self.transformer:
+            r_cnn_cache = torch.cat(r_cnn_cache, dim=1)  # concat on layers
+
+        # <---------forward_chunk END--------->
+
+        # log_ctc_probs = self.ctc.log_softmax(chunk_out)
+        log_ctc_probs = self.ctc.linear(chunk_out)
+        log_probs, log_probs_idx = torch.topk(log_ctc_probs,
+                                              self.beam_size,
+                                              dim=2)
+        log_probs = log_probs.to(chunk_xs.dtype)
+
+        r_offset = offset + chunk_out.shape[1]
+        # the below ops not supported in Tensorrt
+        # chunk_out_lens = torch.div(chunk_lens, subsampling_rate,
+        #                   rounding_mode='floor')
+        chunk_out_lens = chunk_lens // self.subsampling_rate
+        r_offset = r_offset.unsqueeze(1)
+        if self.return_ctc_logprobs:
+            return (
+                log_ctc_probs,
+                chunk_out,
+                chunk_out_lens,
+                r_offset,
+                r_att_cache,
+                r_cnn_cache,
+                r_cache_mask,
+            )
+        else:
+            return (
+                log_probs,
+                log_probs_idx,
+                chunk_out,
+                chunk_out_lens,
+                r_offset,
+                r_att_cache,
+                r_cnn_cache,
+                r_cache_mask,
+            )
+
+
+class StreamingSqueezeformerEncoder(torch.nn.Module):
+
+    def __init__(self, model, required_cache_size, beam_size):
+        super().__init__()
+        self.ctc = model.ctc
+        self.subsampling_rate = model.encoder.embed.subsampling_rate
+        self.embed = model.encoder.embed
+        self.global_cmvn = model.encoder.global_cmvn
+        self.required_cache_size = required_cache_size
+        self.beam_size = beam_size
+        self.encoder = model.encoder
+        self.reduce_idx = model.encoder.reduce_idx
+        self.recover_idx = model.encoder.recover_idx
+        if self.reduce_idx is None:
+            self.time_reduce = None
+        else:
+            if self.recover_idx is None:
+                self.time_reduce = "normal"  # no recovery at the end
+            else:
+                self.time_reduce = "recover"  # recovery at the end
+                assert len(self.reduce_idx) == len(self.recover_idx)
+
+    def calculate_downsampling_factor(self, i: int) -> int:
+        if self.reduce_idx is None:
+            return 1
+        else:
+            reduce_exp, recover_exp = 0, 0
+            for exp, rd_idx in enumerate(self.reduce_idx):
+                if i >= rd_idx:
+                    reduce_exp = exp + 1
+            if self.recover_idx is not None:
+                for exp, rc_idx in enumerate(self.recover_idx):
+                    if i >= rc_idx:
+                        recover_exp = exp + 1
+            return int(2**(reduce_exp - recover_exp))
+
+    def forward(self, chunk_xs, chunk_lens, offset, att_cache, cnn_cache,
+                cache_mask):
+        """Streaming Encoder
+        Args:
+            xs (torch.Tensor): chunk input, with shape (b, time, mel-dim),
+                where `time == (chunk_size - 1) * subsample_rate + \
+                        subsample.right_context + 1`
+            offset (torch.Tensor): offset with shape (b, 1)
+                        1 is retained for triton deployment
+            required_cache_size (int): cache size required for next chunk
+                compuation
+                > 0: actual cache size
+                <= 0: not allowed in streaming gpu encoder                   `
+            att_cache (torch.Tensor): cache tensor for KEY & VALUE in
+                transformer/conformer attention, with shape
+                (b, elayers, head, cache_t1, d_k * 2), where
+                `head * d_k == hidden-dim` and
+                `cache_t1 == chunk_size * num_decoding_left_chunks`.
+            cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer,
+                (b, elayers, b, hidden-dim, cache_t2), where
+                `cache_t2 == cnn.lorder - 1`
+            cache_mask: (torch.Tensor): cache mask with shape (b, required_cache_size)
+                 in a batch of request, each request may have different
+                 history cache. Cache mask is used to indidate the effective
+                 cache for each request
+        Returns:
+            torch.Tensor: log probabilities of ctc output and cutoff by beam size
+                with shape (b, chunk_size, beam)
+            torch.Tensor: index of top beam size probabilities for each timestep
+                with shape (b, chunk_size, beam)
+            torch.Tensor: output of current input xs,
+                with shape (b, chunk_size, hidden-dim).
+            torch.Tensor: new attention cache required for next chunk, with
+                same shape (b, elayers, head, cache_t1, d_k * 2)
+                as the original att_cache
+            torch.Tensor: new conformer cnn cache required for next chunk, with
+                same shape as the original cnn_cache.
+            torch.Tensor: new cache mask, with same shape as the original
+                cache mask
+        """
+        offset = offset.squeeze(1)
+        T = chunk_xs.size(1)
+        chunk_mask = ~make_pad_mask(chunk_lens, T).unsqueeze(1)
+        # B X 1 X T
+        chunk_mask = chunk_mask.to(chunk_xs.dtype)
+        # transpose batch & num_layers dim
+        att_cache = torch.transpose(att_cache, 0, 1)
+        cnn_cache = torch.transpose(cnn_cache, 0, 1)
+
+        # rewrite encoder.forward_chunk
+        # <---------forward_chunk START--------->
+        xs = self.global_cmvn(chunk_xs)
+        # chunk mask is important for batch inferencing since
+        # different sequence in a batch has different length
+        xs, pos_emb, chunk_mask = self.embed(xs, chunk_mask, offset)
+        elayers, cache_size = att_cache.size(0), att_cache.size(3)
+        att_mask = torch.cat((cache_mask, chunk_mask), dim=2)
+        index = offset - cache_size
+
+        pos_emb = self.embed.position_encoding(index, cache_size + xs.size(1))
+        pos_emb = pos_emb.to(dtype=xs.dtype)
+
+        next_cache_start = -self.required_cache_size
+        r_cache_mask = att_mask[:, :, next_cache_start:]
+
+        r_att_cache = []
+        r_cnn_cache = []
+        mask_pad = torch.ones(1,
+                              xs.size(1),
+                              device=xs.device,
+                              dtype=torch.bool)
+        mask_pad = mask_pad.unsqueeze(1)
+        max_att_len: int = 0
+        recover_activations: List[Tuple[torch.Tensor, torch.Tensor,
+                                        torch.Tensor, torch.Tensor]] = []
+        index = 0
+        xs_lens = torch.tensor([xs.size(1)], device=xs.device, dtype=torch.int)
+        xs = self.encoder.preln(xs)
+        for i, layer in enumerate(self.encoder.encoders):
+            if self.reduce_idx is not None:
+                if self.time_reduce is not None and i in self.reduce_idx:
+                    recover_activations.append(
+                        (xs, att_mask, pos_emb, mask_pad))
+                    (
+                        xs,
+                        xs_lens,
+                        att_mask,
+                        mask_pad,
+                    ) = self.encoder.time_reduction_layer(
+                        xs, xs_lens, att_mask, mask_pad)
+                    pos_emb = pos_emb[:, ::2, :]
+                    if self.encoder.pos_enc_layer_type == "rel_pos_repaired":
+                        pos_emb = pos_emb[:, :xs.size(1) * 2 - 1, :]
+                    index += 1
+
+            if self.recover_idx is not None:
+                if self.time_reduce == "recover" and i in self.recover_idx:
+                    index -= 1
+                    (
+                        recover_tensor,
+                        recover_att_mask,
+                        recover_pos_emb,
+                        recover_mask_pad,
+                    ) = recover_activations[index]
+                    # recover output length for ctc decode
+                    xs = xs.unsqueeze(2).repeat(1, 1, 2, 1).flatten(1, 2)
+                    xs = self.encoder.time_recover_layer(xs)
+                    recoverd_t = recover_tensor.size(1)
+                    xs = recover_tensor + xs[:, :recoverd_t, :].contiguous()
+                    att_mask = recover_att_mask
+                    pos_emb = recover_pos_emb
+                    mask_pad = recover_mask_pad
+
+            factor = self.calculate_downsampling_factor(i)
+
+            xs, _, new_att_cache, new_cnn_cache = layer(
+                xs,
+                att_mask,
+                pos_emb,
+                att_cache=att_cache[i][:, :, ::factor, :]
+                [:, :, :pos_emb.size(1) - xs.size(1), :]
+                if elayers > 0 else att_cache[:, :, ::factor, :],
+                cnn_cache=cnn_cache[i] if cnn_cache.size(0) > 0 else cnn_cache,
+            )
+            cached_att = new_att_cache[:, :, next_cache_start // factor:, :]
+            cached_cnn = new_cnn_cache.unsqueeze(1)
+            cached_att = (cached_att.unsqueeze(3).repeat(1, 1, 1, factor,
+                                                         1).flatten(2, 3))
+            if i == 0:
+                # record length for the first block as max length
+                max_att_len = cached_att.size(2)
+            r_att_cache.append(cached_att[:, :, :max_att_len, :].unsqueeze(1))
+            r_cnn_cache.append(cached_cnn)
+
+        chunk_out = xs
+        r_att_cache = torch.cat(r_att_cache, dim=1)  # concat on layers idx
+        r_cnn_cache = torch.cat(r_cnn_cache, dim=1)  # concat on layers
+
+        # <---------forward_chunk END--------->
+
+        # log_ctc_probs = self.ctc.log_softmax(chunk_out)
+        log_ctc_probs = self.ctc.linear(chunk_out)
+        log_probs, log_probs_idx = torch.topk(log_ctc_probs,
+                                              self.beam_size,
+                                              dim=2)
+        log_probs = log_probs.to(chunk_xs.dtype)
+
+        r_offset = offset + chunk_out.shape[1]
+        # the below ops not supported in Tensorrt
+        # chunk_out_lens = torch.div(chunk_lens, subsampling_rate,
+        #                   rounding_mode='floor')
+        chunk_out_lens = chunk_lens // self.subsampling_rate
+        r_offset = r_offset.unsqueeze(1)
+
+        return (
+            log_probs,
+            log_probs_idx,
+            chunk_out,
+            chunk_out_lens,
+            r_offset,
+            r_att_cache,
+            r_cnn_cache,
+            r_cache_mask,
+        )
+
+
+class StreamingEfficientConformerEncoder(torch.nn.Module):
+
+    def __init__(self, model, required_cache_size, beam_size):
+        super().__init__()
+        self.ctc = model.ctc
+        self.subsampling_rate = model.encoder.embed.subsampling_rate
+        self.embed = model.encoder.embed
+        self.global_cmvn = model.encoder.global_cmvn
+        self.required_cache_size = required_cache_size
+        self.beam_size = beam_size
+        self.encoder = model.encoder
+
+        # Efficient Conformer
+        self.stride_layer_idx = model.encoder.stride_layer_idx
+        self.stride = model.encoder.stride
+        self.num_blocks = model.encoder.num_blocks
+        self.cnn_module_kernel = model.encoder.cnn_module_kernel
+
+    def calculate_downsampling_factor(self, i: int) -> int:
+        factor = 1
+        for idx, stride_idx in enumerate(self.stride_layer_idx):
+            if i > stride_idx:
+                factor *= self.stride[idx]
+        return factor
+
+    def forward(self, chunk_xs, chunk_lens, offset, att_cache, cnn_cache,
+                cache_mask):
+        """Streaming Encoder
+        Args:
+            chunk_xs (torch.Tensor): chunk input, with shape (b, time, mel-dim),
+                where `time == (chunk_size - 1) * subsample_rate + \
+                        subsample.right_context + 1`
+            chunk_lens (torch.Tensor):
+            offset (torch.Tensor): offset with shape (b, 1)
+                        1 is retained for triton deployment
+            att_cache (torch.Tensor): cache tensor for KEY & VALUE in
+                transformer/conformer attention, with shape
+                (b, elayers, head, cache_t1, d_k * 2), where
+                `head * d_k == hidden-dim` and
+                `cache_t1 == chunk_size * num_decoding_left_chunks`.
+            cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer,
+                (b, elayers, hidden-dim, cache_t2), where
+                `cache_t2 == cnn.lorder - 1`
+            cache_mask: (torch.Tensor): cache mask with shape (b, required_cache_size)
+                 in a batch of request, each request may have different
+                 history cache. Cache mask is used to indidate the effective
+                 cache for each request
+        Returns:
+            torch.Tensor: log probabilities of ctc output and cutoff by beam size
+                with shape (b, chunk_size, beam)
+            torch.Tensor: index of top beam size probabilities for each timestep
+                with shape (b, chunk_size, beam)
+            torch.Tensor: output of current input xs,
+                with shape (b, chunk_size, hidden-dim).
+            torch.Tensor: new attention cache required for next chunk, with
+                same shape (b, elayers, head, cache_t1, d_k * 2)
+                as the original att_cache
+            torch.Tensor: new conformer cnn cache required for next chunk, with
+                same shape as the original cnn_cache.
+            torch.Tensor: new cache mask, with same shape as the original
+                cache mask
+        """
+        offset = offset.squeeze(1)  # (b, )
+        offset *= self.calculate_downsampling_factor(self.num_blocks + 1)
+
+        T = chunk_xs.size(1)
+        chunk_mask = ~make_pad_mask(chunk_lens, T).unsqueeze(1)  # (b, 1, T)
+        # B X 1 X T
+        chunk_mask = chunk_mask.to(chunk_xs.dtype)
+        # transpose batch & num_layers dim
+        #   Shape(att_cache): (elayers, b, head, cache_t1, d_k * 2)
+        #   Shape(cnn_cache): (elayers, b, outsize, cnn_kernel)
+        att_cache = torch.transpose(att_cache, 0, 1)
+        cnn_cache = torch.transpose(cnn_cache, 0, 1)
+
+        # rewrite encoder.forward_chunk
+        # <---------forward_chunk START--------->
+        xs = self.global_cmvn(chunk_xs)
+        # chunk mask is important for batch inferencing since
+        # different sequence in a batch has different length
+        xs, pos_emb, chunk_mask = self.embed(xs, chunk_mask, offset)
+        cache_size = att_cache.size(3)  # required cache size
+        masks = torch.cat((cache_mask, chunk_mask), dim=2)
+        att_mask = torch.cat((cache_mask, chunk_mask), dim=2)
+        index = offset - cache_size
+
+        pos_emb = self.embed.position_encoding(index, cache_size + xs.size(1))
+        pos_emb = pos_emb.to(dtype=xs.dtype)
+
+        next_cache_start = -self.required_cache_size
+        r_cache_mask = masks[:, :, next_cache_start:]
+
+        r_att_cache = []
+        r_cnn_cache = []
+        mask_pad = chunk_mask.to(torch.bool)
+        max_att_len, max_cnn_len = (
+            0,
+            0,
+        )  # for repeat_interleave of new_att_cache
+        for i, layer in enumerate(self.encoder.encoders):
+            factor = self.calculate_downsampling_factor(i)
+            # NOTE(xcsong): Before layer.forward
+            #   shape(att_cache[i:i + 1]) is (b, head, cache_t1, d_k * 2),
+            #   shape(cnn_cache[i])       is (b=1, hidden-dim, cache_t2)
+            # shape(new_att_cache) = [ batch, head, time2, outdim//head * 2 ]
+            att_cache_trunc = 0
+            if xs.size(1) + att_cache.size(3) / factor > pos_emb.size(1):
+                # The time step is not divisible by the downsampling multiple
+                # We propose to double the chunk_size.
+                att_cache_trunc = (xs.size(1) + att_cache.size(3) // factor -
+                                   pos_emb.size(1) + 1)
+            xs, _, new_att_cache, new_cnn_cache = layer(
+                xs,
+                att_mask,
+                pos_emb,
+                mask_pad=mask_pad,
+                att_cache=att_cache[i][:, :, ::factor, :][:, :,
+                                                          att_cache_trunc:, :],
+                cnn_cache=cnn_cache[i, :, :, :]
+                if cnn_cache.size(0) > 0 else cnn_cache,
+            )
+
+            if i in self.stride_layer_idx:
+                # compute time dimension for next block
+                efficient_index = self.stride_layer_idx.index(i)
+                att_mask = att_mask[:, ::self.stride[efficient_index], ::self.
+                                    stride[efficient_index], ]
+                mask_pad = mask_pad[:, ::self.stride[efficient_index], ::self.
+                                    stride[efficient_index], ]
+                pos_emb = pos_emb[:, ::self.stride[efficient_index], :]
+
+            # shape(new_att_cache) = [batch, head, time2, outdim]
+            new_att_cache = new_att_cache[:, :, next_cache_start // factor:, :]
+            # shape(new_cnn_cache) = [batch, 1, outdim, cache_t2]
+            new_cnn_cache = new_cnn_cache.unsqueeze(1)  # shape(1):layerID
+
+            # use repeat_interleave to new_att_cache
+            # new_att_cache = new_att_cache.repeat_interleave(repeats=factor, dim=2)
+            new_att_cache = (new_att_cache.unsqueeze(3).repeat(
+                1, 1, 1, factor, 1).flatten(2, 3))
+            # padding new_cnn_cache to cnn.lorder for casual convolution
+            new_cnn_cache = F.pad(
+                new_cnn_cache,
+                (self.cnn_module_kernel - 1 - new_cnn_cache.size(3), 0),
+            )
+
+            if i == 0:
+                # record length for the first block as max length
+                max_att_len = new_att_cache.size(2)
+                max_cnn_len = new_cnn_cache.size(3)
+
+            # update real shape of att_cache and cnn_cache
+            r_att_cache.append(new_att_cache[:, :,
+                                             -max_att_len:, :].unsqueeze(1))
+            r_cnn_cache.append(new_cnn_cache[:, :, :, -max_cnn_len:])
+
+        if self.encoder.normalize_before:
+            chunk_out = self.encoder.after_norm(xs)
+        else:
+            chunk_out = xs
+
+        # shape of r_att_cache: (b, elayers, head, time2, outdim)
+        r_att_cache = torch.cat(r_att_cache, dim=1)  # concat on layers idx
+        # shape of r_cnn_cache: (b, elayers, outdim, cache_t2)
+        r_cnn_cache = torch.cat(r_cnn_cache, dim=1)  # concat on layers
+
+        # <---------forward_chunk END--------->
+
+        # log_ctc_probs = self.ctc.log_softmax(chunk_out)
+        log_ctc_probs = self.ctc.linear(chunk_out)
+        log_probs, log_probs_idx = torch.topk(log_ctc_probs,
+                                              self.beam_size,
+                                              dim=2)
+        log_probs = log_probs.to(chunk_xs.dtype)
+
+        r_offset = offset + chunk_out.shape[1]
+        # the below ops not supported in Tensorrt
+        # chunk_out_lens = torch.div(chunk_lens, subsampling_rate,
+        #                   rounding_mode='floor')
+        chunk_out_lens = (
+            chunk_lens // self.subsampling_rate //
+            self.calculate_downsampling_factor(self.num_blocks + 1))
+        chunk_out_lens += 1
+        r_offset = r_offset.unsqueeze(1)
+
+        return (
+            log_probs,
+            log_probs_idx,
+            chunk_out,
+            chunk_out_lens,
+            r_offset,
+            r_att_cache,
+            r_cnn_cache,
+            r_cache_mask,
+        )
+
+
+class Decoder(torch.nn.Module):
+
+    def __init__(
+        self,
+        decoder: TransformerDecoder,
+        ctc_weight: float = 0.5,
+        reverse_weight: float = 0.0,
+        beam_size: int = 10,
+        decoder_fastertransformer: bool = False,
+    ):
+        super().__init__()
+        self.decoder = decoder
+        self.ctc_weight = ctc_weight
+        self.reverse_weight = reverse_weight
+        self.beam_size = beam_size
+        self.decoder_fastertransformer = decoder_fastertransformer
+
+    def forward(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_lens: torch.Tensor,
+        hyps_pad_sos_eos: torch.Tensor,
+        hyps_lens_sos: torch.Tensor,
+        r_hyps_pad_sos_eos: torch.Tensor,
+        ctc_score: torch.Tensor,
+    ):
+        """Encoder
+        Args:
+            encoder_out: B x T x F
+            encoder_lens: B
+            hyps_pad_sos_eos: B x beam x (T2+1),
+                        hyps with sos & eos and padded by ignore id
+            hyps_lens_sos: B x beam, length for each hyp with sos
+            r_hyps_pad_sos_eos: B x beam x (T2+1),
+                    reversed hyps with sos & eos and padded by ignore id
+            ctc_score: B x beam, ctc score for each hyp
+        Returns:
+            decoder_out: B x beam x T2 x V
+            r_decoder_out: B x beam x T2 x V
+            best_index: B
+        """
+        B, T, F = encoder_out.shape
+        bz = self.beam_size
+        B2 = B * bz
+        encoder_out = encoder_out.repeat(1, bz, 1).view(B2, T, F)
+        encoder_mask = ~make_pad_mask(encoder_lens, T).unsqueeze(1)
+        encoder_mask = encoder_mask.repeat(1, bz, 1).view(B2, 1, T)
+        T2 = hyps_pad_sos_eos.shape[2] - 1
+        hyps_pad = hyps_pad_sos_eos.view(B2, T2 + 1)
+        hyps_lens = hyps_lens_sos.view(B2, )
+        hyps_pad_sos = hyps_pad[:, :-1].contiguous()
+        hyps_pad_eos = hyps_pad[:, 1:].contiguous()
+
+        r_hyps_pad = r_hyps_pad_sos_eos.view(B2, T2 + 1)
+        r_hyps_pad_sos = r_hyps_pad[:, :-1].contiguous()
+        r_hyps_pad_eos = r_hyps_pad[:, 1:].contiguous()
+
+        decoder_out, r_decoder_out, _ = self.decoder(
+            encoder_out,
+            encoder_mask,
+            hyps_pad_sos,
+            hyps_lens,
+            r_hyps_pad_sos,
+            self.reverse_weight,
+        )
+        # decoder_out = torch.nn.functional.log_softmax(decoder_out, dim=-1)
+        V = decoder_out.shape[-1]
+        decoder_out = decoder_out.view(B2, T2, V)
+        mask = ~make_pad_mask(hyps_lens, T2)  # B2 x T2
+        # mask index, remove ignore id
+        index = torch.unsqueeze(hyps_pad_eos * mask, 2).to(torch.long)
+        score = decoder_out.gather(2, index).squeeze(2)  # B2 X T2
+        # mask padded part
+        score = score * mask
+        decoder_out = decoder_out.view(B, bz, T2, V)
+        if self.reverse_weight > 0:
+            # r_decoder_out = torch.nn.functional.log_softmax(r_decoder_out,
+            #                                                 dim=-1)
+            r_decoder_out = r_decoder_out.view(B2, T2, V)
+            index = torch.unsqueeze(r_hyps_pad_eos * mask, 2).to(torch.long)
+            r_score = r_decoder_out.gather(2, index).squeeze(2)
+            r_score = r_score * mask
+            score = (score * (1 - self.reverse_weight) +
+                     self.reverse_weight * r_score)
+            r_decoder_out = r_decoder_out.view(B, bz, T2, V)
+        score = torch.sum(score, axis=1)  # B2
+        score = torch.reshape(score, (B, bz)) + self.ctc_weight * ctc_score
+        best_index = torch.argmax(score, dim=1)
+        if self.decoder_fastertransformer:
+            return decoder_out, best_index
+        else:
+            return best_index
+
+
+def to_numpy(tensors):
+    out = []
+    if type(tensors) == torch.tensor:
+        tensors = [tensors]
+    for tensor in tensors:
+        if tensor.requires_grad:
+            tensor = tensor.detach().cpu().numpy()
+        else:
+            tensor = tensor.cpu().numpy()
+        out.append(tensor)
+    return out
+
+
+def test(xlist, blist, rtol=1e-3, atol=1e-5, tolerate_small_mismatch=True):
+    for a, b in zip(xlist, blist):
+        try:
+            torch.testing.assert_allclose(a, b, rtol=rtol, atol=atol)
+        except AssertionError as error:
+            if tolerate_small_mismatch:
+                print(error)
+            else:
+                raise
+
+
+def export_offline_encoder(model, configs, args, logger, encoder_onnx_path):
+    bz = 1
+    seq_len = 1024
+    beam_size = args.beam_size
+    feature_size = configs["input_dim"]
+
+    speech = torch.randn(bz, seq_len, feature_size, dtype=torch.float32)
+    speech_lens = torch.randint(low=10,
+                                high=seq_len,
+                                size=(bz, ),
+                                dtype=torch.int32)
+    encoder = Encoder(model.encoder, model.ctc, beam_size)
+    encoder.eval()
+
+    torch.onnx.export(
+        encoder,
+        (speech, speech_lens),
+        encoder_onnx_path,
+        export_params=True,
+        opset_version=13,
+        do_constant_folding=True,
+        input_names=["speech", "speech_lengths"],
+        output_names=[
+            "encoder_out",
+            "encoder_out_lens",
+            "ctc_log_probs",
+            "beam_log_probs",
+            "beam_log_probs_idx",
+        ],
+        dynamic_axes=None,
+        # dynamic_axes={
+        #     "speech": {
+        #         0: "B",
+        #         1: "T"
+        #     },
+        #     "speech_lengths": {
+        #         0: "B"
+        #     },
+        #     "encoder_out": {
+        #         0: "B",
+        #         1: "T_OUT"
+        #     },
+        #     "encoder_out_lens": {
+        #         0: "B"
+        #     },
+        #     "ctc_log_probs": {
+        #         0: "B",
+        #         1: "T_OUT"
+        #     },
+        #     "beam_log_probs": {
+        #         0: "B",
+        #         1: "T_OUT"
+        #     },
+        #     "beam_log_probs_idx": {
+        #         0: "B",
+        #         1: "T_OUT"
+        #     },
+        # },
+        verbose=False,
+        dynamo=False,
+    )
+    fold_static_pulsar2_subgraphs(encoder_onnx_path)
+    simplify_pulsar2_onnx(encoder_onnx_path)
+    rewrite_pulsar2_bool_not(encoder_onnx_path)
+
+    with torch.no_grad():
+        o0, o1, o2, o3, o4 = encoder(speech, speech_lens)
+
+    providers = ["CPUExecutionProvider"]
+    ort_session = onnxruntime.InferenceSession(encoder_onnx_path,
+                                               providers=providers)
+    ort_inputs = {
+        "speech": to_numpy(speech),
+        "speech_lengths": to_numpy(speech_lens),
+    }
+    ort_outs = ort_session.run(None, ort_inputs)
+
+    # check encoder output
+    test(to_numpy([o0, o1, o2, o3, o4]), ort_outs)
+    logger.info("export offline onnx encoder succeed!")
+    onnx_config = {
+        "beam_size": args.beam_size,
+        "reverse_weight": configs["model_conf"]["reverse_weight"],
+        "ctc_weight": configs["model_conf"]["ctc_weight"],
+    }
+    return onnx_config
+
+
+def export_online_encoder(model, configs, args, logger, encoder_onnx_path):
+    decoding_chunk_size = args.decoding_chunk_size
+    subsampling = model.encoder.embed.subsampling_rate
+    context = model.encoder.embed.right_context + 1
+    decoding_window = (decoding_chunk_size - 1) * subsampling + context
+    batch_size = 1
+    audio_len = decoding_window
+    feature_size = configs["input_dim"]
+    output_size = configs["encoder_conf"]["output_size"]
+    num_layers = configs["encoder_conf"]["num_blocks"]
+    # in transformer the cnn module will not be available
+    transformer = False
+    cnn_module_kernel = configs["encoder_conf"].get("cnn_module_kernel", 1) - 1
+    if not cnn_module_kernel:
+        transformer = True
+    num_decoding_left_chunks = args.num_decoding_left_chunks
+    required_cache_size = decoding_chunk_size * num_decoding_left_chunks
+    if configs["encoder"] == "squeezeformer":
+        encoder = StreamingSqueezeformerEncoder(model, required_cache_size,
+                                                args.beam_size)
+    elif configs["encoder"] == "efficientConformer":
+        encoder = StreamingEfficientConformerEncoder(model,
+                                                     required_cache_size,
+                                                     args.beam_size)
+    else:
+        encoder = StreamingEncoder(
+            model,
+            required_cache_size,
+            args.beam_size,
+            transformer,
+            args.return_ctc_logprobs,
+        )
+    encoder.eval()
+
+    # begin to export encoder
+    chunk_xs = torch.randn(batch_size,
+                           audio_len,
+                           feature_size,
+                           dtype=torch.float32)
+    chunk_lens = torch.ones(batch_size, dtype=torch.int32) * audio_len
+
+    offset = torch.arange(0, batch_size, dtype=torch.int32).unsqueeze(1)
+    #  (elayers, b, head, cache_t1, d_k * 2)
+    head = configs["encoder_conf"]["attention_heads"]
+    d_k = configs["encoder_conf"]["output_size"] // head
+    att_cache = torch.randn(
+        batch_size,
+        num_layers,
+        head,
+        required_cache_size,
+        d_k * 2,
+        dtype=torch.float32,
+    )
+    cnn_cache = torch.randn(
+        batch_size,
+        num_layers,
+        output_size,
+        cnn_module_kernel,
+        dtype=torch.float32,
+    )
+
+    cache_mask = torch.ones(batch_size,
+                            1,
+                            required_cache_size,
+                            dtype=torch.float32)
+    input_names = [
+        "chunk_xs",
+        "chunk_lens",
+        "offset",
+        "att_cache",
+        "cnn_cache",
+        "cache_mask",
+    ]
+    output_names = [
+        "log_probs",
+        "log_probs_idx",
+        "chunk_out",
+        "chunk_out_lens",
+        "r_offset",
+        "r_att_cache",
+        "r_cnn_cache",
+        "r_cache_mask",
+    ]
+    if args.return_ctc_logprobs:
+        output_names = [
+            "ctc_log_probs",
+            "chunk_out",
+            "chunk_out_lens",
+            "r_offset",
+            "r_att_cache",
+            "r_cnn_cache",
+            "r_cache_mask",
+        ]
+    input_tensors = (
+        chunk_xs,
+        chunk_lens,
+        offset,
+        att_cache,
+        cnn_cache,
+        cache_mask,
+    )
+    if transformer:
+        assert (args.return_ctc_logprobs is
+                False), "return_ctc_logprobs is not supported in transformer"
+        output_names.pop(6)
+
+    all_names = input_names + output_names
+    dynamic_axes = {}
+    for name in all_names:
+        # only the first dimension is dynamic
+        # all other dimension is fixed
+        dynamic_axes[name] = {0: "B"}
+
+    torch.onnx.export(
+        encoder,
+        input_tensors,
+        encoder_onnx_path,
+        export_params=True,
+        opset_version=14,
+        do_constant_folding=True,
+        input_names=input_names,
+        output_names=output_names,
+        # dynamic_axes=dynamic_axes,
+        dynamic_axes=None,
+        verbose=False,
+        dynamo=False,
+    )
+    fold_static_pulsar2_subgraphs(encoder_onnx_path)
+    simplify_pulsar2_onnx(encoder_onnx_path)
+    rewrite_pulsar2_bool_not(encoder_onnx_path)
+
+    with torch.no_grad():
+        torch_outs = encoder(chunk_xs, chunk_lens, offset, att_cache,
+                             cnn_cache, cache_mask)
+    if transformer:
+        torch_outs = list(torch_outs).pop(6)
+    ort_session = onnxruntime.InferenceSession(
+        encoder_onnx_path, providers=["CPUExecutionProvider"])
+    ort_inputs = {}
+
+    input_tensors = to_numpy(input_tensors)
+    for idx, name in enumerate(input_names):
+        ort_inputs[name] = input_tensors[idx]
+    if transformer:
+        del ort_inputs["cnn_cache"]
+    ort_outs = ort_session.run(None, ort_inputs)
+    test(to_numpy(torch_outs), ort_outs, rtol=1e-03, atol=1e-05)
+    logger.info("export to onnx streaming encoder succeed!")
+    onnx_config = {
+        "subsampling_rate": subsampling,
+        "context": context,
+        "decoding_chunk_size": decoding_chunk_size,
+        "num_decoding_left_chunks": num_decoding_left_chunks,
+        "beam_size": args.beam_size,
+        "feat_size": feature_size,
+        "decoding_window": decoding_window,
+        "cnn_module_kernel_cache": cnn_module_kernel,
+        "return_ctc_logprobs": args.return_ctc_logprobs,
+    }
+    return onnx_config
+
+
+def export_rescoring_decoder(model, configs, args, logger, decoder_onnx_path,
+                             decoder_fastertransformer):
+    bz, seq_len = 1, 32
+    beam_size = args.beam_size
+    decoder = Decoder(
+        model.decoder,
+        model.ctc_weight,
+        model.reverse_weight,
+        beam_size,
+        decoder_fastertransformer,
+    )
+    decoder.eval()
+
+    hyps_pad_sos_eos = torch.randint(low=3,
+                                     high=1000,
+                                     size=(bz, beam_size, seq_len),
+                                     dtype=torch.int32)
+    hyps_lens_sos = torch.randint(low=3,
+                                  high=seq_len,
+                                  size=(bz, beam_size),
+                                  dtype=torch.int32)
+    r_hyps_pad_sos_eos = torch.randint(low=3,
+                                       high=1000,
+                                       size=(bz, beam_size, seq_len),
+                                       dtype=torch.int32)
+
+    output_size = configs["encoder_conf"]["output_size"]
+    encoder_out = torch.randn(bz, seq_len, output_size, dtype=torch.float32)
+    encoder_out_lens = torch.randint(low=3,
+                                     high=seq_len,
+                                     size=(bz, ),
+                                     dtype=torch.int32)
+    ctc_score = torch.randn(bz, beam_size, dtype=torch.float32)
+
+    input_names = [
+        "encoder_out",
+        "encoder_out_lens",
+        "hyps_pad_sos_eos",
+        "hyps_lens_sos",
+        "r_hyps_pad_sos_eos",
+        "ctc_score",
+    ]
+    output_names = ["best_index"]
+    if decoder_fastertransformer:
+        output_names.insert(0, "decoder_out")
+
+    torch.onnx.export(
+        decoder,
+        (
+            encoder_out,
+            encoder_out_lens,
+            hyps_pad_sos_eos,
+            hyps_lens_sos,
+            r_hyps_pad_sos_eos,
+            ctc_score,
+        ),
+        decoder_onnx_path,
+        export_params=True,
+        opset_version=13,
+        do_constant_folding=True,
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes=None,
+        # dynamic_axes={
+        #     "encoder_out": {
+        #         0: "B",
+        #         1: "T"
+        #     },
+        #     "encoder_out_lens": {
+        #         0: "B"
+        #     },
+        #     "hyps_pad_sos_eos": {
+        #         0: "B",
+        #         2: "T2"
+        #     },
+        #     "hyps_lens_sos": {
+        #         0: "B"
+        #     },
+        #     "r_hyps_pad_sos_eos": {
+        #         0: "B",
+        #         2: "T2"
+        #     },
+        #     "ctc_score": {
+        #         0: "B"
+        #     },
+        #     "best_index": {
+        #         0: "B"
+        #     },
+        # },
+        verbose=False,
+        dynamo=False,
+    )
+    fold_static_pulsar2_subgraphs(decoder_onnx_path)
+    simplify_pulsar2_onnx(decoder_onnx_path)
+    rewrite_pulsar2_bool_not(decoder_onnx_path)
+    rewrite_pulsar2_bool_and(decoder_onnx_path)
+    with torch.no_grad():
+        o0 = decoder(
+            encoder_out,
+            encoder_out_lens,
+            hyps_pad_sos_eos,
+            hyps_lens_sos,
+            r_hyps_pad_sos_eos,
+            ctc_score,
+        )
+    providers = ["CPUExecutionProvider"]
+    ort_session = onnxruntime.InferenceSession(decoder_onnx_path,
+                                               providers=providers)
+
+    input_tensors = [
+        encoder_out,
+        encoder_out_lens,
+        hyps_pad_sos_eos,
+        hyps_lens_sos,
+        r_hyps_pad_sos_eos,
+        ctc_score,
+    ]
+    ort_inputs = {}
+    input_tensors = to_numpy(input_tensors)
+    for idx, name in enumerate(input_names):
+        ort_inputs[name] = input_tensors[idx]
+
+    # if model.reverse weight == 0,
+    # the r_hyps_pad will be removed
+    # from the onnx decoder since it doen't play any role
+    if model.reverse_weight == 0:
+        del ort_inputs["r_hyps_pad_sos_eos"]
+    ort_outs = ort_session.run(None, ort_inputs)
+
+    # check decoder output
+    if decoder_fastertransformer:
+        test(to_numpy(o0), ort_outs, rtol=1e-03, atol=1e-05)
+    else:
+        test(to_numpy([o0]), ort_outs, rtol=1e-03, atol=1e-05)
+    logger.info("export to onnx decoder succeed!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="export x86_gpu model")
+    parser.add_argument(
+        "--pretrained_model_dir",
+        default=DEFAULT_PRETRAINED_MODEL_DIR,
+        help=("pretrained model directory containing train.yaml, final.pt, "
+              "and global_cmvn"),
+    )
+    parser.add_argument(
+        "--pretrained_model_url",
+        default=DEFAULT_PRETRAINED_MODEL_URL,
+        help="pretrained model tar.gz URL used when pretrained_model_dir is missing",
+    )
+    parser.add_argument(
+        "--reverse_weight",
+        default=-1.0,
+        type=float,
+        required=False,
+        help="reverse weight for bitransformer," +
+        "default value is in config file",
+    )
+    parser.add_argument(
+        "--ctc_weight",
+        default=-1.0,
+        type=float,
+        required=False,
+        help="ctc weight, default value is in config file",
+    )
+    parser.add_argument(
+        "--beam_size",
+        default=10,
+        type=int,
+        required=False,
+        help="beam size would be ctc output size",
+    )
+    parser.add_argument(
+        "--output_onnx_dir",
+        default="onnx_model",
+        help="output onnx encoder and decoder directory",
+    )
+    # arguments for streaming encoder
+    # parser.add_argument(
+    #     "--streaming",
+    #     action="store_true",
+    #     help="whether to export streaming encoder, default false",
+    # )
+    parser.add_argument(
+        "--decoding_chunk_size",
+        default=16,
+        type=int,
+        required=False,
+        help="the decoding chunk size, <=0 is not supported",
+    )
+    parser.add_argument(
+        "--num_decoding_left_chunks",
+        default=5,
+        type=int,
+        required=False,
+        help="number of left chunks, <= 0 is not supported",
+    )
+    parser.add_argument(
+        "--decoder_fastertransformer",
+        action="store_true",
+        help="return decoder_out and best_index for ft",
+    )
+    parser.add_argument(
+        "--return_ctc_logprobs",
+        action="store_true",
+        help="return full ctc_log_probs for TLG streaming encoder",
+    )
+    args = parser.parse_args()
+    prepare_pretrained_model(args)
+
+    torch.manual_seed(0)
+    torch.set_printoptions(precision=10)
+
+    with open(args.config, "r") as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+    if os.path.exists(args.cmvn_file):
+        if 'cmvn' not in configs:
+            configs['cmvn'] = "global_cmvn"
+            configs['cmvn_conf'] = {}
+        else:
+            assert configs['cmvn'] == "global_cmvn"
+            assert configs['cmvn_conf'] is not None
+        configs['cmvn_conf']["cmvn_file"] = args.cmvn_file
+        configs['cmvn_conf'].setdefault(
+            "is_json_cmvn", configs.get("is_json_cmvn", True))
+    elif configs.get('cmvn', None) == 'global_cmvn':
+        raise FileNotFoundError(
+            f"Expected global_cmvn in pretrained model dir: {args.cmvn_file}")
+    if (args.reverse_weight != -1.0
+            and "reverse_weight" in configs["model_conf"]):
+        configs["model_conf"]["reverse_weight"] = args.reverse_weight
+        print("Update reverse weight to", args.reverse_weight)
+    if args.ctc_weight != -1:
+        print("Update ctc weight to ", args.ctc_weight)
+        configs["model_conf"]["ctc_weight"] = args.ctc_weight
+    configs["encoder_conf"]["use_dynamic_chunk"] = False
+
+    model, configs = init_model(args, configs)
+    model.eval()
+
+    if not os.path.exists(args.output_onnx_dir):
+        os.mkdir(args.output_onnx_dir)
+    
+    export_enc_func = None
+    # if args.streaming:
+    assert args.decoding_chunk_size > 0
+    assert args.num_decoding_left_chunks > 0
+    export_enc_func = export_online_encoder
+    encoder_onnx_path = os.path.join(args.output_onnx_dir, "encoder_online.onnx")
+    onnx_config = export_enc_func(model, configs, args, logger,
+                                encoder_onnx_path)
+    
+    # else
+    export_enc_func = export_offline_encoder
+    encoder_onnx_path = os.path.join(args.output_onnx_dir, "encoder_offline.onnx")
+    onnx_config = export_enc_func(model, configs, args, logger,
+                                encoder_onnx_path)
+    
+
+    decoder_onnx_path = os.path.join(args.output_onnx_dir, "decoder.onnx")
+    export_rescoring_decoder(
+        model,
+        configs,
+        args,
+        logger,
+        decoder_onnx_path,
+        args.decoder_fastertransformer,
+    )
+
+    config_dir = os.path.join(args.output_onnx_dir, "config.yaml")
+    with open(config_dir, "w") as out:
+        yaml.dump(onnx_config, out)

generate_data.py

@@ -0,0 +1,122 @@
+import argparse
+import os
+import shutil
+
+from tqdm import tqdm
+
+from ort_common import WenetONNXRunner, pack_calibration_dataset
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description="Generate calibration_dataset for exported ONNX models")
+    parser.add_argument("--input",
+                        "-i",
+                        nargs="+",
+                        required=True,
+                        help="Input wav file(s) or directory/directories")
+    parser.add_argument("--config",
+                        required=True,
+                        help="yaml file in checkpoint path")
+    parser.add_argument(
+        "--vocab",
+        required=True,
+        help="pretrained units.txt, for example pretrained/<model>/units.txt",
+    )
+    parser.add_argument("--onnx_dir",
+                        default="onnx_model",
+                        help="directory containing exported ONNX models")
+    parser.add_argument("--calib_data_path",
+                        default="calibration_dataset",
+                        help="output calibration dataset directory")
+    parser.add_argument("--parts",
+                        nargs="+",
+                        choices=["all", "offline", "online", "decoder"],
+                        default=["all"],
+                        help="which model inputs to generate")
+    parser.add_argument("--offline_seq_len", type=int, default=1024)
+    parser.add_argument("--decoder_len", type=int, default=32)
+    parser.add_argument("--decoding_chunk_size", type=int, default=16)
+    parser.add_argument("--num_decoding_left_chunks", type=int, default=5)
+    parser.add_argument("--max_num",
+                        type=int,
+                        default=100,
+                        help="maximum number of audio files used for calibration; set <= 0 to use all")
+    parser.add_argument("--keep_existing",
+                        action="store_true",
+                        help="append to an existing calibration directory")
+    return parser.parse_args()
+
+
+def expand_audio_inputs(inputs):
+    audio_exts = {".wav", ".flac", ".mp3", ".m4a", ".ogg"}
+    audio_files = []
+    for path in inputs:
+        if os.path.isdir(path):
+            for root, _, files in os.walk(path):
+                for filename in files:
+                    if os.path.splitext(filename)[1].lower() in audio_exts:
+                        audio_files.append(os.path.join(root, filename))
+        else:
+            audio_files.append(path)
+    audio_files = sorted(audio_files)
+    if not audio_files:
+        raise FileNotFoundError("No audio files found")
+    return audio_files
+
+
+def normalize_parts(parts):
+    if "all" in parts:
+        return {"offline", "online", "decoder"}
+    return set(parts)
+
+
+def limit_audio_files(audio_files, max_num):
+    if max_num is None or max_num <= 0:
+        return audio_files
+    return audio_files[:max_num]
+
+
+def main():
+    args = get_args()
+    parts = normalize_parts(args.parts)
+    audio_files = limit_audio_files(expand_audio_inputs(args.input),
+                                    args.max_num)
+
+    if os.path.exists(args.calib_data_path) and not args.keep_existing:
+        shutil.rmtree(args.calib_data_path)
+    os.makedirs(args.calib_data_path, exist_ok=True)
+
+    runner = WenetONNXRunner(
+        args.config,
+        args.vocab,
+        onnx_dir=args.onnx_dir,
+        offline_seq_len=args.offline_seq_len,
+        decoder_len=args.decoder_len,
+        decoding_chunk_size=args.decoding_chunk_size,
+        num_decoding_left_chunks=args.num_decoding_left_chunks,
+    )
+
+    counts = {"offline": 0, "online": 0, "decoder": 0}
+    progress = tqdm(audio_files,
+                    desc="Generating calibration data",
+                    unit="wav")
+    for audio_idx, audio_file in enumerate(progress):
+        sample_counts = runner.save_calibration_for_audio(
+            audio_file, parts, args.calib_data_path, audio_idx)
+        for key, value in sample_counts.items():
+            counts[key] += value
+        progress.set_postfix(offline=counts["offline"],
+                             online=counts["online"],
+                             decoder=counts["decoder"])
+
+    print("Packing calibration dataset...")
+    pack_calibration_dataset(args.calib_data_path)
+    print(f"Generated calibration data in {args.calib_data_path}")
+    print(f"offline samples: {counts['offline']}")
+    print(f"online samples: {counts['online']}")
+    print(f"decoder samples: {counts['decoder']}")
+
+
+if __name__ == "__main__":
+    main()

ort_common.py

@@ -0,0 +1,567 @@
+import multiprocessing
+import os
+import tarfile as tf
+
+import numpy as np
+import torch
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+import yaml
+
+try:
+    from swig_decoders import (map_batch, ctc_beam_search_decoder_batch,
+                               TrieVector, PathTrie)
+except ModuleNotFoundError:
+    class PathTrie:
+        pass
+
+    class TrieVector(list):
+        pass
+
+    def _log_add(*values):
+        values = [value for value in values if value != -float("inf")]
+        if not values:
+            return -float("inf")
+        max_value = max(values)
+        return max_value + np.log(sum(np.exp(value - max_value)
+                                      for value in values))
+
+    def _map_sentence(sent, vocabulary, greedy=False, blank_id=0):
+        mapped = []
+        prev = None
+        for token in sent:
+            token = int(token)
+            if greedy and token == prev:
+                prev = token
+                continue
+            prev = token
+            if token == blank_id or token < 0 or token >= len(vocabulary):
+                continue
+            piece = vocabulary[token]
+            if piece.startswith("<") and piece.endswith(">"):
+                continue
+            mapped.append(piece)
+        return "".join(mapped)
+
+    def map_batch(batch_sents, vocabulary, num_processes, greedy=False,
+                  blank_id=0):
+        del num_processes
+        return [_map_sentence(sent, vocabulary, greedy, blank_id)
+                for sent in batch_sents]
+
+    def _ctc_prefix_beam_search(log_probs_seq, log_probs_idx, beam_size,
+                                blank_id):
+        beam = {(): (0.0, -float("inf"))}
+        for frame_probs, frame_ids in zip(log_probs_seq, log_probs_idx):
+            next_beam = {}
+            for prefix, (prob_blank, prob_non_blank) in beam.items():
+                for prob, token in zip(frame_probs, frame_ids):
+                    token = int(token)
+                    prob = float(prob)
+                    next_prob_blank, next_prob_non_blank = next_beam.get(
+                        prefix, (-float("inf"), -float("inf")))
+                    if token == blank_id:
+                        next_beam[prefix] = (
+                            _log_add(next_prob_blank, prob_blank + prob,
+                                     prob_non_blank + prob),
+                            next_prob_non_blank,
+                        )
+                        continue
+
+                    last = prefix[-1] if prefix else None
+                    if token == last:
+                        next_beam[prefix] = (
+                            next_prob_blank,
+                            _log_add(next_prob_non_blank,
+                                     prob_non_blank + prob),
+                        )
+                        new_prefix = prefix + (token, )
+                        nb_blank, nb_non_blank = next_beam.get(
+                            new_prefix, (-float("inf"), -float("inf")))
+                        next_beam[new_prefix] = (
+                            nb_blank,
+                            _log_add(nb_non_blank, prob_blank + prob),
+                        )
+                    else:
+                        new_prefix = prefix + (token, )
+                        nb_blank, nb_non_blank = next_beam.get(
+                            new_prefix, (-float("inf"), -float("inf")))
+                        next_beam[new_prefix] = (
+                            nb_blank,
+                            _log_add(nb_non_blank, prob_blank + prob,
+                                     prob_non_blank + prob),
+                        )
+            beam = dict(sorted(
+                next_beam.items(),
+                key=lambda item: _log_add(item[1][0], item[1][1]),
+                reverse=True)[:beam_size])
+        return [(_log_add(prob_blank, prob_non_blank), prefix)
+                for prefix, (prob_blank, prob_non_blank) in sorted(
+                    beam.items(),
+                    key=lambda item: _log_add(item[1][0], item[1][1]),
+                    reverse=True)]
+
+    def ctc_beam_search_decoder_batch(batch_log_probs_seq,
+                                      batch_log_probs_idx,
+                                      batch_root_trie,
+                                      batch_start,
+                                      beam_size,
+                                      num_processes,
+                                      blank_id=0,
+                                      space_id=-1,
+                                      cutoff_prob=0.999,
+                                      ext_scorer=None):
+        del batch_root_trie, batch_start, num_processes, space_id
+        del cutoff_prob, ext_scorer
+        return [
+            _ctc_prefix_beam_search(log_probs_seq, log_probs_idx, beam_size,
+                                    blank_id)
+            for log_probs_seq, log_probs_idx in zip(batch_log_probs_seq,
+                                                    batch_log_probs_idx)
+        ]
+
+
+def load_config(config_path):
+    with open(config_path, "r") as fin:
+        return yaml.load(fin, Loader=yaml.FullLoader)
+
+
+def load_vocab(vocab_path):
+    vocabulary = []
+    char_dict = {}
+    with open(vocab_path, "r") as fin:
+        for line in fin:
+            arr = line.strip().split()
+            assert len(arr) == 2
+            char_dict[int(arr[1])] = arr[0]
+            vocabulary.append(arr[0])
+    return vocabulary, char_dict
+
+
+def compute_feats(audio_file: str, sr=16000) -> np.ndarray:
+    try:
+        import soundfile as sf
+        waveform, sample_rate = sf.read(audio_file, dtype="int16",
+                                        always_2d=True)
+        waveform = torch.from_numpy(waveform.T).to(torch.float)
+    except ModuleNotFoundError:
+        waveform, sample_rate = torchaudio.load(audio_file, normalize=True)
+        waveform = waveform.to(torch.float)
+    if sample_rate != sr:
+        waveform = torchaudio.transforms.Resample(
+            orig_freq=sample_rate, new_freq=sr)(waveform)
+    feats = kaldi.fbank(waveform,
+                        num_mel_bins=80,
+                        frame_length=25,
+                        frame_shift=10,
+                        energy_floor=0.0,
+                        sample_frequency=sr)
+    return feats.unsqueeze(0).numpy()
+
+
+def pad_array_along_axis(array, pad_width, axis, mode="constant", **kwargs):
+    if array.shape[axis] >= pad_width:
+        return array
+    full_pad_width = [(0, 0)] * array.ndim
+    full_pad_width[axis] = (0, pad_width - array.shape[axis])
+    return np.pad(array, pad_width=full_pad_width, mode=mode, **kwargs)
+
+
+def ctc_decoding(beam_log_probs,
+                 beam_log_probs_idx,
+                 encoder_out_lens,
+                 vocabulary,
+                 mode="ctc_prefix_beam_search"):
+    beam_size = beam_log_probs.shape[-1]
+    batch_size = beam_log_probs.shape[0]
+    num_processes = min(multiprocessing.cpu_count(), batch_size)
+    hyps = []
+    score_hyps = []
+
+    if mode == "ctc_greedy_search":
+        log_probs_idx = beam_log_probs_idx[:, :, 0]
+        batch_sents = []
+        for idx, seq in enumerate(log_probs_idx):
+            batch_sents.append(seq[0:encoder_out_lens[idx]].tolist())
+        hyps = map_batch(batch_sents, vocabulary, num_processes, True, 0)
+    elif mode in ("ctc_prefix_beam_search", "attention_rescoring"):
+        batch_log_probs_seq_list = beam_log_probs.tolist()
+        batch_log_probs_idx_list = beam_log_probs_idx.tolist()
+        batch_len_list = encoder_out_lens.tolist()
+        batch_log_probs_seq = []
+        batch_log_probs_ids = []
+        batch_start = []
+        batch_root = TrieVector()
+        root_dict = {}
+        for i in range(len(batch_len_list)):
+            num_sent = batch_len_list[i]
+            batch_log_probs_seq.append(batch_log_probs_seq_list[i][0:num_sent])
+            batch_log_probs_ids.append(batch_log_probs_idx_list[i][0:num_sent])
+            root_dict[i] = PathTrie()
+            batch_root.append(root_dict[i])
+            batch_start.append(True)
+        score_hyps = ctc_beam_search_decoder_batch(batch_log_probs_seq,
+                                                   batch_log_probs_ids,
+                                                   batch_root,
+                                                   batch_start,
+                                                   beam_size,
+                                                   num_processes,
+                                                   0, -2, 0.99999)
+        if mode == "ctc_prefix_beam_search":
+            for cand_hyps in score_hyps:
+                hyps.append(cand_hyps[0][1])
+            hyps = map_batch(hyps, vocabulary, num_processes, False, 0)
+    return hyps, score_hyps
+
+
+def make_decoder_inputs(encoder_out,
+                        encoder_out_lens,
+                        beam_log_probs,
+                        beam_log_probs_idx,
+                        vocabulary,
+                        sos,
+                        eos,
+                        decoder_len):
+    _, score_hyps = ctc_decoding(beam_log_probs, beam_log_probs_idx,
+                                 encoder_out_lens, vocabulary,
+                                 "attention_rescoring")
+    ignore_id = -1
+    beam_size = beam_log_probs.shape[-1]
+    batch_size = beam_log_probs.shape[0]
+    ctc_score, all_hyps = [], []
+    for hyps in score_hyps:
+        cur_len = len(hyps)
+        if len(hyps) < beam_size:
+            hyps += (beam_size - cur_len) * [(-float("inf"), (0,))]
+        cur_ctc_score = []
+        for hyp in hyps:
+            cur_ctc_score.append(hyp[0])
+            all_hyps.append(list(hyp[1]))
+        ctc_score.append(cur_ctc_score)
+    ctc_score = np.array(ctc_score, dtype=np.float32)
+
+    max_len = decoder_len - 2
+    hyps_pad_sos_eos = np.ones((batch_size, beam_size, max_len + 2),
+                               dtype=np.int64) * ignore_id
+    r_hyps_pad_sos_eos = np.ones((batch_size, beam_size, max_len + 2),
+                                 dtype=np.int64) * ignore_id
+    hyps_lens_sos = np.ones((batch_size, beam_size), dtype=np.int32)
+    k = 0
+    for i in range(batch_size):
+        for j in range(beam_size):
+            cand = all_hyps[k][:max_len]
+            length = len(cand) + 2
+            hyps_pad_sos_eos[i][j][0:length] = [sos] + cand + [eos]
+            r_hyps_pad_sos_eos[i][j][0:length] = [sos] + cand[::-1] + [eos]
+            hyps_lens_sos[i][j] = len(cand) + 1
+            k += 1
+
+    if decoder_len > encoder_out.shape[1]:
+        encoder_out = np.pad(encoder_out,
+                             [(0, 0),
+                              (0, decoder_len - encoder_out.shape[1]),
+                              (0, 0)],
+                             mode="constant",
+                             constant_values=0)
+    elif decoder_len < encoder_out.shape[1]:
+        encoder_out = encoder_out[:, :decoder_len, :]
+
+    return {
+        "encoder_out": encoder_out,
+        "encoder_out_lens": np.full(batch_size,
+                                    fill_value=decoder_len,
+                                    dtype=np.int32),
+        "hyps_pad_sos_eos": hyps_pad_sos_eos.astype(np.int32),
+        "hyps_lens_sos": hyps_lens_sos,
+        "r_hyps_pad_sos_eos": r_hyps_pad_sos_eos.astype(np.int32),
+        "ctc_score": ctc_score,
+    }, all_hyps
+
+
+def make_offline_inputs(feats, seq_len):
+    feats = feats[:, :seq_len, :]
+    speech_lengths = np.array([feats.shape[1]], dtype=np.int32)
+    if feats.shape[1] < seq_len:
+        feats = pad_array_along_axis(feats, pad_width=seq_len, axis=1)
+    return {"speech": feats, "speech_lengths": speech_lengths}
+
+
+def make_online_initial_state(configs,
+                              batch_size=1,
+                              decoding_chunk_size=16,
+                              num_decoding_left_chunks=5):
+    subsampling = 4
+    context = 7
+    stride = subsampling * decoding_chunk_size
+    decoding_window = (decoding_chunk_size - 1) * subsampling + context
+    required_cache_size = decoding_chunk_size * num_decoding_left_chunks
+
+    output_size = configs["encoder_conf"]["output_size"]
+    num_layers = configs["encoder_conf"]["num_blocks"]
+    cnn_module_kernel = configs["encoder_conf"].get("cnn_module_kernel", 1) - 1
+    head = configs["encoder_conf"]["attention_heads"]
+    d_k = configs["encoder_conf"]["output_size"] // head
+
+    state = {
+        "att_cache": np.zeros((batch_size, num_layers, head,
+                               required_cache_size, d_k * 2),
+                              dtype=np.float32),
+        "cnn_cache": np.zeros((batch_size, num_layers, output_size,
+                               cnn_module_kernel),
+                              dtype=np.float32),
+        "cache_mask": np.zeros((batch_size, 1, required_cache_size),
+                               dtype=np.float32),
+        "offset": np.zeros((batch_size, 1), dtype=np.int32),
+    }
+    params = {
+        "batch_size": batch_size,
+        "context": context,
+        "stride": stride,
+        "decoding_window": decoding_window,
+    }
+    return state, params
+
+
+def make_online_encoder_input(feats, cur, params, state):
+    batch_size = params["batch_size"]
+    decoding_window = params["decoding_window"]
+    end = min(cur + decoding_window, feats.shape[1])
+    chunk_xs = feats[:, cur:end, :]
+    if chunk_xs.shape[1] < decoding_window:
+        chunk_xs = pad_array_along_axis(chunk_xs,
+                                        pad_width=decoding_window,
+                                        axis=1)
+    chunk_xs = chunk_xs.astype(np.float32)
+    chunk_lens = np.full(batch_size,
+                         fill_value=chunk_xs.shape[1],
+                         dtype=np.int32)
+    return {
+        "chunk_xs": chunk_xs,
+        "chunk_lens": chunk_lens,
+        "offset": state["offset"],
+        "att_cache": state["att_cache"],
+        "cnn_cache": state["cnn_cache"],
+        "cache_mask": state["cache_mask"],
+    }
+
+
+def update_online_state(state, outputs):
+    state["offset"] = outputs[4]
+    state["att_cache"] = outputs[5]
+    state["cnn_cache"] = outputs[6]
+    state["cache_mask"] = outputs[7]
+
+
+def save_calibration_inputs(calib_data_path, inputs, sample_id):
+    for input_name, data in inputs.items():
+        data_path = os.path.join(calib_data_path, input_name)
+        os.makedirs(data_path, exist_ok=True)
+        np.save(os.path.join(data_path, f"{sample_id}.npy"), data)
+
+
+def pack_calibration_dataset(calib_data_path):
+    for input_name in sorted(os.listdir(calib_data_path)):
+        data_path = os.path.join(calib_data_path, input_name)
+        if not os.path.isdir(data_path):
+            continue
+        tar_path = os.path.join(calib_data_path, input_name + ".tar.gz")
+        with tf.open(tar_path, "w:gz") as tf_file:
+            tf_file.add(data_path, arcname=input_name)
+
+
+class WenetONNXRunner:
+
+    def __init__(self,
+                 config_path,
+                 vocab_path,
+                 onnx_dir="onnx_model",
+                 offline_seq_len=1024,
+                 decoder_len=32,
+                 decoding_chunk_size=16,
+                 num_decoding_left_chunks=5,
+                 batch_size=1,
+                 providers=None):
+        self.config_path = config_path
+        self.vocab_path = vocab_path
+        self.onnx_dir = onnx_dir
+        self.offline_seq_len = offline_seq_len
+        self.decoder_len = decoder_len
+        self.decoding_chunk_size = decoding_chunk_size
+        self.num_decoding_left_chunks = num_decoding_left_chunks
+        self.batch_size = batch_size
+        self.providers = providers or ["CPUExecutionProvider"]
+
+        self.configs = load_config(config_path)
+        self.vocabulary, self.char_dict = load_vocab(vocab_path)
+        self.eos = self.sos = len(self.char_dict) - 1
+
+        self._offline_encoder = None
+        self._online_encoder = None
+        self._decoder = None
+
+    @property
+    def offline_encoder(self):
+        if self._offline_encoder is None:
+            self._offline_encoder = self._new_session("encoder_offline.onnx")
+        return self._offline_encoder
+
+    @property
+    def online_encoder(self):
+        if self._online_encoder is None:
+            self._online_encoder = self._new_session("encoder_online.onnx")
+        return self._online_encoder
+
+    @property
+    def decoder(self):
+        if self._decoder is None:
+            self._decoder = self._new_session("decoder.onnx")
+        return self._decoder
+
+    def _new_session(self, filename):
+        import onnxruntime as ort
+        return ort.InferenceSession(os.path.join(self.onnx_dir, filename),
+                                    providers=self.providers)
+
+    def compute_feats(self, audio_file):
+        return compute_feats(audio_file)
+
+    def run_offline_encoder(self,
+                            feats,
+                            calib_data_path=None,
+                            sample_id=0):
+        encoder_input = make_offline_inputs(feats, self.offline_seq_len)
+        if calib_data_path:
+            save_calibration_inputs(calib_data_path, encoder_input, sample_id)
+        encoder_out, encoder_out_lens, ctc_log_probs, beam_log_probs, beam_log_probs_idx = (
+            self.offline_encoder.run(None, encoder_input))
+        return {
+            "encoder_out": encoder_out,
+            "encoder_out_lens": encoder_out_lens,
+            "ctc_log_probs": ctc_log_probs,
+            "beam_log_probs": beam_log_probs,
+            "beam_log_probs_idx": beam_log_probs_idx,
+        }
+
+    def run_online_encoder(self,
+                           feats,
+                           calib_data_path=None,
+                           sample_prefix=0):
+        state, online_params = make_online_initial_state(
+            self.configs, self.batch_size, self.decoding_chunk_size,
+            self.num_decoding_left_chunks)
+        encoder_out = []
+        beam_log_probs = []
+        beam_log_probs_idx = []
+        num_frames = feats.shape[1]
+
+        for cur in range(0, num_frames - online_params["context"] + 1,
+                         online_params["stride"]):
+            encoder_input = make_online_encoder_input(feats, cur,
+                                                      online_params, state)
+            if calib_data_path:
+                save_calibration_inputs(calib_data_path, encoder_input,
+                                        f"{sample_prefix}_{cur}")
+            outputs = self.online_encoder.run(None, encoder_input)
+            chunk_log_probs, chunk_log_probs_idx, chunk_out, chunk_out_lens = outputs[:4]
+            update_online_state(state, outputs)
+            del chunk_out_lens
+            encoder_out.append(chunk_out)
+            beam_log_probs.append(chunk_log_probs)
+            beam_log_probs_idx.append(chunk_log_probs_idx.astype(np.int32))
+
+        return {
+            "encoder_out": np.concatenate(encoder_out, axis=1),
+            "encoder_out_lens": np.full(self.batch_size,
+                                        fill_value=sum(
+                                            out.shape[1]
+                                            for out in encoder_out),
+                                        dtype=np.int32),
+            "beam_log_probs": np.concatenate(beam_log_probs, axis=1),
+            "beam_log_probs_idx": np.concatenate(beam_log_probs_idx, axis=1),
+            "num_chunks": len(encoder_out),
+        }
+
+    def ctc_decode(self, encoder_outputs, mode):
+        return ctc_decoding(encoder_outputs["beam_log_probs"],
+                            encoder_outputs["beam_log_probs_idx"],
+                            encoder_outputs["encoder_out_lens"],
+                            self.vocabulary, mode)
+
+    def run_decoder(self,
+                    encoder_outputs,
+                    calib_data_path=None,
+                    sample_id=0):
+        decoder_input, all_hyps = make_decoder_inputs(
+            encoder_outputs["encoder_out"],
+            encoder_outputs["encoder_out_lens"],
+            encoder_outputs["beam_log_probs"],
+            encoder_outputs["beam_log_probs_idx"],
+            self.vocabulary,
+            self.sos,
+            self.eos,
+            self.decoder_len,
+        )
+        if calib_data_path:
+            save_calibration_inputs(calib_data_path, decoder_input, sample_id)
+
+        best_index = self.decoder.run(None, decoder_input)[0].astype(np.int32)
+        beam_size = encoder_outputs["beam_log_probs"].shape[-1]
+        num_processes = min(multiprocessing.cpu_count(), best_index.shape[0])
+        best_sents = []
+        k = 0
+        for idx in best_index:
+            best_sents.append(all_hyps[k:k + beam_size][idx])
+            k += beam_size
+        hyps = map_batch(best_sents, self.vocabulary, num_processes)
+        return "".join(hyps)
+
+    def transcribe(self,
+                   audio_file,
+                   online=False,
+                   mode="ctc_prefix_beam_search",
+                   calib_data_path=None):
+        feats = self.compute_feats(audio_file)
+        if online:
+            encoder_outputs = self.run_online_encoder(
+                feats, calib_data_path=calib_data_path, sample_prefix=0)
+        else:
+            encoder_outputs = self.run_offline_encoder(
+                feats, calib_data_path=calib_data_path, sample_id=0)
+
+        if mode == "attention_rescoring":
+            result = self.run_decoder(encoder_outputs, calib_data_path, 0)
+        else:
+            hyps, _ = self.ctc_decode(encoder_outputs, mode)
+            result = "".join(hyps) if hyps else ""
+
+        if calib_data_path:
+            pack_calibration_dataset(calib_data_path)
+        return result
+
+    def save_calibration_for_audio(self,
+                                   audio_file,
+                                   parts,
+                                   calib_data_path,
+                                   sample_id):
+        counts = {"offline": 0, "online": 0, "decoder": 0}
+        feats = self.compute_feats(audio_file)
+        offline_outputs = None
+        if "offline" in parts or "decoder" in parts:
+            offline_outputs = self.run_offline_encoder(
+                feats,
+                calib_data_path=calib_data_path if "offline" in parts else None,
+                sample_id=sample_id)
+            if "offline" in parts:
+                counts["offline"] += 1
+        if "online" in parts:
+            online_outputs = self.run_online_encoder(
+                feats,
+                calib_data_path=calib_data_path,
+                sample_prefix=sample_id)
+            counts["online"] += online_outputs["num_chunks"]
+        if "decoder" in parts:
+            self.run_decoder(offline_outputs,
+                             calib_data_path=calib_data_path,
+                             sample_id=sample_id)
+            counts["decoder"] += 1
+        return counts

requirements_ax.txt

@@ -0,0 +1,3 @@
+https://github.com/AXERA-TECH/pyaxengine/releases/download/0.1.3.rc3/axengine-0.1.3-py3-none-any.whl
+pyyaml
+numpy

requirements_x86.txt

@@ -0,0 +1,12 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+
+pyyaml>=5.1
+torch==2.11.0+cpu
+torchaudio==2.11.0+cpu
+tqdm
+soundfile
+openai-whisper==20250625
+onnx
+onnxruntime
+onnxscript
+onnxsim

run_ax.py

@@ -0,0 +1,65 @@
+import argparse
+
+from ax_common import WenetAXRunner
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", "-i", type=str, required=True, help="Input audio file")
+    parser.add_argument("--online", action="store_true")
+    parser.add_argument("--config",
+                        type=str,
+                        default="pretrained/aishell_u2pp_conformer_exp/train.yaml",
+                        help="yaml file in checkpoint path")
+    parser.add_argument("--vocab",
+                        type=str,
+                        default="pretrained/aishell_u2pp_conformer_exp/units.txt",
+                        help="pretrained units.txt")
+    parser.add_argument("--encoder_online",
+                        type=str,
+                        default="axmodel/encoder_online/encoder_online.axmodel")
+    parser.add_argument("--encoder_offline",
+                        type=str,
+                        default="axmodel/encoder_offline/encoder_offline.axmodel")
+    parser.add_argument("--decoder",
+                        type=str,
+                        default="axmodel/decoder/decoder.axmodel")
+    parser.add_argument("--offline_seq_len", type=int, default=1024)
+    parser.add_argument("--online_seq_len", type=int, default=67)
+    parser.add_argument("--decoder_len", type=int, default=32)
+    parser.add_argument("--decoding_chunk_size", type=int, default=16)
+    parser.add_argument("--num_decoding_left_chunks", type=int, default=5)
+    parser.add_argument("--mode",
+                        choices=[
+                            "ctc_greedy_search", "ctc_prefix_beam_search",
+                            "attention_rescoring"
+                        ],
+                        default="ctc_prefix_beam_search",
+                        help="decoding mode")
+    return parser.parse_args()
+
+
+def main():
+    args = get_args()
+    print(f"online: {args.online}")
+    print(f"mode: {args.mode}")
+
+    runner = WenetAXRunner(
+        args.config,
+        args.vocab,
+        encoder_offline_path=args.encoder_offline,
+        encoder_online_path=args.encoder_online,
+        decoder_path=args.decoder,
+        offline_seq_len=args.offline_seq_len,
+        decoder_len=args.decoder_len,
+        decoding_chunk_size=args.decoding_chunk_size,
+        num_decoding_left_chunks=args.num_decoding_left_chunks,
+    )
+    result = runner.transcribe(args.input,
+                               online=args.online,
+                               mode=args.mode)
+    print(f"ASR Result: {result}")
+
+
+if __name__ == "__main__":
+    main()

run_ort.py

@@ -0,0 +1,55 @@
+import argparse
+import os
+
+from ort_common import WenetONNXRunner
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", "-i", type=str, required=True, help="Input audio file")
+    parser.add_argument("--config", type=str, required=True, help="yaml file in checkpoint path")
+    parser.add_argument(
+        "--vocab",
+        type=str,
+        required=True,
+        help="pretrained units.txt, for example pretrained/<model>/units.txt",
+    )
+    parser.add_argument("--online", action="store_true")
+    parser.add_argument("--offline_seq_len", type=int, default=1024)
+    parser.add_argument("--online_seq_len", type=int, default=67)
+    parser.add_argument("--decoder_len", type=int, default=32)
+    parser.add_argument("--onnx_dir", type=str, default="onnx_model")
+    parser.add_argument('--mode',
+                        choices=[
+                            'ctc_greedy_search', 'ctc_prefix_beam_search', 'attention_rescoring'],
+                        default='ctc_prefix_beam_search',
+                        help='decoding mode')
+    parser.add_argument("--calib_data_path", type=str, default="calibration_dataset", help="Generated calibration data path")
+    return parser.parse_args()
+
+
+def run_ort():
+    args = get_args()
+    print(f"online: {args.online}")
+    print(f"mode: {args.mode}")
+    print(f"calib_data_path: {args.calib_data_path}")
+    calib_data_path = args.calib_data_path or None
+    if calib_data_path:
+        os.makedirs(calib_data_path, exist_ok=True)
+
+    runner = WenetONNXRunner(
+        args.config,
+        args.vocab,
+        onnx_dir=args.onnx_dir,
+        offline_seq_len=args.offline_seq_len,
+        decoder_len=args.decoder_len,
+    )
+    result = runner.transcribe(args.input,
+                               online=args.online,
+                               mode=args.mode,
+                               calib_data_path=calib_data_path)
+
+    print(f"ASR Result: {result}")
+
+if __name__ == "__main__":
+    run_ort()

wenet/LLM/causallm_model.py

@@ -0,0 +1,207 @@
+from typing import Dict, List, Optional, Union
+import torch
+from wenet.LLM.decoder import DecoderOnly
+from wenet.LLM.sampler import sampler
+from wenet.utils.common import IGNORE_ID, th_accuracy
+from wenet.utils.mask import make_pad_mask, subsequent_mask
+
+
+class CausalLM(torch.nn.Module):
+
+    def __init__(
+        self,
+        vocab_size: int,
+        decoder: DecoderOnly,
+        special_tokens: dict,
+        tie_word_embedding: bool = False,
+        linear_bias: bool = False,
+        ignore_id: int = IGNORE_ID,
+        lsm_weight: float = 0.0,
+        reduction: str = 'mean',
+    ) -> None:
+        super().__init__()
+        del special_tokens
+
+        self.embed = torch.nn.Embedding(vocab_size, decoder.hidden_size)
+        self.out = torch.nn.Linear(decoder.hidden_size,
+                                   vocab_size,
+                                   bias=linear_bias)
+
+        self.decoder = decoder
+        self.vocab_size = vocab_size
+        self.criterion_att = torch.nn.CrossEntropyLoss(
+            ignore_index=ignore_id,
+            label_smoothing=lsm_weight,
+            reduction=reduction,
+        )
+        self.tie_word_embedding = tie_word_embedding
+        self.ignore_id = ignore_id
+
+    @torch.jit.unused
+    def forward(
+        self,
+        batch: dict,
+        device: torch.device,
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        """ Forward for training
+        """
+        text = batch['feats'].to(device)
+        target = batch['target'].to(device)
+        text_length = batch['feats_lengths'].to(device)
+
+        mask = ~make_pad_mask(text_length, max_len=text.size(1)).unsqueeze(
+            1)  # (B,1,L)
+        causal_mask = subsequent_mask(
+            mask.size(-1), device=mask.device).unsqueeze(0)  # (1,L,L)
+        att_mask = causal_mask & mask  # (B, L, L)
+
+        embeding = self.embed(text)
+        decoder_out = self.out(self.decoder(embeding,
+                                            att_mask)[0])  # (B, L, vocab_size)
+        loss = self.criterion_att(decoder_out.view(-1, self.vocab_size),
+                                  target.view(-1))
+        acc = th_accuracy(decoder_out.view(-1, self.vocab_size),
+                          target,
+                          ignore_label=self.ignore_id)
+
+        return {
+            "loss": loss,
+            "ppl": torch.exp(loss.detach()),
+            "th_accuracy": acc
+        }
+
+    def tie_or_clone_weights(self, jit_mode: bool):
+        if not self.tie_word_embedding:
+            return
+        if jit_mode:
+            self.out.weight = torch.nn.Parameter(self.embed.weight.clone())
+        else:
+            self.out.weight = self.embed.weight
+            # TODO(Mddct): whether to deal bias for other llm model
+
+    @torch.jit.unused
+    @torch.inference_mode()
+    def generate(
+        self,
+        prompts_tokens: List[List[int]],
+        device: torch.device,
+        stop_tokens: List[int],
+        dtype: torch.dtype = torch.float32,
+        output_len: int = 100,
+        temperature: Union[float, None] = 0.95,
+        top_p: float = 1.0,
+        top_k: int = 100,
+    ) -> List[List[int]]:
+        """Generates responses for given prompts using Gemma model."""
+        # If a single prompt is provided, treat it as a batch of 1.
+        batch_size = len(prompts_tokens)
+        min_prompt_len = min(len(p) for p in prompts_tokens)
+        max_prompt_len = max(len(p) for p in prompts_tokens)
+        max_seq_len = max_prompt_len + output_len
+        assert max_seq_len <= self.decoder.pos_enc.max_len
+
+        # build KV caches
+        kv_caches = []
+        for _ in range(len(self.decoder.decoders)):
+            size = (batch_size, 0, self.decoder.n_kv_head,
+                    self.decoder.head_dim)
+            k_cache = torch.zeros(size=size, dtype=dtype, device=device)
+            v_cache = torch.zeros(size=size, dtype=dtype, device=device)
+            kv_caches.append((k_cache, v_cache))
+
+        # prepare inputs
+        token_ids_tensor = torch.full((batch_size, max_seq_len),
+                                      IGNORE_ID,
+                                      dtype=torch.int64,
+                                      device=device)
+        input_token_ids_tensor = torch.full((batch_size, min_prompt_len),
+                                            IGNORE_ID,
+                                            dtype=torch.int64,
+                                            device=device)
+        # right padding
+        for i, p in enumerate(prompts_tokens):
+            token_ids_tensor[i, :len(p)] = torch.tensor(p)
+            input_token_ids_tensor[i, :min_prompt_len] = torch.tensor(
+                p[:min_prompt_len])
+
+        prompt_mask_tensor = token_ids_tensor != IGNORE_ID
+        input_positions_tensor = torch.arange(0,
+                                              min_prompt_len,
+                                              dtype=torch.int64).to(device)
+        mask_tensor = torch.ones((1, 1, max_seq_len, max_seq_len),
+                                 dtype=torch.bool)
+        mask_tensor = torch.tril(mask_tensor).to(device)
+        curr_mask_tensor = mask_tensor.index_select(2, input_positions_tensor)
+        att_mask = curr_mask_tensor.squeeze(
+            1)[:, :min_prompt_len, :min_prompt_len]
+        output_positions_tensor = torch.LongTensor([min_prompt_len - 1
+                                                    ]).to(device)
+        temperatures_tensor = None if not temperature else torch.FloatTensor(
+            [temperature] * batch_size).to(device)
+        top_ps_tensor = torch.FloatTensor([top_p] * batch_size).to(device)
+        top_ks_tensor = torch.LongTensor([top_k] * batch_size).to(device)
+        output_index = torch.tensor(min_prompt_len,
+                                    dtype=torch.int64).to(device)
+
+        input_token_embeding = self.embed(input_token_ids_tensor)
+        offset = torch.tensor([0] * len(prompts_tokens)).to(device)
+        input_offset = offset
+
+        stop_tokens_tensor = torch.tensor(stop_tokens, device=device)
+        # Prefill up to min_prompt_len tokens, then treat other prefill as
+        # decode and ignore output.
+        for i in range(max_seq_len - min_prompt_len):
+            decoder_out, kv_caches, = self.decoder(
+                input_token_embeding,
+                att_mask,
+                input_offset,
+                kv_caches,
+            )
+            decoder_out = self.out(decoder_out)
+            decoder_out = decoder_out.index_select(1, output_positions_tensor)
+            next_token_ids = sampler(
+                decoder_out,
+                temperatures_tensor,
+                top_ps_tensor,
+                top_ks_tensor,
+            )
+            curr_prompt_mask = prompt_mask_tensor.index_select(
+                1, output_index).squeeze(dim=1)
+            curr_token_ids = token_ids_tensor.index_select(
+                1, output_index).squeeze(dim=1)
+            output_token_ids = torch.where(curr_prompt_mask, curr_token_ids,
+                                           next_token_ids).unsqueeze(dim=1)
+            token_ids_tensor.index_copy_(1, output_index, output_token_ids)
+
+            input_token_ids_tensor = output_token_ids
+            input_token_embeding = self.embed(input_token_ids_tensor)
+
+            input_positions_tensor = output_index.unsqueeze(dim=-1)
+            curr_mask_tensor = mask_tensor.index_select(
+                2, input_positions_tensor)
+            att_mask = curr_mask_tensor.squeeze(1)[:, :output_index +
+                                                   1, :output_index + 1]
+
+            output_positions_tensor = torch.tensor(
+                0, dtype=torch.int64).to(device)
+            input_offset = offset + output_index.unsqueeze(-1)
+            output_index = output_index + 1
+
+            if all(torch.isin(next_token_ids, stop_tokens_tensor)):
+                break
+
+        token_ids = token_ids_tensor.tolist()
+        results = []
+        for i, tokens in enumerate(token_ids):
+            trimmed_output = tokens[len(prompts_tokens[i]
+                                        ):len(prompts_tokens[i]) + output_len]
+            for stop_token in stop_tokens:
+                try:
+                    eos_index = trimmed_output.index(stop_token)
+                    trimmed_output = trimmed_output[:eos_index]
+                    break
+                except Exception:
+                    continue
+            results.append(trimmed_output)
+
+        return results

wenet/LLM/decoder.py

@@ -0,0 +1,161 @@
+from functools import partial
+from typing import List, Optional, Tuple, Union
+import torch
+import torch.utils.checkpoint as ckpt
+from wenet.transformer.attention import T_CACHE
+
+from wenet.transformer.encoder_layer import TransformerEncoderLayer
+from wenet.utils.class_utils import (WENET_ACTIVATION_CLASSES,
+                                     WENET_ATTENTION_CLASSES,
+                                     WENET_EMB_CLASSES, WENET_MLP_CLASSES,
+                                     WENET_NORM_CLASSES)
+from wenet.utils.common import mask_to_bias
+
+
+class DecoderOnly(torch.nn.Module):
+
+    def __init__(
+        self,
+        n_kv_head: int,
+        head_dim: int,
+        hidden_size: int,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        normalize_before: bool = True,
+        query_bias: bool = False,
+        key_bias: bool = False,
+        value_bias: bool = False,
+        mlp_bias: bool = False,
+        activation_type: str = "gelu",
+        gelu_approximate: Union[str, None] = None,
+        max_position_embeding: int = 8192,
+        mlp_type: str = 'gated',
+        layer_norm_type: str = 'rms_norm',
+        norm_eps: float = 1e-5,
+        rms_norm_offset: bool = True,
+        selfattention_layer_type: str = "rope_abs_selfattn",
+        use_sdpa: bool = False,
+        gradient_checkpointing: bool = False,
+        rope_theta: float = 10000.0,
+        rope_style: str = 'google',
+        scale_embed: bool = True,
+    ) -> None:
+        super().__init__()
+
+        assert selfattention_layer_type in ['rope_abs_selfattn']
+        self.pos_enc = WENET_EMB_CLASSES["rope_pos"](
+            hidden_size,
+            head_dim,
+            max_len=max_position_embeding,
+            dropout_rate=positional_dropout_rate,
+            rope_theta=rope_theta,
+            scale=scale_embed)
+        if activation_type == "gelu" and gelu_approximate is not None:
+            activation = WENET_ACTIVATION_CLASSES['gelu'](
+                approximate=gelu_approximate)
+        else:
+            activation = WENET_ACTIVATION_CLASSES[activation_type]()
+
+        mlp_class = WENET_MLP_CLASSES[mlp_type]
+        self.num_blocks = num_blocks
+        # TODO: support lora & refactor lora
+        self.decoders = torch.nn.ModuleList([
+            TransformerEncoderLayer(
+                hidden_size,
+                WENET_ATTENTION_CLASSES[selfattention_layer_type](
+                    attention_heads,
+                    hidden_size,
+                    attention_dropout_rate,
+                    query_bias,
+                    key_bias,
+                    value_bias,
+                    use_sdpa,
+                    n_kv_head,
+                    head_dim,
+                    style=rope_style),
+                mlp_class(hidden_size, linear_units, dropout_rate, activation,
+                          mlp_bias),
+                dropout_rate,
+                normalize_before,
+                layer_norm_type=layer_norm_type,
+                norm_eps=norm_eps,
+                rms_norm_offset=rms_norm_offset,
+            ) for _ in range(self.num_blocks)
+        ])
+        self.pre_norm = normalize_before
+        self.final_norm: Optional[torch.nn.Module] = None
+        if self.pre_norm:
+            norm_class = WENET_NORM_CLASSES[layer_norm_type]
+            if layer_norm_type == "rms_norm":
+                norm_class = partial(
+                    norm_class,
+                    add_unit_offset=rms_norm_offset,
+                )
+            self.final_norm = norm_class(hidden_size, eps=norm_eps)
+
+        self.n_kv_head = n_kv_head
+        self.head_dim = head_dim
+        self._hidden_size = hidden_size
+        self.use_sdpa = use_sdpa
+        self.gradient_checkpointing = gradient_checkpointing
+
+    def forward(
+        self,
+        input: torch.Tensor,
+        att_mask: torch.Tensor,
+        input_position: Union[int, torch.Tensor] = 0,
+        kv_caches: Optional[List[T_CACHE]] = None,
+    ) -> Tuple[torch.Tensor, Union[List[T_CACHE], None]]:
+        xs, pos_emb = self.pos_enc(input, offset=input_position)
+        if self.use_sdpa:
+            att_mask = mask_to_bias(att_mask, xs.dtype)
+
+        if self.gradient_checkpointing and self.training:
+            xs = self.forward_layers_checkpointed(xs, att_mask, pos_emb)
+        else:
+            xs, kv_caches = self.forward_layers(xs, att_mask, pos_emb,
+                                                kv_caches)
+        if self.pre_norm and self.final_norm is not None:
+            xs = self.final_norm(xs)
+        return xs, kv_caches
+
+    def forward_layers(
+        self,
+        xs: torch.Tensor,
+        att_mask: torch.Tensor,
+        pos_emb: torch.Tensor,
+        kv_caches: Optional[List[T_CACHE]] = None,
+    ) -> Tuple[torch.Tensor, Union[List[T_CACHE], None]]:
+        if self.training:
+            for (i, layer) in enumerate(self.decoders):
+                xs, _, _, _ = layer(xs, att_mask, pos_emb)
+            new_kv_caches = kv_caches
+        else:
+            assert kv_caches is not None
+            new_kv_caches = []
+            for (i, layer) in enumerate(self.decoders):
+                xs, _, new_kv_cache, _ = layer(xs,
+                                               att_mask,
+                                               pos_emb,
+                                               att_cache=(kv_caches[i][0],
+                                                          kv_caches[i][1]))
+                new_kv_caches.append(new_kv_cache)
+
+        return xs, new_kv_caches
+
+    @torch.jit.ignore(drop=True)
+    def forward_layers_checkpointed(self, xs: torch.Tensor,
+                                    att_mask: torch.Tensor,
+                                    pos_emb: torch.Tensor) -> torch.Tensor:
+        for layer in self.decoders:
+            xs, _, _, _ = ckpt.checkpoint(layer.__call__, xs, att_mask,
+                                          pos_emb)
+        return xs
+
+    @property
+    def hidden_size(self):
+        return self._hidden_size

wenet/LLM/sampler.py

@@ -0,0 +1,43 @@
+from typing import Union
+import torch
+
+
+# modified from https://github.com/google/gemma_pytorch/blob/main/gemma/model.py#L26
+@torch.no_grad()
+def sampler(
+    logits: torch.Tensor,
+    temperatures: Union[torch.Tensor, None],
+    top_ps: torch.Tensor,
+    top_ks: torch.Tensor,
+) -> torch.Tensor:
+    assert logits.size(1) == 1
+    logits = logits.squeeze(1)  # (batch_size, vocab_size)
+    if temperatures is None:
+        return torch.argmax(logits, dim=-1).squeeze(dim=-1)
+
+    # Apply temperature scaling.
+    logits.div_(temperatures.unsqueeze(dim=1))
+
+    # Calculate probabilities with softmax.
+    probs = torch.softmax(logits, dim=-1, dtype=torch.float)
+    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+
+    # Apply top-p, top-k.
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    top_ps_mask = (probs_sum - probs_sort) > top_ps.unsqueeze(dim=1)
+    probs_sort = torch.where(top_ps_mask, 0, probs_sort)
+
+    top_ks_mask = torch.arange(probs_idx.shape[-1], device=probs_idx.device)
+    top_ks_mask = top_ks_mask.expand(probs_idx.shape[0], -1)
+    top_ks_mask = top_ks_mask >= top_ks.unsqueeze(dim=1)
+    probs_sort = torch.where(top_ks_mask, 0, probs_sort)
+
+    # Re-normalization.
+    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+    probs = torch.gather(probs_sort,
+                         dim=-1,
+                         index=torch.argsort(probs_idx, dim=-1))
+
+    next_token_ids = torch.multinomial(probs, num_samples=1,
+                                       replacement=True).squeeze(dim=-1)
+    return next_token_ids

wenet/README.md

@@ -0,0 +1,29 @@
+# Module Introduction
+
+Here is a brief introduction of each module(directory).
+
+* `bin`: training and recognition binaries
+* `dataset`: IO design
+* `utils`: common utils
+* `transformer`: the core of `WeNet`, in which the standard transformer/conformer is implemented. It contains the common blocks(backbone) of speech transformers.
+  * transformer/attention.py: Standard multi head attention
+  * transformer/embedding.py: Standard position encoding
+  * transformer/positionwise_feed_forward.py: Standard feed forward in transformer
+  * transformer/convolution.py: ConvolutionModule in Conformer model
+  * transformer/subsampling.py: Subsampling implementation for speech task
+* `transducer`: transducer implementation
+* `squeezeformer`: squeezeformer implementation, please refer [paper](https://arxiv.org/pdf/2206.00888.pdf)
+* `efficient_conformer`: efficient conformer implementation, please refer [paper](https://arxiv.org/pdf/2109.01163.pdf)
+* `paraformer`: paraformer implementation, please refer [paper](https://arxiv.org/pdf/1905.11235.pdf)
+   * `paraformer/cif.py`: Continuous Integrate-and-Fire implemented, please refer [paper](https://arxiv.org/pdf/1905.11235.pdf)
+* `branchformer`: branchformer implementation, please refer [paper](https://arxiv.org/abs/2207.02971)
+* `whisper`: whisper implementation, please refer [paper](https://arxiv.org/abs/2212.04356)
+* `ssl`: Self-supervised speech model implementation. e.g. wav2vec2, bestrq, w2vbert.
+* `ctl_model`: Enhancing the Unified Streaming and Non-streaming Model with  with Contrastive Learning implementation [paper](https://arxiv.org/abs/2306.00755)
+
+`transducer`, `squeezeformer`, `efficient_conformer`, `branchformer` and `cif` are all based on `transformer`,
+they resue a lot of the common blocks of `tranformer`.
+
+**If you want to contribute your own x-former, please reuse the current code as much as possible**.
+
+

wenet/__init__.py

@@ -0,0 +1 @@
+from wenet.cli.model import load_model, load_model_pt  # noqa

wenet/bin/alignment.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2021 Mobvoi Inc. (authors: Di Wu)
+#               2022 Tinnove Inc (authors: Wei Ren)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import argparse
+import copy
+import logging
+import os
+import sys
+
+import torch
+import yaml
+from torch.utils.data import DataLoader
+from textgrid import TextGrid, IntervalTier
+import math
+
+from wenet.dataset.dataset import Dataset
+from wenet.utils.ctc_utils import force_align
+from wenet.utils.common import get_subsample
+from wenet.utils.init_model import init_model
+from wenet.utils.init_tokenizer import init_tokenizer
+
+
+def generator_textgrid(maxtime, lines, output):
+    # Download Praat: https://www.fon.hum.uva.nl/praat/
+    interval = maxtime / (len(lines) + 1)
+    margin = 0.0001
+
+    tg = TextGrid(maxTime=maxtime)
+    linetier = IntervalTier(name="line", maxTime=maxtime)
+
+    i = 0
+    for l in lines:
+        s, e, w = l.split()
+        linetier.add(minTime=float(s) + margin, maxTime=float(e), mark=w)
+
+    tg.append(linetier)
+    print("successfully generator {}".format(output))
+    tg.write(output)
+
+
+def get_frames_timestamp(alignment,
+                         prob,
+                         blank_thres=0.999,
+                         thres=0.0000000001):
+    # convert alignment to a praat format, which is a doing phonetics
+    # by computer and helps analyzing alignment
+    timestamp = []
+    # get frames level duration for each token
+    start = 0
+    end = 0
+    local_start = 0
+    while end < len(alignment):
+        while end < len(alignment) and alignment[end] == 0:
+            end += 1
+        if end == len(alignment):
+            timestamp[-1] += alignment[start:]
+            break
+        end += 1
+        while end < len(alignment) and alignment[end - 1] == alignment[end]:
+            end += 1
+        local_start = end - 1
+        # find the possible front border for current token
+        while local_start >= start and (
+                prob[local_start][0] < math.log(blank_thres)
+                or prob[local_start][alignment[end - 1]] > math.log(thres)):
+            alignment[local_start] = alignment[end - 1]
+            local_start -= 1
+        cur_alignment = alignment[start:end]
+        timestamp.append(cur_alignment)
+        start = end
+    return timestamp
+
+
+def get_labformat(timestamp, subsample):
+    begin = 0
+    begin_time = 0
+    duration = 0
+    labformat = []
+    for idx, t in enumerate(timestamp):
+        # 25ms frame_length,10ms hop_length, 1/subsample
+        subsample = get_subsample(configs)
+        # time duration
+        i = 0
+        while t[i] == 0:
+            i += 1
+        begin = i
+        dur = 0
+        while i < len(t) and t[i] != 0:
+            i += 1
+            dur += 1
+        begin = begin_time + begin * 0.01 * subsample
+        duration = dur * 0.01 * subsample
+        if idx < len(timestamp) - 1:
+            print("{:.2f} {:.2f} {}".format(begin, begin + duration,
+                                            char_dict[t[-1]]))
+            labformat.append("{:.2f} {:.2f} {}\n".format(
+                begin, begin + duration, char_dict[t[-1]]))
+        else:  # last token
+            non_blank = 0
+            for i in t:
+                if i != 0:
+                    token = i
+                    break
+            print("{:.2f} {:.2f} {}".format(begin, begin + duration,
+                                            char_dict[token]))
+            labformat.append("{:.2f} {:.2f} {}\n".format(
+                begin, begin + duration, char_dict[token]))
+        begin_time += len(t) * 0.01 * subsample
+    return labformat
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='use ctc to generate alignment')
+    parser.add_argument('--config', required=True, help='config file')
+    parser.add_argument('--input_file', required=True, help='format data file')
+    parser.add_argument('--data_type',
+                        default='raw',
+                        choices=['raw', 'shard'],
+                        help='train and cv data type')
+    parser.add_argument('--gpu',
+                        type=int,
+                        default=-1,
+                        help='gpu id for this rank, -1 for cpu')
+    parser.add_argument('--device',
+                        type=str,
+                        default="cpu",
+                        choices=["cpu", "npu", "cuda"],
+                        help='accelerator to use')
+    parser.add_argument('--blank_thres',
+                        default=0.999999,
+                        type=float,
+                        help='ctc blank thes')
+    parser.add_argument('--thres',
+                        default=0.000001,
+                        type=float,
+                        help='ctc non blank thes')
+    parser.add_argument('--checkpoint', required=True, help='checkpoint model')
+    parser.add_argument('--dict', required=True, help='dict file')
+    parser.add_argument(
+        '--non_lang_syms',
+        help="non-linguistic symbol file. One symbol per line.")
+    parser.add_argument('--result_file',
+                        required=True,
+                        help='alignment result file')
+    parser.add_argument('--batch_size', type=int, default=1, help='batch size')
+    parser.add_argument('--gen_praat',
+                        action='store_true',
+                        help='convert alignment to a praat format')
+    parser.add_argument('--bpe_model',
+                        default=None,
+                        type=str,
+                        help='bpe model for english part')
+
+    args = parser.parse_args()
+    print(args)
+    logging.basicConfig(level=logging.DEBUG,
+                        format='%(asctime)s %(levelname)s %(message)s')
+    if args.gpu != -1:
+        # remain the original usage of gpu
+        args.device = "cuda"
+    if "cuda" in args.device:
+        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
+
+    if args.batch_size > 1:
+        logging.fatal('alignment mode must be running with batch_size == 1')
+        sys.exit(1)
+
+    with open(args.config, 'r') as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+
+    # Load dict
+    char_dict = {}
+    with open(args.dict, 'r') as fin:
+        for line in fin:
+            arr = line.strip().split()
+            assert len(arr) == 2
+            char_dict[int(arr[1])] = arr[0]
+    eos = len(char_dict) - 1
+
+    # Init dataset and data loader
+    ali_conf = copy.deepcopy(configs['dataset_conf'])
+
+    ali_conf['filter_conf']['max_length'] = 102400
+    ali_conf['filter_conf']['min_length'] = 0
+    ali_conf['filter_conf']['token_max_length'] = 102400
+    ali_conf['filter_conf']['token_min_length'] = 0
+    ali_conf['filter_conf']['max_output_input_ratio'] = 102400
+    ali_conf['filter_conf']['min_output_input_ratio'] = 0
+    ali_conf['speed_perturb'] = False
+    ali_conf['spec_aug'] = False
+    ali_conf['spec_trim'] = False
+    ali_conf['shuffle'] = False
+    ali_conf['sort'] = False
+    ali_conf['fbank_conf']['dither'] = 0.0
+    ali_conf['batch_conf']['batch_type'] = "static"
+    ali_conf['batch_conf']['batch_size'] = args.batch_size
+
+    tokenizer = init_tokenizer(configs)
+    ali_dataset = Dataset(args.data_type,
+                          args.input_file,
+                          tokenizer,
+                          ali_conf,
+                          partition=False)
+
+    ali_data_loader = DataLoader(ali_dataset, batch_size=None, num_workers=0)
+
+    # Init asr model from configs
+    model, configs = init_model(args, configs)
+
+    device = torch.device(args.device)
+    model = model.to(device)
+
+    model.eval()
+    with torch.no_grad(), open(args.result_file, 'w',
+                               encoding='utf-8') as fout:
+        for batch_idx, batch in enumerate(ali_data_loader):
+            print("#" * 80)
+            key, feat, target, feats_length, target_length = batch
+
+            feat = feat.to(device)
+            target = target.to(device)
+            feats_length = feats_length.to(device)
+            target_length = target_length.to(device)
+            # Let's assume B = batch_size and N = beam_size
+            # 1. Encoder
+            encoder_out, encoder_mask = model._forward_encoder(
+                feat, feats_length)  # (B, maxlen, encoder_dim)
+            maxlen = encoder_out.size(1)
+            ctc_probs = model.ctc.log_softmax(
+                encoder_out)  # (1, maxlen, vocab_size)
+            # print(ctc_probs.size(1))
+            ctc_probs = ctc_probs.squeeze(0)
+            target = target.squeeze(0)
+            alignment = force_align(ctc_probs, target)
+            fout.write('{} {}\n'.format(key[0], alignment))
+
+            if args.gen_praat:
+                timestamp = get_frames_timestamp(alignment, ctc_probs,
+                                                 args.blank_thres, args.thres)
+                subsample = get_subsample(configs)
+                labformat = get_labformat(timestamp, subsample)
+
+                lab_path = os.path.join(os.path.dirname(args.result_file),
+                                        key[0] + ".lab")
+                with open(lab_path, 'w', encoding='utf-8') as f:
+                    f.writelines(labformat)
+
+                textgrid_path = os.path.join(os.path.dirname(args.result_file),
+                                             key[0] + ".TextGrid")
+                generator_textgrid(maxtime=(len(alignment) + 1) * 0.01 *
+                                   subsample,
+                                   lines=labformat,
+                                   output=textgrid_path)

wenet/bin/average_model.py

@@ -0,0 +1,125 @@
+# Copyright (c) 2020 Mobvoi Inc (Di Wu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import argparse
+import glob
+import sys
+
+import yaml
+import torch
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='average model')
+    parser.add_argument('--dst_model', required=True, help='averaged model')
+    parser.add_argument('--src_path',
+                        required=True,
+                        help='src model path for average')
+    parser.add_argument('--val_best',
+                        action="store_true",
+                        help='averaged model')
+    parser.add_argument('--num',
+                        default=5,
+                        type=int,
+                        help='nums for averaged model')
+    parser.add_argument('--min_epoch',
+                        default=0,
+                        type=int,
+                        help='min epoch used for averaging model')
+    parser.add_argument('--max_epoch',
+                        default=sys.maxsize,
+                        type=int,
+                        help='max epoch used for averaging model')
+    parser.add_argument('--min_step',
+                        default=0,
+                        type=int,
+                        help='min step used for averaging model')
+    parser.add_argument('--max_step',
+                        default=sys.maxsize,
+                        type=int,
+                        help='max step used for averaging model')
+    parser.add_argument('--mode',
+                        default="hybrid",
+                        choices=["hybrid", "epoch", "step"],
+                        type=str,
+                        help='average mode')
+
+    args = parser.parse_args()
+    print(args)
+    return args
+
+
+def main():
+    args = get_args()
+    checkpoints = []
+    val_scores = []
+    if args.val_best:
+        if args.mode == "hybrid":
+            yamls = glob.glob('{}/*.yaml'.format(args.src_path))
+            yamls = [
+                f for f in yamls
+                if not (os.path.basename(f).startswith('train')
+                        or os.path.basename(f).startswith('init'))
+            ]
+        elif args.mode == "step":
+            yamls = glob.glob('{}/step_*.yaml'.format(args.src_path))
+        else:
+            yamls = glob.glob('{}/epoch_*.yaml'.format(args.src_path))
+        for y in yamls:
+            with open(y, 'r') as f:
+                dic_yaml = yaml.load(f, Loader=yaml.FullLoader)
+                loss = dic_yaml['loss_dict']['loss']
+                epoch = dic_yaml['epoch']
+                step = dic_yaml['step']
+                tag = dic_yaml['tag']
+                if epoch >= args.min_epoch and epoch <= args.max_epoch \
+                        and step >= args.min_step and step <= args.max_step:
+                    val_scores += [[epoch, step, loss, tag]]
+        sorted_val_scores = sorted(val_scores,
+                                   key=lambda x: x[2],
+                                   reverse=False)
+        print("best val (epoch, step, loss, tag) = " +
+              str(sorted_val_scores[:args.num]))
+        path_list = [
+            args.src_path + '/{}.pt'.format(score[-1])
+            for score in sorted_val_scores[:args.num]
+        ]
+    else:
+        path_list = glob.glob('{}/[!init]*.pt'.format(args.src_path))
+        path_list = sorted(path_list, key=os.path.getmtime)
+        path_list = path_list[-args.num:]
+    print(path_list)
+    avg = {}
+    num = args.num
+    assert num == len(path_list)
+    for path in path_list:
+        print('Processing {}'.format(path))
+        states = torch.load(path, map_location=torch.device('cpu'))
+        for k in states.keys():
+            if k not in avg.keys():
+                avg[k] = states[k].clone()
+            else:
+                avg[k] += states[k]
+    # average
+    for k in avg.keys():
+        if avg[k] is not None:
+            # pytorch 1.6 use true_divide instead of /=
+            avg[k] = torch.true_divide(avg[k], num)
+    print('Saving to {}'.format(args.dst_model))
+    torch.save(avg, args.dst_model)
+
+
+if __name__ == '__main__':
+    main()

wenet/bin/export_ipex.py

@@ -0,0 +1,95 @@
+# Copyright (C) 2021-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import print_function
+
+import argparse
+import logging
+import os
+
+import torch
+import yaml
+
+from wenet.utils.init_model import init_model
+import intel_extension_for_pytorch as ipex
+from intel_extension_for_pytorch.quantization import prepare, convert
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='export your script model')
+    parser.add_argument('--config', required=True, help='config file')
+    parser.add_argument('--checkpoint', required=True, help='checkpoint model')
+    parser.add_argument('--output_file', default=None, help='output file')
+    parser.add_argument('--dtype',
+                        default="fp32",
+                        help='choose the dtype to run:[fp32,bf16]')
+    parser.add_argument('--output_quant_file',
+                        default=None,
+                        help='output quantized model file')
+    args = parser.parse_args()
+    return args
+
+
+def scripting(model):
+    with torch.inference_mode():
+        script_model = torch.jit.script(model)
+        script_model = torch.jit.freeze(
+            script_model,
+            preserved_attrs=[
+                "forward_encoder_chunk", "ctc_activation",
+                "forward_attention_decoder", "subsampling_rate",
+                "right_context", "sos_symbol", "eos_symbol",
+                "is_bidirectional_decoder"
+            ])
+    return script_model
+
+
+def main():
+    args = get_args()
+    logging.basicConfig(level=logging.DEBUG,
+                        format='%(asctime)s %(levelname)s %(message)s')
+    # No need gpu for model export
+    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
+    with open(args.config, 'r') as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+    model, configs = init_model(args, configs)
+    print(model)
+
+    # Apply IPEX optimization
+    model.eval()
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    model.to(memory_format=torch.channels_last)
+    if args.dtype == "fp32":
+        ipex_model = ipex.optimize(model)
+    elif args.dtype == "bf16":  # For Intel 4th generation Xeon (SPR)
+        ipex_model = ipex.optimize(model,
+                                   dtype=torch.bfloat16,
+                                   weights_prepack=False)
+
+    # Export jit torch script model
+    if args.output_file:
+        if args.dtype == "fp32":
+            script_model = scripting(ipex_model)
+        elif args.dtype == "bf16":
+            torch._C._jit_set_autocast_mode(True)
+            with torch.cpu.amp.autocast():
+                script_model = scripting(ipex_model)
+        script_model.save(args.output_file)
+        print('Export model successfully, see {}'.format(args.output_file))
+
+    # Export quantized jit torch script model
+    if args.output_quant_file:
+        dynamic_qconfig = ipex.quantization.default_dynamic_qconfig
+        dummy_data = (torch.zeros(1, 67, 80), 16, -16,
+                      torch.zeros(12, 4, 32, 128), torch.zeros(12, 1, 256, 7))
+        model = prepare(model, dynamic_qconfig, dummy_data)
+        model = convert(model)
+        script_quant_model = scripting(model)
+        script_quant_model.save(args.output_quant_file)
+        print('Export quantized model successfully, '
+              'see {}'.format(args.output_quant_file))
+
+
+if __name__ == '__main__':
+    main()

wenet/bin/export_jit.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import argparse
+import logging
+import os
+
+import torch
+import yaml
+
+from wenet.utils.init_model import init_model
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='export your script model')
+    parser.add_argument('--config', required=True, help='config file')
+    parser.add_argument('--checkpoint', required=True, help='checkpoint model')
+    parser.add_argument('--output_file', default=None, help='output file')
+    parser.add_argument('--output_quant_file',
+                        default=None,
+                        help='output quantized model file')
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+    args.jit = True
+    logging.basicConfig(level=logging.DEBUG,
+                        format='%(asctime)s %(levelname)s %(message)s')
+    # No need gpu for model export
+    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
+    with open(args.config, 'r') as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+    model, configs = init_model(args, configs)
+    model.eval()
+    print(model)
+    # Export jit torch script model
+
+    if args.output_file:
+        script_model = torch.jit.script(model)
+        script_model.save(args.output_file)
+        print('Export model successfully, see {}'.format(args.output_file))
+
+    # Export quantized jit torch script model
+    if args.output_quant_file:
+        quantized_model = torch.quantization.quantize_dynamic(
+            model, {torch.nn.Linear}, dtype=torch.qint8)
+        print(quantized_model)
+        script_quant_model = torch.jit.script(quantized_model)
+        script_quant_model.save(args.output_quant_file)
+        print('Export quantized model successfully, '
+              'see {}'.format(args.output_quant_file))
+
+
+if __name__ == '__main__':
+    main()

wenet/bin/export_onnx_bpu.py

@@ -0,0 +1,1065 @@
+# Copyright (c) 2022, Horizon Inc. Xingchen Song (sxc19@tsinghua.org.cn)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""NOTE(xcsong): Currently, we only support
+1. specific conformer encoder architecture, see:
+    encoder: conformer
+    encoder_conf:
+      activation_type: **must be** relu
+      attention_heads: 2 or 4 or 8 or any number divisible by output_size
+      causal: **must be** true
+      cnn_module_kernel: 1 ~ 7
+      cnn_module_norm: **must be** batch_norm
+      input_layer: **must be** conv2d8
+      linear_units: 1 ~ 2048
+      normalize_before: **must be** true
+      num_blocks: 1 ~ 12
+      output_size: 1 ~ 512
+      pos_enc_layer_type: **must be** no_pos
+      selfattention_layer_type: **must be** selfattn
+      use_cnn_module: **must be** true
+      use_dynamic_chunk: **must be** true
+      use_dynamic_left_chunk: **must be** true
+
+2. specific decoding method: ctc_greedy_search
+"""
+
+from __future__ import print_function
+
+import os
+import sys
+import copy
+import math
+import yaml
+import logging
+from typing import Tuple
+
+import torch
+import numpy as np
+
+from wenet.transformer.embedding import NoPositionalEncoding
+from wenet.utils.init_model import init_model
+from wenet.bin.export_onnx_cpu import (get_args, to_numpy,
+                                       print_input_output_info)
+
+try:
+    import onnx
+    import onnxruntime
+except ImportError:
+    print('Please install onnx and onnxruntime!')
+    sys.exit(1)
+
+logger = logging.getLogger(__file__)
+logger.setLevel(logging.INFO)
+
+
+class BPULayerNorm(torch.nn.Module):
+    """Refactor torch.nn.LayerNorm to meet 4-D dataflow."""
+
+    def __init__(self, module, chunk_size=8, run_on_bpu=False):
+        super().__init__()
+        original = copy.deepcopy(module)
+        self.hidden = module.weight.size(0)
+        self.chunk_size = chunk_size
+        self.run_on_bpu = run_on_bpu
+
+        if self.run_on_bpu:
+            self.weight = torch.nn.Parameter(
+                module.weight.reshape(1, self.hidden, 1,
+                                      1).repeat(1, 1, 1, chunk_size))
+            self.bias = torch.nn.Parameter(
+                module.bias.reshape(1, self.hidden, 1,
+                                    1).repeat(1, 1, 1, chunk_size))
+            self.negtive = torch.nn.Parameter(
+                torch.ones((1, self.hidden, 1, chunk_size)) * -1.0)
+            self.eps = torch.nn.Parameter(
+                torch.zeros((1, self.hidden, 1, chunk_size)) + module.eps)
+            self.mean_conv_1 = torch.nn.Conv2d(self.hidden, 1, 1, bias=False)
+            self.mean_conv_1.weight = torch.nn.Parameter(
+                torch.ones(self.hidden, self.hidden, 1, 1) /
+                (1.0 * self.hidden))
+            self.mean_conv_2 = torch.nn.Conv2d(self.hidden, 1, 1, bias=False)
+            self.mean_conv_2.weight = torch.nn.Parameter(
+                torch.ones(self.hidden, self.hidden, 1, 1) /
+                (1.0 * self.hidden))
+        else:
+            self.norm = module
+
+        self.check_equal(original)
+
+    def check_equal(self, module):
+        random_data = torch.randn(1, self.chunk_size, self.hidden)
+        orig_out = module(random_data)
+        new_out = self.forward(random_data.transpose(1, 2).unsqueeze(2))
+        np.testing.assert_allclose(to_numpy(orig_out),
+                                   to_numpy(
+                                       new_out.squeeze(2).transpose(1, 2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.run_on_bpu:
+            u = self.mean_conv_1(x)  # (1, h, 1, c)
+            numerator = x + u * self.negtive  # (1, h, 1, c)
+            s = torch.pow(numerator, 2)  # (1, h, 1, c)
+            s = self.mean_conv_2(s)  # (1, h, 1, c)
+            denominator = torch.sqrt(s + self.eps)  # (1, h, 1, c)
+            x = torch.div(numerator, denominator)  # (1, h, 1, c)
+            x = x * self.weight + self.bias
+        else:
+            x = x.squeeze(2).transpose(1, 2).contiguous()
+            x = self.norm(x)
+            x = x.transpose(1, 2).contiguous().unsqueeze(2)
+        return x
+
+
+class BPUIdentity(torch.nn.Module):
+    """Refactor torch.nn.Identity().
+       For inserting BPU node whose input == output.
+    """
+
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.identity_conv = torch.nn.Conv2d(channels,
+                                             channels,
+                                             1,
+                                             groups=channels,
+                                             bias=False)
+        torch.nn.init.dirac_(self.identity_conv.weight.data, groups=channels)
+
+        self.check_equal()
+
+    def check_equal(self):
+        random_data = torch.randn(1, self.channels, 1, 10)
+        result = self.forward(random_data)
+        np.testing.assert_allclose(to_numpy(random_data),
+                                   to_numpy(result),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Identity with 4-D dataflow, input == output.
+        Args:
+            x (torch.Tensor): (batch, in_channel, 1, time)
+
+        Returns:
+            (torch.Tensor): (batch, in_channel, 1, time).
+        """
+        return self.identity_conv(x)
+
+
+class BPULinear(torch.nn.Module):
+    """Refactor torch.nn.Linear or pointwise_conv"""
+
+    def __init__(self, module, is_pointwise_conv=False):
+        super().__init__()
+        # Unchanged submodules and attributes
+        original = copy.deepcopy(module)
+        self.idim = module.weight.size(1)
+        self.odim = module.weight.size(0)
+        self.is_pointwise_conv = is_pointwise_conv
+
+        # Modify weight & bias
+        self.linear = torch.nn.Conv2d(self.idim, self.odim, 1, 1)
+        if is_pointwise_conv:
+            # (odim, idim, kernel=1) -> (odim, idim, 1, 1)
+            self.linear.weight = torch.nn.Parameter(
+                module.weight.unsqueeze(-1))
+        else:
+            # (odim, idim) -> (odim, idim, 1, 1)
+            self.linear.weight = torch.nn.Parameter(
+                module.weight.unsqueeze(2).unsqueeze(3))
+        self.linear.bias = module.bias
+
+        self.check_equal(original)
+
+    def check_equal(self, module):
+        random_data = torch.randn(1, 8, self.idim)
+        if self.is_pointwise_conv:
+            random_data = random_data.transpose(1, 2)
+        original_result = module(random_data)
+        if self.is_pointwise_conv:
+            random_data = random_data.transpose(1, 2)
+            original_result = original_result.transpose(1, 2)
+        random_data = random_data.transpose(1, 2).unsqueeze(2)
+        new_result = self.forward(random_data)
+        np.testing.assert_allclose(to_numpy(original_result),
+                                   to_numpy(
+                                       new_result.squeeze(2).transpose(1, 2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Linear with 4-D dataflow.
+        Args:
+            x (torch.Tensor): (batch, in_channel, 1, time)
+        Returns:
+            (torch.Tensor): (batch, out_channel, 1, time).
+        """
+        return self.linear(x)
+
+
+class BPUGlobalCMVN(torch.nn.Module):
+    """Refactor wenet/transformer/cmvn.py::GlobalCMVN"""
+
+    def __init__(self, module):
+        super().__init__()
+        # Unchanged submodules and attributes
+        self.norm_var = module.norm_var
+
+        # NOTE(xcsong): Expand to 4-D tensor, (mel_dim) -> (1, 1, mel_dim, 1)
+        self.mean = module.mean.unsqueeze(-1).unsqueeze(0).unsqueeze(0)
+        self.istd = module.istd.unsqueeze(-1).unsqueeze(0).unsqueeze(0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """CMVN with 4-D dataflow.
+        Args:
+            x (torch.Tensor): (batch, 1, mel_dim, time)
+        Returns:
+            (torch.Tensor): normalized feature with same shape.
+        """
+        x = x - self.mean
+        if self.norm_var:
+            x = x * self.istd
+        return x
+
+
+class BPUConv2dSubsampling8(torch.nn.Module):
+    """Refactor wenet/transformer/subsampling.py::Conv2dSubsampling8
+
+    NOTE(xcsong): Only support pos_enc_class == NoPositionalEncoding
+    """
+
+    def __init__(self, module):
+        super().__init__()
+        # Unchanged submodules and attributes
+        original = copy.deepcopy(module)
+        self.right_context = module.right_context
+        self.subsampling_rate = module.subsampling_rate
+        assert isinstance(module.pos_enc, NoPositionalEncoding)
+
+        # 1. Modify self.conv
+        # NOTE(xcsong): We change input shape from (1, 1, frames, mel_dim)
+        #   to (1, 1, mel_dim, frames) for more efficient computation.
+        self.conv = module.conv
+        for idx in [0, 2, 4]:
+            self.conv[idx].weight = torch.nn.Parameter(
+                module.conv[idx].weight.transpose(2, 3))
+
+        # 2. Modify self.linear
+        # NOTE(xcsong): Split final projection to meet the requirment of
+        #   maximum kernel_size (7 for XJ3)
+        self.linear = torch.nn.ModuleList()
+        odim = module.linear.weight.size(0)  # 512, in this case
+        freq = module.linear.weight.size(1) // odim  # 4608 // 512 == 9
+        self.odim, self.freq = odim, freq
+        weight = module.linear.weight.reshape(
+            odim, odim, freq,
+            1)  # (odim, odim * freq) -> (odim, odim, freq, 1)
+        self.split_size = []
+        num_split = (freq - 1) // 7 + 1  # XJ3 requires kernel_size <= 7
+        slice_begin = 0
+        for idx in range(num_split):
+            kernel_size = min(freq, (idx + 1) * 7) - idx * 7
+            conv_ele = torch.nn.Conv2d(odim, odim, (kernel_size, 1),
+                                       (kernel_size, 1))
+            conv_ele.weight = torch.nn.Parameter(
+                weight[:, :, slice_begin:slice_begin + kernel_size, :])
+            conv_ele.bias = torch.nn.Parameter(torch.zeros_like(conv_ele.bias))
+            self.linear.append(conv_ele)
+            self.split_size.append(kernel_size)
+            slice_begin += kernel_size
+        self.linear[0].bias = torch.nn.Parameter(module.linear.bias)
+
+        self.check_equal(original)
+
+    def check_equal(self, module):
+        random_data = torch.randn(1, 67, 80)
+        mask = torch.zeros(1, 1, 67)
+        original_result, _, _ = module(random_data, mask)  # (1, 8, 512)
+        random_data = random_data.transpose(1,
+                                            2).unsqueeze(0)  # (1, 1, 80, 67)
+        new_result = self.forward(random_data)  # (1, 512, 1, 8)
+        np.testing.assert_allclose(to_numpy(original_result),
+                                   to_numpy(
+                                       new_result.squeeze(2).transpose(1, 2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Subsample x with 4-D dataflow.
+        Args:
+            x (torch.Tensor): Input tensor (#batch, 1, mel_dim, time).
+
+        Returns:
+            torch.Tensor: Subsampled tensor (#batch, odim, 1, time'),
+                where time' = time // 8.
+        """
+        x = self.conv(x)  # (1, odim, freq, time')
+        x_out = torch.zeros(x.size(0), self.odim, 1, x.size(3))
+        x = torch.split(x, self.split_size, dim=2)
+        for idx, (x_part, layer) in enumerate(zip(x, self.linear)):
+            x_out += layer(x_part)
+        return x_out
+
+
+class BPUMultiHeadedAttention(torch.nn.Module):
+    """Refactor wenet/transformer/attention.py::MultiHeadedAttention
+
+    NOTE(xcsong): Only support attention_class == MultiHeadedAttention,
+        we do not consider RelPositionMultiHeadedAttention currently.
+    """
+
+    def __init__(self, module, chunk_size, left_chunks):
+        super().__init__()
+        # Unchanged submodules and attributes
+        original = copy.deepcopy(module)
+        self.d_k = module.d_k
+        self.h = module.h
+        n_feat = self.d_k * self.h
+        self.chunk_size = chunk_size
+        self.left_chunks = left_chunks
+        self.time = chunk_size * (left_chunks + 1)
+        self.activation = torch.nn.Softmax(dim=-1)
+
+        # 1. Modify self.linear_x
+        self.linear_q = BPULinear(module.linear_q)
+        self.linear_k = BPULinear(module.linear_k)
+        self.linear_v = BPULinear(module.linear_v)
+        self.linear_out = BPULinear(module.linear_out)
+        # 2. denom
+        self.register_buffer(
+            "denom", torch.full((1, self.h, 1, 1), 1.0 / math.sqrt(self.d_k)))
+
+        self.check_equal(original)
+
+    def check_equal(self, module):
+        random_data = torch.randn(1, self.chunk_size, self.d_k * self.h)
+        mask = torch.ones((1, self.h, self.chunk_size, self.time),
+                          dtype=torch.bool)
+        cache = torch.zeros(1, self.h, self.chunk_size * self.left_chunks,
+                            self.d_k * 2)
+        original_out, original_cache = module(random_data, random_data,
+                                              random_data, mask[:, 0, :, :],
+                                              torch.empty(0), cache)
+        random_data = random_data.transpose(1, 2).unsqueeze(2)
+        cache = cache.reshape(1, self.h, self.d_k * 2,
+                              self.chunk_size * self.left_chunks)
+        new_out, new_cache = self.forward(random_data, random_data,
+                                          random_data, mask, cache)
+        np.testing.assert_allclose(to_numpy(original_out),
+                                   to_numpy(
+                                       new_out.squeeze(2).transpose(1, 2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+        np.testing.assert_allclose(to_numpy(original_cache),
+                                   to_numpy(new_cache.transpose(2, 3)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        mask: torch.Tensor,
+        cache: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute scaled dot product attention.
+
+        Args:
+            q (torch.Tensor): Query tensor (#batch, size, 1, chunk_size).
+            k (torch.Tensor): Key tensor (#batch, size, 1, chunk_size).
+            v (torch.Tensor): Value tensor (#batch, size, 1, chunk_size).
+            mask (torch.Tensor): Mask tensor,
+                (#batch, head, chunk_size, cache_t + chunk_size).
+            cache (torch.Tensor): Cache tensor
+                (1, head, d_k * 2, cache_t),
+                where `cache_t == chunk_size * left_chunks`.
+
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, size, 1, chunk_size).
+            torch.Tensor: Cache tensor
+                (1, head, d_k * 2, cache_t + chunk_size)
+                where `cache_t == chunk_size * left_chunks`
+        """
+        # 1. Forward QKV
+        q = self.linear_q(q)  # (1, d, 1, c) d == size, c == chunk_size
+        k = self.linear_k(k)  # (1, d, 1, c)
+        v = self.linear_v(v)  # (1, d, 1, c)
+        q = q.view(1, self.h, self.d_k, self.chunk_size)
+        k = k.view(1, self.h, self.d_k, self.chunk_size)
+        v = v.view(1, self.h, self.d_k, self.chunk_size)
+        q = q.transpose(2, 3)  # (batch, head, time1, d_k)
+        k_cache, v_cache = torch.split(cache, cache.size(2) // 2, dim=2)
+        k = torch.cat((k_cache, k), dim=3)
+        v = torch.cat((v_cache, v), dim=3)
+        new_cache = torch.cat((k, v), dim=2)
+        # 2. (Q^T)K
+        scores = torch.matmul(q, k) * self.denom  # (#b, n_head, time1, time2)
+        # 3. Forward attention
+        mask = mask.eq(0)
+        scores = scores.masked_fill(mask, -float('inf'))
+        attn = self.activation(scores).masked_fill(mask, 0.0)
+        attn = attn.transpose(2, 3)
+        x = torch.matmul(v, attn)
+        x = x.view(1, self.d_k * self.h, 1, self.chunk_size)
+        x_out = self.linear_out(x)
+        return x_out, new_cache
+
+
+class BPUConvolution(torch.nn.Module):
+    """Refactor wenet/transformer/convolution.py::ConvolutionModule
+
+    NOTE(xcsong): Only suport use_layer_norm == False
+    """
+
+    def __init__(self, module):
+        super().__init__()
+        # Unchanged submodules and attributes
+        original = copy.deepcopy(module)
+        self.lorder = module.lorder
+        self.use_layer_norm = False
+        self.activation = module.activation
+        channels = module.pointwise_conv1.weight.size(1)
+        self.channels = channels
+        kernel_size = module.depthwise_conv.weight.size(2)
+        assert module.use_layer_norm is False
+
+        # 1. Modify self.pointwise_conv1
+        self.pointwise_conv1 = BPULinear(module.pointwise_conv1, True)
+
+        # 2. Modify self.depthwise_conv
+        self.depthwise_conv = torch.nn.Conv2d(channels,
+                                              channels, (1, kernel_size),
+                                              stride=1,
+                                              groups=channels)
+        self.depthwise_conv.weight = torch.nn.Parameter(
+            module.depthwise_conv.weight.unsqueeze(-2))
+        self.depthwise_conv.bias = torch.nn.Parameter(
+            module.depthwise_conv.bias)
+
+        # 3. Modify self.norm, Only support batchnorm2d
+        self.norm = torch.nn.BatchNorm2d(channels)
+        self.norm.training = False
+        self.norm.num_features = module.norm.num_features
+        self.norm.eps = module.norm.eps
+        self.norm.momentum = module.norm.momentum
+        self.norm.weight = torch.nn.Parameter(module.norm.weight)
+        self.norm.bias = torch.nn.Parameter(module.norm.bias)
+        self.norm.running_mean = module.norm.running_mean
+        self.norm.running_var = module.norm.running_var
+
+        # 4. Modify self.pointwise_conv2
+        self.pointwise_conv2 = BPULinear(module.pointwise_conv2, True)
+
+        # 5. Identity conv, for running `concat` on BPU
+        self.identity = BPUIdentity(channels)
+
+        self.check_equal(original)
+
+    def check_equal(self, module):
+        random_data = torch.randn(1, 8, self.channels)
+        cache = torch.zeros((1, self.channels, self.lorder))
+        original_out, original_cache = module(random_data, cache=cache)
+        random_data = random_data.transpose(1, 2).unsqueeze(2)
+        cache = cache.unsqueeze(2)
+        new_out, new_cache = self.forward(random_data, cache)
+        np.testing.assert_allclose(to_numpy(original_out),
+                                   to_numpy(
+                                       new_out.squeeze(2).transpose(1, 2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+        np.testing.assert_allclose(to_numpy(original_cache),
+                                   to_numpy(new_cache.squeeze(2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+
+    def forward(self, x: torch.Tensor,
+                cache: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute convolution module.
+        Args:
+            x (torch.Tensor): Input tensor (#batch, channels, 1, chunk_size).
+            cache (torch.Tensor): left context cache, it is only
+                used in causal convolution (#batch, channels, 1, cache_t).
+        Returns:
+            torch.Tensor: Output tensor (#batch, channels, 1, chunk_size).
+            torch.Tensor: Cache tensor (#batch, channels, 1, cache_t).
+        """
+        # Concat cache
+        x = torch.cat((self.identity(cache), self.identity(x)), dim=3)
+        new_cache = x[:, :, :, -self.lorder:]
+
+        # GLU mechanism
+        x = self.pointwise_conv1(x)  # (batch, 2*channel, 1, dim)
+        x = torch.nn.functional.glu(x, dim=1)  # (b, channel, 1, dim)
+
+        # Depthwise Conv
+        x = self.depthwise_conv(x)
+        x = self.activation(self.norm(x))
+        x = self.pointwise_conv2(x)
+        return x, new_cache
+
+
+class BPUFFN(torch.nn.Module):
+    """Refactor wenet/transformer/positionwise_feed_forward.py::PositionwiseFeedForward
+    """
+
+    def __init__(self, module):
+        super().__init__()
+        # Unchanged submodules and attributes
+        original = copy.deepcopy(module)
+        self.activation = module.activation
+
+        # 1. Modify self.w_x
+        self.w_1 = BPULinear(module.w_1)
+        self.w_2 = BPULinear(module.w_2)
+
+        self.check_equal(original)
+
+    def check_equal(self, module):
+        random_data = torch.randn(1, 8, self.w_1.idim)
+        original_out = module(random_data)
+        random_data = random_data.transpose(1, 2).unsqueeze(2)
+        new_out = self.forward(random_data)
+        np.testing.assert_allclose(to_numpy(original_out),
+                                   to_numpy(
+                                       new_out.squeeze(2).transpose(1, 2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward function.
+
+        Args:
+            xs: input tensor (B, D, 1, L)
+        Returns:
+            output tensor, (B, D, 1, L)
+        """
+        return self.w_2(self.activation(self.w_1(x)))
+
+
+class BPUConformerEncoderLayer(torch.nn.Module):
+    """Refactor wenet/transformer/encoder_layer.py::ConformerEncoderLayer
+    """
+
+    def __init__(self, module, chunk_size, left_chunks, ln_run_on_bpu=False):
+        super().__init__()
+        # Unchanged submodules and attributes
+        original = copy.deepcopy(module)
+        self.size = module.size
+        assert module.normalize_before is True
+        assert module.concat_after is False
+
+        # 1. Modify submodules
+        self.feed_forward_macaron = BPUFFN(module.feed_forward_macaron)
+        self.self_attn = BPUMultiHeadedAttention(module.self_attn, chunk_size,
+                                                 left_chunks)
+        self.conv_module = BPUConvolution(module.conv_module)
+        self.feed_forward = BPUFFN(module.feed_forward)
+
+        # 2. Modify norms
+        self.norm_ff = BPULayerNorm(module.norm_ff, chunk_size, ln_run_on_bpu)
+        self.norm_mha = BPULayerNorm(module.norm_mha, chunk_size,
+                                     ln_run_on_bpu)
+        self.norm_ff_macron = BPULayerNorm(module.norm_ff_macaron, chunk_size,
+                                           ln_run_on_bpu)
+        self.norm_conv = BPULayerNorm(module.norm_conv, chunk_size,
+                                      ln_run_on_bpu)
+        self.norm_final = BPULayerNorm(module.norm_final, chunk_size,
+                                       ln_run_on_bpu)
+
+        # 3. 4-D ff_scale
+        self.register_buffer("ff_scale",
+                             torch.full((1, self.size, 1, 1), module.ff_scale))
+
+        self.check_equal(original)
+
+    def check_equal(self, module):
+        time1 = self.self_attn.chunk_size
+        time2 = self.self_attn.time
+        h, d_k = self.self_attn.h, self.self_attn.d_k
+        random_x = torch.randn(1, time1, self.size)
+        att_mask = torch.ones(1, h, time1, time2)
+        att_cache = torch.zeros(1, h, time2 - time1, d_k * 2)
+        cnn_cache = torch.zeros(1, self.size, self.conv_module.lorder)
+        original_x, _, original_att_cache, original_cnn_cache = module(
+            random_x,
+            att_mask[:, 0, :, :],
+            torch.empty(0),
+            att_cache=att_cache,
+            cnn_cache=cnn_cache)
+        random_x = random_x.transpose(1, 2).unsqueeze(2)
+        att_cache = att_cache.reshape(1, h, d_k * 2, time2 - time1)
+        cnn_cache = cnn_cache.unsqueeze(2)
+        new_x, new_att_cache, new_cnn_cache = self.forward(
+            random_x, att_mask, att_cache, cnn_cache)
+        np.testing.assert_allclose(to_numpy(original_att_cache),
+                                   to_numpy(new_att_cache.transpose(2, 3)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+        np.testing.assert_allclose(to_numpy(original_x),
+                                   to_numpy(new_x.squeeze(2).transpose(1, 2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+        np.testing.assert_allclose(to_numpy(original_cnn_cache),
+                                   to_numpy(new_cnn_cache.squeeze(2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+
+    def forward(
+        self, x: torch.Tensor, att_mask: torch.Tensor, att_cache: torch.Tensor,
+        cnn_cache: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Compute encoded features.
+
+        Args:
+            x (torch.Tensor): (#batch, size, 1, chunk_size)
+            att_mask (torch.Tensor): Mask tensor for the input
+                (#batch, head, chunk_size, cache_t1 + chunk_size),
+            att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
+                (#batch=1, head, d_k * 2, cache_t1), head * d_k == size.
+            cnn_cache (torch.Tensor): Convolution cache in conformer layer
+                (#batch=1, size, 1, cache_t2)
+        Returns:
+            torch.Tensor: Output tensor (#batch, size, 1, chunk_size).
+            torch.Tensor: att_cache tensor,
+                (1, head, d_k * 2, cache_t1 + chunk_size).
+            torch.Tensor: cnn_cahce tensor (#batch, size, 1, cache_t2).
+        """
+        # 1. ffn_macaron
+        residual = x
+        x = self.norm_ff_macron(x)
+        x = residual + self.ff_scale * self.feed_forward_macaron(x)
+
+        # 2. attention
+        residual = x
+        x = self.norm_mha(x)
+        x_att, new_att_cache = self.self_attn(x, x, x, att_mask, att_cache)
+        x = residual + x_att
+
+        # 3. convolution
+        residual = x
+        x = self.norm_conv(x)
+        x, new_cnn_cache = self.conv_module(x, cnn_cache)
+        x = residual + x
+
+        # 4. ffn
+        residual = x
+        x = self.norm_ff(x)
+        x = residual + self.ff_scale * self.feed_forward(x)
+
+        # 5. final post-norm
+        x = self.norm_final(x)
+
+        return x, new_att_cache, new_cnn_cache
+
+
+class BPUConformerEncoder(torch.nn.Module):
+    """Refactor wenet/transformer/encoder.py::ConformerEncoder
+    """
+
+    def __init__(self, module, chunk_size, left_chunks, ln_run_on_bpu=False):
+        super().__init__()
+        # Unchanged submodules and attributes
+        original = copy.deepcopy(module)
+        output_size = module.output_size()
+        self._output_size = module.output_size()
+        self.after_norm = module.after_norm
+        self.chunk_size = chunk_size
+        self.left_chunks = left_chunks
+        self.head = module.encoders[0].self_attn.h
+        self.layers = len(module.encoders)
+
+        # 1. Modify submodules
+        self.global_cmvn = BPUGlobalCMVN(module.global_cmvn)
+        self.embed = BPUConv2dSubsampling8(module.embed)
+        self.encoders = torch.nn.ModuleList()
+        for layer in module.encoders:
+            self.encoders.append(
+                BPUConformerEncoderLayer(layer, chunk_size, left_chunks,
+                                         ln_run_on_bpu))
+
+        # 2. Auxiliary conv
+        self.identity_cnncache = BPUIdentity(output_size)
+
+        self.check_equal(original)
+
+    def check_equal(self, module):
+        time1 = self.encoders[0].self_attn.chunk_size
+        time2 = self.encoders[0].self_attn.time
+        layers = self.layers
+        h, d_k = self.head, self.encoders[0].self_attn.d_k
+        decoding_window = (self.chunk_size - 1) * \
+            module.embed.subsampling_rate + \
+            module.embed.right_context + 1
+        lorder = self.encoders[0].conv_module.lorder
+        random_x = torch.randn(1, decoding_window, 80)
+        att_mask = torch.ones(1, h, time1, time2)
+        att_cache = torch.zeros(layers, h, time2 - time1, d_k * 2)
+        cnn_cache = torch.zeros(layers, 1, self._output_size, lorder)
+        orig_x, orig_att_cache, orig_cnn_cache = module.forward_chunk(
+            random_x,
+            0,
+            time2 - time1,
+            att_mask=att_mask[:, 0, :, :],
+            att_cache=att_cache,
+            cnn_cache=cnn_cache)
+        random_x = random_x.unsqueeze(0)
+        att_cache = att_cache.reshape(1, h * layers, d_k * 2, time2 - time1)
+        cnn_cache = cnn_cache.reshape(1, self._output_size, layers, lorder)
+        new_x, new_att_cache, new_cnn_cache = self.forward(
+            random_x, att_cache, cnn_cache, att_mask)
+        caches = torch.split(new_att_cache, h, dim=1)
+        caches = [c.transpose(2, 3) for c in caches]
+        np.testing.assert_allclose(to_numpy(orig_att_cache),
+                                   to_numpy(torch.cat(caches, dim=0)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+        np.testing.assert_allclose(to_numpy(orig_x),
+                                   to_numpy(new_x.squeeze(2).transpose(1, 2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+        np.testing.assert_allclose(
+            to_numpy(orig_cnn_cache),
+            to_numpy(new_cnn_cache.transpose(0, 2).transpose(1, 2)),
+            rtol=1e-02,
+            atol=1e-03)
+
+    def forward(
+        self, xs: torch.Tensor, att_cache: torch.Tensor,
+        cnn_cache: torch.Tensor, att_mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """ Forward just one chunk
+
+        Args:
+            xs (torch.Tensor): chunk input, with shape (b=1, 1, time, mel-dim),
+                where `time == (chunk_size - 1) * subsample_rate + \
+                        subsample.right_context + 1`
+            att_cache (torch.Tensor): cache tensor for KEY & VALUE in
+                transformer/conformer attention, with shape
+                (1, head * elayers, d_k * 2, cache_t1), where
+                `head * d_k == hidden-dim` and
+                `cache_t1 == chunk_size * left_chunks`.
+            cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer,
+                (1, hidden-dim, elayers, cache_t2), where
+                `cache_t2 == cnn.lorder - 1`
+            att_mask (torch.Tensor): Mask tensor for the input
+                (#batch, head, chunk_size, cache_t1 + chunk_size),
+
+        Returns:
+            torch.Tensor: output of current input xs,
+                with shape (b=1, hidden-dim, 1, chunk_size).
+            torch.Tensor: new attention cache required for next chunk, with
+                same shape as the original att_cache.
+            torch.Tensor: new conformer cnn cache required for next chunk, with
+                same shape as the original cnn_cache.
+        """
+        # xs: (B, 1, time, mel_dim) -> (B, 1, mel_dim, time)
+        xs = xs.transpose(2, 3)
+        xs = self.global_cmvn(xs)
+        # xs: (B, 1, mel_dim, time) -> (B, hidden_dim, 1, chunk_size)
+        xs = self.embed(xs)
+
+        att_cache = torch.split(att_cache, self.head, dim=1)
+        cnn_cache = self.identity_cnncache(cnn_cache)
+        cnn_cache = torch.split(cnn_cache, 1, dim=2)
+        r_att_cache = []
+        r_cnn_cache = []
+        for i, layer in enumerate(self.encoders):
+            xs, new_att_cache, new_cnn_cache = layer(xs,
+                                                     att_mask,
+                                                     att_cache=att_cache[i],
+                                                     cnn_cache=cnn_cache[i])
+            r_att_cache.append(new_att_cache[:, :, :, self.chunk_size:])
+            r_cnn_cache.append(new_cnn_cache)
+        r_att_cache = torch.cat(r_att_cache, dim=1)
+        r_cnn_cache = self.identity_cnncache(torch.cat(r_cnn_cache, dim=2))
+
+        xs = xs.squeeze(2).transpose(1, 2).contiguous()
+        xs = self.after_norm(xs)
+        # NOTE(xcsong): 4D in, 4D out to meet the requirment of CTC input.
+        xs = xs.transpose(1, 2).contiguous().unsqueeze(2)  # (B, C, 1, T)
+
+        return (xs, r_att_cache, r_cnn_cache)
+
+
+class BPUCTC(torch.nn.Module):
+    """Refactor wenet/transformer/ctc.py::CTC
+    """
+
+    def __init__(self, module):
+        super().__init__()
+        # Unchanged submodules and attributes
+        original = copy.deepcopy(module)
+        self.idim = module.ctc_lo.weight.size(1)
+        num_class = module.ctc_lo.weight.size(0)
+
+        # 1. Modify self.ctc_lo, Split final projection to meet the
+        #   requirment of maximum in/out channels (2048 for XJ3)
+        self.ctc_lo = torch.nn.ModuleList()
+        self.split_size = []
+        num_split = (num_class - 1) // 2048 + 1
+        for idx in range(num_split):
+            out_channel = min(num_class, (idx + 1) * 2048) - idx * 2048
+            conv_ele = torch.nn.Conv2d(self.idim, out_channel, 1, 1)
+            self.ctc_lo.append(conv_ele)
+            self.split_size.append(out_channel)
+        orig_weight = torch.split(module.ctc_lo.weight, self.split_size, dim=0)
+        orig_bias = torch.split(module.ctc_lo.bias, self.split_size, dim=0)
+        for i, (w, b) in enumerate(zip(orig_weight, orig_bias)):
+            w = w.unsqueeze(2).unsqueeze(3)
+            self.ctc_lo[i].weight = torch.nn.Parameter(w)
+            self.ctc_lo[i].bias = torch.nn.Parameter(b)
+
+        self.check_equal(original)
+
+    def check_equal(self, module):
+        random_data = torch.randn(1, 100, self.idim)
+        original_result = module.ctc_lo(random_data)
+        random_data = random_data.transpose(1, 2).unsqueeze(2)
+        new_result = self.forward(random_data)
+        np.testing.assert_allclose(to_numpy(original_result),
+                                   to_numpy(
+                                       new_result.squeeze(2).transpose(1, 2)),
+                                   rtol=1e-02,
+                                   atol=1e-03)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """frame activations, without softmax.
+
+        Args:
+            Tensor x: 4d tensor (B, hidden_dim, 1, chunk_size)
+        Returns:
+            torch.Tensor: (B, num_class, 1, chunk_size)
+        """
+        out = []
+        for i, layer in enumerate(self.ctc_lo):
+            out.append(layer(x))
+        out = torch.cat(out, dim=1)
+        return out
+
+
+def export_encoder(asr_model, args):
+    logger.info("Stage-1: export encoder")
+    decode_window, mel_dim = args.decoding_window, args.feature_size
+    encoder = BPUConformerEncoder(asr_model.encoder, args.chunk_size,
+                                  args.num_decoding_left_chunks,
+                                  args.ln_run_on_bpu)
+    encoder.eval()
+    encoder_outpath = os.path.join(args.output_dir, 'encoder.onnx')
+
+    logger.info("Stage-1.1: prepare inputs for encoder")
+    chunk = torch.randn((1, 1, decode_window, mel_dim))
+    required_cache_size = encoder.chunk_size * encoder.left_chunks
+    kv_time = required_cache_size + encoder.chunk_size
+    hidden, layers = encoder._output_size, len(encoder.encoders)
+    head = encoder.encoders[0].self_attn.h
+    d_k = hidden // head
+    lorder = encoder.encoders[0].conv_module.lorder
+    att_cache = torch.zeros(1, layers * head, d_k * 2, required_cache_size)
+    att_mask = torch.ones((1, head, encoder.chunk_size, kv_time))
+    att_mask[:, :, :, :required_cache_size] = 0
+    cnn_cache = torch.zeros((1, hidden, layers, lorder))
+    inputs = (chunk, att_cache, cnn_cache, att_mask)
+    logger.info("chunk.size(): {} att_cache.size(): {} "
+                "cnn_cache.size(): {} att_mask.size(): {}".format(
+                    list(chunk.size()), list(att_cache.size()),
+                    list(cnn_cache.size()), list(att_mask.size())))
+
+    logger.info("Stage-1.2: torch.onnx.export")
+    # NOTE(xcsong): Below attributes will be used in
+    #   onnx2horizonbin.py::generate_config()
+    attributes = {}
+    attributes['input_name'] = "chunk;att_cache;cnn_cache;att_mask"
+    attributes['output_name'] = "output;r_att_cache;r_cnn_cache"
+    attributes['input_type'] = "featuremap;featuremap;featuremap;featuremap"
+    attributes['norm_type'] = \
+        "no_preprocess;no_preprocess;no_preprocess;no_preprocess"
+    attributes['input_layout_train'] = "NCHW;NCHW;NCHW;NCHW"
+    attributes['input_layout_rt'] = "NCHW;NCHW;NCHW;NCHW"
+    attributes['input_shape'] = \
+        "{}x{}x{}x{};{}x{}x{}x{};{}x{}x{}x{};{}x{}x{}x{}".format(
+        chunk.size(0), chunk.size(1), chunk.size(2), chunk.size(3),
+        att_cache.size(0), att_cache.size(1), att_cache.size(2),
+        att_cache.size(3), cnn_cache.size(0), cnn_cache.size(1),
+        cnn_cache.size(2), cnn_cache.size(3), att_mask.size(0),
+        att_mask.size(1), att_mask.size(2), att_mask.size(3)
+    )
+    torch.onnx.export(  # NOTE(xcsong): only support opset==11
+        encoder,
+        inputs,
+        encoder_outpath,
+        opset_version=11,
+        export_params=True,
+        do_constant_folding=True,
+        input_names=attributes['input_name'].split(';'),
+        output_names=attributes['output_name'].split(';'),
+        dynamic_axes=None,
+        verbose=False)
+    onnx_encoder = onnx.load(encoder_outpath)
+    for k in vars(args):
+        meta = onnx_encoder.metadata_props.add()
+        meta.key, meta.value = str(k), str(getattr(args, k))
+    for k in attributes:
+        meta = onnx_encoder.metadata_props.add()
+        meta.key, meta.value = str(k), str(attributes[k])
+    onnx.checker.check_model(onnx_encoder)
+    onnx.helper.printable_graph(onnx_encoder.graph)
+    onnx.save(onnx_encoder, encoder_outpath)
+    print_input_output_info(onnx_encoder, "onnx_encoder")
+    logger.info('Export onnx_encoder, done! see {}'.format(encoder_outpath))
+
+    logger.info("Stage-1.3: check onnx_encoder and torch_encoder")
+    torch_output = []
+    torch_chunk, torch_att_mask = copy.deepcopy(chunk), copy.deepcopy(att_mask)
+    torch_att_cache = copy.deepcopy(att_cache)
+    torch_cnn_cache = copy.deepcopy(cnn_cache)
+    for i in range(10):
+        logger.info("torch chunk-{}: {}, att_cache: {}, cnn_cache: {}"
+                    ", att_mask: {}".format(i, list(torch_chunk.size()),
+                                            list(torch_att_cache.size()),
+                                            list(torch_cnn_cache.size()),
+                                            list(torch_att_mask.size())))
+        torch_att_mask[:, :, :, -(encoder.chunk_size * (i + 1)):] = 1
+        out, torch_att_cache, torch_cnn_cache = encoder(
+            torch_chunk, torch_att_cache, torch_cnn_cache, torch_att_mask)
+        torch_output.append(out)
+    torch_output = torch.cat(torch_output, dim=-1)
+
+    onnx_output = []
+    onnx_chunk, onnx_att_mask = to_numpy(chunk), to_numpy(att_mask)
+    onnx_att_cache = to_numpy(att_cache)
+    onnx_cnn_cache = to_numpy(cnn_cache)
+    ort_session = onnxruntime.InferenceSession(encoder_outpath)
+    input_names = [node.name for node in onnx_encoder.graph.input]
+    for i in range(10):
+        logger.info("onnx  chunk-{}: {}, att_cache: {}, cnn_cache: {},"
+                    " att_mask: {}".format(i, onnx_chunk.shape,
+                                           onnx_att_cache.shape,
+                                           onnx_cnn_cache.shape,
+                                           onnx_att_mask.shape))
+        onnx_att_mask[:, :, :, -(encoder.chunk_size * (i + 1)):] = 1
+        ort_inputs = {
+            'chunk': onnx_chunk,
+            'att_cache': onnx_att_cache,
+            'cnn_cache': onnx_cnn_cache,
+            'att_mask': onnx_att_mask,
+        }
+        ort_outs = ort_session.run(None, ort_inputs)
+        onnx_att_cache, onnx_cnn_cache = ort_outs[1], ort_outs[2]
+        onnx_output.append(ort_outs[0])
+    onnx_output = np.concatenate(onnx_output, axis=-1)
+
+    np.testing.assert_allclose(to_numpy(torch_output),
+                               onnx_output,
+                               rtol=1e-03,
+                               atol=1e-04)
+    meta = ort_session.get_modelmeta()
+    logger.info("custom_metadata_map={}".format(meta.custom_metadata_map))
+    logger.info("Check onnx_encoder, pass!")
+    return encoder, ort_session
+
+
+def export_ctc(asr_model, args):
+    logger.info("Stage-2: export ctc")
+    ctc = BPUCTC(asr_model.ctc).eval()
+    ctc_outpath = os.path.join(args.output_dir, 'ctc.onnx')
+
+    logger.info("Stage-2.1: prepare inputs for ctc")
+    hidden = torch.randn((1, args.output_size, 1, args.chunk_size))
+
+    logger.info("Stage-2.2: torch.onnx.export")
+    # NOTE(xcsong): Below attributes will be used in
+    #   onnx2horizonbin.py::generate_config()
+    attributes = {}
+    attributes['input_name'], attributes['input_type'] = "hidden", "featuremap"
+    attributes['norm_type'] = "no_preprocess"
+    attributes['input_layout_train'] = "NCHW"
+    attributes['input_layout_rt'] = "NCHW"
+    attributes['input_shape'] = "{}x{}x{}x{}".format(
+        hidden.size(0),
+        hidden.size(1),
+        hidden.size(2),
+        hidden.size(3),
+    )
+    torch.onnx.export(ctc,
+                      hidden,
+                      ctc_outpath,
+                      opset_version=11,
+                      export_params=True,
+                      do_constant_folding=True,
+                      input_names=['hidden'],
+                      output_names=['probs'],
+                      dynamic_axes=None,
+                      verbose=False)
+    onnx_ctc = onnx.load(ctc_outpath)
+    for k in vars(args):
+        meta = onnx_ctc.metadata_props.add()
+        meta.key, meta.value = str(k), str(getattr(args, k))
+    for k in attributes:
+        meta = onnx_ctc.metadata_props.add()
+        meta.key, meta.value = str(k), str(attributes[k])
+    onnx.checker.check_model(onnx_ctc)
+    onnx.helper.printable_graph(onnx_ctc.graph)
+    onnx.save(onnx_ctc, ctc_outpath)
+    print_input_output_info(onnx_ctc, "onnx_ctc")
+    logger.info('Export onnx_ctc, done! see {}'.format(ctc_outpath))
+
+    logger.info("Stage-2.3: check onnx_ctc and torch_ctc")
+    torch_output = ctc(hidden)
+    ort_session = onnxruntime.InferenceSession(ctc_outpath)
+    onnx_output = ort_session.run(None, {'hidden': to_numpy(hidden)})
+
+    np.testing.assert_allclose(to_numpy(torch_output),
+                               onnx_output[0],
+                               rtol=1e-03,
+                               atol=1e-04)
+    meta = ort_session.get_modelmeta()
+    logger.info("custom_metadata_map={}".format(meta.custom_metadata_map))
+    logger.info("Check onnx_ctc, pass!")
+    return ctc, ort_session
+
+
+def export_decoder(asr_model, args):
+    logger.info("Currently, Decoder is not supported.")
+
+
+if __name__ == '__main__':
+    torch.manual_seed(777)
+    args = get_args()
+    args.ln_run_on_bpu = False
+    # NOTE(xcsong): XJ3 BPU only support static shapes
+    assert args.chunk_size > 0
+    assert args.num_decoding_left_chunks > 0
+    os.system("mkdir -p " + args.output_dir)
+    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
+    with open(args.config, 'r') as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+
+    model, configs = init_model(args, configs)
+    model.eval()
+    print(model)
+
+    args.feature_size = configs['input_dim']
+    args.output_size = model.encoder.output_size()
+    args.decoding_window = (args.chunk_size - 1) * \
+        model.encoder.embed.subsampling_rate + \
+        model.encoder.embed.right_context + 1
+
+    export_encoder(model, args)
+    export_ctc(model, args)
+    export_decoder(model, args)

wenet/bin/export_onnx_cpu.py

@@ -0,0 +1,470 @@
+# Copyright (c) 2022, Xingchen Song (sxc19@mails.tsinghua.edu.cn)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import argparse
+import logging
+import os
+import copy
+import sys
+
+import torch
+import yaml
+import numpy as np
+
+from wenet.utils.init_model import init_model
+
+try:
+    import onnx
+    import onnxruntime
+    from onnxruntime.quantization import quantize_dynamic, QuantType
+except ImportError:
+    print('Please install onnx and onnxruntime!')
+    sys.exit(1)
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='export your script model')
+    parser.add_argument('--config', required=True, help='config file')
+    parser.add_argument('--checkpoint', required=True, help='checkpoint model')
+    parser.add_argument('--output_dir', required=True, help='output directory')
+    parser.add_argument('--chunk_size',
+                        required=True,
+                        type=int,
+                        help='decoding chunk size')
+    parser.add_argument('--num_decoding_left_chunks',
+                        required=True,
+                        type=int,
+                        help='cache chunks')
+    parser.add_argument('--reverse_weight',
+                        default=0.5,
+                        type=float,
+                        help='reverse_weight in attention_rescoing')
+    args = parser.parse_args()
+    return args
+
+
+def to_numpy(tensor):
+    if tensor.requires_grad:
+        return tensor.detach().cpu().numpy()
+    else:
+        return tensor.cpu().numpy()
+
+
+def print_input_output_info(onnx_model, name, prefix="\t\t"):
+    input_names = [node.name for node in onnx_model.graph.input]
+    input_shapes = [[d.dim_value for d in node.type.tensor_type.shape.dim]
+                    for node in onnx_model.graph.input]
+    output_names = [node.name for node in onnx_model.graph.output]
+    output_shapes = [[d.dim_value for d in node.type.tensor_type.shape.dim]
+                     for node in onnx_model.graph.output]
+    print("{}{} inputs : {}".format(prefix, name, input_names))
+    print("{}{} input shapes : {}".format(prefix, name, input_shapes))
+    print("{}{} outputs: {}".format(prefix, name, output_names))
+    print("{}{} output shapes : {}".format(prefix, name, output_shapes))
+
+
+def export_encoder(asr_model, args):
+    print("Stage-1: export encoder")
+    encoder = asr_model.encoder
+    encoder.forward = encoder.forward_chunk
+    encoder_outpath = os.path.join(args['output_dir'], 'encoder.onnx')
+
+    print("\tStage-1.1: prepare inputs for encoder")
+    chunk = torch.randn(
+        (args['batch'], args['decoding_window'], args['feature_size']))
+    offset = 0
+    # NOTE(xcsong): The uncertainty of `next_cache_start` only appears
+    #   in the first few chunks, this is caused by dynamic att_cache shape, i,e
+    #   (0, 0, 0, 0) for 1st chunk and (elayers, head, ?, d_k*2) for subsequent
+    #   chunks. One way to ease the ONNX export is to keep `next_cache_start`
+    #   as a fixed value. To do this, for the **first** chunk, if
+    #   left_chunks > 0, we feed real cache & real mask to the model, otherwise
+    #   fake cache & fake mask. In this way, we get:
+    #   1. 16/-1 mode: next_cache_start == 0 for all chunks
+    #   2. 16/4  mode: next_cache_start == chunk_size for all chunks
+    #   3. 16/0  mode: next_cache_start == chunk_size for all chunks
+    #   4. -1/-1 mode: next_cache_start == 0 for all chunks
+    #   NO MORE DYNAMIC CHANGES!!
+    #
+    # NOTE(Mddct): We retain the current design for the convenience of supporting some
+    #   inference frameworks without dynamic shapes. If you're interested in all-in-one
+    #   model that supports different chunks please see:
+    #   https://github.com/wenet-e2e/wenet/pull/1174
+
+    if args['left_chunks'] > 0:  # 16/4
+        required_cache_size = args['chunk_size'] * args['left_chunks']
+        offset = required_cache_size
+        # Real cache
+        att_cache = torch.zeros(
+            (args['num_blocks'], args['head'], required_cache_size,
+             args['output_size'] // args['head'] * 2))
+        # Real mask
+        att_mask = torch.ones(
+            (args['batch'], 1, required_cache_size + args['chunk_size']),
+            dtype=torch.bool)
+        att_mask[:, :, :required_cache_size] = 0
+    elif args['left_chunks'] <= 0:  # 16/-1, -1/-1, 16/0
+        required_cache_size = -1 if args['left_chunks'] < 0 else 0
+        # Fake cache
+        att_cache = torch.zeros((args['num_blocks'], args['head'], 0,
+                                 args['output_size'] // args['head'] * 2))
+        # Fake mask
+        att_mask = torch.ones((0, 0, 0), dtype=torch.bool)
+    cnn_cache = torch.zeros(
+        (args['num_blocks'], args['batch'], args['output_size'],
+         args['cnn_module_kernel'] - 1))
+    inputs = (chunk, offset, required_cache_size, att_cache, cnn_cache,
+              att_mask)
+    print("\t\tchunk.size(): {}\n".format(chunk.size()),
+          "\t\toffset: {}\n".format(offset),
+          "\t\trequired_cache: {}\n".format(required_cache_size),
+          "\t\tatt_cache.size(): {}\n".format(att_cache.size()),
+          "\t\tcnn_cache.size(): {}\n".format(cnn_cache.size()),
+          "\t\tatt_mask.size(): {}\n".format(att_mask.size()))
+
+    print("\tStage-1.2: torch.onnx.export")
+    dynamic_axes = {
+        'chunk': {
+            1: 'T'
+        },
+        'att_cache': {
+            2: 'T_CACHE'
+        },
+        'att_mask': {
+            2: 'T_ADD_T_CACHE'
+        },
+        'output': {
+            1: 'T'
+        },
+        'r_att_cache': {
+            2: 'T_CACHE'
+        },
+    }
+    # NOTE(xcsong): We keep dynamic axes even if in 16/4 mode, this is
+    #   to avoid padding the last chunk (which usually contains less
+    #   frames than required). For users who want static axes, just pop
+    #   out specific axis.
+    # if args['chunk_size'] > 0:  # 16/4, 16/-1, 16/0
+    #     dynamic_axes.pop('chunk')
+    #     dynamic_axes.pop('output')
+    # if args['left_chunks'] >= 0:  # 16/4, 16/0
+    #     # NOTE(xsong): since we feed real cache & real mask into the
+    #     #   model when left_chunks > 0, the shape of cache will never
+    #     #   be changed.
+    #     dynamic_axes.pop('att_cache')
+    #     dynamic_axes.pop('r_att_cache')
+    torch.onnx.export(encoder,
+                      inputs,
+                      encoder_outpath,
+                      opset_version=13,
+                      export_params=True,
+                      do_constant_folding=True,
+                      input_names=[
+                          'chunk', 'offset', 'required_cache_size',
+                          'att_cache', 'cnn_cache', 'att_mask'
+                      ],
+                      output_names=['output', 'r_att_cache', 'r_cnn_cache'],
+                      dynamic_axes=dynamic_axes,
+                      verbose=False)
+    onnx_encoder = onnx.load(encoder_outpath)
+    for (k, v) in args.items():
+        meta = onnx_encoder.metadata_props.add()
+        meta.key, meta.value = str(k), str(v)
+    onnx.checker.check_model(onnx_encoder)
+    onnx.helper.printable_graph(onnx_encoder.graph)
+    # NOTE(xcsong): to add those metadatas we need to reopen
+    #   the file and resave it.
+    onnx.save(onnx_encoder, encoder_outpath)
+    print_input_output_info(onnx_encoder, "onnx_encoder")
+    # Dynamic quantization
+    model_fp32 = encoder_outpath
+    model_quant = os.path.join(args['output_dir'], 'encoder.quant.onnx')
+    quantize_dynamic(model_fp32, model_quant, weight_type=QuantType.QUInt8)
+    print('\t\tExport onnx_encoder, done! see {}'.format(encoder_outpath))
+
+    print("\tStage-1.3: check onnx_encoder and torch_encoder")
+    torch_output = []
+    torch_chunk = copy.deepcopy(chunk)
+    torch_offset = copy.deepcopy(offset)
+    torch_required_cache_size = copy.deepcopy(required_cache_size)
+    torch_att_cache = copy.deepcopy(att_cache)
+    torch_cnn_cache = copy.deepcopy(cnn_cache)
+    torch_att_mask = copy.deepcopy(att_mask)
+    for i in range(10):
+        print("\t\ttorch chunk-{}: {}, offset: {}, att_cache: {},"
+              " cnn_cache: {}, att_mask: {}".format(
+                  i, list(torch_chunk.size()), torch_offset,
+                  list(torch_att_cache.size()), list(torch_cnn_cache.size()),
+                  list(torch_att_mask.size())))
+        # NOTE(xsong): att_mask of the first few batches need changes if
+        #   we use 16/4 mode.
+        if args['left_chunks'] > 0:  # 16/4
+            torch_att_mask[:, :, -(args['chunk_size'] * (i + 1)):] = 1
+        out, torch_att_cache, torch_cnn_cache = encoder(
+            torch_chunk, torch_offset, torch_required_cache_size,
+            torch_att_cache, torch_cnn_cache, torch_att_mask)
+        torch_output.append(out)
+        torch_offset += out.size(1)
+    torch_output = torch.cat(torch_output, dim=1)
+
+    onnx_output = []
+    onnx_chunk = to_numpy(chunk)
+    onnx_offset = np.array((offset)).astype(np.int64)
+    onnx_required_cache_size = np.array((required_cache_size)).astype(np.int64)
+    onnx_att_cache = to_numpy(att_cache)
+    onnx_cnn_cache = to_numpy(cnn_cache)
+    onnx_att_mask = to_numpy(att_mask)
+    ort_session = onnxruntime.InferenceSession(
+        encoder_outpath, providers=['CPUExecutionProvider'])
+    input_names = [node.name for node in onnx_encoder.graph.input]
+    for i in range(10):
+        print("\t\tonnx  chunk-{}: {}, offset: {}, att_cache: {},"
+              " cnn_cache: {}, att_mask: {}".format(i, onnx_chunk.shape,
+                                                    onnx_offset,
+                                                    onnx_att_cache.shape,
+                                                    onnx_cnn_cache.shape,
+                                                    onnx_att_mask.shape))
+        # NOTE(xsong): att_mask of the first few batches need changes if
+        #   we use 16/4 mode.
+        if args['left_chunks'] > 0:  # 16/4
+            onnx_att_mask[:, :, -(args['chunk_size'] * (i + 1)):] = 1
+        ort_inputs = {
+            'chunk': onnx_chunk,
+            'offset': onnx_offset,
+            'required_cache_size': onnx_required_cache_size,
+            'att_cache': onnx_att_cache,
+            'cnn_cache': onnx_cnn_cache,
+            'att_mask': onnx_att_mask
+        }
+        # NOTE(xcsong): If we use 16/-1, -1/-1 or 16/0 mode, `next_cache_start`
+        #   will be hardcoded to 0 or chunk_size by ONNX, thus
+        #   required_cache_size and att_mask are no more needed and they will
+        #   be removed by ONNX automatically.
+        for k in list(ort_inputs):
+            if k not in input_names:
+                ort_inputs.pop(k)
+        ort_outs = ort_session.run(None, ort_inputs)
+        onnx_att_cache, onnx_cnn_cache = ort_outs[1], ort_outs[2]
+        onnx_output.append(ort_outs[0])
+        onnx_offset += ort_outs[0].shape[1]
+    onnx_output = np.concatenate(onnx_output, axis=1)
+
+    np.testing.assert_allclose(to_numpy(torch_output),
+                               onnx_output,
+                               rtol=1e-03,
+                               atol=1e-05)
+    meta = ort_session.get_modelmeta()
+    print("\t\tcustom_metadata_map={}".format(meta.custom_metadata_map))
+    print("\t\tCheck onnx_encoder, pass!")
+
+
+def export_ctc(asr_model, args):
+    print("Stage-2: export ctc")
+    ctc = asr_model.ctc
+    ctc.forward = ctc.log_softmax
+    ctc_outpath = os.path.join(args['output_dir'], 'ctc.onnx')
+
+    print("\tStage-2.1: prepare inputs for ctc")
+    hidden = torch.randn(
+        (args['batch'], args['chunk_size'] if args['chunk_size'] > 0 else 16,
+         args['output_size']))
+
+    print("\tStage-2.2: torch.onnx.export")
+    dynamic_axes = {'hidden': {1: 'T'}, 'probs': {1: 'T'}}
+    torch.onnx.export(ctc,
+                      hidden,
+                      ctc_outpath,
+                      opset_version=13,
+                      export_params=True,
+                      do_constant_folding=True,
+                      input_names=['hidden'],
+                      output_names=['probs'],
+                      dynamic_axes=dynamic_axes,
+                      verbose=False)
+    onnx_ctc = onnx.load(ctc_outpath)
+    for (k, v) in args.items():
+        meta = onnx_ctc.metadata_props.add()
+        meta.key, meta.value = str(k), str(v)
+    onnx.checker.check_model(onnx_ctc)
+    onnx.helper.printable_graph(onnx_ctc.graph)
+    onnx.save(onnx_ctc, ctc_outpath)
+    print_input_output_info(onnx_ctc, "onnx_ctc")
+    # Dynamic quantization
+    model_fp32 = ctc_outpath
+    model_quant = os.path.join(args['output_dir'], 'ctc.quant.onnx')
+    quantize_dynamic(model_fp32, model_quant, weight_type=QuantType.QUInt8)
+    print('\t\tExport onnx_ctc, done! see {}'.format(ctc_outpath))
+
+    print("\tStage-2.3: check onnx_ctc and torch_ctc")
+    torch_output = ctc(hidden)
+    ort_session = onnxruntime.InferenceSession(
+        ctc_outpath, providers=['CPUExecutionProvider'])
+    onnx_output = ort_session.run(None, {'hidden': to_numpy(hidden)})
+
+    np.testing.assert_allclose(to_numpy(torch_output),
+                               onnx_output[0],
+                               rtol=1e-03,
+                               atol=1e-05)
+    print("\t\tCheck onnx_ctc, pass!")
+
+
+def export_decoder(asr_model, args):
+    print("Stage-3: export decoder")
+    decoder = asr_model
+    # NOTE(lzhin): parameters of encoder will be automatically removed
+    #   since they are not used during rescoring.
+    decoder.forward = decoder.forward_attention_decoder
+    decoder_outpath = os.path.join(args['output_dir'], 'decoder.onnx')
+
+    print("\tStage-3.1: prepare inputs for decoder")
+    # hardcode time->200 nbest->10 len->20, they are dynamic axes.
+    encoder_out = torch.randn((1, 200, args['output_size']))
+    hyps = torch.randint(low=0, high=args['vocab_size'], size=[10, 20])
+    hyps[:, 0] = args['vocab_size'] - 1  # <sos>
+    hyps_lens = torch.randint(low=15, high=21, size=[10])
+
+    print("\tStage-3.2: torch.onnx.export")
+    dynamic_axes = {
+        'hyps': {
+            0: 'NBEST',
+            1: 'L'
+        },
+        'hyps_lens': {
+            0: 'NBEST'
+        },
+        'encoder_out': {
+            1: 'T'
+        },
+        'score': {
+            0: 'NBEST',
+            1: 'L'
+        },
+        'r_score': {
+            0: 'NBEST',
+            1: 'L'
+        }
+    }
+    inputs = (hyps, hyps_lens, encoder_out, args['reverse_weight'])
+    torch.onnx.export(
+        decoder,
+        inputs,
+        decoder_outpath,
+        opset_version=13,
+        export_params=True,
+        do_constant_folding=True,
+        input_names=['hyps', 'hyps_lens', 'encoder_out', 'reverse_weight'],
+        output_names=['score', 'r_score'],
+        dynamic_axes=dynamic_axes,
+        verbose=False)
+    onnx_decoder = onnx.load(decoder_outpath)
+    for (k, v) in args.items():
+        meta = onnx_decoder.metadata_props.add()
+        meta.key, meta.value = str(k), str(v)
+    onnx.checker.check_model(onnx_decoder)
+    onnx.helper.printable_graph(onnx_decoder.graph)
+    onnx.save(onnx_decoder, decoder_outpath)
+    print_input_output_info(onnx_decoder, "onnx_decoder")
+    model_fp32 = decoder_outpath
+    model_quant = os.path.join(args['output_dir'], 'decoder.quant.onnx')
+    quantize_dynamic(model_fp32, model_quant, weight_type=QuantType.QUInt8)
+    print('\t\tExport onnx_decoder, done! see {}'.format(decoder_outpath))
+
+    print("\tStage-3.3: check onnx_decoder and torch_decoder")
+    torch_score, torch_r_score = decoder(hyps, hyps_lens, encoder_out,
+                                         args['reverse_weight'])
+    ort_session = onnxruntime.InferenceSession(
+        decoder_outpath, providers=['CPUExecutionProvider'])
+    input_names = [node.name for node in onnx_decoder.graph.input]
+    ort_inputs = {
+        'hyps': to_numpy(hyps),
+        'hyps_lens': to_numpy(hyps_lens),
+        'encoder_out': to_numpy(encoder_out),
+        'reverse_weight': np.array((args['reverse_weight'])),
+    }
+    for k in list(ort_inputs):
+        if k not in input_names:
+            ort_inputs.pop(k)
+    onnx_output = ort_session.run(None, ort_inputs)
+
+    np.testing.assert_allclose(to_numpy(torch_score),
+                               onnx_output[0],
+                               rtol=1e-03,
+                               atol=1e-05)
+    if args['is_bidirectional_decoder'] and args['reverse_weight'] > 0.0:
+        np.testing.assert_allclose(to_numpy(torch_r_score),
+                                   onnx_output[1],
+                                   rtol=1e-03,
+                                   atol=1e-05)
+    print("\t\tCheck onnx_decoder, pass!")
+
+
+def main():
+    torch.manual_seed(777)
+    args = get_args()
+    logging.basicConfig(level=logging.DEBUG,
+                        format='%(asctime)s %(levelname)s %(message)s')
+    output_dir = args.output_dir
+    os.system("mkdir -p " + output_dir)
+    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
+    with open(args.config, 'r') as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+
+    model, configs = init_model(args, configs)
+    model.eval()
+    print(model)
+
+    arguments = {}
+    arguments['output_dir'] = output_dir
+    arguments['batch'] = 1
+    arguments['chunk_size'] = args.chunk_size
+    arguments['left_chunks'] = args.num_decoding_left_chunks
+    arguments['reverse_weight'] = args.reverse_weight
+    arguments['output_size'] = configs['encoder_conf']['output_size']
+    arguments['num_blocks'] = configs['encoder_conf']['num_blocks']
+    arguments['cnn_module_kernel'] = configs['encoder_conf'].get(
+        'cnn_module_kernel', 1)
+    arguments['head'] = configs['encoder_conf']['attention_heads']
+    arguments['feature_size'] = configs['input_dim']
+    arguments['vocab_size'] = configs['output_dim']
+    # NOTE(xcsong): if chunk_size == -1, hardcode to 67
+    arguments['decoding_window'] = (args.chunk_size - 1) * \
+        model.encoder.embed.subsampling_rate + \
+        model.encoder.embed.right_context + 1 if args.chunk_size > 0 else 67
+    arguments['encoder'] = configs['encoder']
+    arguments['decoder'] = configs['decoder']
+    arguments['subsampling_rate'] = model.subsampling_rate()
+    arguments['right_context'] = model.right_context()
+    arguments['sos_symbol'] = model.sos_symbol()
+    arguments['eos_symbol'] = model.eos_symbol()
+    arguments['is_bidirectional_decoder'] = 1 \
+        if model.is_bidirectional_decoder() else 0
+
+    # NOTE(xcsong): Please note that -1/-1 means non-streaming model! It is
+    #   not a [16/4 16/-1 16/0] all-in-one model and it should not be used in
+    #   streaming mode (i.e., setting chunk_size=16 in `decoder_main`). If you
+    #   want to use 16/-1 or any other streaming mode in `decoder_main`,
+    #   please export onnx in the same config.
+    if arguments['left_chunks'] > 0:
+        assert arguments['chunk_size'] > 0  # -1/4 not supported
+
+    export_encoder(model, arguments)
+    export_ctc(model, arguments)
+    export_decoder(model, arguments)
+
+
+if __name__ == '__main__':
+    main()

wenet/bin/export_onnx_gpu.py

@@ -0,0 +1,1266 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import argparse
+import logging
+import os
+import sys
+
+import torch
+import torch.nn.functional as F
+import yaml
+from wenet.transformer.ctc import CTC
+from wenet.transformer.decoder import TransformerDecoder
+from wenet.transformer.encoder import BaseEncoder
+from wenet.utils.init_model import init_model
+from wenet.utils.mask import make_pad_mask
+
+try:
+    import onnxruntime
+except ImportError:
+    print("Please install onnxruntime-gpu!")
+    sys.exit(1)
+
+logger = logging.getLogger(__file__)
+logger.setLevel(logging.INFO)
+
+
+class Encoder(torch.nn.Module):
+
+    def __init__(self, encoder: BaseEncoder, ctc: CTC, beam_size: int = 10):
+        super().__init__()
+        self.encoder = encoder
+        self.ctc = ctc
+        self.beam_size = beam_size
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+    ):
+        """Encoder
+        Args:
+            speech: (Batch, Length, ...)
+            speech_lengths: (Batch, )
+        Returns:
+            encoder_out: B x T x F
+            encoder_out_lens: B
+            ctc_log_probs: B x T x V
+            beam_log_probs: B x T x beam_size
+            beam_log_probs_idx: B x T x beam_size
+        """
+        encoder_out, encoder_mask = self.encoder(speech, speech_lengths, -1,
+                                                 -1)
+        encoder_out_lens = encoder_mask.squeeze(1).sum(1)
+        ctc_log_probs = self.ctc.log_softmax(encoder_out)
+        encoder_out_lens = encoder_out_lens.int()
+        beam_log_probs, beam_log_probs_idx = torch.topk(ctc_log_probs,
+                                                        self.beam_size,
+                                                        dim=2)
+        return (
+            encoder_out,
+            encoder_out_lens,
+            ctc_log_probs,
+            beam_log_probs,
+            beam_log_probs_idx,
+        )
+
+
+class StreamingEncoder(torch.nn.Module):
+
+    def __init__(
+        self,
+        model,
+        required_cache_size,
+        beam_size,
+        transformer=False,
+        return_ctc_logprobs=False,
+    ):
+        super().__init__()
+        self.ctc = model.ctc
+        self.subsampling_rate = model.encoder.embed.subsampling_rate
+        self.embed = model.encoder.embed
+        self.global_cmvn = model.encoder.global_cmvn
+        self.required_cache_size = required_cache_size
+        self.beam_size = beam_size
+        self.encoder = model.encoder
+        self.transformer = transformer
+        self.return_ctc_logprobs = return_ctc_logprobs
+
+    def forward(self, chunk_xs, chunk_lens, offset, att_cache, cnn_cache,
+                cache_mask):
+        """Streaming Encoder
+        Args:
+            xs (torch.Tensor): chunk input, with shape (b, time, mel-dim),
+                where `time == (chunk_size - 1) * subsample_rate + \
+                        subsample.right_context + 1`
+            offset (torch.Tensor): offset with shape (b, 1)
+                        1 is retained for triton deployment
+            required_cache_size (int): cache size required for next chunk
+                compuation
+                > 0: actual cache size
+                <= 0: not allowed in streaming gpu encoder                   `
+            att_cache (torch.Tensor): cache tensor for KEY & VALUE in
+                transformer/conformer attention, with shape
+                (b, elayers, head, cache_t1, d_k * 2), where
+                `head * d_k == hidden-dim` and
+                `cache_t1 == chunk_size * num_decoding_left_chunks`.
+            cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer,
+                (b, elayers, b, hidden-dim, cache_t2), where
+                `cache_t2 == cnn.lorder - 1`
+            cache_mask: (torch.Tensor): cache mask with shape (b, required_cache_size)
+                 in a batch of request, each request may have different
+                 history cache. Cache mask is used to indidate the effective
+                 cache for each request
+        Returns:
+            torch.Tensor: log probabilities of ctc output and cutoff by beam size
+                with shape (b, chunk_size, beam)
+            torch.Tensor: index of top beam size probabilities for each timestep
+                with shape (b, chunk_size, beam)
+            torch.Tensor: output of current input xs,
+                with shape (b, chunk_size, hidden-dim).
+            torch.Tensor: new attention cache required for next chunk, with
+                same shape (b, elayers, head, cache_t1, d_k * 2)
+                as the original att_cache
+            torch.Tensor: new conformer cnn cache required for next chunk, with
+                same shape as the original cnn_cache.
+            torch.Tensor: new cache mask, with same shape as the original
+                cache mask
+        """
+        offset = offset.squeeze(1)
+        T = chunk_xs.size(1)
+        chunk_mask = ~make_pad_mask(chunk_lens, T).unsqueeze(1)
+        # B X 1 X T
+        chunk_mask = chunk_mask.to(chunk_xs.dtype)
+        # transpose batch & num_layers dim
+        att_cache = torch.transpose(att_cache, 0, 1)
+        cnn_cache = torch.transpose(cnn_cache, 0, 1)
+
+        # rewrite encoder.forward_chunk
+        # <---------forward_chunk START--------->
+        xs = self.global_cmvn(chunk_xs)
+        # chunk mask is important for batch inferencing since
+        # different sequence in a batch has different length
+        xs, pos_emb, chunk_mask = self.embed(xs, chunk_mask, offset)
+        cache_size = att_cache.size(3)  # required cache size
+        masks = torch.cat((cache_mask, chunk_mask), dim=2)
+        index = offset - cache_size
+
+        pos_emb = self.embed.position_encoding(index, cache_size + xs.size(1))
+        pos_emb = pos_emb.to(dtype=xs.dtype)
+
+        next_cache_start = -self.required_cache_size
+        r_cache_mask = masks[:, :, next_cache_start:]
+
+        r_att_cache = []
+        r_cnn_cache = []
+        for i, layer in enumerate(self.encoder.encoders):
+            i_kv_cache = att_cache[i]
+            size = att_cache.size(-1) // 2
+            kv_cache = (i_kv_cache[:, :, :, :size], i_kv_cache[:, :, :, size:])
+            xs, _, new_kv_cache, new_cnn_cache = layer(
+                xs,
+                masks,
+                pos_emb,
+                att_cache=kv_cache,
+                cnn_cache=cnn_cache[i],
+            )
+            #   shape(new_att_cache) is (B, head, attention_key_size, d_k * 2),
+            #   shape(new_cnn_cache) is (B, hidden-dim, cache_t2)
+            new_att_cache = torch.cat(new_kv_cache, dim=-1)
+            r_att_cache.append(
+                new_att_cache[:, :, next_cache_start:, :].unsqueeze(1))
+            if not self.transformer:
+                r_cnn_cache.append(new_cnn_cache.unsqueeze(1))
+        if self.encoder.normalize_before:
+            chunk_out = self.encoder.after_norm(xs)
+        else:
+            chunk_out = xs
+
+        r_att_cache = torch.cat(r_att_cache, dim=1)  # concat on layers idx
+        if not self.transformer:
+            r_cnn_cache = torch.cat(r_cnn_cache, dim=1)  # concat on layers
+
+        # <---------forward_chunk END--------->
+
+        log_ctc_probs = self.ctc.log_softmax(chunk_out)
+        log_probs, log_probs_idx = torch.topk(log_ctc_probs,
+                                              self.beam_size,
+                                              dim=2)
+        log_probs = log_probs.to(chunk_xs.dtype)
+
+        r_offset = offset + chunk_out.shape[1]
+        # the below ops not supported in Tensorrt
+        # chunk_out_lens = torch.div(chunk_lens, subsampling_rate,
+        #                   rounding_mode='floor')
+        chunk_out_lens = chunk_lens // self.subsampling_rate
+        r_offset = r_offset.unsqueeze(1)
+        if self.return_ctc_logprobs:
+            return (
+                log_ctc_probs,
+                chunk_out,
+                chunk_out_lens,
+                r_offset,
+                r_att_cache,
+                r_cnn_cache,
+                r_cache_mask,
+            )
+        else:
+            return (
+                log_probs,
+                log_probs_idx,
+                chunk_out,
+                chunk_out_lens,
+                r_offset,
+                r_att_cache,
+                r_cnn_cache,
+                r_cache_mask,
+            )
+
+
+class StreamingSqueezeformerEncoder(torch.nn.Module):
+
+    def __init__(self, model, required_cache_size, beam_size):
+        super().__init__()
+        self.ctc = model.ctc
+        self.subsampling_rate = model.encoder.embed.subsampling_rate
+        self.embed = model.encoder.embed
+        self.global_cmvn = model.encoder.global_cmvn
+        self.required_cache_size = required_cache_size
+        self.beam_size = beam_size
+        self.encoder = model.encoder
+        self.reduce_idx = model.encoder.reduce_idx
+        self.recover_idx = model.encoder.recover_idx
+        if self.reduce_idx is None:
+            self.time_reduce = None
+        else:
+            if self.recover_idx is None:
+                self.time_reduce = "normal"  # no recovery at the end
+            else:
+                self.time_reduce = "recover"  # recovery at the end
+                assert len(self.reduce_idx) == len(self.recover_idx)
+
+    def calculate_downsampling_factor(self, i: int) -> int:
+        if self.reduce_idx is None:
+            return 1
+        else:
+            reduce_exp, recover_exp = 0, 0
+            for exp, rd_idx in enumerate(self.reduce_idx):
+                if i >= rd_idx:
+                    reduce_exp = exp + 1
+            if self.recover_idx is not None:
+                for exp, rc_idx in enumerate(self.recover_idx):
+                    if i >= rc_idx:
+                        recover_exp = exp + 1
+            return int(2**(reduce_exp - recover_exp))
+
+    def forward(self, chunk_xs, chunk_lens, offset, att_cache, cnn_cache,
+                cache_mask):
+        """Streaming Encoder
+        Args:
+            xs (torch.Tensor): chunk input, with shape (b, time, mel-dim),
+                where `time == (chunk_size - 1) * subsample_rate + \
+                        subsample.right_context + 1`
+            offset (torch.Tensor): offset with shape (b, 1)
+                        1 is retained for triton deployment
+            required_cache_size (int): cache size required for next chunk
+                compuation
+                > 0: actual cache size
+                <= 0: not allowed in streaming gpu encoder                   `
+            att_cache (torch.Tensor): cache tensor for KEY & VALUE in
+                transformer/conformer attention, with shape
+                (b, elayers, head, cache_t1, d_k * 2), where
+                `head * d_k == hidden-dim` and
+                `cache_t1 == chunk_size * num_decoding_left_chunks`.
+            cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer,
+                (b, elayers, b, hidden-dim, cache_t2), where
+                `cache_t2 == cnn.lorder - 1`
+            cache_mask: (torch.Tensor): cache mask with shape (b, required_cache_size)
+                 in a batch of request, each request may have different
+                 history cache. Cache mask is used to indidate the effective
+                 cache for each request
+        Returns:
+            torch.Tensor: log probabilities of ctc output and cutoff by beam size
+                with shape (b, chunk_size, beam)
+            torch.Tensor: index of top beam size probabilities for each timestep
+                with shape (b, chunk_size, beam)
+            torch.Tensor: output of current input xs,
+                with shape (b, chunk_size, hidden-dim).
+            torch.Tensor: new attention cache required for next chunk, with
+                same shape (b, elayers, head, cache_t1, d_k * 2)
+                as the original att_cache
+            torch.Tensor: new conformer cnn cache required for next chunk, with
+                same shape as the original cnn_cache.
+            torch.Tensor: new cache mask, with same shape as the original
+                cache mask
+        """
+        offset = offset.squeeze(1)
+        T = chunk_xs.size(1)
+        chunk_mask = ~make_pad_mask(chunk_lens, T).unsqueeze(1)
+        # B X 1 X T
+        chunk_mask = chunk_mask.to(chunk_xs.dtype)
+        # transpose batch & num_layers dim
+        att_cache = torch.transpose(att_cache, 0, 1)
+        cnn_cache = torch.transpose(cnn_cache, 0, 1)
+
+        # rewrite encoder.forward_chunk
+        # <---------forward_chunk START--------->
+        xs = self.global_cmvn(chunk_xs)
+        # chunk mask is important for batch inferencing since
+        # different sequence in a batch has different length
+        xs, pos_emb, chunk_mask = self.embed(xs, chunk_mask, offset)
+        elayers, cache_size = att_cache.size(0), att_cache.size(3)
+        att_mask = torch.cat((cache_mask, chunk_mask), dim=2)
+        index = offset - cache_size
+
+        pos_emb = self.embed.position_encoding(index, cache_size + xs.size(1))
+        pos_emb = pos_emb.to(dtype=xs.dtype)
+
+        next_cache_start = -self.required_cache_size
+        r_cache_mask = att_mask[:, :, next_cache_start:]
+
+        r_att_cache = []
+        r_cnn_cache = []
+        mask_pad = torch.ones(1,
+                              xs.size(1),
+                              device=xs.device,
+                              dtype=torch.bool)
+        mask_pad = mask_pad.unsqueeze(1)
+        max_att_len: int = 0
+        recover_activations: List[Tuple[torch.Tensor, torch.Tensor,
+                                        torch.Tensor, torch.Tensor]] = []
+        index = 0
+        xs_lens = torch.tensor([xs.size(1)], device=xs.device, dtype=torch.int)
+        xs = self.encoder.preln(xs)
+        for i, layer in enumerate(self.encoder.encoders):
+            if self.reduce_idx is not None:
+                if self.time_reduce is not None and i in self.reduce_idx:
+                    recover_activations.append(
+                        (xs, att_mask, pos_emb, mask_pad))
+                    (
+                        xs,
+                        xs_lens,
+                        att_mask,
+                        mask_pad,
+                    ) = self.encoder.time_reduction_layer(
+                        xs, xs_lens, att_mask, mask_pad)
+                    pos_emb = pos_emb[:, ::2, :]
+                    if self.encoder.pos_enc_layer_type == "rel_pos_repaired":
+                        pos_emb = pos_emb[:, :xs.size(1) * 2 - 1, :]
+                    index += 1
+
+            if self.recover_idx is not None:
+                if self.time_reduce == "recover" and i in self.recover_idx:
+                    index -= 1
+                    (
+                        recover_tensor,
+                        recover_att_mask,
+                        recover_pos_emb,
+                        recover_mask_pad,
+                    ) = recover_activations[index]
+                    # recover output length for ctc decode
+                    xs = xs.unsqueeze(2).repeat(1, 1, 2, 1).flatten(1, 2)
+                    xs = self.encoder.time_recover_layer(xs)
+                    recoverd_t = recover_tensor.size(1)
+                    xs = recover_tensor + xs[:, :recoverd_t, :].contiguous()
+                    att_mask = recover_att_mask
+                    pos_emb = recover_pos_emb
+                    mask_pad = recover_mask_pad
+
+            factor = self.calculate_downsampling_factor(i)
+
+            xs, _, new_att_cache, new_cnn_cache = layer(
+                xs,
+                att_mask,
+                pos_emb,
+                att_cache=att_cache[i][:, :, ::factor, :]
+                [:, :, :pos_emb.size(1) - xs.size(1), :]
+                if elayers > 0 else att_cache[:, :, ::factor, :],
+                cnn_cache=cnn_cache[i] if cnn_cache.size(0) > 0 else cnn_cache,
+            )
+            cached_att = new_att_cache[:, :, next_cache_start // factor:, :]
+            cached_cnn = new_cnn_cache.unsqueeze(1)
+            cached_att = (cached_att.unsqueeze(3).repeat(1, 1, 1, factor,
+                                                         1).flatten(2, 3))
+            if i == 0:
+                # record length for the first block as max length
+                max_att_len = cached_att.size(2)
+            r_att_cache.append(cached_att[:, :, :max_att_len, :].unsqueeze(1))
+            r_cnn_cache.append(cached_cnn)
+
+        chunk_out = xs
+        r_att_cache = torch.cat(r_att_cache, dim=1)  # concat on layers idx
+        r_cnn_cache = torch.cat(r_cnn_cache, dim=1)  # concat on layers
+
+        # <---------forward_chunk END--------->
+
+        log_ctc_probs = self.ctc.log_softmax(chunk_out)
+        log_probs, log_probs_idx = torch.topk(log_ctc_probs,
+                                              self.beam_size,
+                                              dim=2)
+        log_probs = log_probs.to(chunk_xs.dtype)
+
+        r_offset = offset + chunk_out.shape[1]
+        # the below ops not supported in Tensorrt
+        # chunk_out_lens = torch.div(chunk_lens, subsampling_rate,
+        #                   rounding_mode='floor')
+        chunk_out_lens = chunk_lens // self.subsampling_rate
+        r_offset = r_offset.unsqueeze(1)
+
+        return (
+            log_probs,
+            log_probs_idx,
+            chunk_out,
+            chunk_out_lens,
+            r_offset,
+            r_att_cache,
+            r_cnn_cache,
+            r_cache_mask,
+        )
+
+
+class StreamingEfficientConformerEncoder(torch.nn.Module):
+
+    def __init__(self, model, required_cache_size, beam_size):
+        super().__init__()
+        self.ctc = model.ctc
+        self.subsampling_rate = model.encoder.embed.subsampling_rate
+        self.embed = model.encoder.embed
+        self.global_cmvn = model.encoder.global_cmvn
+        self.required_cache_size = required_cache_size
+        self.beam_size = beam_size
+        self.encoder = model.encoder
+
+        # Efficient Conformer
+        self.stride_layer_idx = model.encoder.stride_layer_idx
+        self.stride = model.encoder.stride
+        self.num_blocks = model.encoder.num_blocks
+        self.cnn_module_kernel = model.encoder.cnn_module_kernel
+
+    def calculate_downsampling_factor(self, i: int) -> int:
+        factor = 1
+        for idx, stride_idx in enumerate(self.stride_layer_idx):
+            if i > stride_idx:
+                factor *= self.stride[idx]
+        return factor
+
+    def forward(self, chunk_xs, chunk_lens, offset, att_cache, cnn_cache,
+                cache_mask):
+        """Streaming Encoder
+        Args:
+            chunk_xs (torch.Tensor): chunk input, with shape (b, time, mel-dim),
+                where `time == (chunk_size - 1) * subsample_rate + \
+                        subsample.right_context + 1`
+            chunk_lens (torch.Tensor):
+            offset (torch.Tensor): offset with shape (b, 1)
+                        1 is retained for triton deployment
+            att_cache (torch.Tensor): cache tensor for KEY & VALUE in
+                transformer/conformer attention, with shape
+                (b, elayers, head, cache_t1, d_k * 2), where
+                `head * d_k == hidden-dim` and
+                `cache_t1 == chunk_size * num_decoding_left_chunks`.
+            cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer,
+                (b, elayers, hidden-dim, cache_t2), where
+                `cache_t2 == cnn.lorder - 1`
+            cache_mask: (torch.Tensor): cache mask with shape (b, required_cache_size)
+                 in a batch of request, each request may have different
+                 history cache. Cache mask is used to indidate the effective
+                 cache for each request
+        Returns:
+            torch.Tensor: log probabilities of ctc output and cutoff by beam size
+                with shape (b, chunk_size, beam)
+            torch.Tensor: index of top beam size probabilities for each timestep
+                with shape (b, chunk_size, beam)
+            torch.Tensor: output of current input xs,
+                with shape (b, chunk_size, hidden-dim).
+            torch.Tensor: new attention cache required for next chunk, with
+                same shape (b, elayers, head, cache_t1, d_k * 2)
+                as the original att_cache
+            torch.Tensor: new conformer cnn cache required for next chunk, with
+                same shape as the original cnn_cache.
+            torch.Tensor: new cache mask, with same shape as the original
+                cache mask
+        """
+        offset = offset.squeeze(1)  # (b, )
+        offset *= self.calculate_downsampling_factor(self.num_blocks + 1)
+
+        T = chunk_xs.size(1)
+        chunk_mask = ~make_pad_mask(chunk_lens, T).unsqueeze(1)  # (b, 1, T)
+        # B X 1 X T
+        chunk_mask = chunk_mask.to(chunk_xs.dtype)
+        # transpose batch & num_layers dim
+        #   Shape(att_cache): (elayers, b, head, cache_t1, d_k * 2)
+        #   Shape(cnn_cache): (elayers, b, outsize, cnn_kernel)
+        att_cache = torch.transpose(att_cache, 0, 1)
+        cnn_cache = torch.transpose(cnn_cache, 0, 1)
+
+        # rewrite encoder.forward_chunk
+        # <---------forward_chunk START--------->
+        xs = self.global_cmvn(chunk_xs)
+        # chunk mask is important for batch inferencing since
+        # different sequence in a batch has different length
+        xs, pos_emb, chunk_mask = self.embed(xs, chunk_mask, offset)
+        cache_size = att_cache.size(3)  # required cache size
+        masks = torch.cat((cache_mask, chunk_mask), dim=2)
+        att_mask = torch.cat((cache_mask, chunk_mask), dim=2)
+        index = offset - cache_size
+
+        pos_emb = self.embed.position_encoding(index, cache_size + xs.size(1))
+        pos_emb = pos_emb.to(dtype=xs.dtype)
+
+        next_cache_start = -self.required_cache_size
+        r_cache_mask = masks[:, :, next_cache_start:]
+
+        r_att_cache = []
+        r_cnn_cache = []
+        mask_pad = chunk_mask.to(torch.bool)
+        max_att_len, max_cnn_len = (
+            0,
+            0,
+        )  # for repeat_interleave of new_att_cache
+        for i, layer in enumerate(self.encoder.encoders):
+            factor = self.calculate_downsampling_factor(i)
+            # NOTE(xcsong): Before layer.forward
+            #   shape(att_cache[i:i + 1]) is (b, head, cache_t1, d_k * 2),
+            #   shape(cnn_cache[i])       is (b=1, hidden-dim, cache_t2)
+            # shape(new_att_cache) = [ batch, head, time2, outdim//head * 2 ]
+            att_cache_trunc = 0
+            if xs.size(1) + att_cache.size(3) / factor > pos_emb.size(1):
+                # The time step is not divisible by the downsampling multiple
+                # We propose to double the chunk_size.
+                att_cache_trunc = (xs.size(1) + att_cache.size(3) // factor -
+                                   pos_emb.size(1) + 1)
+            xs, _, new_att_cache, new_cnn_cache = layer(
+                xs,
+                att_mask,
+                pos_emb,
+                mask_pad=mask_pad,
+                att_cache=att_cache[i][:, :, ::factor, :][:, :,
+                                                          att_cache_trunc:, :],
+                cnn_cache=cnn_cache[i, :, :, :]
+                if cnn_cache.size(0) > 0 else cnn_cache,
+            )
+
+            if i in self.stride_layer_idx:
+                # compute time dimension for next block
+                efficient_index = self.stride_layer_idx.index(i)
+                att_mask = att_mask[:, ::self.stride[efficient_index], ::self.
+                                    stride[efficient_index], ]
+                mask_pad = mask_pad[:, ::self.stride[efficient_index], ::self.
+                                    stride[efficient_index], ]
+                pos_emb = pos_emb[:, ::self.stride[efficient_index], :]
+
+            # shape(new_att_cache) = [batch, head, time2, outdim]
+            new_att_cache = new_att_cache[:, :, next_cache_start // factor:, :]
+            # shape(new_cnn_cache) = [batch, 1, outdim, cache_t2]
+            new_cnn_cache = new_cnn_cache.unsqueeze(1)  # shape(1):layerID
+
+            # use repeat_interleave to new_att_cache
+            # new_att_cache = new_att_cache.repeat_interleave(repeats=factor, dim=2)
+            new_att_cache = (new_att_cache.unsqueeze(3).repeat(
+                1, 1, 1, factor, 1).flatten(2, 3))
+            # padding new_cnn_cache to cnn.lorder for casual convolution
+            new_cnn_cache = F.pad(
+                new_cnn_cache,
+                (self.cnn_module_kernel - 1 - new_cnn_cache.size(3), 0),
+            )
+
+            if i == 0:
+                # record length for the first block as max length
+                max_att_len = new_att_cache.size(2)
+                max_cnn_len = new_cnn_cache.size(3)
+
+            # update real shape of att_cache and cnn_cache
+            r_att_cache.append(new_att_cache[:, :,
+                                             -max_att_len:, :].unsqueeze(1))
+            r_cnn_cache.append(new_cnn_cache[:, :, :, -max_cnn_len:])
+
+        if self.encoder.normalize_before:
+            chunk_out = self.encoder.after_norm(xs)
+        else:
+            chunk_out = xs
+
+        # shape of r_att_cache: (b, elayers, head, time2, outdim)
+        r_att_cache = torch.cat(r_att_cache, dim=1)  # concat on layers idx
+        # shape of r_cnn_cache: (b, elayers, outdim, cache_t2)
+        r_cnn_cache = torch.cat(r_cnn_cache, dim=1)  # concat on layers
+
+        # <---------forward_chunk END--------->
+
+        log_ctc_probs = self.ctc.log_softmax(chunk_out)
+        log_probs, log_probs_idx = torch.topk(log_ctc_probs,
+                                              self.beam_size,
+                                              dim=2)
+        log_probs = log_probs.to(chunk_xs.dtype)
+
+        r_offset = offset + chunk_out.shape[1]
+        # the below ops not supported in Tensorrt
+        # chunk_out_lens = torch.div(chunk_lens, subsampling_rate,
+        #                   rounding_mode='floor')
+        chunk_out_lens = (
+            chunk_lens // self.subsampling_rate //
+            self.calculate_downsampling_factor(self.num_blocks + 1))
+        chunk_out_lens += 1
+        r_offset = r_offset.unsqueeze(1)
+
+        return (
+            log_probs,
+            log_probs_idx,
+            chunk_out,
+            chunk_out_lens,
+            r_offset,
+            r_att_cache,
+            r_cnn_cache,
+            r_cache_mask,
+        )
+
+
+class Decoder(torch.nn.Module):
+
+    def __init__(
+        self,
+        decoder: TransformerDecoder,
+        ctc_weight: float = 0.5,
+        reverse_weight: float = 0.0,
+        beam_size: int = 10,
+        decoder_fastertransformer: bool = False,
+    ):
+        super().__init__()
+        self.decoder = decoder
+        self.ctc_weight = ctc_weight
+        self.reverse_weight = reverse_weight
+        self.beam_size = beam_size
+        self.decoder_fastertransformer = decoder_fastertransformer
+
+    def forward(
+        self,
+        encoder_out: torch.Tensor,
+        encoder_lens: torch.Tensor,
+        hyps_pad_sos_eos: torch.Tensor,
+        hyps_lens_sos: torch.Tensor,
+        r_hyps_pad_sos_eos: torch.Tensor,
+        ctc_score: torch.Tensor,
+    ):
+        """Encoder
+        Args:
+            encoder_out: B x T x F
+            encoder_lens: B
+            hyps_pad_sos_eos: B x beam x (T2+1),
+                        hyps with sos & eos and padded by ignore id
+            hyps_lens_sos: B x beam, length for each hyp with sos
+            r_hyps_pad_sos_eos: B x beam x (T2+1),
+                    reversed hyps with sos & eos and padded by ignore id
+            ctc_score: B x beam, ctc score for each hyp
+        Returns:
+            decoder_out: B x beam x T2 x V
+            r_decoder_out: B x beam x T2 x V
+            best_index: B
+        """
+        B, T, F = encoder_out.shape
+        bz = self.beam_size
+        B2 = B * bz
+        encoder_out = encoder_out.repeat(1, bz, 1).view(B2, T, F)
+        encoder_mask = ~make_pad_mask(encoder_lens, T).unsqueeze(1)
+        encoder_mask = encoder_mask.repeat(1, bz, 1).view(B2, 1, T)
+        T2 = hyps_pad_sos_eos.shape[2] - 1
+        hyps_pad = hyps_pad_sos_eos.view(B2, T2 + 1)
+        hyps_lens = hyps_lens_sos.view(B2, )
+        hyps_pad_sos = hyps_pad[:, :-1].contiguous()
+        hyps_pad_eos = hyps_pad[:, 1:].contiguous()
+
+        r_hyps_pad = r_hyps_pad_sos_eos.view(B2, T2 + 1)
+        r_hyps_pad_sos = r_hyps_pad[:, :-1].contiguous()
+        r_hyps_pad_eos = r_hyps_pad[:, 1:].contiguous()
+
+        decoder_out, r_decoder_out, _ = self.decoder(
+            encoder_out,
+            encoder_mask,
+            hyps_pad_sos,
+            hyps_lens,
+            r_hyps_pad_sos,
+            self.reverse_weight,
+        )
+        decoder_out = torch.nn.functional.log_softmax(decoder_out, dim=-1)
+        V = decoder_out.shape[-1]
+        decoder_out = decoder_out.view(B2, T2, V)
+        mask = ~make_pad_mask(hyps_lens, T2)  # B2 x T2
+        # mask index, remove ignore id
+        index = torch.unsqueeze(hyps_pad_eos * mask, 2)
+        score = decoder_out.gather(2, index).squeeze(2)  # B2 X T2
+        # mask padded part
+        score = score * mask
+        decoder_out = decoder_out.view(B, bz, T2, V)
+        if self.reverse_weight > 0:
+            r_decoder_out = torch.nn.functional.log_softmax(r_decoder_out,
+                                                            dim=-1)
+            r_decoder_out = r_decoder_out.view(B2, T2, V)
+            index = torch.unsqueeze(r_hyps_pad_eos * mask, 2)
+            r_score = r_decoder_out.gather(2, index).squeeze(2)
+            r_score = r_score * mask
+            score = (score * (1 - self.reverse_weight) +
+                     self.reverse_weight * r_score)
+            r_decoder_out = r_decoder_out.view(B, bz, T2, V)
+        score = torch.sum(score, axis=1)  # B2
+        score = torch.reshape(score, (B, bz)) + self.ctc_weight * ctc_score
+        best_index = torch.argmax(score, dim=1)
+        if self.decoder_fastertransformer:
+            return decoder_out, best_index
+        else:
+            return best_index
+
+
+def to_numpy(tensors):
+    out = []
+    if type(tensors) == torch.tensor:
+        tensors = [tensors]
+    for tensor in tensors:
+        if tensor.requires_grad:
+            tensor = tensor.detach().cpu().numpy()
+        else:
+            tensor = tensor.cpu().numpy()
+        out.append(tensor)
+    return out
+
+
+def test(xlist, blist, rtol=1e-3, atol=1e-5, tolerate_small_mismatch=True):
+    for a, b in zip(xlist, blist):
+        try:
+            torch.testing.assert_allclose(a, b, rtol=rtol, atol=atol)
+        except AssertionError as error:
+            if tolerate_small_mismatch:
+                print(error)
+            else:
+                raise
+
+
+def export_offline_encoder(model, configs, args, logger, encoder_onnx_path):
+    bz = 32
+    seq_len = 100
+    beam_size = args.beam_size
+    feature_size = configs["input_dim"]
+
+    speech = torch.randn(bz, seq_len, feature_size, dtype=torch.float32)
+    speech_lens = torch.randint(low=10,
+                                high=seq_len,
+                                size=(bz, ),
+                                dtype=torch.int32)
+    encoder = Encoder(model.encoder, model.ctc, beam_size)
+    encoder.eval()
+
+    torch.onnx.export(
+        encoder,
+        (speech, speech_lens),
+        encoder_onnx_path,
+        export_params=True,
+        opset_version=13,
+        do_constant_folding=True,
+        input_names=["speech", "speech_lengths"],
+        output_names=[
+            "encoder_out",
+            "encoder_out_lens",
+            "ctc_log_probs",
+            "beam_log_probs",
+            "beam_log_probs_idx",
+        ],
+        dynamic_axes={
+            "speech": {
+                0: "B",
+                1: "T"
+            },
+            "speech_lengths": {
+                0: "B"
+            },
+            "encoder_out": {
+                0: "B",
+                1: "T_OUT"
+            },
+            "encoder_out_lens": {
+                0: "B"
+            },
+            "ctc_log_probs": {
+                0: "B",
+                1: "T_OUT"
+            },
+            "beam_log_probs": {
+                0: "B",
+                1: "T_OUT"
+            },
+            "beam_log_probs_idx": {
+                0: "B",
+                1: "T_OUT"
+            },
+        },
+        verbose=False,
+    )
+
+    with torch.no_grad():
+        o0, o1, o2, o3, o4 = encoder(speech, speech_lens)
+
+    providers = ["CUDAExecutionProvider"]
+    ort_session = onnxruntime.InferenceSession(encoder_onnx_path,
+                                               providers=providers)
+    ort_inputs = {
+        "speech": to_numpy(speech),
+        "speech_lengths": to_numpy(speech_lens),
+    }
+    ort_outs = ort_session.run(None, ort_inputs)
+
+    # check encoder output
+    test(to_numpy([o0, o1, o2, o3, o4]), ort_outs)
+    logger.info("export offline onnx encoder succeed!")
+    onnx_config = {
+        "beam_size": args.beam_size,
+        "reverse_weight": args.reverse_weight,
+        "ctc_weight": args.ctc_weight,
+        "fp16": args.fp16,
+    }
+    return onnx_config
+
+
+def export_online_encoder(model, configs, args, logger, encoder_onnx_path):
+    decoding_chunk_size = args.decoding_chunk_size
+    subsampling = model.encoder.embed.subsampling_rate
+    context = model.encoder.embed.right_context + 1
+    decoding_window = (decoding_chunk_size - 1) * subsampling + context
+    batch_size = 32
+    audio_len = decoding_window
+    feature_size = configs["input_dim"]
+    output_size = configs["encoder_conf"]["output_size"]
+    num_layers = configs["encoder_conf"]["num_blocks"]
+    # in transformer the cnn module will not be available
+    transformer = False
+    cnn_module_kernel = configs["encoder_conf"].get("cnn_module_kernel", 1) - 1
+    if not cnn_module_kernel:
+        transformer = True
+    num_decoding_left_chunks = args.num_decoding_left_chunks
+    required_cache_size = decoding_chunk_size * num_decoding_left_chunks
+    if configs["encoder"] == "squeezeformer":
+        encoder = StreamingSqueezeformerEncoder(model, required_cache_size,
+                                                args.beam_size)
+    elif configs["encoder"] == "efficientConformer":
+        encoder = StreamingEfficientConformerEncoder(model,
+                                                     required_cache_size,
+                                                     args.beam_size)
+    else:
+        encoder = StreamingEncoder(
+            model,
+            required_cache_size,
+            args.beam_size,
+            transformer,
+            args.return_ctc_logprobs,
+        )
+    encoder.eval()
+
+    # begin to export encoder
+    chunk_xs = torch.randn(batch_size,
+                           audio_len,
+                           feature_size,
+                           dtype=torch.float32)
+    chunk_lens = torch.ones(batch_size, dtype=torch.int32) * audio_len
+
+    offset = torch.arange(0, batch_size).unsqueeze(1)
+    #  (elayers, b, head, cache_t1, d_k * 2)
+    head = configs["encoder_conf"]["attention_heads"]
+    d_k = configs["encoder_conf"]["output_size"] // head
+    att_cache = torch.randn(
+        batch_size,
+        num_layers,
+        head,
+        required_cache_size,
+        d_k * 2,
+        dtype=torch.float32,
+    )
+    cnn_cache = torch.randn(
+        batch_size,
+        num_layers,
+        output_size,
+        cnn_module_kernel,
+        dtype=torch.float32,
+    )
+
+    cache_mask = torch.ones(batch_size,
+                            1,
+                            required_cache_size,
+                            dtype=torch.float32)
+    input_names = [
+        "chunk_xs",
+        "chunk_lens",
+        "offset",
+        "att_cache",
+        "cnn_cache",
+        "cache_mask",
+    ]
+    output_names = [
+        "log_probs",
+        "log_probs_idx",
+        "chunk_out",
+        "chunk_out_lens",
+        "r_offset",
+        "r_att_cache",
+        "r_cnn_cache",
+        "r_cache_mask",
+    ]
+    if args.return_ctc_logprobs:
+        output_names = [
+            "ctc_log_probs",
+            "chunk_out",
+            "chunk_out_lens",
+            "r_offset",
+            "r_att_cache",
+            "r_cnn_cache",
+            "r_cache_mask",
+        ]
+    input_tensors = (
+        chunk_xs,
+        chunk_lens,
+        offset,
+        att_cache,
+        cnn_cache,
+        cache_mask,
+    )
+    if transformer:
+        assert (args.return_ctc_logprobs is
+                False), "return_ctc_logprobs is not supported in transformer"
+        output_names.pop(6)
+
+    all_names = input_names + output_names
+    dynamic_axes = {}
+    for name in all_names:
+        # only the first dimension is dynamic
+        # all other dimension is fixed
+        dynamic_axes[name] = {0: "B"}
+
+    torch.onnx.export(
+        encoder,
+        input_tensors,
+        encoder_onnx_path,
+        export_params=True,
+        opset_version=14,
+        do_constant_folding=True,
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes=dynamic_axes,
+        verbose=False,
+    )
+
+    with torch.no_grad():
+        torch_outs = encoder(chunk_xs, chunk_lens, offset, att_cache,
+                             cnn_cache, cache_mask)
+    if transformer:
+        torch_outs = list(torch_outs).pop(6)
+    ort_session = onnxruntime.InferenceSession(
+        encoder_onnx_path, providers=["CUDAExecutionProvider"])
+    ort_inputs = {}
+
+    input_tensors = to_numpy(input_tensors)
+    for idx, name in enumerate(input_names):
+        ort_inputs[name] = input_tensors[idx]
+    if transformer:
+        del ort_inputs["cnn_cache"]
+    ort_outs = ort_session.run(None, ort_inputs)
+    test(to_numpy(torch_outs), ort_outs, rtol=1e-03, atol=1e-05)
+    logger.info("export to onnx streaming encoder succeed!")
+    onnx_config = {
+        "subsampling_rate": subsampling,
+        "context": context,
+        "decoding_chunk_size": decoding_chunk_size,
+        "num_decoding_left_chunks": num_decoding_left_chunks,
+        "beam_size": args.beam_size,
+        "fp16": args.fp16,
+        "feat_size": feature_size,
+        "decoding_window": decoding_window,
+        "cnn_module_kernel_cache": cnn_module_kernel,
+        "return_ctc_logprobs": args.return_ctc_logprobs,
+    }
+    return onnx_config
+
+
+def export_rescoring_decoder(model, configs, args, logger, decoder_onnx_path,
+                             decoder_fastertransformer):
+    bz, seq_len = 32, 100
+    beam_size = args.beam_size
+    decoder = Decoder(
+        model.decoder,
+        model.ctc_weight,
+        model.reverse_weight,
+        beam_size,
+        decoder_fastertransformer,
+    )
+    decoder.eval()
+
+    hyps_pad_sos_eos = torch.randint(low=3,
+                                     high=1000,
+                                     size=(bz, beam_size, seq_len))
+    hyps_lens_sos = torch.randint(low=3,
+                                  high=seq_len,
+                                  size=(bz, beam_size),
+                                  dtype=torch.int32)
+    r_hyps_pad_sos_eos = torch.randint(low=3,
+                                       high=1000,
+                                       size=(bz, beam_size, seq_len))
+
+    output_size = configs["encoder_conf"]["output_size"]
+    encoder_out = torch.randn(bz, seq_len, output_size, dtype=torch.float32)
+    encoder_out_lens = torch.randint(low=3,
+                                     high=seq_len,
+                                     size=(bz, ),
+                                     dtype=torch.int32)
+    ctc_score = torch.randn(bz, beam_size, dtype=torch.float32)
+
+    input_names = [
+        "encoder_out",
+        "encoder_out_lens",
+        "hyps_pad_sos_eos",
+        "hyps_lens_sos",
+        "r_hyps_pad_sos_eos",
+        "ctc_score",
+    ]
+    output_names = ["best_index"]
+    if decoder_fastertransformer:
+        output_names.insert(0, "decoder_out")
+
+    torch.onnx.export(
+        decoder,
+        (
+            encoder_out,
+            encoder_out_lens,
+            hyps_pad_sos_eos,
+            hyps_lens_sos,
+            r_hyps_pad_sos_eos,
+            ctc_score,
+        ),
+        decoder_onnx_path,
+        export_params=True,
+        opset_version=13,
+        do_constant_folding=True,
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes={
+            "encoder_out": {
+                0: "B",
+                1: "T"
+            },
+            "encoder_out_lens": {
+                0: "B"
+            },
+            "hyps_pad_sos_eos": {
+                0: "B",
+                2: "T2"
+            },
+            "hyps_lens_sos": {
+                0: "B"
+            },
+            "r_hyps_pad_sos_eos": {
+                0: "B",
+                2: "T2"
+            },
+            "ctc_score": {
+                0: "B"
+            },
+            "best_index": {
+                0: "B"
+            },
+        },
+        verbose=False,
+    )
+    with torch.no_grad():
+        o0 = decoder(
+            encoder_out,
+            encoder_out_lens,
+            hyps_pad_sos_eos,
+            hyps_lens_sos,
+            r_hyps_pad_sos_eos,
+            ctc_score,
+        )
+    providers = ["CUDAExecutionProvider"]
+    ort_session = onnxruntime.InferenceSession(decoder_onnx_path,
+                                               providers=providers)
+
+    input_tensors = [
+        encoder_out,
+        encoder_out_lens,
+        hyps_pad_sos_eos,
+        hyps_lens_sos,
+        r_hyps_pad_sos_eos,
+        ctc_score,
+    ]
+    ort_inputs = {}
+    input_tensors = to_numpy(input_tensors)
+    for idx, name in enumerate(input_names):
+        ort_inputs[name] = input_tensors[idx]
+
+    # if model.reverse weight == 0,
+    # the r_hyps_pad will be removed
+    # from the onnx decoder since it doen't play any role
+    if model.reverse_weight == 0:
+        del ort_inputs["r_hyps_pad_sos_eos"]
+    ort_outs = ort_session.run(None, ort_inputs)
+
+    # check decoder output
+    if decoder_fastertransformer:
+        test(to_numpy(o0), ort_outs, rtol=1e-03, atol=1e-05)
+    else:
+        test(to_numpy([o0]), ort_outs, rtol=1e-03, atol=1e-05)
+    logger.info("export to onnx decoder succeed!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="export x86_gpu model")
+    parser.add_argument("--config", required=True, help="config file")
+    parser.add_argument("--checkpoint", required=True, help="checkpoint model")
+    parser.add_argument(
+        "--cmvn_file",
+        required=False,
+        default="",
+        type=str,
+        help="global_cmvn file, default path is in config file",
+    )
+    parser.add_argument(
+        "--reverse_weight",
+        default=-1.0,
+        type=float,
+        required=False,
+        help="reverse weight for bitransformer," +
+        "default value is in config file",
+    )
+    parser.add_argument(
+        "--ctc_weight",
+        default=-1.0,
+        type=float,
+        required=False,
+        help="ctc weight, default value is in config file",
+    )
+    parser.add_argument(
+        "--beam_size",
+        default=10,
+        type=int,
+        required=False,
+        help="beam size would be ctc output size",
+    )
+    parser.add_argument(
+        "--output_onnx_dir",
+        default="onnx_model",
+        help="output onnx encoder and decoder directory",
+    )
+    parser.add_argument(
+        "--fp16",
+        action="store_true",
+        help="whether to export fp16 model, default false",
+    )
+    # arguments for streaming encoder
+    parser.add_argument(
+        "--streaming",
+        action="store_true",
+        help="whether to export streaming encoder, default false",
+    )
+    parser.add_argument(
+        "--decoding_chunk_size",
+        default=16,
+        type=int,
+        required=False,
+        help="the decoding chunk size, <=0 is not supported",
+    )
+    parser.add_argument(
+        "--num_decoding_left_chunks",
+        default=5,
+        type=int,
+        required=False,
+        help="number of left chunks, <= 0 is not supported",
+    )
+    parser.add_argument(
+        "--decoder_fastertransformer",
+        action="store_true",
+        help="return decoder_out and best_index for ft",
+    )
+    parser.add_argument(
+        "--return_ctc_logprobs",
+        action="store_true",
+        help="return full ctc_log_probs for TLG streaming encoder",
+    )
+    args = parser.parse_args()
+
+    torch.manual_seed(0)
+    torch.set_printoptions(precision=10)
+
+    with open(args.config, "r") as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+    if args.cmvn_file and os.path.exists(args.cmvn_file):
+        if 'cmvn' not in configs:
+            configs['cmvn'] = "global_cmvn"
+            configs['cmvn_conf'] = {}
+        else:
+            assert configs['cmvn'] == "global_cmvn"
+            assert configs['cmvn_conf'] is not None
+        configs['cmvn_conf']["cmvn_file"] = args.cmvn_file
+    if (args.reverse_weight != -1.0
+            and "reverse_weight" in configs["model_conf"]):
+        configs["model_conf"]["reverse_weight"] = args.reverse_weight
+        print("Update reverse weight to", args.reverse_weight)
+    if args.ctc_weight != -1:
+        print("Update ctc weight to ", args.ctc_weight)
+        configs["model_conf"]["ctc_weight"] = args.ctc_weight
+    configs["encoder_conf"]["use_dynamic_chunk"] = False
+
+    model, configs = init_model(args, configs)
+    model.eval()
+
+    if not os.path.exists(args.output_onnx_dir):
+        os.mkdir(args.output_onnx_dir)
+    encoder_onnx_path = os.path.join(args.output_onnx_dir, "encoder.onnx")
+    export_enc_func = None
+    if args.streaming:
+        assert args.decoding_chunk_size > 0
+        assert args.num_decoding_left_chunks > 0
+        export_enc_func = export_online_encoder
+    else:
+        export_enc_func = export_offline_encoder
+
+    onnx_config = export_enc_func(model, configs, args, logger,
+                                  encoder_onnx_path)
+
+    decoder_onnx_path = os.path.join(args.output_onnx_dir, "decoder.onnx")
+    export_rescoring_decoder(
+        model,
+        configs,
+        args,
+        logger,
+        decoder_onnx_path,
+        args.decoder_fastertransformer,
+    )
+
+    if args.fp16:
+        try:
+            import onnxmltools
+            from onnxmltools.utils.float16_converter import \
+                convert_float_to_float16
+        except ImportError:
+            print("Please install onnxmltools!")
+            sys.exit(1)
+        encoder_onnx_model = onnxmltools.utils.load_model(encoder_onnx_path)
+        encoder_onnx_model = convert_float_to_float16(encoder_onnx_model)
+        encoder_onnx_path = os.path.join(args.output_onnx_dir,
+                                         "encoder_fp16.onnx")
+        onnxmltools.utils.save_model(encoder_onnx_model, encoder_onnx_path)
+        decoder_onnx_model = onnxmltools.utils.load_model(decoder_onnx_path)
+        decoder_onnx_model = convert_float_to_float16(decoder_onnx_model)
+        decoder_onnx_path = os.path.join(args.output_onnx_dir,
+                                         "decoder_fp16.onnx")
+        onnxmltools.utils.save_model(decoder_onnx_model, decoder_onnx_path)
+    # dump configurations
+
+    config_dir = os.path.join(args.output_onnx_dir, "config.yaml")
+    with open(config_dir, "w") as out:
+        yaml.dump(onnx_config, out)

wenet/bin/recognize.py

@@ -0,0 +1,317 @@
+# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Xiaoyu Chen, Di Wu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import argparse
+import copy
+import logging
+import os
+
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+from wenet.dataset.dataset import Dataset
+from wenet.utils.config import override_config
+from wenet.utils.init_model import init_model
+from wenet.utils.init_tokenizer import init_tokenizer
+from wenet.utils.context_graph import ContextGraph
+from wenet.utils.ctc_utils import get_blank_id
+from wenet.utils.common import TORCH_NPU_AVAILABLE  # noqa just ensure to check torch-npu
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='recognize with your model')
+    parser.add_argument('--config', required=True, help='config file')
+    parser.add_argument('--test_data', required=True, help='test data file')
+    parser.add_argument('--data_type',
+                        default='raw',
+                        choices=['raw', 'shard'],
+                        help='train and cv data type')
+    parser.add_argument('--gpu',
+                        type=int,
+                        default=-1,
+                        help='gpu id for this rank, -1 for cpu')
+    parser.add_argument('--device',
+                        type=str,
+                        default="cpu",
+                        choices=["cpu", "npu", "cuda"],
+                        help='accelerator to use')
+    parser.add_argument('--dtype',
+                        type=str,
+                        default='fp32',
+                        choices=['fp16', 'fp32', 'bf16'],
+                        help='model\'s dtype')
+    parser.add_argument('--num_workers',
+                        default=0,
+                        type=int,
+                        help='num of subprocess workers for reading')
+    parser.add_argument('--checkpoint', required=True, help='checkpoint model')
+    parser.add_argument('--beam_size',
+                        type=int,
+                        default=10,
+                        help='beam size for search')
+    parser.add_argument('--length_penalty',
+                        type=float,
+                        default=0.0,
+                        help='length penalty')
+    parser.add_argument('--blank_penalty',
+                        type=float,
+                        default=0.0,
+                        help='blank penalty')
+    parser.add_argument('--result_dir', required=True, help='asr result file')
+    parser.add_argument('--batch_size',
+                        type=int,
+                        default=16,
+                        help='asr result file')
+    parser.add_argument('--modes',
+                        nargs='+',
+                        help="""decoding mode, support the following:
+                             attention
+                             ctc_greedy_search
+                             ctc_prefix_beam_search
+                             attention_rescoring
+                             rnnt_greedy_search
+                             rnnt_beam_search
+                             rnnt_beam_attn_rescoring
+                             ctc_beam_td_attn_rescoring
+                             hlg_onebest
+                             hlg_rescore
+                             paraformer_greedy_search
+                             paraformer_beam_search""")
+    parser.add_argument('--search_ctc_weight',
+                        type=float,
+                        default=1.0,
+                        help='ctc weight for nbest generation')
+    parser.add_argument('--search_transducer_weight',
+                        type=float,
+                        default=0.0,
+                        help='transducer weight for nbest generation')
+    parser.add_argument('--ctc_weight',
+                        type=float,
+                        default=0.0,
+                        help='ctc weight for rescoring weight in \
+                                  attention rescoring decode mode \
+                              ctc weight for rescoring weight in \
+                                  transducer attention rescore decode mode')
+
+    parser.add_argument('--transducer_weight',
+                        type=float,
+                        default=0.0,
+                        help='transducer weight for rescoring weight in '
+                        'transducer attention rescore mode')
+    parser.add_argument('--attn_weight',
+                        type=float,
+                        default=0.0,
+                        help='attention weight for rescoring weight in '
+                        'transducer attention rescore mode')
+    parser.add_argument('--decoding_chunk_size',
+                        type=int,
+                        default=-1,
+                        help='''decoding chunk size,
+                                <0: for decoding, use full chunk.
+                                >0: for decoding, use fixed chunk size as set.
+                                0: used for training, it's prohibited here''')
+    parser.add_argument('--num_decoding_left_chunks',
+                        type=int,
+                        default=-1,
+                        help='number of left chunks for decoding')
+    parser.add_argument('--simulate_streaming',
+                        action='store_true',
+                        help='simulate streaming inference')
+    parser.add_argument('--reverse_weight',
+                        type=float,
+                        default=0.0,
+                        help='''right to left weight for attention rescoring
+                                decode mode''')
+    parser.add_argument('--override_config',
+                        action='append',
+                        default=[],
+                        help="override yaml config")
+
+    parser.add_argument('--word',
+                        default='',
+                        type=str,
+                        help='word file, only used for hlg decode')
+    parser.add_argument('--hlg',
+                        default='',
+                        type=str,
+                        help='hlg file, only used for hlg decode')
+    parser.add_argument('--lm_scale',
+                        type=float,
+                        default=0.0,
+                        help='lm scale for hlg attention rescore decode')
+    parser.add_argument('--decoder_scale',
+                        type=float,
+                        default=0.0,
+                        help='lm scale for hlg attention rescore decode')
+    parser.add_argument('--r_decoder_scale',
+                        type=float,
+                        default=0.0,
+                        help='lm scale for hlg attention rescore decode')
+
+    parser.add_argument(
+        '--context_bias_mode',
+        type=str,
+        default='',
+        help='''Context bias mode, selectable from the following
+                                option: decoding-graph, deep-biasing''')
+    parser.add_argument('--context_list_path',
+                        type=str,
+                        default='',
+                        help='Context list path')
+    parser.add_argument('--context_graph_score',
+                        type=float,
+                        default=0.0,
+                        help='''The higher the score, the greater the degree of
+                                bias using decoding-graph for biasing''')
+
+    parser.add_argument('--use_lora',
+                        type=bool,
+                        default=False,
+                        help='''Whether to use lora for biasing''')
+    parser.add_argument("--lora_ckpt_path",
+                        default=None,
+                        type=str,
+                        help="lora checkpoint path.")
+    args = parser.parse_args()
+    print(args)
+    return args
+
+
+def main():
+    args = get_args()
+    logging.basicConfig(level=logging.DEBUG,
+                        format='%(asctime)s %(levelname)s %(message)s')
+    if args.gpu != -1:
+        # remain the original usage of gpu
+        args.device = "cuda"
+    if "cuda" in args.device:
+        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
+
+    with open(args.config, 'r') as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+    if len(args.override_config) > 0:
+        configs = override_config(configs, args.override_config)
+
+    test_conf = copy.deepcopy(configs['dataset_conf'])
+
+    test_conf['filter_conf']['max_length'] = 102400
+    test_conf['filter_conf']['min_length'] = 0
+    test_conf['filter_conf']['token_max_length'] = 102400
+    test_conf['filter_conf']['token_min_length'] = 0
+    test_conf['filter_conf']['max_output_input_ratio'] = 102400
+    test_conf['filter_conf']['min_output_input_ratio'] = 0
+    test_conf['speed_perturb'] = False
+    test_conf['spec_aug'] = False
+    test_conf['spec_sub'] = False
+    test_conf['spec_trim'] = False
+    test_conf['shuffle'] = False
+    test_conf['sort'] = False
+    test_conf['cycle'] = 1
+    test_conf['list_shuffle'] = False
+    if 'fbank_conf' in test_conf:
+        test_conf['fbank_conf']['dither'] = 0.0
+    elif 'mfcc_conf' in test_conf:
+        test_conf['mfcc_conf']['dither'] = 0.0
+    test_conf['batch_conf']['batch_type'] = "static"
+    test_conf['batch_conf']['batch_size'] = args.batch_size
+
+    tokenizer = init_tokenizer(configs)
+    test_dataset = Dataset(args.data_type,
+                           args.test_data,
+                           tokenizer,
+                           test_conf,
+                           partition=False)
+
+    test_data_loader = DataLoader(test_dataset,
+                                  batch_size=None,
+                                  num_workers=args.num_workers)
+
+    # Init asr model from configs
+    args.jit = False
+    model, configs = init_model(args, configs)
+
+    device = torch.device(args.device)
+    model = model.to(device)
+    model.eval()
+    dtype = torch.float32
+    if args.dtype == 'fp16':
+        dtype = torch.float16
+    elif args.dtype == 'bf16':
+        dtype = torch.bfloat16
+    logging.info("compute dtype is {}".format(dtype))
+
+    context_graph = None
+    if 'decoding-graph' in args.context_bias_mode:
+        context_graph = ContextGraph(args.context_list_path,
+                                     tokenizer.symbol_table,
+                                     configs['tokenizer_conf']['bpe_path'],
+                                     args.context_graph_score)
+
+    _, blank_id = get_blank_id(configs, tokenizer.symbol_table)
+    logging.info("blank_id is {}".format(blank_id))
+
+    # TODO(Dinghao Zhou): Support RNN-T related decoding
+    # TODO(Lv Xiang): Support k2 related decoding
+    # TODO(Kaixun Huang): Support context graph
+    files = {}
+    for mode in args.modes:
+        dir_name = os.path.join(args.result_dir, mode)
+        os.makedirs(dir_name, exist_ok=True)
+        file_name = os.path.join(dir_name, 'text')
+        files[mode] = open(file_name, 'w')
+    max_format_len = max([len(mode) for mode in args.modes])
+
+    with torch.cuda.amp.autocast(enabled=True,
+                                 dtype=dtype,
+                                 cache_enabled=False):
+        with torch.no_grad():
+            for batch_idx, batch in enumerate(test_data_loader):
+                keys = batch["keys"]
+                feats = batch["feats"].to(device)
+                target = batch["target"].to(device)
+                feats_lengths = batch["feats_lengths"].to(device)
+                target_lengths = batch["target_lengths"].to(device)
+                infos = {"tasks": batch["tasks"], "langs": batch["langs"]}
+                results = model.decode(
+                    args.modes,
+                    feats,
+                    feats_lengths,
+                    args.beam_size,
+                    decoding_chunk_size=args.decoding_chunk_size,
+                    num_decoding_left_chunks=args.num_decoding_left_chunks,
+                    ctc_weight=args.ctc_weight,
+                    simulate_streaming=args.simulate_streaming,
+                    reverse_weight=args.reverse_weight,
+                    context_graph=context_graph,
+                    blank_id=blank_id,
+                    blank_penalty=args.blank_penalty,
+                    length_penalty=args.length_penalty,
+                    infos=infos)
+                for i, key in enumerate(keys):
+                    for mode, hyps in results.items():
+                        tokens = hyps[i].tokens
+                        line = '{} {}'.format(key,
+                                              tokenizer.detokenize(tokens)[0])
+                        logging.info('{} {}'.format(mode.ljust(max_format_len),
+                                                    line))
+                        files[mode].write(line + '\n')
+        for mode, f in files.items():
+            f.close()
+
+
+if __name__ == '__main__':
+    main()

wenet/bin/recognize_onnx_gpu.py

@@ -0,0 +1,297 @@
+# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Xiaoyu Chen, Di Wu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This script is for testing exported onnx encoder and decoder from
+export_onnx_gpu.py. The exported onnx models only support batch offline ASR inference.
+It requires a python wrapped c++ ctc decoder.
+Please install it by following:
+https://github.com/Slyne/ctc_decoder.git
+"""
+from __future__ import print_function
+
+import argparse
+import copy
+import logging
+import os
+import sys
+
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+from wenet.dataset.dataset import Dataset
+from wenet.utils.common import IGNORE_ID
+from wenet.utils.config import override_config
+from wenet.utils.init_tokenizer import init_tokenizer
+
+import onnxruntime as rt
+import multiprocessing
+import numpy as np
+
+try:
+    from swig_decoders import map_batch, \
+        ctc_beam_search_decoder_batch, \
+        TrieVector, PathTrie
+except ImportError:
+    print('Please install ctc decoders first by refering to\n' +
+          'https://github.com/Slyne/ctc_decoder.git')
+    sys.exit(1)
+
+def get_args():
+    parser = argparse.ArgumentParser(description='recognize with your model')
+    parser.add_argument('--config', required=True, help='config file')
+    parser.add_argument('--test_data', required=True, help='test data file')
+    parser.add_argument('--data_type',
+                        default='raw',
+                        choices=['raw', 'shard'],
+                        help='train and cv data type')
+    parser.add_argument('--gpu',
+                        type=int,
+                        default=-1,
+                        help='gpu id for this rank, -1 for cpu')
+    parser.add_argument('--dict', required=True, help='dict file')
+    parser.add_argument('--encoder_onnx',
+                        required=True,
+                        help='encoder onnx file')
+    parser.add_argument('--decoder_onnx',
+                        required=True,
+                        help='decoder onnx file')
+    parser.add_argument('--result_file', required=True, help='asr result file')
+    parser.add_argument('--batch_size',
+                        type=int,
+                        default=32,
+                        help='asr result file')
+    parser.add_argument('--mode',
+                        choices=[
+                            'ctc_greedy_search', 'ctc_prefix_beam_search',
+                            'attention_rescoring'
+                        ],
+                        default='attention_rescoring',
+                        help='decoding mode')
+    parser.add_argument('--bpe_model',
+                        default=None,
+                        type=str,
+                        help='bpe model for english part')
+    parser.add_argument('--override_config',
+                        action='append',
+                        default=[],
+                        help="override yaml config")
+    parser.add_argument('--fp16',
+                        action='store_true',
+                        help='whether to export fp16 model, default false')
+    args = parser.parse_args()
+    return args
+
+def main():
+    args = get_args()
+    logging.basicConfig(level=logging.DEBUG,
+                        format='%(asctime)s %(levelname)s %(message)s')
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
+
+    with open(args.config, 'r') as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+    if len(args.override_config) > 0:
+        configs = override_config(configs, args.override_config)
+
+    reverse_weight = configs["model_conf"].get("reverse_weight", 0.0)
+    special_tokens = configs.get('tokenizer_conf', {}).get('special_tokens', None)
+    test_conf = copy.deepcopy(configs['dataset_conf'])
+    test_conf['filter_conf']['max_length'] = 102400
+    test_conf['filter_conf']['min_length'] = 0
+    test_conf['filter_conf']['token_max_length'] = 102400
+    test_conf['filter_conf']['token_min_length'] = 0
+    test_conf['filter_conf']['max_output_input_ratio'] = 102400
+    test_conf['filter_conf']['min_output_input_ratio'] = 0
+    test_conf['speed_perturb'] = False
+    test_conf['spec_aug'] = False
+    test_conf['spec_sub'] = False
+    test_conf['spec_trim'] = False
+    test_conf['shuffle'] = False
+    test_conf['sort'] = False
+    test_conf['fbank_conf']['dither'] = 0.0
+    test_conf['batch_conf']['batch_type'] = "static"
+    test_conf['batch_conf']['batch_size'] = args.batch_size
+
+    tokenizer = init_tokenizer(configs)
+    test_dataset = Dataset(args.data_type,
+                           args.test_data,
+                           tokenizer,
+                           test_conf,
+                           partition=False)
+    test_data_loader = DataLoader(test_dataset, batch_size=None, num_workers=0)
+
+    # Init asr model from configs
+    use_cuda = args.gpu >= 0 and torch.cuda.is_available()
+    if use_cuda:
+        EP_list = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+    else:
+        EP_list = ['CPUExecutionProvider']
+
+    encoder_ort_session = rt.InferenceSession(args.encoder_onnx,
+                                              providers=EP_list)
+    decoder_ort_session = None
+    if args.mode == "attention_rescoring":
+        decoder_ort_session = rt.InferenceSession(args.decoder_onnx,
+                                                  providers=EP_list)
+
+    # Load dict
+    vocabulary = []
+    char_dict = {}
+    with open(args.dict, 'r') as fin:
+        for line in fin:
+            arr = line.strip().split()
+            assert len(arr) == 2
+            char_dict[int(arr[1])] = arr[0]
+            vocabulary.append(arr[0])
+
+    vocab_size = len(char_dict)
+    sos = (vocab_size - 1 if special_tokens is None else
+           special_tokens.get("<sos>", vocab_size - 1))
+    eos = (vocab_size - 1 if special_tokens is None else
+           special_tokens.get("<eos>", vocab_size - 1))
+
+    with torch.no_grad(), open(args.result_file, 'w') as fout:
+        for _, batch in enumerate(test_data_loader):
+            keys = batch['keys']
+            feats = batch['feats']
+            feats_lengths = batch['feats_lengths']
+            feats, feats_lengths = feats.numpy(), feats_lengths.numpy()
+            if args.fp16:
+                feats = feats.astype(np.float16)
+            ort_inputs = {
+                encoder_ort_session.get_inputs()[0].name: feats,
+                encoder_ort_session.get_inputs()[1].name: feats_lengths
+            }
+            ort_outs = encoder_ort_session.run(None, ort_inputs)
+            encoder_out, encoder_out_lens, ctc_log_probs, \
+                beam_log_probs, beam_log_probs_idx = ort_outs
+            beam_size = beam_log_probs.shape[-1]
+            batch_size = beam_log_probs.shape[0]
+            num_processes = min(multiprocessing.cpu_count(), batch_size)
+            if args.mode == 'ctc_greedy_search':
+                if beam_size != 1:
+                    log_probs_idx = beam_log_probs_idx[:, :, 0]
+                batch_sents = []
+                for idx, seq in enumerate(log_probs_idx):
+                    batch_sents.append(seq[0:encoder_out_lens[idx]].tolist())
+                hyps = map_batch(batch_sents, vocabulary, num_processes, True,
+                                 0)
+            elif args.mode in ('ctc_prefix_beam_search',
+                               "attention_rescoring"):
+                batch_log_probs_seq_list = beam_log_probs.tolist()
+                batch_log_probs_idx_list = beam_log_probs_idx.tolist()
+                batch_len_list = encoder_out_lens.tolist()
+                batch_log_probs_seq = []
+                batch_log_probs_ids = []
+                batch_start = []  # only effective in streaming deployment
+                batch_root = TrieVector()
+                root_dict = {}
+                for i in range(len(batch_len_list)):
+                    num_sent = batch_len_list[i]
+                    batch_log_probs_seq.append(
+                        batch_log_probs_seq_list[i][0:num_sent])
+                    batch_log_probs_ids.append(
+                        batch_log_probs_idx_list[i][0:num_sent])
+                    root_dict[i] = PathTrie()
+                    batch_root.append(root_dict[i])
+                    batch_start.append(True)
+                score_hyps = ctc_beam_search_decoder_batch(
+                    batch_log_probs_seq, batch_log_probs_ids, batch_root,
+                    batch_start, beam_size, num_processes, 0, -2, 0.99999)
+                if args.mode == 'ctc_prefix_beam_search':
+                    hyps = []
+                    for cand_hyps in score_hyps:
+                        hyps.append(cand_hyps[0][1])
+                    hyps = map_batch(hyps, vocabulary, num_processes, False, 0)
+            if args.mode == 'attention_rescoring':
+                ctc_score, all_hyps = [], []
+                max_len = 0
+                for hyps in score_hyps:
+                    cur_len = len(hyps)
+                    if len(hyps) < beam_size:
+                        hyps += (beam_size - cur_len) * [(-float("INF"),
+                                                          (0, ))]
+                    cur_ctc_score = []
+                    for hyp in hyps:
+                        cur_ctc_score.append(hyp[0])
+                        all_hyps.append(list(hyp[1]))
+                        if len(hyp[1]) > max_len:
+                            max_len = len(hyp[1])
+                    ctc_score.append(cur_ctc_score)
+                if args.fp16:
+                    ctc_score = np.array(ctc_score, dtype=np.float16)
+                else:
+                    ctc_score = np.array(ctc_score, dtype=np.float32)
+                hyps_pad_sos_eos = np.ones(
+                    (batch_size, beam_size, max_len + 2),
+                    dtype=np.int64) * IGNORE_ID
+                r_hyps_pad_sos_eos = np.ones(
+                    (batch_size, beam_size, max_len + 2),
+                    dtype=np.int64) * IGNORE_ID
+                hyps_lens_sos = np.ones((batch_size, beam_size),
+                                        dtype=np.int32)
+                k = 0
+                for i in range(batch_size):
+                    for j in range(beam_size):
+                        cand = all_hyps[k]
+                        l = len(cand) + 2
+                        hyps_pad_sos_eos[i][j][0:l] = [sos] + cand + [eos]
+                        r_hyps_pad_sos_eos[i][j][0:l] = [sos] + cand[::-1] + [
+                            eos
+                        ]
+                        hyps_lens_sos[i][j] = len(cand) + 1
+                        k += 1
+                decoder_ort_inputs = {
+                    decoder_ort_session.get_inputs()[0].name: encoder_out,
+                    decoder_ort_session.get_inputs()[1].name: encoder_out_lens,
+                    decoder_ort_session.get_inputs()[2].name: hyps_pad_sos_eos,
+                    decoder_ort_session.get_inputs()[3].name: hyps_lens_sos,
+                    decoder_ort_session.get_inputs()[-1].name: ctc_score
+                }
+                if reverse_weight > 0:
+                    r_hyps_pad_sos_eos_name = decoder_ort_session.get_inputs(
+                    )[4].name
+                    decoder_ort_inputs[
+                        r_hyps_pad_sos_eos_name] = r_hyps_pad_sos_eos
+                best_index = decoder_ort_session.run(None,
+                                                     decoder_ort_inputs)[0]
+                best_sents = []
+                k = 0
+                for idx in best_index:
+                    cur_best_sent = all_hyps[k:k + beam_size][idx]
+                    best_sents.append(cur_best_sent)
+                    k += beam_size
+                hyps = map_batch(best_sents, vocabulary, num_processes)
+
+            for i, key in enumerate(keys):
+                content = hyps[i]
+                logging.info('{} {}'.format(key, content))
+                fout.write('{} {}\n'.format(key, content))
+
+if __name__ == '__main__':
+    main()

wenet/bin/train.py

@@ -0,0 +1,190 @@
+# Copyright (c) 2021 Mobvoi Inc. (authors: Binbin Zhang)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import argparse
+import datetime
+import logging
+import os
+import torch
+import yaml
+
+import torch.distributed as dist
+
+from torch.distributed.elastic.multiprocessing.errors import record
+from wenet.utils.common import lrs_to_str, TORCH_NPU_AVAILABLE  # noqa just ensure to check torch-npu
+
+from wenet.utils.executor import Executor
+from wenet.utils.config import override_config
+from wenet.utils.init_model import init_model
+from wenet.utils.init_tokenizer import init_tokenizer
+from wenet.utils.train_utils import (
+    add_fsdp_args, add_model_args, add_dataset_args, add_ddp_args,
+    add_deepspeed_args, add_trace_args, init_distributed,
+    init_dataset_and_dataloader, check_modify_and_save_config,
+    init_optimizer_and_scheduler, init_scaler, trace_and_print_model,
+    wrap_cuda_model, init_summarywriter, save_model, log_per_epoch,
+    add_lora_args, reinit_lora)
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='training your network')
+    parser.add_argument('--train_engine',
+                        default='torch_ddp',
+                        choices=['torch_ddp', 'torch_fsdp', 'deepspeed'],
+                        help='Engine for paralleled training')
+    # set default value of device to "cuda", avoiding the modify of original scripts
+    parser.add_argument('--device',
+                        type=str,
+                        default='cuda',
+                        choices=["cpu", "npu", "cuda"],
+                        help='accelerator for training')
+    parser = add_model_args(parser)
+    parser = add_dataset_args(parser)
+    parser = add_ddp_args(parser)
+    parser = add_lora_args(parser)
+    parser = add_deepspeed_args(parser)
+    parser = add_fsdp_args(parser)
+    parser = add_trace_args(parser)
+    args = parser.parse_args()
+    if args.train_engine == "deepspeed":
+        args.deepspeed = True
+        assert args.deepspeed_config is not None
+    return args
+
+
+# NOTE(xcsong): On worker errors, this recod tool will summarize the
+#   details of the error (e.g. time, rank, host, pid, traceback, etc).
+@record
+def main():
+    args = get_args()
+    logging.basicConfig(level=logging.DEBUG,
+                        format='%(asctime)s %(levelname)s %(message)s')
+
+    # Set random seed
+    torch.manual_seed(777)
+
+    # Read config
+    with open(args.config, 'r') as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+    if len(args.override_config) > 0:
+        configs = override_config(configs, args.override_config)
+
+    # init tokenizer
+    tokenizer = init_tokenizer(configs)
+
+    # Init env for ddp OR deepspeed
+    _, _, rank = init_distributed(args)
+
+    # Get dataset & dataloader
+    train_dataset, cv_dataset, train_data_loader, cv_data_loader = \
+        init_dataset_and_dataloader(args, configs, tokenizer)
+
+    # Do some sanity checks and save config to arsg.model_dir
+    configs = check_modify_and_save_config(args, configs,
+                                           tokenizer.symbol_table)
+
+    # Init asr model from configs
+    model, configs = init_model(args, configs)
+
+    if hasattr(args, 'lora_reinit') and args.lora_reinit:
+        reinit_lora(model, args, configs, tokenizer)
+
+    # Check model is jitable & print model archtectures
+    trace_and_print_model(args, model)
+
+    # Tensorboard summary
+    writer = init_summarywriter(args)
+
+    # Dispatch model from cpu to gpu
+    model, device = wrap_cuda_model(args, model, configs)
+
+    # Get optimizer & scheduler
+    model, optimizer, scheduler = init_optimizer_and_scheduler(
+        args, configs, model)
+
+    # Save checkpoints
+    save_model(model,
+               info_dict={
+                   "save_time":
+                   datetime.datetime.now().strftime('%d/%m/%Y %H:%M:%S'),
+                   "tag":
+                   "init",
+                   **configs
+               })
+
+    # Get executor
+    tag = configs["init_infos"].get("tag", "init")
+    executor = Executor(global_step=configs["init_infos"].get('step', -1),
+                        device=device)
+
+    # Init scaler, used for pytorch amp mixed precision training
+    scaler = init_scaler(args)
+
+    # Start training loop
+    start_epoch = configs["init_infos"].get('epoch', 0) + int("epoch_" in tag)
+    # if save_interval in configs, steps mode else epoch mode
+    end_epoch = configs.get('max_epoch',
+                            100) if "save_interval" not in configs else 1
+    assert start_epoch <= end_epoch
+    configs.pop("init_infos", None)
+    final_epoch = None
+    for epoch in range(start_epoch, end_epoch):
+        configs['epoch'] = epoch
+
+        lrs = [group['lr'] for group in optimizer.param_groups]
+        logging.info('Epoch {} Step {} TRAIN info lr {} rank {}'.format(
+            epoch, executor.step, lrs_to_str(lrs), rank))
+
+        dist.barrier(
+        )  # NOTE(xcsong): Ensure all ranks start Train at the same time.
+        # NOTE(xcsong): Why we need a new group? see `train_utils.py::wenet_join`
+        group_join = dist.new_group(
+            backend="gloo", timeout=datetime.timedelta(seconds=args.timeout))
+        executor.train(model, optimizer, scheduler, train_data_loader,
+                       cv_data_loader, writer, configs, scaler, group_join)
+        dist.destroy_process_group(group_join)
+
+        dist.barrier(
+        )  # NOTE(xcsong): Ensure all ranks start CV at the same time.
+        loss_dict = executor.cv(model, cv_data_loader, configs)
+        info_dict = {
+            'epoch': epoch,
+            'lrs': [group['lr'] for group in optimizer.param_groups],
+            'step': executor.step,
+            'save_time': datetime.datetime.now().strftime('%d/%m/%Y %H:%M:%S'),
+            'tag': "epoch_{}".format(epoch),
+            'loss_dict': loss_dict,
+            **configs
+        }
+        # epoch cv: tensorboard && log
+        log_per_epoch(writer, info_dict=info_dict)
+        save_model(model, info_dict=info_dict)
+
+        final_epoch = epoch
+
+    if final_epoch is not None and rank == 0:
+        final_model_path = os.path.join(args.model_dir, 'final.pt')
+        os.remove(final_model_path) if os.path.exists(
+            final_model_path) else None
+        os.symlink('{}.pt'.format(final_epoch), final_model_path)
+        writer.close()
+    dist.barrier(
+    )  # NOTE(yktian): Ensure all ranks end Train before destroy process group.
+    dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    main()

wenet/branchformer/cgmlp.py

@@ -0,0 +1,194 @@
+# Copyright (c) 2022 Yifan Peng (Carnegie Mellon University)
+#               2023 Voicecomm Inc (Kai Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+"""MLP with convolutional gating (cgMLP) definition.
+
+References:
+    https://openreview.net/forum?id=RA-zVvZLYIy
+    https://arxiv.org/abs/2105.08050
+
+"""
+
+from typing import Tuple
+import torch
+import torch.nn as nn
+from wenet.utils.class_utils import WENET_ACTIVATION_CLASSES
+
+
+class ConvolutionalSpatialGatingUnit(torch.nn.Module):
+    """Convolutional Spatial Gating Unit (CSGU)."""
+
+    def __init__(
+        self,
+        size: int,
+        kernel_size: int,
+        dropout_rate: float,
+        use_linear_after_conv: bool,
+        gate_activation: str,
+        causal: bool = True,
+    ):
+        super().__init__()
+
+        # split input channels
+        n_channels = size // 2
+        self.norm = nn.LayerNorm(n_channels)
+        # self.lorder is used to distinguish if it's a causal convolution,
+        # if self.lorder > 0: it's a causal convolution, the input will be
+        #    padded with self.lorder frames on the left in forward.
+        # else: it's a symmetrical convolution
+        if causal:
+            padding = 0
+            self.lorder = kernel_size - 1
+        else:
+            # kernel_size should be an odd number for none causal convolution
+            assert (kernel_size - 1) % 2 == 0
+            padding = (kernel_size - 1) // 2
+            self.lorder = 0
+        self.conv = torch.nn.Conv1d(
+            n_channels,
+            n_channels,
+            kernel_size,
+            1,
+            padding,
+            groups=n_channels,
+        )
+        if use_linear_after_conv:
+            self.linear = torch.nn.Linear(n_channels, n_channels)
+        else:
+            self.linear = None
+
+        if gate_activation == "identity":
+            self.act = torch.nn.Identity()
+        else:
+            self.act = WENET_ACTIVATION_CLASSES[gate_activation]()
+
+        self.dropout = torch.nn.Dropout(dropout_rate)
+
+    def espnet_initialization_fn(self):
+        torch.nn.init.normal_(self.conv.weight, std=1e-6)
+        torch.nn.init.ones_(self.conv.bias)
+        if self.linear is not None:
+            torch.nn.init.normal_(self.linear.weight, std=1e-6)
+            torch.nn.init.ones_(self.linear.bias)
+
+    def forward(
+        self, x: torch.Tensor, cache: torch.Tensor = torch.zeros((0, 0, 0))
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward method
+
+        Args:
+            x (torch.Tensor): (batch, time, channels)
+            cache (torch.Tensor): left context cache, it is only
+                used in causal convolution (#batch, channels, cache_t),
+                (0, 0, 0) meas fake cache.
+
+        Returns:
+            out (torch.Tensor): (batch, time, channels/2)
+        """
+
+        x_r, x_g = x.chunk(2, dim=-1)
+        # exchange the temporal dimension and the feature dimension
+        x_g = x_g.transpose(1, 2)  # (#batch, channels, time)
+
+        if self.lorder > 0:
+            if cache.size(2) == 0:  # cache_t == 0
+                x_g = nn.functional.pad(x_g, (self.lorder, 0), 'constant', 0.0)
+            else:
+                assert cache.size(0) == x_g.size(0)  # equal batch
+                assert cache.size(1) == x_g.size(1)  # equal channel
+                x_g = torch.cat((cache, x_g), dim=2)
+            assert (x_g.size(2) > self.lorder)
+            new_cache = x_g[:, :, -self.lorder:]
+        else:
+            # It's better we just return None if no cache is required,
+            # However, for JIT export, here we just fake one tensor instead of
+            # None.
+            new_cache = torch.zeros((0, 0, 0),
+                                    dtype=x_g.dtype,
+                                    device=x_g.device)
+
+        x_g = x_g.transpose(1, 2)
+        x_g = self.norm(x_g)  # (N, T, D/2)
+        x_g = self.conv(x_g.transpose(1, 2)).transpose(1, 2)  # (N, T, D/2)
+        if self.linear is not None:
+            x_g = self.linear(x_g)
+
+        x_g = self.act(x_g)
+        out = x_r * x_g  # (N, T, D/2)
+        out = self.dropout(out)
+        return out, new_cache
+
+
+class ConvolutionalGatingMLP(torch.nn.Module):
+    """Convolutional Gating MLP (cgMLP)."""
+
+    def __init__(
+        self,
+        size: int,
+        linear_units: int,
+        kernel_size: int,
+        dropout_rate: float,
+        use_linear_after_conv: bool,
+        gate_activation: str,
+        causal: bool = True,
+    ):
+        super().__init__()
+
+        self.channel_proj1 = torch.nn.Sequential(
+            torch.nn.Linear(size, linear_units), torch.nn.GELU())
+        self.csgu = ConvolutionalSpatialGatingUnit(
+            size=linear_units,
+            kernel_size=kernel_size,
+            dropout_rate=dropout_rate,
+            use_linear_after_conv=use_linear_after_conv,
+            gate_activation=gate_activation,
+            causal=causal,
+        )
+        self.channel_proj2 = torch.nn.Linear(linear_units // 2, size)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        cache: torch.Tensor = torch.zeros((0, 0, 0))
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward method
+
+        Args:
+            x (torch.Tensor): (batch, time, channels)
+            mask_pad (torch.Tensor): used for batch padding (#batch, 1, time),
+                (0, 0, 0) means fake mask. Not used yet
+            cache (torch.Tensor): left context cache, it is only
+                used in causal convolution (#batch, channels, cache_t),
+                (0, 0, 0) meas fake cache.
+
+        Returns:
+            out (torch.Tensor): (batch, time, channels/2)
+        """
+
+        xs_pad = x
+
+        # size -> linear_units
+        xs_pad = self.channel_proj1(xs_pad)
+
+        # linear_units -> linear_units/2
+        xs_pad, new_cnn_cache = self.csgu(xs_pad, cache)
+
+        # linear_units/2 -> size
+        xs_pad = self.channel_proj2(xs_pad)
+
+        out = xs_pad
+
+        return out, new_cnn_cache

wenet/branchformer/encoder.py

@@ -0,0 +1,177 @@
+# Copyright (c) 2022 Yifan Peng (Carnegie Mellon University)
+#               2023 Voicecomm Inc (Kai Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+"""Encoder definition."""
+
+import torch
+
+from typing import List, Optional, Union
+
+from wenet.branchformer.encoder_layer import BranchformerEncoderLayer
+from wenet.branchformer.cgmlp import ConvolutionalGatingMLP
+from wenet.transformer.encoder import BaseEncoder
+from wenet.utils.class_utils import (
+    WENET_ATTENTION_CLASSES, )
+
+
+class BranchformerEncoder(BaseEncoder):
+    """Branchformer encoder module."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 256,
+        use_attn: bool = True,
+        attention_heads: int = 4,
+        selfattention_layer_type: str = "rel_selfattn",
+        pos_enc_layer_type: str = "rel_pos",
+        use_cgmlp: bool = True,
+        cgmlp_linear_units: int = 2048,
+        cgmlp_conv_kernel: int = 31,
+        use_linear_after_conv: bool = False,
+        gate_activation: str = "identity",
+        merge_method: str = "concat",
+        cgmlp_weight: Union[float, List[float]] = 0.5,
+        attn_branch_drop_rate: Union[float, List[float]] = 0.0,
+        num_blocks: int = 12,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        input_layer: str = "conv2d",
+        stochastic_depth_rate: Union[float, List[float]] = 0.0,
+        static_chunk_size: int = 0,
+        use_dynamic_chunk: bool = False,
+        global_cmvn: torch.nn.Module = None,
+        use_dynamic_left_chunk: bool = False,
+        causal: bool = False,
+        query_bias: bool = True,
+        key_bias: bool = True,
+        value_bias: bool = True,
+        gradient_checkpointing: bool = False,
+        use_sdpa: bool = False,
+        layer_norm_type: str = 'layer_norm',
+        norm_eps: float = 1e-5,
+        n_kv_head: Optional[int] = None,
+        head_dim: Optional[int] = None,
+    ):
+        super().__init__(input_size, output_size, attention_heads,
+                         cgmlp_linear_units, num_blocks, dropout_rate,
+                         positional_dropout_rate, attention_dropout_rate,
+                         input_layer, pos_enc_layer_type, True,
+                         static_chunk_size, use_dynamic_chunk, global_cmvn,
+                         use_dynamic_left_chunk, gradient_checkpointing,
+                         use_sdpa, layer_norm_type, norm_eps)
+
+        encoder_selfattn_layer_args = (
+            attention_heads,
+            output_size,
+            attention_dropout_rate,
+            query_bias,
+            key_bias,
+            value_bias,
+            use_sdpa,
+            n_kv_head,
+            head_dim,
+        )
+
+        cgmlp_layer = ConvolutionalGatingMLP
+        cgmlp_layer_args = (
+            output_size,
+            cgmlp_linear_units,
+            cgmlp_conv_kernel,
+            dropout_rate,
+            use_linear_after_conv,
+            gate_activation,
+            causal,
+        )
+
+        if isinstance(stochastic_depth_rate, float):
+            stochastic_depth_rate = [stochastic_depth_rate] * num_blocks
+        if len(stochastic_depth_rate) != num_blocks:
+            raise ValueError(
+                f"Length of stochastic_depth_rate ({len(stochastic_depth_rate)}) "
+                f"should be equal to num_blocks ({num_blocks})")
+
+        if isinstance(cgmlp_weight, float):
+            cgmlp_weight = [cgmlp_weight] * num_blocks
+        if len(cgmlp_weight) != num_blocks:
+            raise ValueError(
+                f"Length of cgmlp_weight ({len(cgmlp_weight)}) should be equal to "
+                f"num_blocks ({num_blocks})")
+
+        if isinstance(attn_branch_drop_rate, float):
+            attn_branch_drop_rate = [attn_branch_drop_rate] * num_blocks
+        if len(attn_branch_drop_rate) != num_blocks:
+            raise ValueError(
+                f"Length of attn_branch_drop_rate ({len(attn_branch_drop_rate)}) "
+                f"should be equal to num_blocks ({num_blocks})")
+
+        self.encoders = LayerDropModuleList(
+            p=stochastic_depth_rate,
+            modules=[
+                BranchformerEncoderLayer(
+                    output_size,
+                    WENET_ATTENTION_CLASSES[selfattention_layer_type](
+                        *encoder_selfattn_layer_args) if use_attn else None,
+                    cgmlp_layer(*cgmlp_layer_args) if use_cgmlp else None,
+                    dropout_rate,
+                    merge_method,
+                    cgmlp_weight[lnum],
+                    attn_branch_drop_rate[lnum],
+                    stochastic_depth_rate[lnum],
+                ) for lnum in range(num_blocks)
+            ])
+
+
+# modify from : https://github.com/facebookresearch/fairseq/blob/main/fairseq/modules/layer_drop.py # noqa
+class LayerDropModuleList(torch.nn.ModuleList):
+    """
+    A LayerDrop implementation based on :class:`torch.nn.ModuleList`.
+
+    We refresh the choice of which layers to drop every time we iterate
+    over the LayerDropModuleList instance. During evaluation we always
+    iterate over all layers.
+
+    Usage::
+
+        layers = LayerDropList(p=0.5, modules=[layer1, layer2, layer3])
+        for layer in layers:  # this might iterate over layers 1 and 3
+            x = layer(x)
+        for layer in layers:  # this might iterate over all layers
+            x = layer(x)
+        for layer in layers:  # this might not iterate over any layers
+            x = layer(x)
+
+    Args:
+        p (float): probability of dropping out each layer
+        modules (iterable, optional): an iterable of modules to add
+
+    Limitations:
+        1 can work with ddp when layer's gradient checkpoint disabled
+        2 can't work with ddp when layer's gradient checkpoint enables
+        3 can work with fsdp
+        4 can work with deepspeed
+    """
+
+    def __init__(self, p: List[float], modules=None):
+        super().__init__(modules)
+        assert len(p) == len(self)
+        self.p = p
+
+    def __iter__(self):
+        dropout_probs = torch.empty(len(self)).uniform_()
+        for i, m in enumerate(super().__iter__()):
+            if not self.training or (dropout_probs[i] > self.p[i]):
+                yield m

wenet/branchformer/encoder_layer.py

@@ -0,0 +1,245 @@
+# Copyright (c) 2022 Yifan Peng (Carnegie Mellon University)
+#               2023 Voicecomm Inc (Kai Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+"""BranchformerEncoderLayer definition."""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple
+
+from wenet.transformer.attention import T_CACHE
+
+
+class BranchformerEncoderLayer(torch.nn.Module):
+    """Branchformer encoder layer module.
+
+    Args:
+        size (int): model dimension
+        attn: standard self-attention or efficient attention, optional
+        cgmlp: ConvolutionalGatingMLP, optional
+        dropout_rate (float): dropout probability
+        merge_method (str): concat, learned_ave, fixed_ave
+        cgmlp_weight (float): weight of the cgmlp branch, between 0 and 1,
+            used if merge_method is fixed_ave
+        attn_branch_drop_rate (float): probability of dropping the attn branch,
+            used if merge_method is learned_ave
+        stochastic_depth_rate (float): stochastic depth probability
+    """
+
+    def __init__(
+        self,
+        size: int,
+        attn: Optional[torch.nn.Module],
+        cgmlp: Optional[torch.nn.Module],
+        dropout_rate: float,
+        merge_method: str,
+        cgmlp_weight: float = 0.5,
+        attn_branch_drop_rate: float = 0.0,
+        stochastic_depth_rate: float = 0.0,
+    ):
+        super().__init__()
+        assert (attn is not None) or (
+            cgmlp is not None), "At least one branch should be valid"
+
+        self.size = size
+        self.attn = attn
+        self.cgmlp = cgmlp
+        self.merge_method = merge_method
+        self.cgmlp_weight = cgmlp_weight
+        self.attn_branch_drop_rate = attn_branch_drop_rate
+        self.stochastic_depth_rate = stochastic_depth_rate
+        self.use_two_branches = (attn is not None) and (cgmlp is not None)
+
+        if attn is not None:
+            self.norm_mha = nn.LayerNorm(size)  # for the MHA module
+        if cgmlp is not None:
+            self.norm_mlp = nn.LayerNorm(size)  # for the MLP module
+        self.norm_final = nn.LayerNorm(
+            size)  # for the final output of the block
+
+        self.dropout = torch.nn.Dropout(dropout_rate)
+
+        # # attention-based pooling for two branches
+        self.pooling_proj1 = torch.nn.Linear(size, 1)
+        self.pooling_proj2 = torch.nn.Linear(size, 1)
+
+        # # linear projections for calculating merging weights
+        self.weight_proj1 = torch.nn.Linear(size, 1)
+        self.weight_proj2 = torch.nn.Linear(size, 1)
+
+        if self.use_two_branches:
+            if self.merge_method == "concat":
+                self.merge_proj = torch.nn.Linear(size + size, size)
+
+            elif self.merge_method == "learned_ave":
+                # linear projection after weighted average
+                self.merge_proj = torch.nn.Linear(size, size)
+
+            elif self.merge_method == "fixed_ave":
+                assert (0.0 <= cgmlp_weight <=
+                        1.0), "cgmlp weight should be between 0.0 and 1.0"
+
+                # remove the other branch if only one branch is used
+                if cgmlp_weight == 0.0:
+                    self.use_two_branches = False
+                    self.cgmlp = None
+                    self.norm_mlp = None
+                elif cgmlp_weight == 1.0:
+                    self.use_two_branches = False
+                    self.attn = None
+                    self.norm_mha = None
+
+                # linear projection after weighted average
+                self.merge_proj = torch.nn.Linear(size, size)
+            else:
+                raise ValueError(f"unknown merge method: {merge_method}")
+        else:
+            self.merge_proj = torch.nn.Identity()
+
+    def _forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        pos_emb: torch.Tensor,
+        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        att_cache: T_CACHE = (torch.zeros(
+            (0, 0, 0, 0)), torch.zeros(0, 0, 0, 0)),
+        cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+        stoch_layer_coeff: float = 1.0
+    ) -> Tuple[torch.Tensor, torch.Tensor, T_CACHE, torch.Tensor]:
+        # Two branches
+        x1 = x
+        x2 = x
+
+        # Branch 1: multi-headed attention module
+        if self.attn is not None:
+            x1 = self.norm_mha(x1)
+            x_att, new_att_cache = self.attn(x1, x1, x1, mask, pos_emb,
+                                             att_cache)
+            x1 = self.dropout(x_att)
+
+        # Branch 2: convolutional gating mlp
+        # Fake new cnn cache here, and then change it in conv_module
+        new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
+        if self.cgmlp is not None:
+            x2 = self.norm_mlp(x2)
+            x2, new_cnn_cache = self.cgmlp(x2, mask_pad, cnn_cache)
+            x2 = self.dropout(x2)
+
+        # Merge two branches
+        if self.use_two_branches:
+            if self.merge_method == "concat":
+                x = x + stoch_layer_coeff * self.dropout(
+                    self.merge_proj(torch.cat([x1, x2], dim=-1)))
+            elif self.merge_method == "learned_ave":
+                if (self.training and self.attn_branch_drop_rate > 0
+                        and torch.rand(1).item() < self.attn_branch_drop_rate):
+                    # Drop the attn branch
+                    w1, w2 = torch.tensor(0.0), torch.tensor(1.0)
+                else:
+                    # branch1
+                    score1 = (self.pooling_proj1(x1).transpose(1, 2) /
+                              self.size**0.5)
+                    score1 = score1.masked_fill(mask_pad.eq(0), -float('inf'))
+                    score1 = torch.softmax(score1, dim=-1).masked_fill(
+                        mask_pad.eq(0), 0.0)
+
+                    pooled1 = torch.matmul(score1,
+                                           x1).squeeze(1)  # (batch, size)
+                    weight1 = self.weight_proj1(pooled1)  # (batch, 1)
+
+                    # branch2
+                    score2 = (self.pooling_proj2(x2).transpose(1, 2) /
+                              self.size**0.5)
+                    score2 = score2.masked_fill(mask_pad.eq(0), -float('inf'))
+                    score2 = torch.softmax(score2, dim=-1).masked_fill(
+                        mask_pad.eq(0), 0.0)
+
+                    pooled2 = torch.matmul(score2,
+                                           x2).squeeze(1)  # (batch, size)
+                    weight2 = self.weight_proj2(pooled2)  # (batch, 1)
+
+                    # normalize weights of two branches
+                    merge_weights = torch.softmax(torch.cat([weight1, weight2],
+                                                            dim=-1),
+                                                  dim=-1)  # (batch, 2)
+                    merge_weights = merge_weights.unsqueeze(-1).unsqueeze(
+                        -1)  # (batch, 2, 1, 1)
+                    w1, w2 = merge_weights[:,
+                                           0], merge_weights[:,
+                                                             1]  # (batch, 1, 1)
+
+                x = x + stoch_layer_coeff * self.dropout(
+                    self.merge_proj(w1 * x1 + w2 * x2))
+            elif self.merge_method == "fixed_ave":
+                x = x + stoch_layer_coeff * self.dropout(
+                    self.merge_proj((1.0 - self.cgmlp_weight) * x1 +
+                                    self.cgmlp_weight * x2))
+            else:
+                raise RuntimeError(
+                    f"unknown merge method: {self.merge_method}")
+        else:
+            if self.attn is None:
+                x = x + stoch_layer_coeff * self.dropout(self.merge_proj(x2))
+            elif self.cgmlp is None:
+                x = x + stoch_layer_coeff * self.dropout(self.merge_proj(x1))
+            else:
+                # This should not happen
+                raise RuntimeError(
+                    "Both branches are not None, which is unexpected.")
+
+        x = self.norm_final(x)
+
+        return x, mask, new_att_cache, new_cnn_cache
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        pos_emb: torch.Tensor,
+        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        att_cache: T_CACHE = (torch.zeros(
+            (0, 0, 0, 0)), torch.zeros(0, 0, 0, 0)),
+        cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+    ) -> Tuple[torch.Tensor, torch.Tensor, T_CACHE, torch.Tensor]:
+        """Compute encoded features.
+
+        Args:
+            x (Union[Tuple, torch.Tensor]): Input tensor  (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time, time).
+            pos_emb (torch.Tensor): positional encoding, must not be None
+                for BranchformerEncoderLayer.
+            mask_pad (torch.Tensor): batch padding mask used for conv module.
+                (#batch, 1，time), (0, 0, 0) means fake mask.
+            att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
+                (#batch=1, head, cache_t1, d_k * 2), head * d_k == size.
+            cnn_cache (torch.Tensor): Convolution cache in cgmlp layer
+                (#batch=1, size, cache_t2)
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time, time.
+            torch.Tensor: att_cache tensor,
+                (#batch=1, head, cache_t1 + time, d_k * 2).
+            torch.Tensor: cnn_cahce tensor (#batch, size, cache_t2).
+        """
+
+        stoch_layer_coeff = 1.0
+        # with stochastic depth, residual connection `x + f(x)` becomes
+        # `x <- x + 1 / (1 - p) * f(x)` at training time.
+        if self.training:
+            stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate)
+        return self._forward(x, mask, pos_emb, mask_pad, att_cache, cnn_cache,
+                             stoch_layer_coeff)

wenet/cli/hub.py

@@ -0,0 +1,118 @@
+# Copyright (c) 2022  Mddct(hamddct@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+import tarfile
+from pathlib import Path
+from urllib.request import urlretrieve
+
+import requests
+import tqdm
+
+
+def download(url: str, dest: str, only_child=True):
+    """ download from url to dest
+    """
+    assert os.path.exists(dest)
+    print('Downloading {} to {}'.format(url, dest))
+
+    def progress_hook(t):
+        last_b = [0]
+
+        def update_to(b=1, bsize=1, tsize=None):
+            if tsize not in (None, -1):
+                t.total = tsize
+            displayed = t.update((b - last_b[0]) * bsize)
+            last_b[0] = b
+            return displayed
+
+        return update_to
+
+    # *.tar.gz
+    name = url.split('?')[0].split('/')[-1]
+    tar_path = os.path.join(dest, name)
+    with tqdm.tqdm(unit='B',
+                   unit_scale=True,
+                   unit_divisor=1024,
+                   miniters=1,
+                   desc=(name)) as t:
+        urlretrieve(url,
+                    filename=tar_path,
+                    reporthook=progress_hook(t),
+                    data=None)
+        t.total = t.n
+
+    with tarfile.open(tar_path) as f:
+        if not only_child:
+            f.extractall(dest)
+        else:
+            for tarinfo in f:
+                if "/" not in tarinfo.name:
+                    continue
+                name = os.path.basename(tarinfo.name)
+                fileobj = f.extractfile(tarinfo)
+                with open(os.path.join(dest, name), "wb") as writer:
+                    writer.write(fileobj.read())
+
+
+class Hub(object):
+    """Hub for wenet pretrain runtime model
+    """
+    # TODO(Mddct): make assets class to support other language
+    Assets = {
+        # wenetspeech
+        "chinese": "wenetspeech_u2pp_conformer_libtorch.tar.gz",
+        # gigaspeech
+        "english": "gigaspeech_u2pp_conformer_libtorch.tar.gz",
+        # paraformer
+        "paraformer": "paraformer.tar.gz",
+        # punc
+        "punc": "punc.tar.gz"
+    }
+
+    def __init__(self) -> None:
+        pass
+
+    @staticmethod
+    def get_model_by_lang(lang: str) -> str:
+        if lang not in Hub.Assets.keys():
+            print('ERROR: Unsupported language {} !!!'.format(lang))
+            sys.exit(1)
+
+        # NOTE(Mddct): model_dir structure
+        # Path.Home()/.wenet
+        # - chs
+        #    - units.txt
+        #    - final.zip
+        # - en
+        #    - units.txt
+        #    - final.zip
+        model = Hub.Assets[lang]
+        model_dir = os.path.join(Path.home(), ".wenet", lang)
+        if not os.path.exists(model_dir):
+            os.makedirs(model_dir)
+        # TODO(Mddct): model metadata
+        if set(["final.zip",
+                "units.txt"]).issubset(set(os.listdir(model_dir))):
+            return model_dir
+        # If not exist, download
+        response = requests.get(
+            "https://modelscope.cn/api/v1/datasets/wenet/wenet_pretrained_models/oss/tree"  # noqa
+        )
+        model_info = next(data for data in response.json()["Data"]
+                          if data["Key"] == model)
+        model_url = model_info['Url']
+        download(model_url, model_dir, only_child=True)
+        return model_dir

wenet/cli/model.py

@@ -0,0 +1,209 @@
+# Copyright (c) 2023 Binbin Zhang (binbzha@qq.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import torch
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+import yaml
+
+from wenet.cli.hub import Hub
+from wenet.utils.ctc_utils import (force_align, gen_ctc_peak_time,
+                                   gen_timestamps_from_peak)
+from wenet.utils.file_utils import read_symbol_table
+from wenet.utils.init_model import init_model
+from wenet.transformer.search import (attention_rescoring,
+                                      ctc_prefix_beam_search, DecodeResult)
+from wenet.utils.context_graph import ContextGraph
+from wenet.utils.common import TORCH_NPU_AVAILABLE  # noqa just ensure to check torch-npu
+
+
+class Model:
+
+    def __init__(self,
+                 model_dir: str,
+                 gpu: int = -1,
+                 beam: int = 5,
+                 context_path: str = None,
+                 context_score: float = 6.0,
+                 resample_rate: int = 16000):
+        model_path = os.path.join(model_dir, 'final.zip')
+        units_path = os.path.join(model_dir, 'units.txt')
+        self.model = torch.jit.load(model_path)
+        self.resample_rate = resample_rate
+        self.model.eval()
+        if gpu >= 0:
+            device = 'cuda:{}'.format(gpu)
+        else:
+            device = 'cpu'
+        self.device = torch.device(device)
+        self.model.to(device)
+        self.symbol_table = read_symbol_table(units_path)
+        self.char_dict = {v: k for k, v in self.symbol_table.items()}
+        self.beam = beam
+        if context_path is not None:
+            self.context_graph = ContextGraph(context_path,
+                                              self.symbol_table,
+                                              context_score=context_score)
+        else:
+            self.context_graph = None
+
+    def compute_feats(self, audio_file: str) -> torch.Tensor:
+        waveform, sample_rate = torchaudio.load(audio_file, normalize=False)
+        waveform = waveform.to(torch.float)
+        if sample_rate != self.resample_rate:
+            waveform = torchaudio.transforms.Resample(
+                orig_freq=sample_rate, new_freq=self.resample_rate)(waveform)
+        # NOTE (MengqingCao): complex dtype not supported in torch_npu.abs() now,
+        # thus, delay placing data on NPU after the calculation of fbank.
+        # revert me after complex dtype is supported.
+        if "npu" not in self.device.__str__():
+            waveform = waveform.to(self.device)
+        feats = kaldi.fbank(waveform,
+                            num_mel_bins=80,
+                            frame_length=25,
+                            frame_shift=10,
+                            energy_floor=0.0,
+                            sample_frequency=self.resample_rate)
+        if "npu" in self.device.__str__():
+            feats = feats.to(self.device)
+        feats = feats.unsqueeze(0)
+        return feats
+
+    @torch.no_grad()
+    def _decode(self,
+                audio_file: str,
+                tokens_info: bool = False,
+                label: str = None) -> dict:
+        feats = self.compute_feats(audio_file)
+        encoder_out, _, _ = self.model.forward_encoder_chunk(feats, 0, -1)
+        encoder_lens = torch.tensor([encoder_out.size(1)],
+                                    dtype=torch.long,
+                                    device=encoder_out.device)
+        ctc_probs = self.model.ctc_activation(encoder_out)
+        if label is None:
+            ctc_prefix_results = ctc_prefix_beam_search(
+                ctc_probs,
+                encoder_lens,
+                self.beam,
+                context_graph=self.context_graph)
+        else:  # force align mode, construct ctc prefix result from alignment
+            label_t = self.tokenize(label)
+            alignment = force_align(ctc_probs.squeeze(0),
+                                    torch.tensor(label_t, dtype=torch.long))
+            peaks = gen_ctc_peak_time(alignment)
+            ctc_prefix_results = [
+                DecodeResult(tokens=label_t,
+                             score=0.0,
+                             times=peaks,
+                             nbest=[label_t],
+                             nbest_scores=[0.0],
+                             nbest_times=[peaks])
+            ]
+        rescoring_results = attention_rescoring(self.model, ctc_prefix_results,
+                                                encoder_out, encoder_lens, 0.3,
+                                                0.5)
+        res = rescoring_results[0]
+        result = {}
+        result['text'] = ''.join([self.char_dict[x] for x in res.tokens])
+        result['confidence'] = res.confidence
+
+        if tokens_info:
+            frame_rate = self.model.subsampling_rate(
+            ) * 0.01  # 0.01 seconds per frame
+            max_duration = encoder_out.size(1) * frame_rate
+            times = gen_timestamps_from_peak(res.times, max_duration,
+                                             frame_rate, 1.0)
+            tokens_info = []
+            for i, x in enumerate(res.tokens):
+                tokens_info.append({
+                    'token': self.char_dict[x],
+                    'start': round(times[i][0], 3),
+                    'end': round(times[i][1], 3),
+                    'confidence': round(res.tokens_confidence[i], 2)
+                })
+            result['tokens'] = tokens_info
+        return result
+
+    def transcribe(self, audio_file: str, tokens_info: bool = False) -> dict:
+        return self._decode(audio_file, tokens_info)
+
+    def tokenize(self, label: str):
+        # TODO(Binbin Zhang): Support BPE
+        tokens = []
+        for c in label:
+            if c == ' ':
+                c = "▁"
+            tokens.append(c)
+        token_list = []
+        for c in tokens:
+            if c in self.symbol_table:
+                token_list.append(self.symbol_table[c])
+            elif '<unk>' in self.symbol_table:
+                token_list.append(self.symbol_table['<unk>'])
+        return token_list
+
+    def align(self, audio_file: str, label: str) -> dict:
+        return self._decode(audio_file, True, label)
+
+
+def load_model(language: str = None,
+               model_dir: str = None,
+               gpu: int = -1,
+               beam: int = 5,
+               context_path: str = None,
+               context_score: float = 6.0,
+               device: str = "cpu") -> Model:
+    if model_dir is None:
+        model_dir = Hub.get_model_by_lang(language)
+
+    if gpu != -1:
+        # remain the original usage of gpu
+        device = "cuda"
+    model = Model(model_dir, gpu, beam, context_path, context_score)
+    model.device = torch.device(device)
+    model.model.to(device)
+    return model
+
+# Load the pytorch pt model which contains all the details compared with jit.
+# And we can use the pt model as a third party pytorch nn.Module for training
+def load_model_pt(model_dir):
+    """ There are the followi files in in `model_dir`
+        * final.pt, required
+        * train.yaml, required
+        * units.txt, required
+        * global_cmvn, optional
+    """
+    # Check required files
+    required_files = ['train.yaml', 'final.pt', 'units.txt']
+    for file in required_files:
+        file_path = os.path.join(model_dir, file)
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(
+                f"Required file {file} not found in {model_dir}")
+    # Read config and override some config
+    config_file = os.path.join(model_dir, 'train.yaml')
+    with open(config_file, 'r') as fin:
+        configs = yaml.load(fin, Loader=yaml.FullLoader)
+    token_file = os.path.join(model_dir, 'units.txt')
+    configs['tokenizer_conf']['symbol_table_path'] = token_file
+    cmvn_file = os.path.join(model_dir, 'global_cmvn')
+    if os.path.exists(cmvn_file):
+        configs['cmvn_conf']['cmvn_file'] = cmvn_file
+    # Read model
+    pt_file = os.path.join(model_dir, 'final.pt')
+    args = {'checkpoint': pt_file}
+    model, configs = init_model(args, configs)
+    return model

wenet/cli/paraformer_model.py

@@ -0,0 +1,110 @@
+import io
+import os
+from typing import Dict, List, Union
+
+import torch
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+from wenet.cli.hub import Hub
+from wenet.paraformer.search import (gen_timestamps_from_peak,
+                                     paraformer_greedy_search)
+from wenet.text.paraformer_tokenizer import ParaformerTokenizer
+
+
+class Paraformer:
+
+    def __init__(self, model_dir: str, resample_rate: int = 16000) -> None:
+
+        model_path = os.path.join(model_dir, 'final.zip')
+        units_path = os.path.join(model_dir, 'units.txt')
+        self.model = torch.jit.load(model_path)
+        self.resample_rate = resample_rate
+        self.device = torch.device("cpu")
+        self.tokenizer = ParaformerTokenizer(symbol_table=units_path)
+
+    @torch.inference_mode()
+    def transcribe_batch(self,
+                         audio_files: List[Union[str, bytes]],
+                         tokens_info: bool = False) -> List[Dict]:
+        feats_lst = []
+        feats_lens_lst = []
+        for audio in audio_files:
+            if isinstance(audio, bytes):
+                with io.BytesIO(audio) as fobj:
+                    waveform, sample_rate = torchaudio.load(fobj,
+                                                            normalize=False)
+            else:
+                waveform, sample_rate = torchaudio.load(audio, normalize=False)
+            if sample_rate != self.resample_rate:
+                waveform = torchaudio.transforms.Resample(
+                    orig_freq=sample_rate,
+                    new_freq=self.resample_rate)(waveform)
+
+            waveform = waveform.to(torch.float)
+            feats = kaldi.fbank(waveform,
+                                num_mel_bins=80,
+                                frame_length=25,
+                                frame_shift=10,
+                                energy_floor=0.0,
+                                sample_frequency=self.resample_rate,
+                                window_type="hamming")
+            feats_lst.append(feats)
+            feats_lens_lst.append(
+                torch.tensor(feats.shape[0], dtype=torch.int64))
+        feats_tensor = torch.nn.utils.rnn.pad_sequence(
+            feats_lst, batch_first=True).to(device=self.device)
+        feats_lens_tensor = torch.tensor(feats_lens_lst, device=self.device)
+
+        decoder_out, token_num, tp_alphas, frames = self.model.forward_paraformer(
+            feats_tensor, feats_lens_tensor)
+        frames = frames.cpu().numpy()
+        cif_peaks = self.model.forward_cif_peaks(tp_alphas, token_num)
+
+        results = paraformer_greedy_search(decoder_out, token_num, cif_peaks)
+
+        r = []
+        for (i, res) in enumerate(results):
+            result = {}
+            result['confidence'] = res.confidence
+            result['text'] = self.tokenizer.detokenize(res.tokens)[0]
+            if tokens_info:
+                tokens_info_l = []
+                times = gen_timestamps_from_peak(res.times,
+                                                 num_frames=frames[i],
+                                                 frame_rate=0.02)
+
+                for i, x in enumerate(res.tokens[:len(times)]):
+                    tokens_info_l.append({
+                        'token':
+                        self.tokenizer.char_dict[x],
+                        'start':
+                        round(times[i][0], 3),
+                        'end':
+                        round(times[i][1], 3),
+                        'confidence':
+                        round(res.tokens_confidence[i], 2)
+                    })
+                    result['tokens'] = tokens_info_l
+            r.append(result)
+        return r
+
+    def transcribe(self, audio_file: str, tokens_info: bool = False) -> dict:
+        result = self.transcribe_batch([audio_file], tokens_info)[0]
+        return result
+
+    def align(self, audio_file: str, label: str) -> dict:
+        raise NotImplementedError("Align is currently not supported")
+
+
+def load_model(model_dir: str = None,
+               gpu: int = -1,
+               device: str = "cpu") -> Paraformer:
+    if model_dir is None:
+        model_dir = Hub.get_model_by_lang('paraformer')
+    if gpu != -1:
+        # remain the original usage of gpu
+        device = "cuda"
+    paraformer = Paraformer(model_dir)
+    paraformer.device = torch.device(device)
+    paraformer.model.to(device)
+    return paraformer

wenet/cli/punc_model.py

@@ -0,0 +1,115 @@
+import os
+from typing import List
+
+import jieba
+import torch
+from wenet.cli.hub import Hub
+from wenet.paraformer.search import _isAllAlpha
+from wenet.text.char_tokenizer import CharTokenizer
+
+
+class PuncModel:
+
+    def __init__(self, model_dir: str) -> None:
+        self.model_dir = model_dir
+        model_path = os.path.join(model_dir, 'final.zip')
+        units_path = os.path.join(model_dir, 'units.txt')
+
+        self.model = torch.jit.load(model_path)
+        self.tokenizer = CharTokenizer(units_path)
+        self.device = torch.device("cpu")
+        self.use_jieba = False
+
+        self.punc_table = ['<unk>', '', '，', '。', '？', '、']
+
+    def split_words(self, text: str):
+        if not self.use_jieba:
+            self.use_jieba = True
+            import logging
+
+            # Disable jieba's logger
+            logging.getLogger('jieba').disabled = True
+            jieba.load_userdict(os.path.join(self.model_dir, 'jieba_usr_dict'))
+
+        result_list = []
+        tokens = text.split()
+        current_language = None
+        buffer = []
+
+        for token in tokens:
+            is_english = token.isascii()
+            if is_english:
+                language = "English"
+            else:
+                language = "Chinese"
+
+            if current_language and language != current_language:
+                if current_language == "Chinese":
+                    result_list.extend(jieba.cut(''.join(buffer), HMM=False))
+                else:
+                    result_list.extend(buffer)
+                buffer = []
+
+            buffer.append(token)
+            current_language = language
+
+        if buffer:
+            if current_language == "Chinese":
+                result_list.extend(jieba.cut(''.join(buffer), HMM=False))
+            else:
+                result_list.extend(buffer)
+
+        return result_list
+
+    def add_punc_batch(self, texts: List[str]):
+        batch_text_words = []
+        batch_text_ids = []
+        batch_text_lens = []
+
+        for text in texts:
+            words = self.split_words(text)
+            ids = self.tokenizer.tokens2ids(words)
+            batch_text_words.append(words)
+            batch_text_ids.append(ids)
+            batch_text_lens.append(len(ids))
+
+        texts_tensor = torch.tensor(batch_text_ids,
+                                    device=self.device,
+                                    dtype=torch.int64)
+        texts_lens_tensor = torch.tensor(batch_text_lens,
+                                         device=self.device,
+                                         dtype=torch.int64)
+
+        log_probs, _ = self.model(texts_tensor, texts_lens_tensor)
+        result = []
+        outs = log_probs.argmax(-1).cpu().numpy()
+        for i, out in enumerate(outs):
+            punc_id = out[:batch_text_lens[i]]
+            sentence = ''
+            for j, word in enumerate(batch_text_words[i]):
+                if _isAllAlpha(word):
+                    word = '▁' + word
+                word += self.punc_table[punc_id[j]]
+                sentence += word
+            result.append(sentence.replace('▁', ' '))
+        return result
+
+    def __call__(self, text: str):
+        if text != '':
+            r = self.add_punc_batch([text])[0]
+            return r
+        return ''
+
+
+def load_model(model_dir: str = None,
+               gpu: int = -1,
+               device: str = "cpu") -> PuncModel:
+    if model_dir is None:
+        model_dir = Hub.get_model_by_lang('punc')
+    if gpu != -1:
+        # remain the original usage of gpu
+        device = "cuda"
+    punc = PuncModel(model_dir)
+    punc.device = torch.device(device)
+    punc.model.to(device)
+    return punc

wenet/cli/transcribe.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2023 Binbin Zhang (binbzha@qq.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+from wenet.cli.model import load_model
+from wenet.cli.paraformer_model import load_model as load_paraformer
+from wenet.cli.punc_model import load_model as load_punc_model
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('audio_file', help='audio file to transcribe')
+    parser.add_argument('-l',
+                        '--language',
+                        choices=[
+                            'chinese',
+                            'english',
+                        ],
+                        default='chinese',
+                        help='language type')
+    parser.add_argument('-m',
+                        '--model_dir',
+                        default=None,
+                        help='specify your own model dir')
+    parser.add_argument('-g',
+                        '--gpu',
+                        type=int,
+                        default='-1',
+                        help='gpu id to decode, default is cpu.')
+    parser.add_argument('--device',
+                        type=str,
+                        default='cpu',
+                        choices=["cpu", "npu", "cuda"],
+                        help='accelerator to use')
+    parser.add_argument('-t',
+                        '--show_tokens_info',
+                        action='store_true',
+                        help='whether to output token(word) level information'
+                        ', such times/confidence')
+    parser.add_argument('--align',
+                        action='store_true',
+                        help='force align the input audio and transcript')
+    parser.add_argument('--label', type=str, help='the input label to align')
+    parser.add_argument('--paraformer',
+                        action='store_true',
+                        help='whether to use the best chinese model')
+    parser.add_argument('--beam', type=int, default=5, help="beam size")
+    parser.add_argument('--context_path',
+                        type=str,
+                        default=None,
+                        help='context list file')
+    parser.add_argument('--context_score',
+                        type=float,
+                        default=6.0,
+                        help='context score')
+    parser.add_argument('--punc', action='store_true', help='context score')
+
+    parser.add_argument('-pm',
+                        '--punc_model_dir',
+                        default=None,
+                        help='specify your own punc model dir')
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    if args.paraformer:
+        model = load_paraformer(args.model_dir, args.gpu, args.device)
+    else:
+        model = load_model(args.language, args.model_dir, args.gpu, args.beam,
+                           args.context_path, args.context_score, args.device)
+    punc_model = None
+    if args.punc:
+        punc_model = load_punc_model(args.punc_model_dir, args.gpu,
+                                     args.device)
+    if args.align:
+        result = model.align(args.audio_file, args.label)
+    else:
+        result = model.transcribe(args.audio_file, args.show_tokens_info)
+        if args.punc:
+            assert punc_model is not None
+            result['text_with_punc'] = punc_model(result['text'])
+    print(result)
+
+
+if __name__ == "__main__":
+    main()

wenet/ctl_model/asr_model_ctl.py

@@ -0,0 +1,277 @@
+# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu)
+#               2023 NetEase Inc. (authors: Yuting Yang)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet) and
+# fairseq(https://github.com/facebookresearch/fairseq)
+
+from typing import Dict, Optional
+
+import torch
+import torch.nn.functional as F
+from wenet.transformer.ctc import CTC
+from wenet.transformer.decoder import TransformerDecoder
+from wenet.ctl_model.encoder import TransformerEncoder
+from wenet.transformer.asr_model import ASRModel
+from wenet.utils.common import IGNORE_ID
+
+
+class CTLModel(ASRModel):
+    """
+        Implementation of Interspeecch 2023 paper:
+        'Enhancing the Unified Streaming and Non-streaming Model
+         with Contrastive Learning'
+        https://arxiv.org/abs/2306.00755
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        encoder: TransformerEncoder,
+        decoder: TransformerDecoder,
+        ctc: CTC,
+        ctc_weight: float = 0.5,
+        ignore_id: int = IGNORE_ID,
+        reverse_weight: float = 0.0,
+        lsm_weight: float = 0.0,
+        length_normalized_loss: bool = False,
+        logit_temp: float = 0.1,
+        n_negatives: int = 0,
+        ctl_weight: float = 1,
+        special_tokens: dict = None,
+    ):
+        assert 0.0 <= ctc_weight <= 1.0, ctc_weight
+        super().__init__(vocab_size,
+                         encoder,
+                         decoder,
+                         ctc,
+                         ctc_weight,
+                         ignore_id,
+                         reverse_weight,
+                         lsm_weight,
+                         length_normalized_loss,
+                         special_tokens=special_tokens)
+
+        # For CTL Loss
+        self.n_negatives = n_negatives
+        self.ctl_weight = ctl_weight
+        self.logit_temp = logit_temp
+
+    @torch.jit.unused
+    def forward(
+        self,
+        batch: dict,
+        device: torch.device,
+    ) -> Dict[str, Optional[torch.Tensor]]:
+
+        speech = batch['feats'].to(device)
+        speech_lengths = batch['feats_lengths'].to(device)
+        text = batch['target'].to(device)
+        text_lengths = batch['target_lengths'].to(device)
+        loss_full, encoder_out_full, _, _ = self.forward_full(
+            speech, speech_lengths, text, text_lengths)
+        loss_chunk, encoder_out, lens_chunk, encoder_mask = self.forward_chunk(
+            speech, speech_lengths, text, text_lengths)
+
+        ctl_loss = 0.0
+        if self.ctl_weight > 0 and self.n_negatives > 0:
+            num = encoder_out_full.size(1)
+            targets = encoder_out_full
+            src = encoder_out
+            negs, negs_idxs = self.sample_negatives(targets,
+                                                    targets.size(1),
+                                                    speech_lengths=lens_chunk)
+            ctl_loss = self.CTL(src, targets, negs, encoder_mask)
+
+        loss = loss_full + loss_chunk + self.ctl_weight * ctl_loss
+        return {
+            "loss": loss,
+            "loss_full": loss_full,
+            "loss_chunk": loss_chunk,
+            "loss_ctl": ctl_loss
+        }
+
+    def forward_full(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+    ):
+        """Full context mode
+        Frontend + Encoder + Decoder + Calc loss
+
+        Args:
+            speech: (Batch, Length, ...)
+            speech_lengths: (Batch, )
+            text: (Batch, Length)
+            text_lengths: (Batch,)
+        """
+
+        assert text_lengths.dim() == 1, text_lengths.shape
+        # Check that batch_size is unified
+        assert (speech.shape[0] == speech_lengths.shape[0] == text.shape[0] ==
+                text_lengths.shape[0]), (speech.shape, speech_lengths.shape,
+                                         text.shape, text_lengths.shape)
+        # 1. Encoder
+        encoder_out, encoder_mask = self.encoder.forward_full(
+            speech, speech_lengths)
+        encoder_out_lens = encoder_mask.squeeze(1).sum(1)
+
+        # 2a. Attention-decoder branch
+        if self.ctc_weight != 1.0:
+            loss_att, acc_att = self._calc_att_loss(encoder_out, encoder_mask,
+                                                    text, text_lengths)
+        else:
+            loss_att = None
+
+        # 2b. CTC branch
+        if self.ctc_weight != 0.0:
+            loss_ctc = self.ctc(encoder_out, encoder_out_lens, text,
+                                text_lengths)
+        else:
+            loss_ctc = None
+
+        if loss_ctc is None:
+            loss = loss_att
+        elif loss_att is None:
+            loss = loss_ctc
+        else:
+            loss = self.ctc_weight * loss_ctc[0] + (1 -
+                                                    self.ctc_weight) * loss_att
+        return loss, encoder_out, encoder_out_lens, encoder_mask
+
+    def forward_chunk(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor,
+        text: torch.Tensor,
+        text_lengths: torch.Tensor,
+    ):
+        """Chunk-based context mode
+        Frontend + Encoder + Decoder + Calc loss
+
+        Args:
+            speech: (Batch, Length, ...)
+            speech_lengths: (Batch, )
+            text: (Batch, Length)
+            text_lengths: (Batch,)
+        """
+
+        assert text_lengths.dim() == 1, text_lengths.shape
+        # Check that batch_size is unified
+        assert (speech.shape[0] == speech_lengths.shape[0] == text.shape[0] ==
+                text_lengths.shape[0]), (speech.shape, speech_lengths.shape,
+                                         text.shape, text_lengths.shape)
+        # 1. Encoder
+        encoder_out, encoder_mask = self.encoder(speech, speech_lengths)
+        encoder_out_lens = encoder_mask.squeeze(1).sum(1)
+
+        # 2a. Attention-decoder branch
+        if self.ctc_weight != 1.0:
+            loss_att, acc_att = self._calc_att_loss(encoder_out, encoder_mask,
+                                                    text, text_lengths)
+        else:
+            loss_att = None
+
+        # 2b. CTC branch
+        if self.ctc_weight != 0.0:
+            loss_ctc = self.ctc(encoder_out, encoder_out_lens, text,
+                                text_lengths)
+        else:
+            loss_ctc = None
+
+        if loss_ctc is None:
+            loss = loss_att
+        elif loss_att is None:
+            loss = loss_ctc
+        else:
+            loss = self.ctc_weight * loss_ctc[0] + (1 -
+                                                    self.ctc_weight) * loss_att
+        return loss, encoder_out, encoder_out_lens, encoder_mask
+
+    def sample_negatives(self, y, num, padding_count=0, speech_lengths=None):
+        if self.n_negatives == 0:
+            return y.new(0)
+        bsz, tsz, fsz = y.shape
+        y = y.reshape(-1, fsz)  # BTC => (BxT)C
+
+        # FIXME: what happens if padding_count is specified?
+        high = tsz - (padding_count or 0)
+        with torch.no_grad():
+            assert high > 1, f"{bsz,tsz,fsz}"
+
+            if self.n_negatives > 0:
+                tszs = (torch.arange(num).unsqueeze(-1).expand(
+                    -1, self.n_negatives).flatten())
+                if speech_lengths is not None:
+                    neg_idxs = [
+                        torch.randint(low=0,
+                                      high=speech_lengths[i].item() - 1,
+                                      size=(1, self.n_negatives * tsz))
+                        for i in range(len(speech_lengths))
+                    ]
+                    neg_idxs = torch.cat(neg_idxs).reshape(
+                        bsz, self.n_negatives * tsz)
+                else:
+                    neg_idxs = torch.randint(low=0,
+                                             high=num - 1,
+                                             size=(bsz,
+                                                   self.n_negatives * tsz))
+                neg_idxs[neg_idxs >= tszs] += 1
+
+        if self.n_negatives > 0:
+            neg_idxs = neg_idxs + (torch.arange(bsz).unsqueeze(1) * high)
+
+        negs = y[neg_idxs.view(-1)]
+        negs = negs.contiguous().view(bsz, num, self.n_negatives,
+                                      fsz).permute(2, 0, 1, 3)  # to NxBxTxC
+        return negs, neg_idxs
+
+    def compute_preds(self, x, y, negatives):
+        neg_is_pos = (y == negatives).all(-1)
+        y = y.unsqueeze(0)
+        targets = torch.cat([y, negatives], dim=0)
+
+        logits = torch.cosine_similarity(x.float(), targets.float(), dim=-1)
+        logits = logits / self.logit_temp
+        logits = logits.type_as(x)
+
+        if neg_is_pos.any():
+            if not hasattr(self, "_inftensor"):
+                self._inftensor = float("-inf")
+            # logits[1:] = index_put(logits[1:], neg_is_pos, self._inftensor)
+            logits[1:][neg_is_pos] = self._inftensor
+        logits = logits.transpose(0, 2)
+        logits = logits.transpose(0, 1)
+        logits = logits.reshape(-1, logits.size(-1))
+        return logits
+
+    def CTL(self, x, y, negs, mask=None):
+        # Step1: compute cosine similarity, shape [B*T, n_negatives+1]
+        logits = self.compute_preds(x, y, negs)
+
+        # Step2: target shape [B*T]
+        target = x.new_zeros(x.size(0) * x.size(1), dtype=torch.long)
+
+        # Step3: compute CTL loss
+        if mask is not None:
+            normalize_length = mask.sum()
+            bz, sz = mask.size(0), mask.size(-1)
+            mask = mask.squeeze(1).reshape(bz * sz).eq(0)
+            ce = F.cross_entropy(logits, target, reduction='none')
+            loss = ce.masked_fill(mask, 0).sum() / normalize_length
+        else:
+            loss = F.cross_entropy(logits, target)
+
+        return loss

wenet/ctl_model/encoder.py

@@ -0,0 +1,172 @@
+# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)
+#               2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn)
+#               2023 NetEase Inc
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Modified from ESPnet(https://github.com/espnet/espnet)
+"""Encoder definition."""
+from typing import Optional, Tuple
+
+import torch
+
+from wenet.utils.mask import make_pad_mask
+from wenet.transformer.encoder import TransformerEncoder, ConformerEncoder
+
+
+class DualTransformerEncoder(TransformerEncoder):
+    """Transformer encoder module."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 256,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        input_layer: str = "conv2d",
+        pos_enc_layer_type: str = "abs_pos",
+        normalize_before: bool = True,
+        static_chunk_size: int = 0,
+        use_dynamic_chunk: bool = False,
+        global_cmvn: torch.nn.Module = None,
+        use_dynamic_left_chunk: bool = False,
+        query_bias: bool = True,
+        key_bias: bool = True,
+        value_bias: bool = True,
+        activation_type: str = "relu",
+        gradient_checkpointing: bool = False,
+        use_sdpa: bool = False,
+        layer_norm_type: str = 'layer_norm',
+        norm_eps: float = 1e-5,
+        n_kv_head: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        selfattention_layer_type: str = "selfattn",
+        mlp_type: str = 'position_wise_feed_forward',
+        mlp_bias: bool = True,
+        n_expert: int = 8,
+        n_expert_activated: int = 2,
+    ):
+        """ Construct DualTransformerEncoder
+        Support both the full context mode and the streaming mode separately
+        """
+        super().__init__(input_size, output_size, attention_heads,
+                         linear_units, num_blocks, dropout_rate,
+                         positional_dropout_rate, attention_dropout_rate,
+                         input_layer, pos_enc_layer_type, normalize_before,
+                         static_chunk_size, use_dynamic_chunk, global_cmvn,
+                         use_dynamic_left_chunk, query_bias, key_bias,
+                         value_bias, activation_type, gradient_checkpointing,
+                         use_sdpa, layer_norm_type, norm_eps, n_kv_head,
+                         head_dim, selfattention_layer_type, mlp_type,
+                         mlp_bias, n_expert, n_expert_activated)
+
+    def forward_full(
+        self,
+        xs: torch.Tensor,
+        xs_lens: torch.Tensor,
+        decoding_chunk_size: int = 0,
+        num_decoding_left_chunks: int = -1,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        T = xs.size(1)
+        masks = ~make_pad_mask(xs_lens, T).unsqueeze(1)  # (B, 1, T)
+        if self.global_cmvn is not None:
+            xs = self.global_cmvn(xs)
+        xs, pos_emb, masks = self.embed(xs, masks)
+        mask_pad = masks  # (B, 1, T/subsample_rate)
+        for layer in self.encoders:
+            xs, masks, _, _ = layer(xs, masks, pos_emb, mask_pad)
+        if self.normalize_before:
+            xs = self.after_norm(xs)
+        return xs, masks
+
+
+class DualConformerEncoder(ConformerEncoder):
+    """Conformer encoder module."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 256,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        input_layer: str = "conv2d",
+        pos_enc_layer_type: str = "rel_pos",
+        normalize_before: bool = True,
+        static_chunk_size: int = 0,
+        use_dynamic_chunk: bool = False,
+        global_cmvn: torch.nn.Module = None,
+        use_dynamic_left_chunk: bool = False,
+        positionwise_conv_kernel_size: int = 1,
+        macaron_style: bool = True,
+        selfattention_layer_type: str = "rel_selfattn",
+        activation_type: str = "swish",
+        use_cnn_module: bool = True,
+        cnn_module_kernel: int = 15,
+        causal: bool = False,
+        cnn_module_norm: str = "batch_norm",
+        query_bias: bool = True,
+        key_bias: bool = True,
+        value_bias: bool = True,
+        conv_bias: bool = True,
+        gradient_checkpointing: bool = False,
+        use_sdpa: bool = False,
+        layer_norm_type: str = 'layer_norm',
+        norm_eps: float = 1e-5,
+        n_kv_head: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        mlp_type: str = 'position_wise_feed_forward',
+        mlp_bias: bool = True,
+        n_expert: int = 8,
+        n_expert_activated: int = 2,
+    ):
+        """ Construct DualConformerEncoder
+        Support both the full context mode and the streaming mode separately
+        """
+        super().__init__(
+            input_size, output_size, attention_heads, linear_units, num_blocks,
+            dropout_rate, positional_dropout_rate, attention_dropout_rate,
+            input_layer, pos_enc_layer_type, normalize_before,
+            static_chunk_size, use_dynamic_chunk, global_cmvn,
+            use_dynamic_left_chunk, positionwise_conv_kernel_size,
+            macaron_style, selfattention_layer_type, activation_type,
+            use_cnn_module, cnn_module_kernel, causal, cnn_module_norm,
+            query_bias, key_bias, value_bias, conv_bias,
+            gradient_checkpointing, use_sdpa, layer_norm_type, norm_eps,
+            n_kv_head, head_dim, mlp_type, mlp_bias, n_expert,
+            n_expert_activated)
+
+    def forward_full(
+        self,
+        xs: torch.Tensor,
+        xs_lens: torch.Tensor,
+        decoding_chunk_size: int = 0,
+        num_decoding_left_chunks: int = -1,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        T = xs.size(1)
+        masks = ~make_pad_mask(xs_lens, T).unsqueeze(1)  # (B, 1, T)
+        if self.global_cmvn is not None:
+            xs = self.global_cmvn(xs)
+        xs, pos_emb, masks = self.embed(xs, masks)
+        mask_pad = masks  # (B, 1, T/subsample_rate)
+        for layer in self.encoders:
+            xs, masks, _, _ = layer(xs, masks, pos_emb, mask_pad)
+        if self.normalize_before:
+            xs = self.after_norm(xs)
+        return xs, masks

wenet/dataset/datapipes.py

@@ -0,0 +1,470 @@
+# Copyright (c) 2023 Wenet Community. (authors: Dinghao Zhou)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import collections
+from collections.abc import Callable
+import copy
+import sys
+import tarfile
+import logging
+from typing import List, Optional
+import numpy as np
+import torch
+from torch.utils.data import IterDataPipe, functional_datapipe
+from torch.utils.data import datapipes
+from torch.utils.data.datapipes.iter import Mapper
+from torch.utils.data.datapipes.iter.sharding import (
+    SHARDING_PRIORITIES, ShardingFilterIterDataPipe)
+from torch.utils.data.datapipes.utils.common import _check_unpickable_fn
+
+from wenet.dataset.processor import parse_url
+
+
+@functional_datapipe("map_ignore_error")
+class MapperIgnoreErrorDataPipe(Mapper):
+
+    def __init__(self,
+                 dataset: IterDataPipe,
+                 fn: Callable,
+                 input_col=None,
+                 output_col=None,
+                 log_error: bool = True) -> None:
+        super().__init__(dataset, fn, input_col, output_col)
+        self._iter = None
+        self.log_error = log_error
+
+    def __iter__(self):
+        if self._iter is None:
+            self._iter = iter(self.datapipe)
+
+        while True:
+            try:
+                elem = next(self._iter)
+                yield self._apply_fn(elem)
+            except StopIteration:
+                self._iter = None
+                return
+            except Exception as ex:
+                if self.log_error:
+                    logging.warning(str(ex))
+
+
+@functional_datapipe('bucket_by_sequence_length')
+class BucketBySequenceLengthDataPipe(IterDataPipe):
+
+    def __init__(
+        self,
+        dataset: IterDataPipe,
+        elem_length_func,
+        bucket_boundaries: List[int],
+        bucket_batch_sizes: List[int],
+        wrapper_class=None,
+    ) -> None:
+        super().__init__()
+        _check_unpickable_fn(elem_length_func)
+        assert len(bucket_batch_sizes) == len(bucket_boundaries) + 1
+        self.bucket_batch_sizes = bucket_batch_sizes
+        self.bucket_boundaries = bucket_boundaries + [sys.maxsize]
+        self.elem_length_func = elem_length_func
+
+        self._group_dp = GroupByWindowDataPipe(dataset,
+                                               self._element_to_bucket_id,
+                                               self._window_size_func,
+                                               wrapper_class=wrapper_class)
+
+    def __iter__(self):
+        yield from self._group_dp
+
+    def _element_to_bucket_id(self, elem):
+        seq_len = self.elem_length_func(elem)
+        bucket_id = 0
+        for (i, b) in enumerate(self.bucket_boundaries):
+            if seq_len < b:
+                bucket_id = i
+                break
+        return bucket_id
+
+    def _window_size_func(self, bucket_id):
+        return self.bucket_batch_sizes[bucket_id]
+
+
+@functional_datapipe("group_by_window")
+class GroupByWindowDataPipe(datapipes.iter.Grouper):
+
+    def __init__(
+        self,
+        dataset: IterDataPipe,
+        key_func,
+        window_size_func,
+        wrapper_class=None,
+    ):
+        super().__init__(dataset,
+                         key_func,
+                         keep_key=False,
+                         group_size=None,
+                         drop_remaining=False)
+        _check_unpickable_fn(window_size_func)
+        self.dp = dataset
+        self.window_size_func = window_size_func
+        if wrapper_class is not None:
+            _check_unpickable_fn(wrapper_class)
+            del self.wrapper_class
+            self.wrapper_class = wrapper_class
+
+    def __iter__(self):
+        for x in self.datapipe:
+            key = self.group_key_fn(x)
+
+            self.buffer_elements[key].append(x)
+            self.curr_buffer_size += 1
+
+            group_size = self.window_size_func(key)
+            if group_size == len(self.buffer_elements[key]):
+                result = self.wrapper_class(self.buffer_elements[key])
+                yield result
+                self.curr_buffer_size -= len(self.buffer_elements[key])
+                del self.buffer_elements[key]
+
+            if self.curr_buffer_size == self.max_buffer_size:
+                result_to_yield = self._remove_biggest_key()
+                if result_to_yield is not None:
+                    result = self.wrapper_class(result_to_yield)
+                    yield result
+
+        for key in tuple(self.buffer_elements.keys()):
+            result = self.wrapper_class(self.buffer_elements.pop(key))
+            self.curr_buffer_size -= len(result)
+            yield result
+
+
+@functional_datapipe("sort")
+class SortDataPipe(IterDataPipe):
+
+    def __init__(self,
+                 dataset: IterDataPipe,
+                 buffer_size: int = 500,
+                 key_func=None,
+                 reverse=False) -> None:
+        if key_func is not None:
+            _check_unpickable_fn(key_func)
+        self.buffer_size = buffer_size
+        super().__init__()
+        self.dp = dataset
+        self._buffer = []
+        self.key_func = key_func
+        self.reverse = reverse
+
+    def __iter__(self):
+        for elem in self.dp:
+            self._buffer.append(elem)
+            if len(self._buffer) >= self.buffer_size:
+                self._buffer.sort(key=self.key_func, reverse=self.reverse)
+                for x in self._buffer:
+                    yield x
+                del self._buffer
+                self._buffer = []
+        # The sample left over
+        self._buffer.sort(key=self.key_func, reverse=self.reverse)
+        for x in self._buffer:
+            yield x
+        del self._buffer
+        self._buffer = []
+
+
+@functional_datapipe("dynamic_batch")
+class DynamicBatchDataPipe(IterDataPipe):
+
+    def __init__(self, dataset: IterDataPipe, window_class,
+                 wrapper_class) -> None:
+        _check_unpickable_fn(window_class)
+        _check_unpickable_fn(wrapper_class)
+        super().__init__()
+        self.dp = dataset
+        assert window_class is not None
+        assert wrapper_class is not None
+        self.window_class = window_class
+        self._buffer = []
+        self._wrappr_class = wrapper_class
+
+    def __iter__(self):
+        for elem in self.dp:
+            if not self.window_class(elem, len(self._buffer)):
+                self._buffer.append(elem)
+            else:
+                if len(self._buffer) > 0:
+                    yield self._wrappr_class(self._buffer)
+                del self._buffer
+                self._buffer = [elem]
+        if len(self._buffer) > 0:
+            yield self._wrappr_class(self._buffer)
+        del self._buffer
+        self._buffer = []
+
+
+@functional_datapipe("prefetch")
+class PrefetchDataPipe(IterDataPipe):
+    """Performs prefetching"""
+
+    def __init__(
+        self,
+        dataset: IterDataPipe,
+        buffer_size: int = 500,
+    ):
+        # TODO(Mddct): support multiprocessing pool with shared-memory to
+        #   prefetch
+        super().__init__()
+        self.dp = dataset
+        self._iter = None
+        self._prefetch_buffer_size = buffer_size
+        self._buffer = None
+        if self._prefetch_buffer_size > 0:
+            self._buffer = collections.deque(maxlen=self._prefetch_buffer_size)
+
+    def __iter__(self):
+        if self._prefetch_buffer_size > 0:
+            if self._iter is None:
+                self._iter = iter(self.dp)
+            assert self._buffer is not None
+
+            while True:
+                if len(self._buffer) <= self._prefetch_buffer_size // 2:
+                    while len(self._buffer) < self._prefetch_buffer_size:
+                        try:
+                            self._buffer.append(next(self._iter))
+                        except StopIteration:
+                            if len(self._buffer) != 0:
+                                while len(self._buffer) > 0:
+                                    yield self._buffer.popleft()
+                            self._iter = None
+                            return
+                while len(self._buffer) > self._prefetch_buffer_size // 2:
+                    elem = self._buffer.popleft()
+                    yield elem
+
+        else:
+            yield from self.dp
+
+
+@functional_datapipe("repeat")
+class RepeatDatapipe(IterDataPipe):
+
+    def __init__(self, dataset: IterDataPipe, count: int = -1):
+        super().__init__()
+        self.dp = dataset
+        self.count = count
+
+    def __iter__(self):
+        if self.count == 1:
+            yield from self.dp
+            return
+        i = 0
+        while self.count < 0 or i < self.count:
+            for elem in self.dp:
+                new_elem = copy.copy(elem)
+                yield new_elem
+            i += 1
+
+
+@functional_datapipe("shard")
+class ShardDataPipe(ShardingFilterIterDataPipe):
+
+    def __init__(self, dataset: IterDataPipe, partition: bool = False):
+        super().__init__(dataset, None)
+        self.partition = partition
+        self.dp = dataset
+
+    def apply_sharding(self, num_of_instances: int, instance_id: int,
+                       sharding_group: SHARDING_PRIORITIES):
+        if self.partition:
+            return super().apply_sharding(num_of_instances, instance_id,
+                                          sharding_group)
+        else:
+            # We can not handle uneven data for CV on DDP, so we don't
+            # sample data by rank, that means every GPU gets the same
+            # and all the CV data
+            info = torch.utils.data.get_worker_info()
+            if info is None:
+                self.num_of_instances = 1
+                self.instance_id = 0
+            else:
+                n_workers_per_device = info.num_workers
+                self.num_of_instances = n_workers_per_device
+                self.instance_id = info.id
+
+
+@functional_datapipe("interleave")
+class InterlaveDataPipe(IterDataPipe):
+
+    def __init__(
+        self,
+        source_datapipes: List[IterDataPipe],
+        weights: Optional[List[float]] = None,
+        seed=2027,
+    ):
+        super().__init__()
+        self.rng = np.random.default_rng(seed)
+        self.source_datapipes = source_datapipes
+        self.weights = weights
+        if weights is None:
+            self.weights = [1 / len(self.source_datapipes)] * len(
+                self.source_datapipes)
+        else:
+            self.weights = [weight / sum(weights) for weight in weights]
+        self.iters = None
+
+    def __iter__(self):
+        weights = copy.deepcopy(self.weights)
+        exhausted = len(self.source_datapipes) * [False]
+        if self.iters is None:
+            self.iters = [(i, iter(d))
+                          for i, d in enumerate(self.source_datapipes)]
+        while True:
+            # TODO(Mddct): rng
+            index_iter = self.rng.choice(self.iters, p=weights)
+            i, ite = index_iter
+            try:
+                elem = next(ite)
+                yield elem
+            except StopIteration:
+                weights[i] = 0.
+                exhausted[i] = True
+                if all(exhausted):
+                    return
+                weights = [weight / sum(weights) for weight in weights]
+
+
+class TextLineDataPipe(IterDataPipe):
+    """ Streamming Text line
+    """
+
+    def __init__(self, filenames, mode='r'):
+        super().__init__()
+        _dp = datapipes.iter.FileLister(filenames)
+        _dp = datapipes.iter.FileOpener(_dp, mode=mode)
+        self.dp = _dp
+
+    def __iter__(self):
+        for fname, stream in self.dp:
+            for line in stream:
+                line = line.strip('\n')
+                yield {"file_name": fname, "line": line}
+            stream.close()
+
+
+@functional_datapipe("tar_file_and_group")
+class TarsDataPipe(IterDataPipe):
+    """ Decode wenet's tar , yield {'txt': "...", "raw": "..."}
+    """
+
+    def __init__(self, dataset: IterDataPipe) -> None:
+        super().__init__()
+        self.dp = dataset
+
+    def __iter__(self):
+        from wenet.dataset.processor import AUDIO_FORMAT_SETS
+        for sample in self.dp:
+            assert 'file_name' in sample
+            assert 'line' in sample
+            assert 'stream' in sample
+            try:
+                with tarfile.open(fileobj=sample['stream'],
+                                  mode="r:*") as stream:
+                    prev_prefix = None
+                    example = {
+                        'file_name': sample['file_name'],
+                        'tar_file_name': sample['line']
+                    }
+                    valid = True
+                    for tarinfo in stream:
+                        name = tarinfo.name
+                        pos = name.rfind('.')
+                        assert pos > 0
+                        prefix, postfix = name[:pos], name[pos + 1:]
+                        if prev_prefix is not None and prefix != prev_prefix:
+                            example['key'] = prev_prefix
+                            if valid:
+                                yield example
+                            example = {
+                                'file_name': sample['file_name'],
+                                'tar_file_name': sample['line']
+                            }
+                            valid = True
+                        with stream.extractfile(tarinfo) as file_obj:
+                            try:
+                                if postfix == 'txt':
+                                    example['txt'] = file_obj.read().decode(
+                                        'utf8').strip()
+                                elif postfix in AUDIO_FORMAT_SETS:
+                                    example['wav'] = file_obj.read()
+                                else:
+                                    example[postfix] = file_obj.read()
+                            except Exception as ex:
+                                valid = False
+                                logging.warning(
+                                    'error to parse {}'.format(name))
+                            prev_prefix = prefix
+                    if prev_prefix is not None:
+                        example['key'] = prev_prefix
+                        yield example
+            except Exception as ex:
+                msg = 'In tar_file_and_group: {} when processing {}'.format(
+                    ex, sample['line'])
+                logging.warning(msg)
+            finally:
+                if 'process' in sample:
+                    sample['process'].communicate()
+                sample['stream'].close()
+
+
+class WenetRawDatasetSource(IterDataPipe):
+
+    def __init__(self,
+                 filenames: str,
+                 prefetch: int = 500,
+                 partition: bool = True,
+                 shuffle: bool = False,
+                 shuffle_size: int = 10000,
+                 cycle: int = 1) -> None:
+        super().__init__()
+        self.dp = TextLineDataPipe(filenames)
+        if shuffle:
+            self.dp = self.dp.shuffle(buffer_size=shuffle_size)
+        self.dp = self.dp.repeat(cycle).prefetch(prefetch)
+        self.dp = self.dp.shard(partition)
+
+    def __iter__(self):
+        for d in self.dp:
+            yield d
+
+
+class WenetTarShardDatasetSource(IterDataPipe):
+
+    def __init__(self,
+                 filenames: str,
+                 prefetch: int = 500,
+                 partition: bool = True,
+                 shuffle: bool = False,
+                 shuffle_size: int = 10000,
+                 cycle: int = 1) -> None:
+        super().__init__()
+        self.dp = TextLineDataPipe(filenames)
+        if shuffle:
+            self.dp = self.dp.shuffle(buffer_size=shuffle_size)
+        self.dp = self.dp.repeat(cycle)
+        self.dp = self.dp.shard(partition).map_ignore_error(
+            parse_url).tar_file_and_group().prefetch(prefetch)
+
+    def __iter__(self):
+        for d in self.dp:
+            yield d

wenet/dataset/dataset.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2021 Wenet Community. (authors: Binbin Zhang)
+#               2023 Wenet Community. (authors: Dinghao Zhou)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import partial
+import sys
+from typing import Optional
+from wenet.dataset import processor
+from wenet.dataset.datapipes import (WenetRawDatasetSource,
+                                     WenetTarShardDatasetSource)
+from wenet.text.base_tokenizer import BaseTokenizer
+from wenet.utils.file_utils import read_symbol_table
+
+
+def Dataset(data_type,
+            data_list_file,
+            tokenizer: Optional[BaseTokenizer] = None,
+            conf=None,
+            partition=True):
+    """ Construct dataset from arguments
+
+        We have two shuffle stage in the Dataset. The first is global
+        shuffle at shards tar/raw file level. The second is global shuffle
+        at training samples level.
+
+        Args:
+            data_type(str): raw/shard
+            tokenizer (BaseTokenizer or None): tokenizer to tokenize
+            partition(bool): whether to do data partition in terms of rank
+    """
+    assert conf is not None
+    assert data_type in ['raw', 'shard']
+    # cycle dataset
+    cycle = conf.get('cycle', 1)
+    # stage1 shuffle: source
+    list_shuffle = conf.get('list_shuffle', True)
+    list_shuffle_size = sys.maxsize
+    if list_shuffle:
+        list_shuffle_conf = conf.get('list_shuffle_conf', {})
+        list_shuffle_size = list_shuffle_conf.get('shuffle_size',
+                                                  list_shuffle_size)
+    if data_type == 'raw':
+        dataset = WenetRawDatasetSource(data_list_file,
+                                        partition=partition,
+                                        shuffle=list_shuffle,
+                                        shuffle_size=list_shuffle_size,
+                                        cycle=cycle)
+        dataset = dataset.map(processor.parse_json)
+    else:
+        dataset = WenetTarShardDatasetSource(data_list_file,
+                                             partition=partition,
+                                             shuffle=list_shuffle,
+                                             shuffle_size=list_shuffle_size,
+                                             cycle=cycle)
+    dataset = dataset.map_ignore_error(processor.decode_wav)
+
+    singal_channel_conf = conf.get('singal_channel_conf', {})
+    dataset = dataset.map(
+        partial(processor.singal_channel, **singal_channel_conf))
+
+    speaker_conf = conf.get('speaker_conf', None)
+    if speaker_conf is not None:
+        assert 'speaker_table_path' in speaker_conf
+        speaker_table = read_symbol_table(speaker_conf['speaker_table_path'])
+        dataset = dataset.map(
+            partial(processor.parse_speaker, speaker_dict=speaker_table))
+
+    if tokenizer is not None:
+        dataset = dataset.map(partial(processor.tokenize, tokenizer=tokenizer))
+
+    filter_conf = conf.get('filter_conf', {})
+    dataset = dataset.filter(partial(processor.filter, **filter_conf))
+
+    resample_conf = conf.get('resample_conf', {})
+    dataset = dataset.map(partial(processor.resample, **resample_conf))
+
+    speed_perturb = conf.get('speed_perturb', False)
+    if speed_perturb:
+        dataset = dataset.map(partial(processor.speed_perturb))
+
+    feats_type = conf.get('feats_type', 'fbank')
+    assert feats_type in ['fbank', 'mfcc', 'log_mel_spectrogram']
+    if feats_type == 'fbank':
+        fbank_conf = conf.get('fbank_conf', {})
+        dataset = dataset.map(partial(processor.compute_fbank, **fbank_conf))
+    elif feats_type == 'mfcc':
+        mfcc_conf = conf.get('mfcc_conf', {})
+        dataset = dataset.map(partial(processor.compute_mfcc, **mfcc_conf))
+    elif feats_type == 'log_mel_spectrogram':
+        log_mel_spectrogram_conf = conf.get('log_mel_spectrogram_conf', {})
+        dataset = dataset.map(
+            partial(processor.compute_log_mel_spectrogram,
+                    **log_mel_spectrogram_conf))
+    spec_aug = conf.get('spec_aug', True)
+    spec_sub = conf.get('spec_sub', False)
+    spec_trim = conf.get('spec_trim', False)
+    if spec_aug:
+        spec_aug_conf = conf.get('spec_aug_conf', {})
+        dataset = dataset.map(partial(processor.spec_aug, **spec_aug_conf))
+    if spec_sub:
+        spec_sub_conf = conf.get('spec_sub_conf', {})
+        dataset = dataset.map(partial(processor.spec_sub, **spec_sub_conf))
+    if spec_trim:
+        spec_trim_conf = conf.get('spec_trim_conf', {})
+        dataset = dataset.map(partial(processor.spec_trim, **spec_trim_conf))
+
+    language_conf = conf.get('language_conf', {"limited_langs": ['zh', 'en']})
+    dataset = dataset.map(partial(processor.detect_language, **language_conf))
+    dataset = dataset.map(processor.detect_task)
+
+    shuffle = conf.get('shuffle', True)
+    if shuffle:
+        shuffle_conf = conf.get('shuffle_conf', {})
+        dataset = dataset.shuffle(buffer_size=shuffle_conf['shuffle_size'])
+
+    sort = conf.get('sort', True)
+    if sort:
+        sort_conf = conf.get('sort_conf', {})
+        dataset = dataset.sort(buffer_size=sort_conf['sort_size'],
+                               key_func=processor.sort_by_feats)
+
+    batch_conf = conf.get('batch_conf', {})
+    batch_type = batch_conf.get('batch_type', 'static')
+    assert batch_type in ['static', 'bucket', 'dynamic']
+    if batch_type == 'static':
+        assert 'batch_size' in batch_conf
+        batch_size = batch_conf.get('batch_size', 16)
+        dataset = dataset.batch(batch_size, wrapper_class=processor.padding)
+    elif batch_type == 'bucket':
+        assert 'bucket_boundaries' in batch_conf
+        assert 'bucket_batch_sizes' in batch_conf
+        dataset = dataset.bucket_by_sequence_length(
+            processor.feats_length_fn,
+            batch_conf['bucket_boundaries'],
+            batch_conf['bucket_batch_sizes'],
+            wrapper_class=processor.padding)
+    else:
+        max_frames_in_batch = batch_conf.get('max_frames_in_batch', 12000)
+        dataset = dataset.dynamic_batch(
+            processor.DynamicBatchWindow(max_frames_in_batch),
+            wrapper_class=processor.padding,
+        )
+
+    return dataset

wenet/dataset/deprecated/dataset.py

@@ -0,0 +1,202 @@
+# Copyright (c) 2021 Mobvoi Inc. (authors: Binbin Zhang)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+
+import torch
+import torch.distributed as dist
+from torch.utils.data import IterableDataset
+
+import wenet.dataset.deprecated.processor as processor
+from wenet.text.base_tokenizer import BaseTokenizer
+from wenet.utils.file_utils import read_lists
+
+
+class Processor(IterableDataset):
+
+    def __init__(self, source, f, *args, **kw):
+        assert callable(f)
+        self.source = source
+        self.f = f
+        self.args = args
+        self.kw = kw
+
+    def set_epoch(self, epoch):
+        self.source.set_epoch(epoch)
+
+    def __iter__(self):
+        """ Return an iterator over the source dataset processed by the
+            given processor.
+        """
+        assert self.source is not None
+        assert callable(self.f)
+        return self.f(iter(self.source), *self.args, **self.kw)
+
+    def apply(self, f):
+        assert callable(f)
+        return Processor(self, f, *self.args, **self.kw)
+
+
+class DistributedSampler:
+
+    def __init__(self, shuffle=True, partition=True):
+        self.epoch = -1
+        self.update()
+        self.shuffle = shuffle
+        self.partition = partition
+
+    def update(self):
+        assert dist.is_available()
+        if dist.is_initialized():
+            self.rank = dist.get_rank()
+            self.world_size = dist.get_world_size()
+        else:
+            self.rank = 0
+            self.world_size = 1
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info is None:
+            self.worker_id = 0
+            self.num_workers = 1
+        else:
+            self.worker_id = worker_info.id
+            self.num_workers = worker_info.num_workers
+        return dict(rank=self.rank,
+                    world_size=self.world_size,
+                    worker_id=self.worker_id,
+                    num_workers=self.num_workers)
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+    def sample(self, data):
+        """ Sample data according to rank/world_size/num_workers
+
+            Args:
+                data(List): input data list
+
+            Returns:
+                List: data list after sample
+        """
+        data = list(range(len(data)))
+        # TODO(Binbin Zhang): fix this
+        # We can not handle uneven data for CV on DDP, so we don't
+        # sample data by rank, that means every GPU gets the same
+        # and all the CV data
+        if self.partition:
+            if self.shuffle:
+                random.Random(self.epoch).shuffle(data)
+            data = data[self.rank::self.world_size]
+        data = data[self.worker_id::self.num_workers]
+        return data
+
+
+class DataList(IterableDataset):
+
+    def __init__(self, lists, shuffle=True, partition=True):
+        self.lists = lists
+        self.sampler = DistributedSampler(shuffle, partition)
+
+    def set_epoch(self, epoch):
+        self.sampler.set_epoch(epoch)
+
+    def __iter__(self):
+        sampler_info = self.sampler.update()
+        indexes = self.sampler.sample(self.lists)
+        for index in indexes:
+            # yield dict(src=src)
+            data = dict(src=self.lists[index])
+            data.update(sampler_info)
+            yield data
+
+
+def Dataset(data_type,
+            data_list_file,
+            tokenizer: BaseTokenizer,
+            conf,
+            partition=True):
+    """ Construct dataset from arguments
+
+        We have two shuffle stage in the Dataset. The first is global
+        shuffle at shards tar/raw file level. The second is global shuffle
+        at training samples level.
+
+        Args:
+            data_type(str): raw/shard
+            bpe_model(str): model for english bpe part
+            partition(bool): whether to do data partition in terms of rank
+    """
+    assert data_type in ['raw', 'shard']
+    lists = read_lists(data_list_file)
+    shuffle = conf.get('shuffle', True)
+    dataset = DataList(lists, shuffle=shuffle, partition=partition)
+    if data_type == 'shard':
+        dataset = Processor(dataset, processor.url_opener)
+        dataset = Processor(dataset, processor.tar_file_and_group)
+    else:
+        dataset = Processor(dataset, processor.parse_raw)
+
+    speaker_conf = conf.get('speaker_conf', None)
+    if speaker_conf is not None:
+        dataset = Processor(dataset, processor.parse_speaker, **speaker_conf)
+
+    dataset = Processor(dataset, processor.tokenize, tokenizer)
+    filter_conf = conf.get('filter_conf', {})
+    dataset = Processor(dataset, processor.filter, **filter_conf)
+
+    resample_conf = conf.get('resample_conf', {})
+    dataset = Processor(dataset, processor.resample, **resample_conf)
+
+    speed_perturb = conf.get('speed_perturb', False)
+    if speed_perturb:
+        dataset = Processor(dataset, processor.speed_perturb)
+
+    feats_type = conf.get('feats_type', 'fbank')
+    assert feats_type in ['fbank', 'mfcc', 'log_mel_spectrogram']
+    if feats_type == 'fbank':
+        fbank_conf = conf.get('fbank_conf', {})
+        dataset = Processor(dataset, processor.compute_fbank, **fbank_conf)
+    elif feats_type == 'mfcc':
+        mfcc_conf = conf.get('mfcc_conf', {})
+        dataset = Processor(dataset, processor.compute_mfcc, **mfcc_conf)
+    elif feats_type == 'log_mel_spectrogram':
+        log_mel_spectrogram_conf = conf.get('log_mel_spectrogram_conf', {})
+        dataset = Processor(dataset, processor.compute_log_mel_spectrogram,
+                            **log_mel_spectrogram_conf)
+
+    spec_aug = conf.get('spec_aug', True)
+    spec_sub = conf.get('spec_sub', False)
+    spec_trim = conf.get('spec_trim', False)
+    if spec_aug:
+        spec_aug_conf = conf.get('spec_aug_conf', {})
+        dataset = Processor(dataset, processor.spec_aug, **spec_aug_conf)
+    if spec_sub:
+        spec_sub_conf = conf.get('spec_sub_conf', {})
+        dataset = Processor(dataset, processor.spec_sub, **spec_sub_conf)
+    if spec_trim:
+        spec_trim_conf = conf.get('spec_trim_conf', {})
+        dataset = Processor(dataset, processor.spec_trim, **spec_trim_conf)
+
+    if shuffle:
+        shuffle_conf = conf.get('shuffle_conf', {})
+        dataset = Processor(dataset, processor.shuffle, **shuffle_conf)
+
+    sort = conf.get('sort', True)
+    if sort:
+        sort_conf = conf.get('sort_conf', {})
+        dataset = Processor(dataset, processor.sort, **sort_conf)
+
+    batch_conf = conf.get('batch_conf', {})
+    dataset = Processor(dataset, processor.batch, **batch_conf)
+    dataset = Processor(dataset, processor.padding)
+    return dataset