[add]tn module

cookbooks/test.py

@@ -1,12 +1,19 @@
 import warnings
 import torch
 from transformers import AutoProcessor
+import os
+import sys
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.append(parent_dir)
 
 from modeling_bailingmm import BailingMMNativeForConditionalGeneration
-
 import random
 import numpy as np
 from loguru import logger
+from sentence_manager.sentence_manager import SentenceNormalizer
+import re
+import yaml
 
 def seed_everything(seed=1895):
     random.seed(seed)
@@ -32,7 +39,30 @@ def __init__(self, model_path, device="cuda:0"):
         self.tokenizer = self.processor.tokenizer
         self.sample_rate = self.processor.audio_processor.sample_rate
         self.patch_size = self.processor.audio_processor.patch_size
+        self.normalizer = self.init_tn_normalizer(tokenizer=self.tokenizer)
+
+    def init_tn_normalizer(self, config_file_path=None, tokenizer=None):
+        if config_file_path is None:
+            default_config_path = os.path.join(
+                os.path.dirname(os.path.dirname(os.path.realpath(__file__))), 
+                "sentence_manager/default_config.yaml"
+            )
+            config_file_path = default_config_path
+
+        with open(config_file_path, 'r') as f:
+            self.sentence_manager_config = yaml.safe_load(f)
 
+        if "split_token" not in self.sentence_manager_config:
+            self.sentence_manager_config["split_token"] = []
+
+        assert isinstance(self.sentence_manager_config["split_token"], list)
+        if tokenizer is not None:
+            self.sentence_manager_config["split_token"].append(re.escape(tokenizer.eos_token))
+
+        normalizer = SentenceNormalizer(self.sentence_manager_config.get("text_norm", {}))
+
+        return normalizer
+
     def speech_understanding(self, messages, lang=None):
         text = self.processor.apply_chat_template(messages, add_generation_prompt=True)
         image_inputs, video_inputs, audio_inputs = self.processor.process_vision_info(messages)
@@ -77,6 +107,7 @@ def speech_generation(
         lang='zh',
         output_wav_path='out.wav'
     ):
+        text = self.normalizer.normalize(text)
         waveform = self.model.generate_tts(
             text=text,
             prompt_wav_path=prompt_wav_path,
@@ -127,7 +158,6 @@ def speech_edit(
 
 if __name__ == "__main__":
     model = MingAudio("inclusionAI/Ming-UniAudio-16B-A3B")
-
     # ASR
     messages = [
         {

sentence_manager/default_config.yaml

@@ -0,0 +1,43 @@
+text_norm:
+  postprocess:
+    - # EN2CN
+      "…" : "。"
+      "!" : "！"
+      "\\?" : "？"
+      ";" : "；"
+      ":" : "："
+      "," : "，"
+      "\\(" : "（"
+      "\\)" : "）"
+    - # EN2CN
+      "……": "。"
+    - # OTHER2CN
+      "﹐" : "，"
+      "﹔" : "；"
+      "｡"  : "。"
+      # CN2CN
+      "；" : "。"
+      "：" : "，"
+      "、" : "，"
+    - # 处理连续句号"。"
+      "。+": "。"   
+    - # 正则后的 "/"
+      "/": "每"
+    - # 处理_
+      "_": " "            
+    - # 处理正则后的[~～]+，根据是否在句尾替换为“。”或“至”
+      "～+": "～"
+      "~+": "~"
+      "[~～]": "。"
+    - # 删除除英文内的“-”, "'"
+      "(?<=[^a-zA-Z])[-']+": "，"
+      "[-']+(?=[^a-zA-Z])": "，"
+    - # 删除除了标准中文标点、英文、-、’、空格、数字、中文外的其他符号
+      "[^。！？，\u4e00-\u4E27\u4E29-\u4E3E\u4E42-\u9fa4a-zA-Z ]": ""
+    - # 处理连续逗号"。"
+      "，+": "，"   
+    - # 处理连续空格"。"
+      " +": " "  
+
+split_token: ["。", "，"]
+split_cn_length: null

sentence_manager/sentence_manager.py

@@ -0,0 +1,111 @@
+# -*- encoding: utf-8 -*-
+# Time: 2024/12/02 19:58:58
+# Desc:
+
+import os
+import re
+import yaml
+from typing import Optional, List
+from enum import Enum
+
+from .text_norm.normalizer import Normalizer
+
+
+CN_REGEX = "\u4e00-\u4E27\u4E29-\u4E3E\u4E42-\u9fa4"  # 跳过笔画\u4E28\u4E3F\u4E40\u4E41
+
+
+class SentencePieceType(str, Enum):
+    END_OF_SENTENCE = "END_OF_SENTENCE"
+
+
+def split_with_separator(regx_sep, text):
+    """difference with `re.split`: include separator
+    """
+    split_list = []
+    start = 0
+    for match in re.finditer(regx_sep, text):
+        end = match.span()[1]
+        assert end > start
+        split_list.append(text[start: end])
+        start = end
+    split_list.append(text[start:])  # could be empty string
+    return split_list
+
+
+def split(text, split_pattern, split_cn_length=None):
+    """
+    Args:
+        text
+    Return:
+        (split_list: List[str], remain: str)
+    """
+    split_list = split_with_separator(split_pattern, text)
+    remain = split_list.pop(-1)  # 最后一项为不完整子句，可能为空字符串
+
+    # 针对末尾的不完整子句，若满足中文字数条件，也添加进split_list
+    if split_cn_length is not None:
+        text_split = re.search(f"^[。！？，{CN_REGEX}]" + "{" + f"{split_cn_length}," + "}", remain)
+        if text_split:
+            text_split = text_split.group()
+            split_list.append(text_split)
+            remain = remain[len(text_split):]
+    return split_list, remain
+
+
+class SentenceNormalizer(Normalizer):
+    def __init__(self, config={}):
+        self.config = config
+
+    def normalize(self, text, context: str = ""):
+        text = self.preprocess(text)
+        text, norm_details = self.normalize_regular(text, is_en=False)
+        text = self.postprocess(text, custom=self.config["postprocess"])
+        text = text[len(context):]
+        return text
+
+
+class SentenceManager:
+    def __init__(self, tokenizer, normalizer, config):
+        """
+        Args:
+            tokenizer: tokenizer为必填参数，因为对英文等特殊符号，tokenize和拼接操作不是可交换的
+        """
+
+        self.split_pattern = "|".join(config["split_token"])
+        self.split_cn_length = config["split_cn_length"]
+
+        self.tokenizer = tokenizer
+        self.normalizer = normalizer
+
+        self.context: Optional[str] = ""
+        self.cache: List[int] = []
+        self.output_queue: List[List[int]] = []
+
+    def put(self, token_id):
+        text = self.tokenizer.decode([*self.cache, token_id])
+        split_list, remain = split(text, self.split_pattern, split_cn_length=self.split_cn_length)
+        assert split_list or remain
+
+        if split_list:
+            normalized_split_list = [
+                self.normalizer.normalize(x) if i < len(split_list) else self.normalizer.normalize(x, self.context) 
+                for i, x in enumerate(split_list)
+            ]
+            if len(normalized_split_list[-1]) == len(split_list[-1]):
+                self.context = split_list[-1]
+            token_ids_list = [self.tokenizer.encode(x) for x in normalized_split_list if x]
+            self.output_queue.extend(token_ids_list)
+            if re.search(f"({self.split_pattern})$", split_list[-1]):
+                self.output_queue.append(SentencePieceType.END_OF_SENTENCE)
+
+        if remain:
+            token_ids_remain = self.tokenizer.encode(remain)
+            self.cache = token_ids_remain
+        else:
+            self.cache = []
+
+    def get(self):
+        if self.output_queue:
+            return self.output_queue.pop(0)
+        else:
+            return None