Reformat codes

SenseVoiceAx.py

@@ -4,7 +4,7 @@ import librosa
 from frontend import WavFrontend
 import os
 import time
-from typing import List, Union
+from typing import List, Union, Optional
 from asr_decoder import CTCDecoder
 from tokenizer import SentencepiecesTokenizer
 from online_fbank import OnlineFbank
@@ -15,93 +15,117 @@ def sequence_mask(lengths, maxlen=None, dtype=np.float32):
     # 如果 maxlen 未指定，则取 lengths 中的最大值
     if maxlen is None:
         maxlen = np.max(lengths)
-    
+
     # 创建一个从 0 到 maxlen-1 的行向量
     row_vector = np.arange(0, maxlen, 1)
-    
+
     # 将 lengths 转换为列向量
     matrix = np.expand_dims(lengths, axis=-1)
-    
+
     # 比较生成掩码
     mask = row_vector < matrix
     if mask.shape[-1] < lengths[0]:
-        mask = np.concatenate([mask, np.zeros((mask.shape[0], lengths[0] - mask.shape[-1]), dtype=np.float32)], axis=-1)
-    
+        mask = np.concatenate(
+            [
+                mask,
+                np.zeros(
+                    (mask.shape[0], lengths[0] - mask.shape[-1]), dtype=np.float32
+                ),
+            ],
+            axis=-1,
+        )
+
     # 返回指定数据类型的掩码
     return mask.astype(dtype)[None, ...]
-    
+
 
 def unique_consecutive_np(arr):
     """
     找出数组中连续的唯一值，模拟 torch.unique_consecutive(yseq, dim=-1)
-    
+
     参数:
     arr: 一维numpy数组
-    
+
     返回:
     unique_values: 去除连续重复值后的数组
     """
     if len(arr) == 0:
         return np.array([])
-    
+
     if len(arr) == 1:
         return arr.copy()
-    
+
     # 找出变化的位置
     diff = np.diff(arr)
     change_positions = np.where(diff != 0)[0] + 1
-    
+
     # 添加起始位置
     start_positions = np.concatenate(([0], change_positions))
-    
+
     # 获取唯一值（每个连续段的第一个值）
     unique_values = arr[start_positions]
-    
-    return unique_values
-
 
-class Tokenizer:
-    def __init__(self, symbol_path):
-        self.symbol_tables = {}
-        with open(symbol_path, 'r') as f:
-            i = 0
-            for line in f:
-                token = line.strip()
-                self.symbol_tables[token] = i
-                i += 1
-        
-    def tokens2text(self, token):
-        return self.symbol_tables[token]
+    return unique_values
 
 
 class SenseVoiceAx:
-    def __init__(self, model_path, 
-                 max_len=256, 
-                 beam_size=3,
-                 language="auto", 
-                 hot_words=Union[List[str], None],
-                 use_itn=True, 
-                 streaming=False):
+    """ SenseVoice axmodel runner """
+    
+    def __init__(
+        self,
+        model_path: str,
+        max_len: int = 256,
+        beam_size: int = 3,
+        language: str = "auto",
+        hot_words: Optional[List[str]] = None,
+        use_itn: bool = True,
+        streaming: bool = False,
+    ):
+        """
+        Initialize SenseVoiceAx
+
+        Args:
+            model_path: Path of axmodel
+            max_len:    Fixed shape of input of axmodel
+            beam_size:  Max number of hypos to hold after each decode step
+            language:   Support auto, zh(Chinese), en(English), yue(Cantonese), ja(Japanese), ko(Korean)
+            hot_words:  Words that may fail to recognize, 
+                        special words/phrases (aka hotwords) like rare words, personalized information etc. 
+            use_itn:    Allow Invert Text Normalization if True, 
+                        ITN converts ASR model output into its written form to improve text readability,
+                        For example, the ITN module replaces “one hundred and twenty-three dollars” transcribed by an ASR model with “$123.”
+            streaming:  Processes audio in small segments or "chunks" sequentially and outputs text on the fly.
+                        Use stream_infer method if streaming is true otherwise infer.              
+
+        """
         model_path_root = os.path.dirname(model_path)
-        emb_path = os.path.join(model_path_root, '../embeddings.npy')
-        cmvn_file = os.path.join(model_path_root, '../am.mvn')
-        bpe_model = os.path.join(model_path_root, '../chn_jpn_yue_eng_ko_spectok.bpe.model')
+        emb_path = os.path.join(model_path_root, "../embeddings.npy")
+        cmvn_file = os.path.join(model_path_root, "../am.mvn")
+        bpe_model = os.path.join(
+            model_path_root, "../chn_jpn_yue_eng_ko_spectok.bpe.model"
+        )
         if streaming:
-            self.position_encoding = np.load(os.path.join(model_path_root, '../pe_streaming.npy'))
+            self.position_encoding = np.load(
+                os.path.join(model_path_root, "../pe_streaming.npy")
+            )
         else:
-            self.position_encoding = np.load(os.path.join(model_path_root, '../pe_nonstream.npy'))
+            self.position_encoding = np.load(
+                os.path.join(model_path_root, "../pe_nonstream.npy")
+            )
 
         self.streaming = streaming
         self.tokenizer = SentencepiecesTokenizer(bpemodel=bpe_model)
 
-        self.frontend = WavFrontend(cmvn_file=cmvn_file,
-                                    fs=16000, 
-                                    window="hamming", 
-                                    n_mels=80, 
-                                    frame_length=25, 
-                                    frame_shift=10,
-                                    lfr_m=7,
-                                    lfr_n=6)
+        self.frontend = WavFrontend(
+            cmvn_file=cmvn_file,
+            fs=16000,
+            window="hamming",
+            n_mels=80,
+            frame_length=25,
+            frame_shift=10,
+            lfr_m=7,
+            lfr_n=6,
+        )
         self.model = axe.InferenceSession(model_path)
         self.sample_rate = 16000
         self.blank_id = 0
@@ -109,11 +133,32 @@ class SenseVoiceAx:
         self.padding = 16
         self.input_size = 560
 
-        self.lid_dict = {"auto": 0, "zh": 3, "en": 4, "yue": 7, "ja": 11, "ko": 12, "nospeech": 13}
-        self.lid_int_dict = {24884: 3, 24885: 4, 24888: 7, 24892: 11, 24896: 12, 24992: 13}
+        self.lid_dict = {
+            "auto": 0,
+            "zh": 3,
+            "en": 4,
+            "yue": 7,
+            "ja": 11,
+            "ko": 12,
+            "nospeech": 13,
+        }
+        self.lid_int_dict = {
+            24884: 3,
+            24885: 4,
+            24888: 7,
+            24892: 11,
+            24896: 12,
+            24992: 13,
+        }
         self.textnorm_dict = {"withitn": 14, "woitn": 15}
         self.textnorm_int_dict = {25016: 14, 25017: 15}
-        self.emo_dict = {"unk": 25009, "happy": 25001, "sad": 25002, "angry": 25003, "neutral": 25004}
+        self.emo_dict = {
+            "unk": 25009,
+            "happy": 25001,
+            "sad": 25002,
+            "angry": 25003,
+            "neutral": 25004,
+        }
 
         self.load_embeddings(emb_path, language, use_itn)
         self.language = language
@@ -135,39 +180,48 @@ class SenseVoiceAx:
             self.caches_shape = (max_len, self.input_size)
             self.caches = np.zeros(self.caches_shape, dtype=np.float32)
             self.zeros = np.zeros((1, self.input_size), dtype=np.float32)
-            self.neg_mean, self.inv_stddev = self.frontend.cmvn[0, :], self.frontend.cmvn[1, :]
+            self.neg_mean, self.inv_stddev = (
+                self.frontend.cmvn[0, :],
+                self.frontend.cmvn[1, :],
+            )
 
             self.fbank = OnlineFbank(window_type="hamming")
-            self.masks = sequence_mask(np.array([self.max_len], dtype=np.int32), maxlen=self.max_len, dtype=np.float32)
-
+            self.masks = sequence_mask(
+                np.array([self.max_len], dtype=np.int32),
+                maxlen=self.max_len,
+                dtype=np.float32,
+            )
 
     @property
     def language_options(self):
         return list(self.lid_dict.keys())
-    
+
     @property
     def textnorm_options(self):
         return list(self.textnorm_dict.keys())
-    
+
     def load_embeddings(self, emb_path, language, use_itn):
         self.embeddings = np.load(emb_path, allow_pickle=True).item()
         self.language_query = self.embeddings[language]
-        self.textnorm_query = self.embeddings['withitn'] if use_itn else self.embeddings['woitn']
-        self.event_emo_query = self.embeddings['event_emo']
-        self.input_query = np.concatenate((self.textnorm_query, self.language_query, self.event_emo_query), axis=1)
+        self.textnorm_query = (
+            self.embeddings["withitn"] if use_itn else self.embeddings["woitn"]
+        )
+        self.event_emo_query = self.embeddings["event_emo"]
+        self.input_query = np.concatenate(
+            (self.textnorm_query, self.language_query, self.event_emo_query), axis=1
+        )
         self.query_num = self.input_query.shape[1]
 
-    
     def choose_language(self, language):
         self.language_query = self.embeddings[language]
-        self.input_query = np.concatenate((self.textnorm_query, self.language_query, self.event_emo_query), axis=1)
+        self.input_query = np.concatenate(
+            (self.textnorm_query, self.language_query, self.event_emo_query), axis=1
+        )
         self.language = language
 
-
     def load_data(self, filepath: str) -> np.ndarray:
         waveform, _ = librosa.load(filepath, sr=self.sample_rate)
         return waveform.flatten()
-    
 
     @staticmethod
     def pad_feats(feats: List[np.ndarray], max_feat_len: int) -> np.ndarray:
@@ -179,7 +233,6 @@ class SenseVoiceAx:
         feats = np.array(feat_res).astype(np.float32)
         return feats
 
-
     def preprocess(self, waveform):
         feats, feats_len = [], []
         for wf in [waveform]:
@@ -191,11 +244,10 @@ class SenseVoiceAx:
         feats = self.pad_feats(feats, np.max(feats_len))
         feats_len = np.array(feats_len).astype(np.int32)
         return feats, feats_len
-    
 
     def postprocess(self, ctc_logits, encoder_out_lens):
         # 提取数据
-        x = ctc_logits[0, 4:encoder_out_lens[0], :]
+        x = ctc_logits[0, 4 : encoder_out_lens[0], :]
 
         # 获取最大值索引
         yseq = np.argmax(x, axis=-1)
@@ -208,7 +260,6 @@ class SenseVoiceAx:
         token_int = yseq[mask].tolist()
 
         return token_int
-    
 
     def infer_waveform(self, waveform: np.ndarray, language="auto"):
         if language != self.language:
@@ -224,32 +275,46 @@ class SenseVoiceAx:
         asr_res = []
         for i in range(slice_num):
             if i == 0:
-                sub_feat = feat[:, i*slice_len:(i+1)*slice_len, :]
+                sub_feat = feat[:, i * slice_len : (i + 1) * slice_len, :]
             else:
-                sub_feat = feat[:, i*slice_len - self.padding:(i+1)*slice_len - self.padding, :]
+                sub_feat = feat[
+                    :,
+                    i * slice_len - self.padding : (i + 1) * slice_len - self.padding,
+                    :,
+                ]
             # concat query
             sub_feat = np.concatenate([self.input_query, sub_feat], axis=1)
             real_len = sub_feat.shape[1]
             if real_len < self.max_len:
-                sub_feat = np.concatenate([
-                        sub_feat, 
-                        np.zeros((1, self.max_len - real_len, sub_feat.shape[-1]), dtype=np.float32)
+                sub_feat = np.concatenate(
+                    [
+                        sub_feat,
+                        np.zeros(
+                            (1, self.max_len - real_len, sub_feat.shape[-1]),
+                            dtype=np.float32,
+                        ),
                     ],
-                    axis=1)
-                
-            masks = sequence_mask(np.array([self.max_len], dtype=np.int32), maxlen=real_len, dtype=np.float32)
+                    axis=1,
+                )
+
+            masks = sequence_mask(
+                np.array([self.max_len], dtype=np.int32),
+                maxlen=real_len,
+                dtype=np.float32,
+            )
 
             # start = time.time()
-            outputs = self.model.run(None, {"speech": sub_feat,
-                                            "masks": masks,
-                                            "position_encoding": self.position_encoding})
+            outputs = self.model.run(
+                None,
+                {
+                    "speech": sub_feat,
+                    "masks": masks,
+                    "position_encoding": self.position_encoding,
+                },
+            )
             ctc_logits, encoder_out_lens = outputs
-            # print(f"ctc_logits.shape: {ctc_logits.shape}")
-            # print(f"Run model take {time.time() - start}s")
 
-            # start = time.time()
             token_int = self.postprocess(ctc_logits, encoder_out_lens)
-            # print(f"Postprocess take {time.time() - start}s")
 
             if self.tokenizer is not None:
                 asr_res.append(self.tokenizer.tokens2text(token_int))
@@ -257,9 +322,12 @@ class SenseVoiceAx:
                 asr_res.append(token_int)
 
         return asr_res
-    
 
-    def infer(self, filepath_or_data: Union[np.ndarray, str], language="auto", print_rtf=True):
+    def infer(
+        self, filepath_or_data: Union[np.ndarray, str], language="auto", print_rtf=False
+    ):
+        assert not self.streaming, "This method is for non-streaming model"
+
         if isinstance(filepath_or_data, str):
             waveform = self.load_data(filepath_or_data)
         else:
@@ -284,22 +352,21 @@ class SenseVoiceAx:
             times_ms.append(step * 60)
         return times_ms, self.tokenizer.decode(tokens)
 
-
     def reset(self):
         self.cur_idx = -1
         self.decoder.reset()
         self.fbank = OnlineFbank(window_type="hamming")
         self.caches = np.zeros(self.caches_shape)
 
-
     def get_size(self):
         effective_size = self.cur_idx + 1 - self.padding
         if effective_size <= 0:
             return 0
         return effective_size % self.chunk_size or self.chunk_size
-    
 
     def stream_infer(self, audio, is_last, language="auto"):
+        assert self.streaming, "This method is for streaming model"
+
         if language != self.language:
             self.choose_language(language)
 
@@ -321,13 +388,18 @@ class SenseVoiceAx:
                 continue
 
             speech = self.caches[None, ...]
-            outputs = self.model.run(None, {"speech": speech,
-                                            "masks": self.masks,
-                                            "position_encoding": self.position_encoding})
+            outputs = self.model.run(
+                None,
+                {
+                    "speech": speech,
+                    "masks": self.masks,
+                    "position_encoding": self.position_encoding,
+                },
+            )
             ctc_logits, encoder_out_lens = outputs
-            probs = ctc_logits[0, 4:encoder_out_lens[0]]
+            probs = ctc_logits[0, 4 : encoder_out_lens[0]]
             probs = torch.from_numpy(probs)
-            
+
             if cur_size != self.chunk_size:
                 probs = probs[self.chunk_size - cur_size :]
             if not is_last:

download_utils.py

@@ -1,4 +1,5 @@
 import os
+
 # Speed up hf download using mirror url
 os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
 from huggingface_hub import snapshot_download
@@ -7,6 +8,7 @@ current_file_path = os.path.dirname(__file__)
 REPO_ROOT = "AXERA-TECH"
 CACHE_PATH = os.path.join(current_file_path, "models")
 
+
 def download_model(model_name: str) -> str:
     """
     Download model from AXERA-TECH's huggingface space.
@@ -23,7 +25,9 @@ def download_model(model_name: str) -> str:
     model_path = os.path.join(CACHE_PATH, model_name)
     if not os.path.exists(model_path):
         print(f"Downloading {model_name}...")
-        snapshot_download(repo_id=f"{REPO_ROOT}/{model_name}", 
-                          local_dir=os.path.join(CACHE_PATH, model_name))
-    
+        snapshot_download(
+            repo_id=f"{REPO_ROOT}/{model_name}",
+            local_dir=os.path.join(CACHE_PATH, model_name),
+        )
+
     return model_path

frontend.py

@@ -96,7 +96,9 @@ class WavFrontend:
         T = T + (lfr_m - 1) // 2
         for i in range(T_lfr):
             if lfr_m <= T - i * lfr_n:
-                LFR_inputs.append((inputs[i * lfr_n : i * lfr_n + lfr_m]).reshape(1, -1))
+                LFR_inputs.append(
+                    (inputs[i * lfr_n : i * lfr_n + lfr_m]).reshape(1, -1)
+                )
             else:
                 # process last LFR frame
                 num_padding = lfr_m - (T - i * lfr_n)
@@ -180,7 +182,9 @@ class WavFrontendOnline(WavFrontend):
         splice_idx = T_lfr
         for i in range(T_lfr):
             if lfr_m <= T - i * lfr_n:
-                LFR_inputs.append((inputs[i * lfr_n : i * lfr_n + lfr_m]).reshape(1, -1))
+                LFR_inputs.append(
+                    (inputs[i * lfr_n : i * lfr_n + lfr_m]).reshape(1, -1)
+                )
             else:  # process last LFR frame
                 if is_final:
                     num_padding = lfr_m - (T - i * lfr_n)
@@ -201,8 +205,12 @@ class WavFrontendOnline(WavFrontend):
     def compute_frame_num(
         sample_length: int, frame_sample_length: int, frame_shift_sample_length: int
     ) -> int:
-        frame_num = int((sample_length - frame_sample_length) / frame_shift_sample_length + 1)
-        return frame_num if frame_num >= 1 and sample_length >= frame_sample_length else 0
+        frame_num = int(
+            (sample_length - frame_sample_length) / frame_shift_sample_length + 1
+        )
+        return (
+            frame_num if frame_num >= 1 and sample_length >= frame_sample_length else 0
+        )
 
     def fbank(
         self, input: np.ndarray, input_lengths: np.ndarray
@@ -238,7 +246,9 @@ class WavFrontendOnline(WavFrontend):
                 )
                 waveform = waveform * (1 << 15)
 
-                self.fbank_fn.accept_waveform(self.opts.frame_opts.samp_freq, waveform.tolist())
+                self.fbank_fn.accept_waveform(
+                    self.opts.frame_opts.samp_freq, waveform.tolist()
+                )
                 frames = self.fbank_fn.num_frames_ready
                 mat = np.empty([frames, self.opts.mel_opts.num_bins])
                 for i in range(frames):
@@ -291,7 +301,9 @@ class WavFrontendOnline(WavFrontend):
         assert (
             batch_size == 1
         ), "we support to extract feature online only when the batch size is equal to 1 now"
-        waveforms, feats, feats_lengths = self.fbank(input, input_lengths)  # input shape: B T D
+        waveforms, feats, feats_lengths = self.fbank(
+            input, input_lengths
+        )  # input shape: B T D
         if feats.shape[0]:
             self.waveforms = (
                 waveforms
@@ -301,7 +313,9 @@ class WavFrontendOnline(WavFrontend):
             if not self.lfr_splice_cache:
                 for i in range(batch_size):
                     self.lfr_splice_cache.append(
-                        np.expand_dims(feats[i][0, :], axis=0).repeat((self.lfr_m - 1) // 2, axis=0)
+                        np.expand_dims(feats[i][0, :], axis=0).repeat(
+                            (self.lfr_m - 1) // 2, axis=0
+                        )
                     )
 
             if feats_lengths[0] + self.lfr_splice_cache[0].shape[0] >= self.lfr_m:
@@ -313,7 +327,9 @@ class WavFrontendOnline(WavFrontend):
                     / self.frame_shift_sample_length
                     + 1
                 )
-                minus_frame = (self.lfr_m - 1) // 2 if self.reserve_waveforms is None else 0
+                minus_frame = (
+                    (self.lfr_m - 1) // 2 if self.reserve_waveforms is None else 0
+                )
                 feats, feats_lengths, lfr_splice_frame_idxs = self.lfr_cmvn(
                     feats, feats_lengths, is_final
                 )
@@ -346,7 +362,9 @@ class WavFrontendOnline(WavFrontend):
         else:
             if is_final:
                 self.waveforms = (
-                    waveforms if self.reserve_waveforms is None else self.reserve_waveforms
+                    waveforms
+                    if self.reserve_waveforms is None
+                    else self.reserve_waveforms
                 )
                 feats = np.stack(self.lfr_splice_cache)
                 feats_lengths = np.zeros(batch_size, dtype=np.int32) + feats.shape[1]
@@ -377,20 +395,33 @@ def load_bytes(input):
     i = np.iinfo(middle_data.dtype)
     abs_max = 2 ** (i.bits - 1)
     offset = i.min + abs_max
-    array = np.frombuffer((middle_data.astype(dtype) - offset) / abs_max, dtype=np.float32)
+    array = np.frombuffer(
+        (middle_data.astype(dtype) - offset) / abs_max, dtype=np.float32
+    )
     return array
 
 
 class SinusoidalPositionEncoderOnline:
     """Streaming Positional encoding."""
 
-    def encode(self, positions: np.ndarray = None, depth: int = None, dtype: np.dtype = np.float32):
+    def encode(
+        self,
+        positions: np.ndarray = None,
+        depth: int = None,
+        dtype: np.dtype = np.float32,
+    ):
         batch_size = positions.shape[0]
         positions = positions.astype(dtype)
-        log_timescale_increment = np.log(np.array([10000], dtype=dtype)) / (depth / 2 - 1)
-        inv_timescales = np.exp(np.arange(depth / 2).astype(dtype) * (-log_timescale_increment))
+        log_timescale_increment = np.log(np.array([10000], dtype=dtype)) / (
+            depth / 2 - 1
+        )
+        inv_timescales = np.exp(
+            np.arange(depth / 2).astype(dtype) * (-log_timescale_increment)
+        )
         inv_timescales = np.reshape(inv_timescales, [batch_size, -1])
-        scaled_time = np.reshape(positions, [1, -1, 1]) * np.reshape(inv_timescales, [1, 1, -1])
+        scaled_time = np.reshape(positions, [1, -1, 1]) * np.reshape(
+            inv_timescales, [1, 1, -1]
+        )
         encoding = np.concatenate((np.sin(scaled_time), np.cos(scaled_time)), axis=2)
         return encoding.astype(dtype)
 

gradio_demo.py

@@ -5,7 +5,7 @@ from tokenizer import SentencepiecesTokenizer
 from print_utils import rich_transcription_postprocess
 from download_utils import download_model
 
-use_itn = True # 标点符号预测
+use_itn = True  # 标点符号预测
 max_len = 256
 
 model_path = os.path.join("sensevoice_ax650", "sensevoice.axmodel")
@@ -14,11 +14,10 @@ bpemodel = "chn_jpn_yue_eng_ko_spectok.bpe.model"
 assert os.path.exists(model_path), f"model {model_path} not exist"
 
 tokenizer = SentencepiecesTokenizer(bpemodel=bpemodel)
-pipeline = SenseVoiceAx(model_path, 
-                        max_len=max_len,
-                        language="auto", 
-                        use_itn=use_itn, 
-                        tokenizer=tokenizer)
+pipeline = SenseVoiceAx(
+    model_path, max_len=max_len, language="auto", use_itn=use_itn, tokenizer=tokenizer
+)
+
 
 def speech_to_text(audio_path, lang):
     """
@@ -27,7 +26,7 @@ def speech_to_text(audio_path, lang):
     """
     if not audio_path:
         return "无音频"
-    
+
     pipeline.choose_language(language=lang)
     asr_res = pipeline.infer(audio_path, print_rtf=True)
     res = " ".join([rich_transcription_postprocess(i) for i in asr_res])
@@ -38,34 +37,26 @@ def speech_to_text(audio_path, lang):
 def main():
     with gr.Blocks() as demo:
         with gr.Row():
-            output_text = gr.Textbox(
-                label="识别结果",
-                lines=5
-            )
-            
+            output_text = gr.Textbox(label="识别结果", lines=5)
 
         with gr.Row():
             audio_input = gr.Audio(
-                sources=["upload"],
-                type="filepath",
-                label="录制或上传音频",
-                format="mp3"
+                sources=["upload"], type="filepath", label="录制或上传音频", format="mp3"
             )
             lang_dropdown = gr.Dropdown(
                 choices=["auto", "zh", "en", "yue", "ja", "ko"],
                 value="auto",
-                label="选择音频语言"
+                label="选择音频语言",
             )
 
         audio_input.change(
-            fn=speech_to_text,
-            inputs=[audio_input, lang_dropdown],
-            outputs=output_text
+            fn=speech_to_text, inputs=[audio_input, lang_dropdown], outputs=output_text
         )
 
     demo.launch(
-            server_name="0.0.0.0",
-        )
+        server_name="0.0.0.0",
+    )
+
 
 if __name__ == "__main__":
-    main()
+    main()

main.py

@@ -8,8 +8,17 @@ import time
 
 def get_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--input", "-i", required=True, type=str, help="Input audio file")
-    parser.add_argument("--language", "-l", required=False, type=str, default="auto", choices=["auto", "zh", "en", "yue", "ja", "ko"])
+    parser.add_argument(
+        "--input", "-i", required=True, type=str, help="Input audio file"
+    )
+    parser.add_argument(
+        "--language",
+        "-l",
+        required=False,
+        type=str,
+        default="auto",
+        choices=["auto", "zh", "en", "yue", "ja", "ko"],
+    )
     parser.add_argument("--streaming", action="store_true")
     return parser.parse_args()
 
@@ -19,7 +28,7 @@ def main():
 
     input_audio = args.input
     language = args.language
-    use_itn = True # 标点符号预测
+    use_itn = True  # 标点符号预测
     if not args.streaming:
         max_len = 256
         model_path = os.path.join("sensevoice_ax650", "sensevoice.axmodel")
@@ -35,14 +44,16 @@ def main():
     print(f"model_path: {model_path}")
     print(f"streaming: {args.streaming}")
 
-    pipeline = SenseVoiceAx(model_path, 
-                            max_len=max_len, 
-                            beam_size=3,
-                            language="auto", 
-                            hot_words=None,
-                            use_itn=True, 
-                            streaming=args.streaming)
-    
+    pipeline = SenseVoiceAx(
+        model_path,
+        max_len=max_len,
+        beam_size=3,
+        language="auto",
+        hot_words=None,
+        use_itn=True,
+        streaming=args.streaming,
+    )
+
     if not args.streaming:
         asr_res = pipeline.infer(input_audio, print_rtf=True)
         print("ASR result: " + asr_res)
@@ -57,11 +68,12 @@ def main():
             is_last = i + step >= len(samples)
             for res in pipeline.stream_infer(samples[i : i + step], is_last):
                 print(res)
-        
+
         end = time.time()
         cost_time = end - start
 
         print(f"RTF: {cost_time / duration}")
 
+
 if __name__ == "__main__":
-    main()
+    main()

print_utils.py

@@ -90,6 +90,7 @@ def format_str_v2(s):
         s = s.replace(emoji + " ", emoji)
     return s.strip()
 
+
 def rich_transcription_postprocess(s):
     def get_emo(s):
         return s[-1] if s[-1] in emo_set else None
@@ -116,6 +117,7 @@ def rich_transcription_postprocess(s):
     new_s = new_s.replace("The.", " ")
     return new_s.strip()
 
+
 def rich_print_asr_res(asr_res, will_print=True, remove_punc=False):
     res = "".join([rich_transcription_postprocess(i) for i in asr_res])
 
@@ -126,4 +128,4 @@ def rich_print_asr_res(asr_res, will_print=True, remove_punc=False):
     if will_print:
         print(res)
 
-    return res
+    return res

server.py

@@ -20,6 +20,7 @@ app = FastAPI(title="ASR Server", description="Automatic Speech Recognition API"
 # 全局变量存储模型
 asr_model = None
 
+
 @app.on_event("startup")
 async def load_model():
     """
@@ -27,11 +28,11 @@ async def load_model():
     """
     global asr_model
     logger.info("Loading ASR model...")
-    
+
     try:
         # 模型加载
         language = "auto"
-        use_itn = True # 标点符号预测
+        use_itn = True  # 标点符号预测
         max_len = 256
 
         model_path = os.path.join("sensevoice_ax650", "sensevoice.axmodel")
@@ -44,63 +45,74 @@ async def load_model():
         print(f"model_path: {model_path}")
 
         tokenizer = SentencepiecesTokenizer(bpemodel=bpemodel)
-        asr_model = SenseVoiceAx(model_path, 
-                                max_len=max_len,
-                                language=language, 
-                                use_itn=use_itn, 
-                                tokenizer=tokenizer)
-        
+        asr_model = SenseVoiceAx(
+            model_path,
+            max_len=max_len,
+            language=language,
+            use_itn=use_itn,
+            tokenizer=tokenizer,
+        )
+
         logger.info("ASR model loaded successfully")
     except Exception as e:
         logger.error(f"Failed to load ASR model: {str(e)}")
         raise
 
+
 def validate_audio_data(audio_data: List[float]) -> np.ndarray:
     """
     验证并转换音频数据为numpy数组
-    
+
     参数:
     - audio_data: 浮点数列表表示的音频数据
-    
+
     返回:
     - 验证后的numpy数组
     """
     try:
         # 转换为numpy数组
         np_array = np.array(audio_data, dtype=np.float32)
-        
+
         # 验证数据有效性
         if np_array.ndim != 1:
             raise ValueError("Audio data must be 1-dimensional")
-            
+
         if len(np_array) == 0:
             raise ValueError("Audio data cannot be empty")
-            
+
         return np_array
     except Exception as e:
         raise ValueError(f"Invalid audio data: {str(e)}")
-    
+
+
 @app.get("/get_language", summary="Get current language")
 async def get_language():
     return JSONResponse(content={"language": asr_model.language})
 
-@app.get("/get_language_options", summary="Get possible language options, possible options include [auto, zh, en, yue, ja, ko]")
+
+@app.get(
+    "/get_language_options",
+    summary="Get possible language options, possible options include [auto, zh, en, yue, ja, ko]",
+)
 async def get_language_options():
     return JSONResponse(content={"language_options": asr_model.language_options})
 
+
 @app.post("/asr", summary="Recognize speech from numpy audio data")
 async def recognize_speech(
-    audio_data: List[float] = Body(..., embed=True, description="Audio data as list of floats"),
+    audio_data: List[float] = Body(
+        ..., embed=True, description="Audio data as list of floats"
+    ),
     sample_rate: Optional[int] = Body(16000, description="Audio sample rate in Hz"),
-    language: Optional[str] = Body("auto", description="Language")
+    language: Optional[str] = Body("auto", description="Language"),
 ):
     """
     接收numpy数组格式的音频数据并返回识别结果
-    
+
     参数:
     - audio_data: 浮点数列表表示的音频数据
     - sample_rate: 音频采样率(默认16000Hz)
-    
+
     返回:
     - JSON包含识别文本
     """
@@ -108,19 +120,19 @@ async def recognize_speech(
         # 检查模型是否已加载
         if asr_model is None:
             raise HTTPException(status_code=503, detail="ASR model not loaded")
-        
+
         logger.info(f"Received audio data with length: {len(audio_data)}")
-        
+
         # 验证并转换数据
         np_audio = validate_audio_data(audio_data)
         if sample_rate != asr_model.sample_rate:
             np_audio = librosa.resample(np_audio, sample_rate, asr_model.sample_rate)
-        
+
         # 调用模型进行识别
         result = asr_model.infer_waveform(np_audio, language)
-        
+
         return JSONResponse(content={"text": result})
-    
+
     except ValueError as e:
         logger.error(f"Validation error: {str(e)}")
         raise HTTPException(status_code=400, detail=str(e))
@@ -128,6 +140,8 @@ async def recognize_speech(
         logger.error(f"Recognition error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
+
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

test_wer.py

@@ -14,35 +14,35 @@ 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'
-    
+    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='w', encoding='utf-8')
+    file_handler = logging.FileHandler(log_file, mode="w", 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
 
 
@@ -50,21 +50,21 @@ class AIShellDataset:
     def __init__(self, gt_path: str):
         """
         初始化数据集
-        
+
         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, "aishell_S0764")
-        
+
         # 检查必要文件和文件夹是否存在
         assert os.path.exists(gt_path), f"gt文件不存在: {gt_path}"
         assert os.path.exists(self.voice_dir), f"aishell_S0764文件夹不存在: {self.voice_dir}"
-        
+
         # 加载数据
         self.data = []
-        with open(gt_path, 'r', encoding='utf-8') as f:
+        with open(gt_path, "r", encoding="utf-8") as f:
             for line in f:
                 line = line.strip()
                 audio_path, gt = line.split(" ")
@@ -74,50 +74,50 @@ class AIShellDataset:
         # 使用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:
+        with open(tsv_path, "r", encoding="utf-8") as f:
             f.readline()
             for line in f:
                 line = line.strip()
@@ -126,79 +126,101 @@ class CommonVoiceDataset:
                 gt = splits[3]
                 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()
-    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("--language", "-l", required=False, type=str, default="auto", choices=["auto", "zh", "en", "yue", "ja", "ko"])
-    parser.add_argument("--max_num", type=int, default=-1, required=False, help="Maximum test data num")
+    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(
+        "--language",
+        "-l",
+        required=False,
+        type=str,
+        default="auto",
+        choices=["auto", "zh", "en", "yue", "ja", "ko"],
+    )
+    parser.add_argument(
+        "--max_num", type=int, default=-1, required=False, help="Maximum test data num"
+    )
     return parser.parse_args()
 
 
 def min_distance(word1: str, word2: str) -> int:
- 
+
     row = len(word1) + 1
     column = len(word2) + 1
- 
-    cache = [ [0]*column for i in range(row) ]
- 
+
+    cache = [[0] * column for i in range(row)]
+
     for i in range(row):
         for j in range(column):
- 
-            if i ==0 and j ==0:
+
+            if i == 0 and j == 0:
                 cache[i][j] = 0
-            elif i == 0 and j!=0:
+            elif i == 0 and j != 0:
                 cache[i][j] = j
-            elif j == 0 and i!=0:
+            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]
+                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
- 
+                    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]
+
+    return cache[row - 1][column - 1]
 
 
 def remove_punctuation(text):
     # 定义正则表达式模式，匹配所有标点符号
     # 这个模式包括常见的标点符号和中文标点
-    pattern = r'[^\w\s]|_'
-    
+    pattern = r"[^\w\s]|_"
+
     # 使用sub方法将所有匹配的标点符号替换为空字符串
-    cleaned_text = re.sub(pattern, '', text)
-    
+    cleaned_text = re.sub(pattern, "", text)
+
     return cleaned_text
 
 
@@ -207,7 +229,7 @@ def main():
     args = get_args()
 
     language = args.language
-    use_itn = False # 标点符号预测
+    use_itn = False  # 标点符号预测
     max_num = args.max_num
 
     dataset_type = args.dataset.lower()
@@ -230,7 +252,9 @@ def main():
     logger.info(f"model_path: {model_path}")
 
     tokenizer = SentencepiecesTokenizer(bpemodel=bpemodel)
-    pipeline = SenseVoiceAx(model_path, language=language, use_itn=use_itn, tokenizer=tokenizer, max_len=256)
+    pipeline = SenseVoiceAx(
+        model_path, language=language, use_itn=use_itn, tokenizer=tokenizer, max_len=256
+    )
 
     # Iterate over dataset
     hyp = []
@@ -242,8 +266,10 @@ def main():
         reference = remove_punctuation(reference).lower()
 
         asr_res = pipeline.infer(audio_path, print_rtf=False)
-        hypothesis = rich_print_asr_res(asr_res, will_print=False, remove_punc=True).lower()
-        hypothesis = emoji.replace_emoji(hypothesis, replace='')
+        hypothesis = rich_print_asr_res(
+            asr_res, will_print=False, remove_punc=True
+        ).lower()
+        hypothesis = emoji.replace_emoji(hypothesis, replace="")
 
         character_error_num = min_distance(reference, hypothesis)
         character_num = len(reference)
@@ -254,7 +280,7 @@ def main():
 
         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}%"
         logger.info(line_content)
 
@@ -265,5 +291,6 @@ def main():
 
     logger.info(f"Total WER: {total_character_error_rate}%")
 
+
 if __name__ == "__main__":
-    main()
+    main()

tokenizer.py

@@ -52,7 +52,9 @@ class BaseTokenizer(ABC):
 
             self.unk_symbol = unk_symbol
             if self.unk_symbol not in self.token2id:
-                raise RuntimeError(f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list")
+                raise RuntimeError(
+                    f"Unknown symbol '{unk_symbol}' doesn't exist in the token_list"
+                )
             self.unk_id = self.token2id[self.unk_symbol]
 
     def encode(self, text, **kwargs):
@@ -84,7 +86,7 @@ class BaseTokenizer(ABC):
     @abstractmethod
     def tokens2text(self, tokens: Iterable[str]) -> str:
         raise NotImplementedError
-    
+
 
 class SentencepiecesTokenizer(BaseTokenizer):
     def __init__(self, bpemodel: Union[Path, str], **kwargs):
@@ -130,4 +132,4 @@ class SentencepiecesTokenizer(BaseTokenizer):
         return self.decode(*args, **kwargs)
 
     def tokens2ids(self, *args, **kwargs):
-        return self.encode(*args, **kwargs)
+        return self.encode(*args, **kwargs)