Update model, no need for decoder_main

axmodel/decoder_loop.axmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6b109e1281135a673b613c1fd92f5d12d64e02d1f3da47561c142bbc57295d5d
-size 446759232
+oid sha256:7c9b3e351557d20846f50d819e18c59d6f10a8adfc40322e5e3034b404b3e038
+size 435136795

axmodel/decoder_loop_u8.axmodel

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

axmodel/decoder_main.axmodel

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

axmodel/decoder_main_u8.axmodel

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

axmodel/encoder.axmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d1ceef85b578ecffec2e6eaee4dc27987c0e342f109b14a375376935121c5a2c
-size 851312087
+oid sha256:6cc674ba54cf0e57f3c7dffa3824cd53700e4e7709827893f8708c4958e116c1
+size 851656147

fireredasr_axmodel.py

@@ -0,0 +1,336 @@
+from fireredasr.data.asr_feat import ASRFeatExtractor
+from fireredasr.tokenizer.aed_tokenizer import ChineseCharEnglishSpmTokenizer
+
+import axengine as axe
+import torch
+import torch.nn.functional as F
+import numpy as np
+from torch import Tensor
+from typing import Tuple, List, Dict
+import os
+import time
+
+INF = 1e10
+
+def to_numpy(tensor):
+    if isinstance(tensor, np.ndarray):
+        return tensor
+    if tensor.requires_grad:
+        return tensor.detach().cpu().numpy()
+    else:
+        return tensor.cpu().numpy()
+    
+    
+def set_finished_beam_score_to_zero(scores, is_finished):
+    NB, B = scores.size()
+    is_finished = is_finished.float()
+    mask_score = torch.tensor([0.0] + [-INF]*(B-1)).float()
+    mask_score = mask_score.view(1, B).repeat(NB, 1)
+    return scores * (1 - is_finished) + mask_score * is_finished
+
+
+def set_finished_beam_y_to_eos(ys, is_finished, eos_id):
+    is_finished = is_finished.long()
+    return ys * (1 - is_finished) + eos_id * is_finished
+
+
+class FireRedASRAxModel:
+    def __init__(
+        self, 
+        encoder_path: str, 
+        decoder_loop_path: str,
+        cmvn_file: str,
+        dict_file: str, 
+        spm_model_path: str,
+        providers=['AxEngineExecutionProvider'],
+        decode_max_len=128,
+        audio_dur=10
+    ):
+        # NOTE: 参考whisper设置的最大的解码长度
+        # FireRedASR-AED 模型支持的最长语音为 60s
+        # ref: https://github.com/FireRedTeam/FireRedASR?tab=readme-ov-file#input-length-limitations
+        self.decode_max_len = decode_max_len
+        
+        self.decoder_hidden_dim = 1280
+        self.audio_dur = audio_dur
+        self.max_feat_len = self.calc_feat_len(audio_dur)
+        self.num_decoder_blocks = 16
+        self.blank_id = 0
+        self.sos_id = 3
+        self.eos_id = 4
+        self.pad_id = 2
+        
+        self.feature_extractor = ASRFeatExtractor(cmvn_file)
+        self.tokenizer = ChineseCharEnglishSpmTokenizer(dict_file, spm_model_path)
+        
+        self.init_encoder(encoder_path, providers)
+        self.init_decoder_loop(decoder_loop_path, providers)
+        self.pe = self.init_pe(decoder_loop_path)
+        
+    def init_encoder(self, encoder_path, providers=None):
+        self.encoder = axe.InferenceSession(
+            encoder_path,
+            providers=providers
+        )
+
+    def init_decoder_loop(self, decoder_path, providers=None):
+        self.decoder_loop = axe.InferenceSession(
+            decoder_path,
+            providers=providers
+        )
+
+    def init_pe(self, decoder_path):
+        decoder_path = os.path.dirname(decoder_path)
+        decoder_path = os.path.join(decoder_path, "pe.npy")
+        
+        return np.load(decoder_path)
+    
+    def run_encoder(self, input: np.ndarray, 
+                    input_length: np.ndarray
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        n_layer_cross_k, n_layer_cross_v, cross_attn_mask = self.encoder.run(
+            None,
+            {
+                "encoder_input": input,
+                "encoder_input_lengths": input_length
+            }
+        )
+        return (
+            n_layer_cross_k,
+            n_layer_cross_v,
+            cross_attn_mask
+        )
+
+    def decode_loop_one_token(
+        self,
+        tokens: np.ndarray,
+        n_layer_self_k_cache: np.ndarray,
+        n_layer_self_v_cache: np.ndarray,
+        n_layer_cross_k_cache: np.ndarray,
+        n_layer_cross_v_cache: np.ndarray,
+        pe: np.ndarray,
+        self_attn_mask: np.ndarray,
+        cross_attn_mask: np.ndarray
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        logits, out_n_layer_self_k_cache, out_n_layer_self_v_cache = self.decoder_loop.run(
+            None,
+            {
+                "tokens": tokens,
+                "in_n_layer_self_k_cache": n_layer_self_k_cache,
+                "in_n_layer_self_v_cache": n_layer_self_v_cache,
+                "n_layer_cross_k": n_layer_cross_k_cache,
+                "n_layer_cross_v": n_layer_cross_v_cache,
+                "pe": pe,
+                "self_attn_mask": self_attn_mask,
+                "cross_attn_mask": cross_attn_mask,
+            }
+        )
+        return (
+            logits,
+            out_n_layer_self_k_cache,
+            out_n_layer_self_v_cache
+        )
+    
+    def run_decoder(
+        self,
+        n_layer_cross_k, 
+        n_layer_cross_v, 
+        cross_attn_mask,
+        beam_size,
+        nbest
+    ):
+        
+        num_layer, batch_size, Ti, encoder_out_dim = n_layer_cross_k.shape
+        encoder_out_length = cross_attn_mask.shape[-1]
+        
+        cross_attn_mask = torch.from_numpy(cross_attn_mask).to(torch.float32)
+        cross_attn_mask = cross_attn_mask.unsqueeze(1).repeat(
+            1, beam_size, 1, 1
+        ).view(beam_size * batch_size, -1, encoder_out_length)
+        
+        n_layer_cross_k = torch.from_numpy(n_layer_cross_k)
+        n_layer_cross_v = torch.from_numpy(n_layer_cross_v)
+        n_layer_cross_k = n_layer_cross_k.unsqueeze(2).repeat(
+            1, 1, beam_size, 1, 1
+        ).view(num_layer, beam_size * batch_size, Ti, encoder_out_dim)
+        n_layer_cross_v = n_layer_cross_v.unsqueeze(2).repeat(
+            1, 1, beam_size, 1, 1
+        ).view(num_layer, beam_size * batch_size, Ti, encoder_out_dim)
+
+        prediction_tokens = torch.ones(
+            beam_size * batch_size, 1).fill_(self.sos_id).long()
+        tokens = prediction_tokens
+        offset = torch.zeros(1, dtype=torch.int64)
+        n_layer_self_k_cache, n_layer_self_v_cache = self.get_initialized_self_cache(
+            batch_size, beam_size
+        )
+        
+        scores = torch.tensor([0.0] + [-INF]*(beam_size - 1)).float()
+        scores = scores.repeat(batch_size).view(batch_size * beam_size, 1)
+        is_finished = torch.zeros_like(scores)
+        
+        self_attn_mask = np.zeros((batch_size * beam_size, 1, 1), dtype=np.float32)
+
+        for i in range(self.decode_max_len):
+
+            tokens = to_numpy(tokens).astype(np.int32)
+            n_layer_self_k_cache = to_numpy(n_layer_self_k_cache)
+            n_layer_self_v_cache = to_numpy(n_layer_self_v_cache)
+            n_layer_cross_k = to_numpy(n_layer_cross_k)
+            n_layer_cross_v = to_numpy(n_layer_cross_v)
+            cross_attn_mask = to_numpy(cross_attn_mask)
+
+            self_attn_mask = np.zeros((batch_size * beam_size, 1, self.decode_max_len), dtype=np.float32)
+            self_attn_mask[:, :, :self.decode_max_len - offset[0] - 1] = -np.inf
+
+            logits, n_layer_self_k_cache, n_layer_self_v_cache = self.decode_loop_one_token(
+                    to_numpy(tokens),
+                    to_numpy(n_layer_self_k_cache),
+                    to_numpy(n_layer_self_v_cache),
+                    to_numpy(n_layer_cross_k),
+                    to_numpy(n_layer_cross_v),
+                    self.pe[offset],
+                    self_attn_mask,
+                    to_numpy(cross_attn_mask)
+                )
+            
+            offset += 1
+            logits = torch.from_numpy(logits)
+            
+            logits = logits.squeeze(1)
+            t_scores = F.log_softmax(logits, dim=-1)
+            t_topB_scores, t_topB_ys = torch.topk(t_scores, k=beam_size, dim=1)
+            t_topB_scores = set_finished_beam_score_to_zero(t_topB_scores, is_finished)
+            t_topB_ys = set_finished_beam_y_to_eos(t_topB_ys, is_finished, self.eos_id)
+            
+            scores = scores + t_topB_scores
+            
+            scores = scores.view(batch_size, beam_size * beam_size)
+            scores, topB_score_ids = torch.topk(scores, k=beam_size, dim=1)
+            scores = scores.view(-1, 1)
+            
+            topB_row_number_in_each_B_rows_of_ys = torch.div(
+                topB_score_ids, beam_size).view(batch_size * beam_size)
+            stride = beam_size * torch.arange(batch_size).view(
+                batch_size, 1).repeat(1, beam_size).view(batch_size * beam_size)
+            topB_row_number_in_ys = topB_row_number_in_each_B_rows_of_ys.long() + stride.long()
+            
+            prediction_tokens = prediction_tokens[topB_row_number_in_ys]
+            t_ys = torch.gather(
+                t_topB_ys.view(batch_size, beam_size * beam_size), 
+                dim=1, index=topB_score_ids
+            ).view(beam_size * batch_size, 1)
+            
+            tokens = t_ys
+        
+            prediction_tokens = torch.cat((prediction_tokens, t_ys), dim=1)
+            
+            n_layer_self_k_cache = torch.from_numpy(n_layer_self_k_cache)
+            n_layer_self_v_cache = torch.from_numpy(n_layer_self_v_cache)
+            
+            for i, self_k_cache in enumerate(n_layer_self_k_cache):
+                n_layer_self_k_cache[i] = n_layer_self_k_cache[i][topB_row_number_in_ys]
+            
+            for i, self_v_cache in enumerate(n_layer_self_v_cache):
+                n_layer_self_v_cache[i] = n_layer_self_v_cache[i][topB_row_number_in_ys]
+                
+            is_finished = t_ys.eq(self.eos_id)
+            if is_finished.sum().item() == beam_size * batch_size:
+                break
+            
+        scores = scores.view(batch_size, beam_size)
+        prediction_valid_token_lengths = torch.sum(
+            torch.ne(
+                prediction_tokens.view(batch_size, beam_size, -1), 
+                self.eos_id),
+            dim=-1
+        ).int()
+       
+        nbest_scores, nbest_ids = torch.topk(scores, k=nbest, dim=1)
+        index = nbest_ids + beam_size * torch.arange(batch_size).view(batch_size, 1).long()
+        nbest_prediction_tokens = prediction_tokens.view(batch_size * beam_size, -1)[index.view(-1)]
+        nbest_prediction_tokens = nbest_prediction_tokens.view(batch_size, nbest_ids.size(1), -1)
+        nbest_prediction_valid_token_lengths = prediction_valid_token_lengths.view(
+            batch_size * beam_size)[index.view(-1)].view(batch_size, -1)
+        nbest_hyps: List[List[Dict[str, torch.Tensor]]] = []
+        for i in range(batch_size):
+            i_best_hyps: List[Dict[str, torch.Tensor]] = []
+            for j, score in enumerate(nbest_scores[i]):
+                hyp = {
+                    "token_ids": nbest_prediction_tokens[i, j, 1:nbest_prediction_valid_token_lengths[i, j]],
+                    "score": score
+                }
+                i_best_hyps.append(hyp)
+            nbest_hyps.append(i_best_hyps)
+
+        return nbest_hyps
+        
+    def get_initialized_self_cache(self, 
+                                   batch_size, 
+                                   beam_size
+                                   ) -> Tuple[Tensor, Tensor]:
+        n_layer_self_k_cache = torch.zeros(
+            self.num_decoder_blocks,
+            batch_size * beam_size,
+            self.decode_max_len,
+            self.decoder_hidden_dim,
+        )
+        n_layer_self_v_cache = torch.zeros(
+            self.num_decoder_blocks,
+            batch_size * beam_size,
+            self.decode_max_len,
+            self.decoder_hidden_dim,
+        )
+        return n_layer_self_k_cache, n_layer_self_v_cache
+    
+    def calc_feat_len(self, audio_dur):
+        import math
+        sample_rate = 16000
+        frame_length = 25 * sample_rate / 1000
+        frame_shift = 10 * sample_rate / 1000
+        length = math.floor((audio_dur * sample_rate - frame_length) / frame_shift) + 1
+        return length
+ 
+    def transcribe(self, 
+                   batch_wav_path: List[str],
+                   beam_size: int = 1,
+                   nbest: int = 1
+                ) -> List[Dict]:
+        feats, lengths, wav_durations = self.feature_extractor(batch_wav_path)
+        # print(f"feats.shape: {feats.shape}")
+        if feats.shape[1] < self.max_feat_len:
+            feats = np.concatenate([feats, np.zeros((1, self.max_feat_len - feats.shape[1], 80), dtype=np.float32)], axis=1)
+        feats = feats[:, :self.max_feat_len, :]
+        lengths = torch.minimum(lengths, torch.tensor(self.max_feat_len))
+
+        feats = to_numpy(feats)
+        lengths = to_numpy(lengths).astype(np.int32)
+
+        start_time = time.time()
+        n_layer_cross_k, n_layer_cross_v, cross_attn_mask = self.run_encoder(
+            to_numpy(feats),
+            to_numpy(lengths)
+        )
+        # print(f"run encoder take {(time.time() - start_time) * 1000}ms")
+        nbest_hyps = self.run_decoder(n_layer_cross_k,
+                                      n_layer_cross_v,
+                                      cross_attn_mask,
+                                      beam_size,
+                                      nbest,
+                                      )
+        transcribe_durations = time.time() - start_time
+        results: List[Dict] = []
+        for wav, hyp in zip(batch_wav_path, nbest_hyps):
+            hyp = hyp[0]
+            hyp_ids = [int(id) for id in hyp["token_ids"].cpu()]
+            score = hyp["score"].item()
+            text = self.tokenizer.detokenize(hyp_ids)
+            results.append(
+                {
+                    "wav": wav,
+                    "text": text,
+                    "score": score
+                }
+            )
+            
+        return results, wav_durations, transcribe_durations   

test_ax_model.py

@@ -1,546 +1,30 @@
-from fireredasr.data.asr_feat import ASRFeatExtractor
-from fireredasr.tokenizer.aed_tokenizer import ChineseCharEnglishSpmTokenizer
-
-import axengine as axe
-import torch
-import torch.nn.functional as F
-import numpy as np
-from torch import Tensor
-from typing import Tuple, List, Dict
 import argparse
 import os
 import time
 import logging
 
+from fireredasr_axmodel import FireRedASRAxModel
+
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
 logger_stream_hander = logging.StreamHandler()
 logger_stream_hander.setLevel("INFO")
 logger.addHandler(logger_stream_hander)
 
-
-INF = 1e10
-
-
-def to_numpy(tensor):
-    if isinstance(tensor, np.ndarray):
-        return tensor
-    if tensor.requires_grad:
-        return tensor.detach().cpu().numpy()
-    else:
-        return tensor.cpu().numpy()
-    
-    
-def set_finished_beam_score_to_zero(scores, is_finished):
-    NB, B = scores.size()
-    is_finished = is_finished.float()
-    mask_score = torch.tensor([0.0] + [-INF]*(B-1)).float()
-    mask_score = mask_score.view(1, B).repeat(NB, 1)
-    return scores * (1 - is_finished) + mask_score * is_finished
-
-
-def set_finished_beam_y_to_eos(ys, is_finished, eos_id):
-    is_finished = is_finished.long()
-    return ys * (1 - is_finished) + eos_id * is_finished
-
-
-class FireRedASROnnxModel:
-    def __init__(
-        self, 
-        encoder_path: str, 
-        decoder_path: str,
-        cmvn_file: str,
-        dict_file: str, 
-        spm_model_path: str,
-        providers=['AXCLRTExecutionProvider', 'AxEngineExecutionProvider'],
-        decode_max_len=128
-    ):
-        # NOTE: 参考whisper设置的最大的解码长度
-        # FireRedASR-AED 模型支持的最长语音为 60s
-        # ref: https://github.com/FireRedTeam/FireRedASR?tab=readme-ov-file#input-length-limitations
-        self.decode_max_len = decode_max_len
-        
-        self.decoder_hidden_dim = 1280
-        self.num_decoder_blocks = 16
-        self.blank_id = 0
-        self.sos_id = 3
-        self.eos_id = 4
-        self.pad_id = 2
-        
-        self.feature_extractor = ASRFeatExtractor(cmvn_file)
-        self.tokenizer = ChineseCharEnglishSpmTokenizer(dict_file, spm_model_path)
-        self.encoder = None
-        self.decoder = None
-        
-        self.init_encoder(encoder_path, providers)
-        self.init_decoder_main(decoder_path, providers)
-        self.init_decoder_loop(decoder_path, providers)
-        self.pe = self.init_pe(decoder_path)
-        
-    def init_encoder(self, encoder_path, providers=None):
-        start_time = time.time()
-        self.encoder = axe.InferenceSession(
-            encoder_path,
-            # sess_options=self.session_opts,
-            providers=providers
-        )
-        end_time = time.time()
-        logger.info(f"load encoder cost {end_time - start_time} seconds")
-
-    def init_decoder_main(self, decoder_path, providers=None):
-        decoder_path = os.path.dirname(decoder_path)
-        decoder_path = os.path.join(decoder_path, "decoder_main.axmodel")
-        start_time = time.time()
-        self.decoder_main = axe.InferenceSession(
-            decoder_path,
-            # sess_options=self.session_opts,
-            providers=providers
-        )
-        end_time = time.time()
-        logger.info(f"load decoder_main cost {end_time - start_time} seconds")
-
-        # input_names = [i.name for i in self.decoder_main.get_inputs()]
-        # print(f"decoder_main.input_names: {input_names}")
-
-    def init_decoder_loop(self, decoder_path, providers=None):
-        decoder_path = os.path.dirname(decoder_path)
-        decoder_path = os.path.join(decoder_path, "decoder_loop.axmodel")
-
-        start_time = time.time()
-        self.decoder_loop = axe.InferenceSession(
-            decoder_path,
-            # sess_options=self.session_opts,
-            providers=providers
-        )
-        end_time = time.time()
-        logger.info(f"load decoder_loop cost {end_time - start_time} seconds")
-
-        # input_names = [i.name for i in self.decoder_loop.get_inputs()]
-        # print(f"decoder_loop.input_names: {input_names}")
-
-    def init_pe(self, decoder_path):
-        decoder_path = os.path.dirname(decoder_path)
-        decoder_path = os.path.join(decoder_path, "pe.npy")
-        
-        return np.load(decoder_path)
-    
-    def run_encoder(self, input: np.ndarray, 
-                    input_length: np.ndarray
-    ) -> Tuple[Tensor, Tensor, Tensor]:
-        n_layer_cross_k, n_layer_cross_v, cross_attn_mask = self.encoder.run(
-            None,
-            {
-                "encoder_input": input,
-                "encoder_input_lengths": input_length
-            }
-        )
-        # n_layer_cross_k, n_layer_cross_v, cross_attn_mask = \
-        #     outputs["n_layer_cross_k"], outputs["n_layer_cross_v"], outputs["cross_attn_mask"]
-        return (
-            n_layer_cross_k,
-            n_layer_cross_v,
-            cross_attn_mask
-        )
-        
-    def decode_one_token(
-        self,
-        tokens: np.ndarray,
-        n_layer_self_k_cache: np.ndarray,
-        n_layer_self_v_cache: np.ndarray,
-        n_layer_cross_k_cache: np.ndarray,
-        n_layer_cross_v_cache: np.ndarray,
-        offset: np.ndarray,
-        self_attn_mask: np.ndarray,
-        cross_attn_mask: np.ndarray
-    ) -> Tuple[Tensor, Tensor, Tensor]:
-        print("decode:")
-        print(f"tokens.shape: {tokens.shape}")
-        print(f"n_layer_self_k_cache.shape: {n_layer_self_k_cache.shape}")
-        print(f"n_layer_self_v_cache.shape: {n_layer_self_v_cache.shape}")
-        print(f"n_layer_cross_k_cache.shape: {n_layer_cross_k_cache.shape}")
-        print(f"n_layer_cross_v_cache.shape: {n_layer_cross_v_cache.shape}")
-        print(f"offset.shape: {offset.shape}")
-        print(f"self_attn_mask.shape: {self_attn_mask.shape}")
-        print(f"cross_attn_mask.shape: {cross_attn_mask.shape}")
-        # print(f"self_attn_mask: {self_attn_mask}")
-
-        logits, out_n_layer_self_k_cache, out_n_layer_self_v_cache = self.decoder.run(
-            None,
-            {
-                self.decoder.get_inputs()[0].name: tokens,
-                self.decoder.get_inputs()[1].name: n_layer_self_k_cache,
-                self.decoder.get_inputs()[2].name: n_layer_self_v_cache,
-                self.decoder.get_inputs()[3].name: n_layer_cross_k_cache,
-                self.decoder.get_inputs()[4].name: n_layer_cross_v_cache,
-                self.decoder.get_inputs()[5].name: offset,
-                self.decoder.get_inputs()[6].name: self_attn_mask,
-                self.decoder.get_inputs()[7].name: cross_attn_mask,
-            }
-        )
-        return (
-            logits,
-            out_n_layer_self_k_cache,
-            out_n_layer_self_v_cache
-        )
-    
-    def decode_main_one_token(
-        self,
-        tokens: np.ndarray,
-        n_layer_self_k_cache: np.ndarray,
-        n_layer_self_v_cache: np.ndarray,
-        n_layer_cross_k_cache: np.ndarray,
-        n_layer_cross_v_cache: np.ndarray,
-        pe: np.ndarray,
-        self_attn_mask: np.ndarray,
-        cross_attn_mask: np.ndarray
-    ) -> Tuple[Tensor, Tensor, Tensor]:
-        # print("decode_main:")
-        # print(f"tokens.shape: {tokens.shape}")
-        # print(f"n_layer_self_k_cache.shape: {n_layer_self_k_cache.shape}")
-        # print(f"n_layer_self_v_cache.shape: {n_layer_self_v_cache.shape}")
-        # print(f"n_layer_cross_k_cache.shape: {n_layer_cross_k_cache.shape}")
-        # print(f"n_layer_cross_v_cache.shape: {n_layer_cross_v_cache.shape}")
-        # print(f"pe.shape: {pe.shape}")
-        # print(f"self_attn_mask.shape: {self_attn_mask.shape}")
-        # print(f"cross_attn_mask.shape: {cross_attn_mask.shape}")
-
-        logits, out_n_layer_self_k_cache, out_n_layer_self_v_cache = self.decoder_main.run(
-            None,
-            {
-                "tokens": tokens,
-                # self.decoder_main.get_inputs()[1].name: n_layer_self_k_cache,
-                "n_layer_cross_k": n_layer_cross_k_cache,
-                "n_layer_cross_v": n_layer_cross_v_cache,
-                # "pe": pe,
-                # "self_attn_mask": self_attn_mask,
-                "cross_attn_mask": cross_attn_mask,
-                # self.decoder_main.get_inputs()[7].name: cross_attn_mask,
-            }
-        )
-        # logits, out_n_layer_self_k_cache, out_n_layer_self_v_cache = \
-        #     outputs["logits"], outputs["out_n_layer_self_k_cache"], outputs["out_n_layer_self_v_cache"]
-        return (
-            logits,
-            out_n_layer_self_k_cache,
-            out_n_layer_self_v_cache
-        )
-    
-    def decode_loop_one_token(
-        self,
-        tokens: np.ndarray,
-        n_layer_self_k_cache: np.ndarray,
-        n_layer_self_v_cache: np.ndarray,
-        n_layer_cross_k_cache: np.ndarray,
-        n_layer_cross_v_cache: np.ndarray,
-        pe: np.ndarray,
-        self_attn_mask: np.ndarray,
-        cross_attn_mask: np.ndarray
-    ) -> Tuple[Tensor, Tensor, Tensor]:
-        # print("decode_loop:")
-        # print(f"tokens.shape: {tokens.shape}")
-        # print(f"n_layer_self_k_cache.shape: {n_layer_self_k_cache.shape}")
-        # print(f"n_layer_self_v_cache.shape: {n_layer_self_v_cache.shape}")
-        # print(f"n_layer_cross_k_cache.shape: {n_layer_cross_k_cache.shape}")
-        # print(f"n_layer_cross_v_cache.shape: {n_layer_cross_v_cache.shape}")
-        # print(f"pe.shape: {pe.shape}")
-        # print(f"self_attn_mask.shape: {self_attn_mask.shape}")
-        # print(f"cross_attn_mask.shape: {cross_attn_mask.shape}")
-
-        logits, out_n_layer_self_k_cache, out_n_layer_self_v_cache = self.decoder_loop.run(
-            None,
-            {
-                "tokens": tokens,
-                "in_n_layer_self_k_cache": n_layer_self_k_cache,
-                "in_n_layer_self_v_cache": n_layer_self_v_cache,
-                "n_layer_cross_k": n_layer_cross_k_cache,
-                "n_layer_cross_v": n_layer_cross_v_cache,
-                "pe": pe,
-                "self_attn_mask": self_attn_mask,
-                "cross_attn_mask": cross_attn_mask,
-            }
-        )
-        # logits, out_n_layer_self_k_cache, out_n_layer_self_v_cache = \
-        #     outputs["logits"], outputs["out_n_layer_self_k_cache"], outputs["out_n_layer_self_v_cache"]
-        return (
-            logits,
-            out_n_layer_self_k_cache,
-            out_n_layer_self_v_cache
-        )
-    
-    def run_decoder(
-        self,
-        n_layer_cross_k, 
-        n_layer_cross_v, 
-        cross_attn_mask,
-        beam_size,
-        nbest
-    ):
-        
-        num_layer, batch_size, Ti, encoder_out_dim = n_layer_cross_k.shape
-        encoder_out_length = cross_attn_mask.shape[-1]
-        
-        cross_attn_mask = torch.from_numpy(cross_attn_mask).to(torch.float32)
-        cross_attn_mask = cross_attn_mask.unsqueeze(1).repeat(
-            1, beam_size, 1, 1
-        ).view(beam_size * batch_size, -1, encoder_out_length)
-        
-        n_layer_cross_k = torch.from_numpy(n_layer_cross_k)
-        n_layer_cross_v = torch.from_numpy(n_layer_cross_v)
-        n_layer_cross_k = n_layer_cross_k.unsqueeze(2).repeat(
-            1, 1, beam_size, 1, 1
-        ).view(num_layer, beam_size * batch_size, Ti, encoder_out_dim)
-        n_layer_cross_v = n_layer_cross_v.unsqueeze(2).repeat(
-            1, 1, beam_size, 1, 1
-        ).view(num_layer, beam_size * batch_size, Ti, encoder_out_dim)
-
-        prediction_tokens = torch.ones(
-            beam_size * batch_size, 1).fill_(self.sos_id).long()
-        tokens = prediction_tokens
-        offset = torch.zeros(1, dtype=torch.int64)
-        n_layer_self_k_cache, n_layer_self_v_cache = self.get_initialized_self_cache(
-            batch_size, beam_size
-        )
-        
-        scores = torch.tensor([0.0] + [-INF]*(beam_size - 1)).float()
-        scores = scores.repeat(batch_size).view(batch_size * beam_size, 1)
-        is_finished = torch.zeros_like(scores)
-        
-        # self_attn_mask = torch.zeros(
-        #     batch_size * beam_size,
-        #     1, 1
-        # )
-        self_attn_mask = np.zeros((batch_size * beam_size, 1, 1), dtype=np.float32)
-
-        results = [self.sos_id]
-        for i in range(self.decode_max_len):
-
-            # self_attn_mask = torch.empty(
-            #     batch_size * beam_size, 
-            #     prediction_tokens.shape[-1], prediction_tokens.shape[-1]
-            # ).fill_(-np.inf).triu_(1)
-            # self_attn_mask = self_attn_mask[:, -1:, :]
-            # self_attn_mask = to_numpy(self_attn_mask)
-
-            # logits, n_layer_self_k_cache, n_layer_self_v_cache = self.decode_one_token(
-            #     to_numpy(tokens),
-            #     to_numpy(n_layer_self_k_cache),
-            #     to_numpy(n_layer_self_v_cache),
-            #     to_numpy(n_layer_cross_k),
-            #     to_numpy(n_layer_cross_v),
-            #     to_numpy(offset),
-            #     to_numpy(self_attn_mask),
-            #     to_numpy(cross_attn_mask)
-            # )
-
-            tokens = to_numpy(tokens).astype(np.int32)
-            n_layer_self_k_cache = to_numpy(n_layer_self_k_cache)
-            n_layer_self_v_cache = to_numpy(n_layer_self_v_cache)
-            n_layer_cross_k = to_numpy(n_layer_cross_k)
-            n_layer_cross_v = to_numpy(n_layer_cross_v)
-            cross_attn_mask = to_numpy(cross_attn_mask)
-
-            self_attn_mask = np.zeros((batch_size * beam_size, 1, self.decode_max_len), dtype=np.float32)
-            self_attn_mask[:, :, :self.decode_max_len - offset[0] - 1] = -np.inf
-
-            # for name, npy in zip(
-            #     ["tokens", "n_layer_self_k_cache", "n_layer_self_v_cache", "n_layer_cross_k", "n_layer_cross_v", "pe", "self_attn_mask", "cross_attn_mask"],
-            #     [tokens, n_layer_self_k_cache, n_layer_self_v_cache, n_layer_cross_k, n_layer_cross_v, self.pe[offset], self_attn_mask, cross_attn_mask]
-            # ):
-            #     file_path = os.path.join(decoder_data_path, name)
-            #     os.makedirs(file_path, exist_ok=True)
-            #     np.save(os.path.join(file_path, f"{i}.npy"), npy)
-
-            if i == 0:
-                start_time = time.time()
-                logits, n_layer_self_k_cache, n_layer_self_v_cache = self.decode_main_one_token(
-                    to_numpy(tokens),
-                    to_numpy(n_layer_self_k_cache),
-                    to_numpy(n_layer_self_v_cache),
-                    to_numpy(n_layer_cross_k),
-                    to_numpy(n_layer_cross_v),
-                    self.pe[offset],
-                    self_attn_mask,
-                    to_numpy(cross_attn_mask)
-                )
-                print(f"run decoder_main take {(time.time() - start_time) * 1000}ms")
-            else:
-                start_time = time.time()
-                logits, n_layer_self_k_cache, n_layer_self_v_cache = self.decode_loop_one_token(
-                    to_numpy(tokens),
-                    to_numpy(n_layer_self_k_cache),
-                    to_numpy(n_layer_self_v_cache),
-                    to_numpy(n_layer_cross_k),
-                    to_numpy(n_layer_cross_v),
-                    self.pe[offset],
-                    self_attn_mask,
-                    to_numpy(cross_attn_mask)
-                )
-                print(f"run decoder_loop take {(time.time() - start_time) * 1000}ms")
-            
-            offset += 1
-            logits = torch.from_numpy(logits)
-            
-            logits = logits.squeeze(1)
-            t_scores = F.log_softmax(logits, dim=-1)
-            t_topB_scores, t_topB_ys = torch.topk(t_scores, k=beam_size, dim=1)
-            t_topB_scores = set_finished_beam_score_to_zero(t_topB_scores, is_finished)
-            t_topB_ys = set_finished_beam_y_to_eos(t_topB_ys, is_finished, self.eos_id)
-            
-            scores = scores + t_topB_scores
-            
-            scores = scores.view(batch_size, beam_size * beam_size)
-            scores, topB_score_ids = torch.topk(scores, k=beam_size, dim=1)
-            scores = scores.view(-1, 1)
-            
-            topB_row_number_in_each_B_rows_of_ys = torch.div(
-                topB_score_ids, beam_size).view(batch_size * beam_size)
-            stride = beam_size * torch.arange(batch_size).view(
-                batch_size, 1).repeat(1, beam_size).view(batch_size * beam_size)
-            topB_row_number_in_ys = topB_row_number_in_each_B_rows_of_ys.long() + stride.long()
-            
-            prediction_tokens = prediction_tokens[topB_row_number_in_ys]
-            t_ys = torch.gather(
-                t_topB_ys.view(batch_size, beam_size * beam_size), 
-                dim=1, index=topB_score_ids
-            ).view(beam_size * batch_size, 1)
-            
-            tokens = t_ys
-        
-            prediction_tokens = torch.cat((prediction_tokens, t_ys), dim=1)
-            
-            n_layer_self_k_cache = torch.from_numpy(n_layer_self_k_cache)
-            n_layer_self_v_cache = torch.from_numpy(n_layer_self_v_cache)
-            
-            for i, self_k_cache in enumerate(n_layer_self_k_cache):
-                n_layer_self_k_cache[i] = n_layer_self_k_cache[i][topB_row_number_in_ys]
-            
-            for i, self_v_cache in enumerate(n_layer_self_v_cache):
-                n_layer_self_v_cache[i] = n_layer_self_v_cache[i][topB_row_number_in_ys]
-                
-            is_finished = t_ys.eq(self.eos_id)
-            if is_finished.sum().item() == beam_size * batch_size:
-                break
-            
-        scores = scores.view(batch_size, beam_size)
-        prediction_valid_token_lengths = torch.sum(
-            torch.ne(
-                prediction_tokens.view(batch_size, beam_size, -1), 
-                self.eos_id),
-            dim=-1
-        ).int()
-       
-        nbest_scores, nbest_ids = torch.topk(scores, k=nbest, dim=1)
-        index = nbest_ids + beam_size * torch.arange(batch_size).view(batch_size, 1).long()
-        nbest_prediction_tokens = prediction_tokens.view(batch_size * beam_size, -1)[index.view(-1)]
-        nbest_prediction_tokens = nbest_prediction_tokens.view(batch_size, nbest_ids.size(1), -1)
-        nbest_prediction_valid_token_lengths = prediction_valid_token_lengths.view(
-            batch_size * beam_size)[index.view(-1)].view(batch_size, -1)
-        nbest_hyps: List[List[Dict[str, torch.Tensor]]] = []
-        for i in range(batch_size):
-            i_best_hyps: List[Dict[str, torch.Tensor]] = []
-            for j, score in enumerate(nbest_scores[i]):
-                hyp = {
-                    "token_ids": nbest_prediction_tokens[i, j, 1:nbest_prediction_valid_token_lengths[i, j]],
-                    "score": score
-                }
-                i_best_hyps.append(hyp)
-            nbest_hyps.append(i_best_hyps)
-
-        return nbest_hyps
-        
-    def get_initialized_self_cache(self, 
-                                   batch_size, 
-                                   beam_size
-                                   ) -> Tuple[Tensor, Tensor]:
-        n_layer_self_k_cache = torch.zeros(
-            self.num_decoder_blocks,
-            batch_size * beam_size,
-            self.decode_max_len,
-            self.decoder_hidden_dim,
-        )
-        n_layer_self_v_cache = torch.zeros(
-            self.num_decoder_blocks,
-            batch_size * beam_size,
-            self.decode_max_len,
-            self.decoder_hidden_dim,
-        )
-        return n_layer_self_k_cache, n_layer_self_v_cache
-    
-    def calc_feat_len(self, audio_dur):
-        import math
-        sample_rate = 16000
-        frame_length = 25 * sample_rate / 1000
-        frame_shift = 10 * sample_rate / 1000
-        length = math.floor((audio_dur * sample_rate - frame_length) / frame_shift) + 1
-        return length
- 
-    def transcribe(self, 
-                   batch_wav_path: List[str],
-                   beam_size: int = 1,
-                   nbest: int = 1
-                ) -> List[Dict]:
-        feats, lengths, wav_durations = self.feature_extractor(batch_wav_path)
-        # print(f"feats.shape: {feats.shape}")
-        maxlen = self.calc_feat_len(10)
-        if feats.shape[1] < maxlen:
-            feats = np.concatenate([feats, np.zeros((1, maxlen - feats.shape[1], 80), dtype=np.float32)], axis=1)
-        feats = feats[:, :maxlen, :]
-
-        # encoder_data_path = os.path.join("calib_dataset", "encoder", os.path.basename(batch_wav_path[0]))
-        # decoder_data_path = os.path.join("calib_dataset", "decoder", os.path.basename(batch_wav_path[0]))
-        # os.makedirs(encoder_data_path, exist_ok=True)
-        # os.makedirs(decoder_data_path, exist_ok=True)
-
-        feats = to_numpy(feats)
-        lengths = to_numpy(lengths).astype(np.int32)
-
-        # for name, npy in zip(["encoder_input", "encoder_input_lengths"], [feats, lengths]):
-        #     file_path = os.path.join(encoder_data_path, name + ".npy")
-        #     np.save(file_path, npy)
-
-        start_time = time.time()
-        n_layer_cross_k, n_layer_cross_v, cross_attn_mask = self.run_encoder(
-            to_numpy(feats),
-            to_numpy(lengths)
-        )
-        print(f"run encoder take {(time.time() - start_time) * 1000}ms")
-        nbest_hyps = self.run_decoder(n_layer_cross_k,
-                                      n_layer_cross_v,
-                                      cross_attn_mask,
-                                      beam_size,
-                                      nbest,
-                                      )
-        transcribe_durations = time.time() - start_time
-        results: List[Dict] = []
-        for wav, hyp in zip(batch_wav_path, nbest_hyps):
-            hyp = hyp[0]
-            hyp_ids = [int(id) for id in hyp["token_ids"].cpu()]
-            score = hyp["score"].item()
-            text = self.tokenizer.detokenize(hyp_ids)
-            results.append(
-                {
-                    "wav": wav,
-                    "text": text,
-                    "score": score
-                }
-            )
-            
-        return results, wav_durations, transcribe_durations   
- 
  
 def parse_args():
-    parser = argparse.ArgumentParser(description="FireRedASROnnxModel Test")
+    parser = argparse.ArgumentParser(description="FireRedASRAxModel Test")
     parser.add_argument(
         "--encoder", 
         type=str, 
         default="axmodel/encoder.axmodel",
-        help="Path to onnx encoder"
+        help="Path to axmodel encoder"
     )
     parser.add_argument(
-        "--decoder", 
+        "--decoder_loop", 
         type=str, 
-        default="axmodel/decoder_main.axmodel",
-        help="Path to onnx decoder"
+        default="axmodel/decoder_loop.axmodel",
+        help="Path to axmodel decoder loop"
     )
     parser.add_argument(
         "--cmvn",
@@ -585,10 +69,16 @@ def parse_args():
         help=""
     )
     parser.add_argument(
-        "--max_len",
+        "--decode_max_len",
         type=int,
         default=128,
-        help=""
+        help="max token len"
+    )
+    parser.add_argument(
+        "--max_dur",
+        type=int,
+        default=10,
+        help="max audio len"
     )
     
     return parser.parse_args()
@@ -611,12 +101,14 @@ def main():
     args = parse_args()
     print(args)
     
-    onnx_model = FireRedASROnnxModel(args.encoder,
-                                     args.decoder,
+    model = FireRedASRAxModel(args.encoder,
+                                     args.decoder_main,
+                                     args.decoder_loop,
                                      args.cmvn,
                                      args.dict,
                                      args.spm_model,
-                                     decode_max_len=args.max_len
+                                     decode_max_len=args.decode_max_len,
+                                     audio_dur=args.max_dur
                                      )
     
     wf = open(args.hypo, "wt")
@@ -626,7 +118,7 @@ def main():
     total_transcribe_durations = 0
     for wav in wavlist:
         batch_wav = [wav]
-        results, wav_durations, transcribe_durations = onnx_model.transcribe(
+        results, wav_durations, transcribe_durations = model.transcribe(
             batch_wav, args.beam_size, args.nbest)
         
         wav_durations = sum(wav_durations)

test_wer.py

@@ -0,0 +1,348 @@
+import argparse
+import os
+import logging
+import re
+from fireredasr_axmodel import FireRedASRAxModel
+
+
+def setup_logging():
+    """配置日志系统，同时输出到控制台和文件"""
+    # 获取脚本所在目录
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    log_file = os.path.join(script_dir, "test_wer.log")
+    
+    # 配置日志格式
+    log_format = '%(asctime)s - %(levelname)s - %(message)s'
+    date_format = '%Y-%m-%d %H:%M:%S'
+    
+    # 创建logger
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    
+    # 清除现有的handler
+    for handler in logger.handlers[:]:
+        logger.removeHandler(handler)
+    
+    # 创建文件handler
+    file_handler = logging.FileHandler(log_file, mode='a', encoding='utf-8')
+    file_handler.setLevel(logging.INFO)
+    file_formatter = logging.Formatter(log_format, date_format)
+    file_handler.setFormatter(file_formatter)
+    
+    # 创建控制台handler
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.INFO)
+    console_formatter = logging.Formatter(log_format, date_format)
+    console_handler.setFormatter(console_formatter)
+    
+    # 添加handler到logger
+    logger.addHandler(file_handler)
+    logger.addHandler(console_handler)
+    
+    return logger
+
+
+class AIShellDataset:
+    def __init__(self, gt_path: str, voice_dir='wav'):
+        """
+        初始化数据集
+        
+        Args:
+            json_path: voice.json文件的路径
+        """
+        self.gt_path = gt_path
+        self.dataset_dir = os.path.dirname(gt_path)
+        self.voice_dir = os.path.join(self.dataset_dir, voice_dir)
+        
+        # 检查必要文件和文件夹是否存在
+        assert os.path.exists(gt_path), f"gt文件不存在: {gt_path}"
+        assert os.path.exists(self.voice_dir), f"文件夹不存在: {self.voice_dir}"
+        
+        # 加载数据
+        self.data = []
+        with open(gt_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                audio_path, gt = line.split(" ")
+                audio_path = os.path.join(self.voice_dir, audio_path + ".wav")
+                self.data.append({"audio_path": audio_path, "gt": gt})
+
+        # 使用logging而不是print
+        logger = logging.getLogger()
+        logger.info(f"加载了 {len(self.data)} 条数据")
+    
+    def __iter__(self):
+        """返回迭代器"""
+        self.index = 0
+        return self
+    
+    def __next__(self):
+        """返回下一个数据项"""
+        if self.index >= len(self.data):
+            raise StopIteration
+        
+        item = self.data[self.index]
+        audio_path = item["audio_path"]
+        ground_truth = item["gt"]
+        
+        self.index += 1
+        return audio_path, ground_truth
+    
+    def __len__(self):
+        """返回数据集大小"""
+        return len(self.data)
+    
+
+class CommonVoiceDataset:
+    """Common Voice数据集解析器"""
+    
+    def __init__(self, tsv_path: str):
+        """
+        初始化数据集
+        
+        Args:
+            json_path: voice.json文件的路径
+        """
+        self.tsv_path = tsv_path
+        self.dataset_dir = os.path.dirname(tsv_path)
+        self.voice_dir = os.path.join(self.dataset_dir, "clips")
+        
+        # 检查必要文件和文件夹是否存在
+        assert os.path.exists(tsv_path), f"{tsv_path}文件不存在: {tsv_path}"
+        assert os.path.exists(self.voice_dir), f"voice文件夹不存在: {self.voice_dir}"
+        
+        # 加载JSON数据
+        self.data = []
+        with open(tsv_path, 'r', encoding='utf-8') as f:
+            f.readline()
+            for line in f:
+                line = line.strip()
+                splits = line.split("\t")
+                audio_path = splits[1]
+                gt = splits[2]
+                audio_path = os.path.join(self.voice_dir, audio_path)
+                self.data.append({"audio_path": audio_path, "gt": gt})
+        
+        # 使用logging而不是print
+        logger = logging.getLogger()
+        logger.info(f"加载了 {len(self.data)} 条数据")
+    
+    def __iter__(self):
+        """返回迭代器"""
+        self.index = 0
+        return self
+    
+    def __next__(self):
+        """返回下一个数据项"""
+        if self.index >= len(self.data):
+            raise StopIteration
+        
+        item = self.data[self.index]
+        audio_path = item["audio_path"]
+        ground_truth = item["gt"]
+        
+        self.index += 1
+        return audio_path, ground_truth
+    
+    def __len__(self):
+        """返回数据集大小"""
+        return len(self.data)
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        prog="whisper",
+        description="Test WER on dataset"
+    )
+    parser.add_argument("--dataset", "-d", type=str, required=True, choices=["aishell", "common_voice"], help="Test dataset")
+    parser.add_argument("--gt_path", "-g", type=str, required=True, help="Test dataset ground truth file")
+    parser.add_argument("--max_num", type=int, default=-1, required=False, help="Maximum test data num")
+    parser.add_argument("--language", "-l", type=str, required=False, default="zh", help="Target language, support en, zh, ja, and others. See languages.py for more options.")
+    parser.add_argument(
+        "--encoder", 
+        type=str, 
+        default="axmodel/encoder.axmodel",
+        help="Path to onnx encoder"
+    )
+    parser.add_argument(
+        "--decoder_main", 
+        type=str, 
+        default="axmodel/decoder_main.axmodel",
+        help="Path to axmodel decoder main"
+    )
+    parser.add_argument(
+        "--decoder_loop", 
+        type=str, 
+        default="axmodel/decoder_loop.axmodel",
+        help="Path to axmodel decoder loop"
+    )
+    parser.add_argument(
+        "--cmvn",
+        type=str,
+        default="axmodel/cmvn.ark",
+        help="Path to cmvn"
+    )
+    parser.add_argument(
+        "--dict",
+        type=str,
+        default="axmodel/dict.txt",
+        help="Path to dict"
+    )
+    parser.add_argument(
+        "--spm_model",
+        type=str,
+        default="axmodel/train_bpe1000.model",
+        help="Path to spm model"
+    )
+    parser.add_argument(
+        "--wavlist",
+        type=str,
+        default="wavlist.txt",
+        help="File to wav path list"
+    )
+    parser.add_argument(
+        "--hypo",
+        type=str,
+        default="hypo_axmodel.txt",
+        help="File of hypos"
+    )
+    parser.add_argument(
+        "--beam_size",
+        type=int,
+        default=3,
+        help=""
+    )
+    parser.add_argument(
+        "--nbest",
+        type=int,
+        default=1,
+        help=""
+    )
+    parser.add_argument(
+        "--max_len",
+        type=int,
+        default=128,
+        help=""
+    )
+    return parser.parse_args()
+
+
+def print_args(args):
+    logger = logging.getLogger()
+    logger.info(f"dataset: {args.dataset}")
+    logger.info(f"gt_path: {args.gt_path}")
+    logger.info(f"max_num: {args.max_num}")
+    logger.info(f"language: {args.language}")
+
+
+def min_distance(word1: str, word2: str) -> int:
+ 
+    row = len(word1) + 1
+    column = len(word2) + 1
+ 
+    cache = [ [0]*column for i in range(row) ]
+ 
+    for i in range(row):
+        for j in range(column):
+ 
+            if i ==0 and j ==0:
+                cache[i][j] = 0
+            elif i == 0 and j!=0:
+                cache[i][j] = j
+            elif j == 0 and i!=0:
+                cache[i][j] = i
+            else:
+                if word1[i-1] == word2[j-1]:
+                    cache[i][j] = cache[i-1][j-1]
+                else:
+                    replace = cache[i-1][j-1] + 1
+                    insert = cache[i][j-1] + 1
+                    remove = cache[i-1][j] + 1
+ 
+                    cache[i][j] = min(replace, insert, remove)
+ 
+    return cache[row-1][column-1]
+
+
+def remove_punctuation(text):
+    # 定义正则表达式模式，匹配所有标点符号
+    # 这个模式包括常见的标点符号和中文标点
+    pattern = r'[^\w\s]|_'
+    
+    # 使用sub方法将所有匹配的标点符号替换为空字符串
+    cleaned_text = re.sub(pattern, '', text)
+    
+    return cleaned_text
+
+
+def main():
+    # 设置日志系统
+    logger = setup_logging()
+
+    args = get_args()
+    print_args(args)
+
+    dataset_type = args.dataset.lower()
+    if dataset_type == "aishell":
+        dataset = AIShellDataset(args.gt_path)
+    elif dataset_type == "common_voice":
+        dataset = CommonVoiceDataset(args.gt_path)
+    else:
+        raise ValueError(f"Unknown dataset type {dataset_type}")
+
+    max_num = args.max_num
+
+    # Load model
+    model = FireRedASRAxModel(args.encoder,
+                            args.decoder_main,
+                            args.decoder_loop,
+                            args.cmvn,
+                            args.dict,
+                            args.spm_model,
+                            decode_max_len=args.max_len,
+                            audio_dur=10
+    )
+    
+
+    # Iterate over dataset
+    references = []
+    hyp = []
+    all_character_error_num = 0
+    all_character_num = 0
+    wer_file = open("wer.txt", "w")
+    max_data_num = max_num if max_num > 0 else len(dataset)
+    for n, (audio_path, reference) in enumerate(dataset):
+        batch_uttid = [os.path.splitext(os.path.basename(audio_path))[0]]
+        batch_wav = [audio_path]
+        results, _, _ = model.transcribe(
+            batch_wav, args.beam_size, args.nbest)
+
+        hypothesis = results[0]['text']
+
+        hypothesis = remove_punctuation(hypothesis)
+        reference = remove_punctuation(reference)
+
+        character_error_num = min_distance(reference, hypothesis)
+        character_num = len(reference)
+        character_error_rate = character_error_num / character_num * 100
+
+        all_character_error_num += character_error_num
+        all_character_num += character_num
+
+        hyp.append(hypothesis)
+        references.append(reference)
+        
+        line_content = f"({n+1}/{max_data_num}) {os.path.basename(audio_path)}  gt: {reference}  predict: {hypothesis}  WER: {character_error_rate}%"
+        wer_file.write(line_content + "\n")
+        logger.info(line_content)
+
+        if n + 1 >= max_data_num:
+            break
+
+    total_character_error_rate = all_character_error_num / all_character_num * 100
+
+    logger.info(f"Total WER: {total_character_error_rate}%")
+    wer_file.write(f"Total WER: {total_character_error_rate}%")
+    wer_file.close()
+
+if __name__ == "__main__":
+    main()