update codes

README.md

@@ -1,9 +1,12 @@
----
-license: mit
----
 # sensevoice.axera
 FunASR SenseVoice on Axera, official repo: https://github.com/FunAudioLLM/SenseVoice
 
+## TODO
+
+- [ ] 支持AX630C
+- [ ] 支持C++
+- [ ] 支持FastAPI
+
 ## 功能
  - 语音识别
  - 自动识别语言(支持中文、英文、粤语、日语、韩语)
@@ -60,6 +63,10 @@ RTF: 0.03026517820946964    Latency: 0.15689468383789062s  Total length: 5.184s
 python test_wer.py -d datasets -l zh
 ```
 
+## 模型转换
+
+参考[model_convert](model_convert/README.md)
+
 ## 技术讨论
 
 - Github issues

SenseVoiceAx.py

@@ -5,6 +5,7 @@ from frontend import WavFrontend
 import os
 import time
 from typing import Any, Dict, Iterable, List, NamedTuple, Set, Tuple, Union
+from print_utils import rich_transcription_postprocess
 
 def sequence_mask(lengths, maxlen=None, dtype=np.float32):
     # 如果 maxlen 未指定，则取 lengths 中的最大值
@@ -19,6 +20,8 @@ def sequence_mask(lengths, maxlen=None, dtype=np.float32):
     
     # 比较生成掩码
     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)
     
     # 返回指定数据类型的掩码
     return mask.astype(dtype)[None, ...]
@@ -67,10 +70,40 @@ def unique_consecutive_np(x, dim=None, return_inverse=False, return_counts=False
     
     return results[0] if len(results) == 1 else results
 
+
+def longest_common_suffix_prefix_with_tolerance(
+    lhs, 
+    rhs, 
+    tolerate: int = 0
+) -> int:
+    """
+    计算两个数组的最长公共子序列，该子序列必须同时满足：
+    - 是 lhs 的后 n 个元素（后缀）
+    - 是 rhs 的前 n 个元素（前缀）
+    并且允许最多 `tolerate` 个元素不匹配。
+
+    参数:
+        lhs: np.ndarray, 第一个数组
+        rhs: np.ndarray, 第二个数组
+        tolerate: int, 允许的不匹配元素数量（默认为 0，即完全匹配）
+
+    返回:
+        int: 最长公共后缀/前缀的长度（如果没有则返回 0）
+    """
+    max_possible_n = min(len(lhs), len(rhs))
+    
+    for n in range(max_possible_n, 0, -1):
+        mismatches = np.sum(lhs[-n:] != rhs[:n])
+        if mismatches <= tolerate:
+            return n
+    
+    return 0
+
 class SenseVoiceAx:
-    def __init__(self, model_path, language="auto", use_itn=True, tokenizer=None):
-        model_path_root = os.path.join(os.path.dirname(model_path), "../embeddings")
-        self.frontend = WavFrontend(cmvn_file="am.mvn",
+    def __init__(self, model_path, max_len=68, language="auto", use_itn=True, tokenizer=None):
+        model_path_root = os.path.join(os.path.dirname(model_path), "..")
+        embedding_root = os.path.join(model_path_root, "embeddings")
+        self.frontend = WavFrontend(cmvn_file=f"{model_path_root}/am.mvn",
                                     fs=16000, 
                                     window="hamming", 
                                     n_mels=80, 
@@ -82,7 +115,8 @@ class SenseVoiceAx:
         self.sample_rate = 16000
         self.tokenizer = tokenizer
         self.blank_id = 0
-        self.max_len = 34
+        self.max_len = max_len
+        self.padding = 16
 
         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}
@@ -90,13 +124,12 @@ class SenseVoiceAx:
         self.textnorm_int_dict = {25016: 14, 25017: 15}
         self.emo_dict = {"unk": 25009, "happy": 25001, "sad": 25002, "angry": 25003, "neutral": 25004}
 
-        self.position_encoding = np.load(f"{model_path_root}/position_encoding.npy")
-        language_query = np.load(f"{model_path_root}/{language}.npy")
-        textnorm_query = np.load(f"{model_path_root}/withitn.npy") if use_itn else np.load(f"{model_path_root}/woitn.npy")
-        event_emo_query = np.load(f"{model_path_root}/event_emo.npy")
+        self.position_encoding = np.load(f"{embedding_root}/position_encoding.npy")
+        language_query = np.load(f"{embedding_root}/{language}.npy")
+        textnorm_query = np.load(f"{embedding_root}/withitn.npy") if use_itn else np.load(f"{embedding_root}/woitn.npy")
+        event_emo_query = np.load(f"{embedding_root}/event_emo.npy")
         self.input_query = np.concatenate((textnorm_query, language_query, event_emo_query), axis=1)
         self.query_num = self.input_query.shape[1]
-        self.masks = sequence_mask(np.array([self.max_len], dtype=np.int32), dtype=np.float32)
 
     def load_data(self, filepath: str) -> np.ndarray:
         waveform, _ = librosa.load(filepath, sr=self.sample_rate)
@@ -147,28 +180,41 @@ class SenseVoiceAx:
         slice_num = int(np.ceil(feat.shape[1] / slice_len))
 
         asr_res = []
+        prev_token_int = None
         for i in range(slice_num):
-            sub_feat = feat[:, i*slice_len:(i+1)*slice_len, :]
+            if i == 0:
+                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, :]
             # concat query
             sub_feat = np.concatenate([self.input_query, sub_feat], axis=1)
-
-            if sub_feat.shape[1] < self.max_len:
+            real_len = sub_feat.shape[1]
+            if real_len < self.max_len:
                 sub_feat = np.concatenate([
                         sub_feat, 
-                        np.zeros((1, self.max_len - sub_feat.shape[1], sub_feat.shape[-1]), dtype=np.float32)
+                        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)
+
             outputs = self.model.run(None, {"speech": sub_feat,
-                                            "masks": self.masks,
+                                            "masks": masks,
                                             "position_encoding": self.position_encoding})
             ctc_logits, encoder_out_lens = outputs
 
             token_int = self.postprocess(ctc_logits, encoder_out_lens)
-            if self.tokenizer is not None:
-                asr_res.append(self.tokenizer.tokens2text(token_int))
-            else:
-                asr_res.append(token_int)
+
+            # common prefix
+            if self.padding > 0 and prev_token_int is not None:
+                # prefix_len = common_prefix_len(prev_token_int, token_int)
+                prefix_len = longest_common_suffix_prefix_with_tolerance(prev_token_int, token_int, 6)
+                common_prefix = rich_transcription_postprocess(self.tokenizer.tokens2text(token_int[:prefix_len]))
+
+                asr_res[-1] = asr_res[-1][:-len(common_prefix)]
+            prev_token_int = np.copy(token_int)
+
+            asr_res.append(self.tokenizer.tokens2text(token_int))
 
         return asr_res
     

download_dataset.sh

@@ -0,0 +1,2 @@
+wget https://github.com/ml-inory/whisper.axera/releases/download/v1.0/datasets.zip
+unzip datasets.zip -d ./

main.py

@@ -19,9 +19,11 @@ def main():
     input_audio = args.input
     language = args.language
     use_itn = True # 标点符号预测
+    max_len = 68
 
-    model_path = os.path.join("sensevoice_ax650", "sensevoice.axmodel")
-    bpemodel = "chn_jpn_yue_eng_ko_spectok.bpe.model"
+    model_path_root = download_model("SenseVoice")
+    model_path = os.path.join(model_path_root, "sensevoice_ax650", "sensevoice.axmodel")
+    bpemodel = os.path.join(model_path_root, "chn_jpn_yue_eng_ko_spectok.bpe.model")
 
     assert os.path.exists(model_path), f"model {model_path} not exist"
 
@@ -31,7 +33,11 @@ def main():
     print(f"model_path: {model_path}")
 
     tokenizer = SentencepiecesTokenizer(bpemodel=bpemodel)
-    pipeline = SenseVoiceAx(model_path, language, use_itn, tokenizer=tokenizer)
+    pipeline = SenseVoiceAx(model_path, 
+                            max_len=max_len,
+                            language=language, 
+                            use_itn=use_itn, 
+                            tokenizer=tokenizer)
     asr_res = pipeline.infer(input_audio, print_rtf=True)
     print([rich_transcription_postprocess(i) for i in asr_res])
     # rich_print_asr_res(asr_res)

print_utils.py

@@ -116,6 +116,14 @@ def rich_transcription_postprocess(s):
     new_s = new_s.replace("The.", " ")
     return new_s.strip()
 
-def rich_print_asr_res(asr_res):
+def rich_print_asr_res(asr_res, will_print=True, remove_punc=False):
     res = "".join([rich_transcription_postprocess(i) for i in asr_res])
-    print(res)
+
+    if remove_punc:
+        res = res.replace("，", "")
+        res = res.replace("。", "")
+
+    if will_print:
+        print(res)
+
+    return res

test_wer.py

@@ -0,0 +1,76 @@
+import os, sys
+import argparse
+from SenseVoiceAx import SenseVoiceAx
+from tokenizer import SentencepiecesTokenizer
+from print_utils import rich_transcription_postprocess, rich_print_asr_res
+from download_utils import download_model
+import jiwer
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", "-d", required=True, type=str, help="Input dataset")
+    parser.add_argument("--language", "-l", required=False, type=str, default="auto", choices=["auto", "zh", "en", "yue", "ja", "ko"])
+    return parser.parse_args()
+
+
+def main():
+    args = get_args()
+
+    dataset = args.dataset
+    language = args.language
+    use_itn = False # 标点符号预测
+
+    model_path_root = download_model("SenseVoice")
+    model_path = os.path.join(model_path_root, "sensevoice_ax650", "sensevoice.axmodel")
+    bpemodel = os.path.join(model_path_root, "chn_jpn_yue_eng_ko_spectok.bpe.model")
+
+    assert os.path.exists(model_path), f"model {model_path} not exist"
+
+    print(f"dataset: {dataset}")
+    print(f"language: {language}")
+    print(f"use_itn: {use_itn}")
+    print(f"model_path: {model_path}")
+
+    tokenizer = SentencepiecesTokenizer(bpemodel=bpemodel)
+    pipeline = SenseVoiceAx(model_path, language, use_itn, tokenizer=tokenizer)
+
+    # Load dataset
+    wav_names = []
+    references = []
+    with open(os.path.join(dataset, "ground_truth.txt"), "r") as f:
+        for line in f:
+            line = line.strip()
+            w, r = line.split(" ")
+            wav_names.append(w)
+            references.append(r)
+
+    # Iterate over dataset
+    hyp = []
+    wer_file = open("wer.txt", "w")
+    for wav_name, reference in zip(wav_names, references):
+        wav_path = os.path.join(dataset, "aishell_S0764", wav_name + ".wav")
+
+        asr_res = pipeline.infer(wav_path, print_rtf=False)
+        hypothesis = rich_print_asr_res(asr_res, will_print=False, remove_punc=True)
+        hyp.append(hypothesis)
+
+        wer = jiwer.cer(
+                    reference,
+                    hypothesis
+                )
+        
+        line_content = f"{wav_name}  reference: {reference}  hypothesis: {hypothesis}  WER: {wer}"
+        wer_file.write(line_content + "\n")
+        print(line_content)
+
+    total_wer = jiwer.cer(
+                    references,
+                    hyp
+                )
+    print(f"Total WER: {total_wer}")
+    wer_file.write(f"Total WER: {total_wer}")
+    wer_file.close()
+
+if __name__ == "__main__":
+    main()