Refactor: add style_bert_vits2/text_processing/bert_models.py to hold loaded BERT models/tokenizer and replace all from_pretrained() to load_model/load_tokenizer

server_editor.py

@@ -28,7 +28,6 @@ from fastapi.responses import JSONResponse, Response
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from scipy.io import wavfile
-from transformers import AutoTokenizer
 
 from common.tts_model import ModelHolder
 from style_bert_vits2.constants import (
@@ -42,6 +41,7 @@ from style_bert_vits2.constants import (
     Languages,
 )
 from style_bert_vits2.logging import logger
+from style_bert_vits2.text_processing import bert_models
 from style_bert_vits2.text_processing.japanese.g2p_utils import g2kata_tone, kata_tone2phone_tone
 from style_bert_vits2.text_processing.japanese.normalizer import normalize_text
 from style_bert_vits2.text_processing.japanese.user_dict import (
@@ -150,8 +150,10 @@ def save_last_download(latest_release):
 # 最初に pyopenjtalk の辞書を更新
 update_dict()
 
-# 単語分割に使う BERT トークナイザーをロード
+# 単語分割に使う BERT モデル/トークナイザーを事前にロードしておく
-tokenizer = AutoTokenizer.from_pretrained("./bert/deberta-v2-large-japanese-char-wwm")
+## server_editor.py は日本語にしか対応していないため、日本語の BERT モデル/トークナイザーのみロードする
+bert_models.load_model(Languages.JP)
+bert_models.load_tokenizer(Languages.JP)
 
 
 class AudioResponse(Response):
@@ -227,7 +229,7 @@ async def read_item(item: TextRequest):
     try:
         # 最初に正規化しないと整合性がとれない
         text = normalize_text(item.text)
-        kata_tone_list = g2kata_tone(text, tokenizer)
+        kata_tone_list = g2kata_tone(text)
     except Exception as e:
         raise HTTPException(
             status_code=400,

style_bert_vits2/constants.py

@@ -1,19 +1,31 @@
-from enum import Enum
+from enum import StrEnum
 from pathlib import Path
 
 
 # Style-Bert-VITS2 のバージョン
 VERSION = "2.3.1"
 
-# Gradio のテーマ
+# Style-Bert-VITS2 のベースディレクトリ
-## Built-in theme: "default", "base", "monochrome", "soft", "glass"
+BASE_DIR = Path(__file__).parent.parent
-## See https://huggingface.co/spaces/gradio/theme-gallery for more themes
+
-GRADIO_THEME = "NoCrypt/miku"
+# 利用可能な言語
+## JP-Extra モデル利用時は JP 以外の言語の音声合成はできない
+class Languages(StrEnum):
+    JP = "JP"
+    EN = "EN"
+    ZH = "ZH"
+
+# 言語ごとのデフォルトの BERT トークナイザーのパス
+DEFAULT_BERT_TOKENIZER_PATHS = {
+    Languages.JP: BASE_DIR / "bert" / "deberta-v2-large-japanese-char-wwm",
+    Languages.EN: BASE_DIR / "bert" / "deberta-v3-large",
+    Languages.ZH: BASE_DIR / "bert" / "chinese-roberta-wwm-ext-large",
+}
 
 # デフォルトのユーザー辞書ディレクトリ
 ## style_bert_vits2.text_processing.japanese.user_dict モジュールのデフォルト値として利用される
 ## ライブラリとしての利用などで外部のユーザー辞書を指定したい場合は、user_dict 以下の各関数の実行時、引数に辞書データファイルのパスを指定する
-DEFAULT_USER_DICT_DIR = Path(__file__).parent.parent / "dict_data"
+DEFAULT_USER_DICT_DIR = BASE_DIR / "dict_data"
 
 # デフォルトの推論パラメータ
 DEFAULT_STYLE = "Neutral"
@@ -27,9 +39,7 @@ DEFAULT_SPLIT_INTERVAL = 0.5
 DEFAULT_ASSIST_TEXT_WEIGHT = 0.7
 DEFAULT_ASSIST_TEXT_WEIGHT = 1.0
 
-# 利用可能な言語
+# Gradio のテーマ
-## JP-Extra モデル利用時は JP 以外の言語の音声合成はできない
+## Built-in theme: "default", "base", "monochrome", "soft", "glass"
-class Languages(str, Enum):
+## See https://huggingface.co/spaces/gradio/theme-gallery for more themes
-    JP = "JP"
+GRADIO_THEME = "NoCrypt/miku"
-    EN = "EN"
-    ZH = "ZH"

style_bert_vits2/models/infer.py

@@ -1,9 +1,8 @@
-from typing import Literal
-
 import torch
 
 import utils
 from text import cleaned_text_to_sequence, get_bert
+from style_bert_vits2.constants import Languages
 from style_bert_vits2.logging import logger
 from style_bert_vits2.models import commons
 from style_bert_vits2.models.models import SynthesizerTrn
@@ -48,7 +47,7 @@ def get_net_g(model_path: str, version: str, device: str, hps):
 
 def get_text(
     text: str,
-    language_str: Literal["JP", "EN", "ZH"],
+    language_str: Languages,
     hps,
     device: str,
     assist_text: str | None = None,
@@ -89,15 +88,15 @@ def get_text(
     del word2ph
     assert bert_ori.shape[-1] == len(phone), phone
 
-    if language_str == "ZH":
+    if language_str == Languages.ZH:
         bert = bert_ori
         ja_bert = torch.zeros(1024, len(phone))
         en_bert = torch.zeros(1024, len(phone))
-    elif language_str == "JP":
+    elif language_str == Languages.JP:
         bert = torch.zeros(1024, len(phone))
         ja_bert = bert_ori
         en_bert = torch.zeros(1024, len(phone))
-    elif language_str == "EN":
+    elif language_str == Languages.EN:
         bert = torch.zeros(1024, len(phone))
         ja_bert = torch.zeros(1024, len(phone))
         en_bert = bert_ori
@@ -122,7 +121,7 @@ def infer(
     noise_scale_w: float,
     length_scale: float,
     sid: int,  # In the original Bert-VITS2, its speaker_name: str, but here it's id
-    language: Literal["JP", "EN", "ZH"],
+    language: Languages,
     hps,
     net_g,
     device: str,
@@ -222,7 +221,7 @@ def infer_multilang(
     noise_scale_w: float,
     length_scale: float,
     sid: int,
-    language: Literal["JP", "EN", "ZH"],
+    language: Languages,
     hps,
     net_g,
     device: str,

style_bert_vits2/text_processing/bert_models.py

@@ -0,0 +1,123 @@
+"""
+Style-Bert-VITS2 の学習・推論に必要な各言語ごとの BERT モデルをロード/取得するためのモジュール。
+
+オリジナルの Bert-VITS2 では各言語ごとの BERT モデルが初回インポート時にハードコードされたパスから「暗黙的に」ロードされているが、
+場合によっては多重にロードされて非効率なほか、BERT モデルのロード元のパスがハードコードされているためライブラリ化ができない。
+
+そこで、ライブラリの利用前に、音声合成に利用する言語の BERT モデルだけを「明示的に」ロードできるようにした。
+一度 load_tokenizer() で当該言語の BERT モデルがロードされていれば、ライブラリ内部のどこからでもロード済みのモデル/トークナイザーを取得できる。
+"""
+
+from typing import cast
+
+from transformers import (
+    AutoModelForMaskedLM,
+    AutoTokenizer,
+    DebertaV2Model,
+    DebertaV2Tokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+)
+
+from style_bert_vits2.constants import DEFAULT_BERT_TOKENIZER_PATHS, Languages
+from style_bert_vits2.logging import logger
+
+
+# 各言語ごとのロード済みの BERT モデルを格納する辞書
+loaded_models: dict[Languages, PreTrainedModel | DebertaV2Model] = {}
+
+# 各言語ごとのロード済みの BERT トークナイザーを格納する辞書
+loaded_tokenizers: dict[Languages, PreTrainedTokenizer | PreTrainedTokenizerFast | DebertaV2Tokenizer] = {}
+
+
+def load_model(
+    language: Languages,
+    pretrained_model_name_or_path: str | None = None,
+) -> PreTrainedModel | DebertaV2Model:
+    """
+    指定された言語の BERT モデルをロードし、ロード済みの BERT モデルを返す
+    一度ロードされていれば、ロード済みの BERT モデルを即座に返す
+    ライブラリ利用時は常に pretrain_model_name_or_path (Hugging Face のリポジトリ名 or ローカルのファイルパス) を指定する必要がある
+    ロードにはそれなりに時間がかかるため、ライブラリ利用前に明示的に pretrained_model_name_or_path を指定してロードしておくべき
+
+    Style-Bert-VITS2 では、BERT モデルに下記の 3 つが利用されている
+    これ以外の BERT モデルを指定した場合は正常に動作しない可能性が高い
+    - 日本語: ku-nlp/deberta-v2-large-japanese-char-wwm
+    - 英語: microsoft/deberta-v3-large
+    - 中国語: hfl/chinese-roberta-wwm-ext-large
+
+    Args:
+        language (Languages): ロードする学習済みモデルの対象言語
+        pretrained_model_name_or_path (str | None): ロードする学習済みモデルの名前またはパス。指定しない場合はデフォルトのパスが利用される (デフォルト: None)
+
+    Returns:
+        PreTrainedModel | DebertaV2Model: ロード済みの BERT モデル
+    """
+
+    # すでにロード済みの場合はそのまま返す
+    if language in loaded_models:
+        return loaded_models[language]
+
+    # pretrained_model_name_or_path が指定されていない場合はデフォルトのパスを利用
+    if pretrained_model_name_or_path is None:
+        assert DEFAULT_BERT_TOKENIZER_PATHS[language].exists(), \
+            f"The default {language} BERT model does not exist on the file system. Please specify the path to the pre-trained model."
+        pretrained_model_name_or_path = str(DEFAULT_BERT_TOKENIZER_PATHS[language])
+
+    # BERT モデルをロードし、辞書に格納して返す
+    ## 英語のみ DebertaV2Model でロードする必要がある
+    if language == Languages.EN:
+        model = cast(DebertaV2Model, DebertaV2Model.from_pretrained(pretrained_model_name_or_path))
+    else:
+        model = AutoModelForMaskedLM.from_pretrained(pretrained_model_name_or_path)
+    loaded_models[language] = model
+    logger.info(f"Loaded the {language} BERT model from {pretrained_model_name_or_path}")
+
+    return model
+
+
+def load_tokenizer(
+    language: Languages,
+    pretrained_model_name_or_path: str | None = None,
+) -> PreTrainedTokenizer | PreTrainedTokenizerFast | DebertaV2Tokenizer:
+    """
+    指定された言語の BERT モデルをロードし、ロード済みの BERT トークナイザーを返す
+    一度ロードされていれば、ロード済みの BERT トークナイザーを即座に返す
+    ライブラリ利用時は常に pretrain_model_name_or_path (Hugging Face のリポジトリ名 or ローカルのファイルパス) を指定する必要がある
+    ロードにはそれなりに時間がかかるため、ライブラリ利用前に明示的に pretrained_model_name_or_path を指定してロードしておくべき
+
+    Style-Bert-VITS2 では、BERT モデルに下記の 3 つが利用されている
+    これ以外の BERT モデルを指定した場合は正常に動作しない可能性が高い
+    - 日本語: ku-nlp/deberta-v2-large-japanese-char-wwm
+    - 英語: microsoft/deberta-v3-large
+    - 中国語: hfl/chinese-roberta-wwm-ext-large
+
+    Args:
+        language (Languages): ロードする学習済みモデルの対象言語
+        pretrained_model_name_or_path (str | None): ロードする学習済みモデルの名前またはパス。指定しない場合はデフォルトのパスが利用される (デフォルト: None)
+
+    Returns:
+        PreTrainedTokenizer | PreTrainedTokenizerFast | DebertaV2Tokenizer: ロード済みの BERT トークナイザー
+    """
+
+    # すでにロード済みの場合はそのまま返す
+    if language in loaded_tokenizers:
+        return loaded_tokenizers[language]
+
+    # pretrained_model_name_or_path が指定されていない場合はデフォルトのパスを利用
+    if pretrained_model_name_or_path is None:
+        assert DEFAULT_BERT_TOKENIZER_PATHS[language].exists(), \
+            f"The default {language} BERT tokenizer does not exist on the file system. Please specify the path to the pre-trained model."
+        pretrained_model_name_or_path = str(DEFAULT_BERT_TOKENIZER_PATHS[language])
+
+    # BERT トークナイザーをロードし、辞書に格納して返す
+    ## 英語のみ DebertaV2Tokenizer でロードする必要がある
+    if language == Languages.EN:
+        tokenizer = DebertaV2Tokenizer.from_pretrained(pretrained_model_name_or_path)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
+    loaded_tokenizers[language] = tokenizer
+    logger.info(f"Loaded the {language} BERT tokenizer from {pretrained_model_name_or_path}")
+
+    return tokenizer

style_bert_vits2/text_processing/cleaner.py

@@ -1,9 +1,9 @@
-from typing import Literal
+from style_bert_vits2.constants import Languages
 
 
 def clean_text(
     text: str,
-    language: Literal["JP", "EN", "ZH"],
+    language: Languages,
     use_jp_extra: bool = True,
     raise_yomi_error: bool = False,
 ) -> tuple[str, list[str], list[int], list[int]]:
@@ -12,7 +12,7 @@ def clean_text(
 
     Args:
         text (str): クリーニングするテキスト
-        language (Literal["JP", "EN", "ZH"]): テキストの言語
+        language (Languages): テキストの言語
         use_jp_extra (bool, optional): テキストが日本語の場合に JP-Extra モデルを利用するかどうか。Defaults to True.
         raise_yomi_error (bool, optional): False の場合、読めない文字が消えたような扱いとして処理される。Defaults to False.
 
@@ -21,22 +21,16 @@ def clean_text(
     """
 
     # Changed to import inside if condition to avoid unnecessary import
-    if language == "JP":
+    if language == Languages.JP:
-        from transformers import AutoTokenizer
         from style_bert_vits2.text_processing.japanese.g2p import g2p
         from style_bert_vits2.text_processing.japanese.normalizer import normalize_text
         norm_text = normalize_text(text)
-        phones, tones, word2ph = g2p(
+        phones, tones, word2ph = g2p(norm_text, use_jp_extra, raise_yomi_error)
-            norm_text,
+    elif language == Languages.EN:
-            tokenizer = AutoTokenizer.from_pretrained("./bert/deberta-v2-large-japanese-char-wwm"),  # 暫定的にここで指定
-            use_jp_extra = use_jp_extra,
-            raise_yomi_error = raise_yomi_error,
-        )
-    elif language == "EN":
         from ...text import english as language_module
         norm_text = language_module.normalize_text(text)
         phones, tones, word2ph = language_module.g2p(norm_text)
-    elif language == "ZH":
+    elif language == Languages.ZH:
         from ...text import chinese as language_module
         norm_text = language_module.normalize_text(text)
         phones, tones, word2ph = language_module.g2p(norm_text)

style_bert_vits2/text_processing/japanese/g2p.py

@@ -1,8 +1,9 @@
 import pyopenjtalk
 import re
-from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
 
+from style_bert_vits2.constants import Languages
 from style_bert_vits2.logging import logger
+from style_bert_vits2.text_processing import bert_models
 from style_bert_vits2.text_processing.japanese.mora_list import MORA_KATA_TO_MORA_PHONEMES
 from style_bert_vits2.text_processing.japanese.normalizer import replace_punctuation
 from style_bert_vits2.text_processing.symbols import PUNCTUATIONS
@@ -10,7 +11,6 @@ from style_bert_vits2.text_processing.symbols import PUNCTUATIONS
 
 def g2p(
     norm_text: str,
-    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
     use_jp_extra: bool = True,
     raise_yomi_error: bool = False
 ) -> tuple[list[str], list[int], list[int]]:
@@ -21,11 +21,9 @@ def g2p(
     - word2ph: 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
     のタプルを返す。
     ただし `phones` と `tones` の最初と終わりに `_` が入り、応じて `word2ph` の最初と最後に 1 が追加される。
-    tokenizer には deberta-v2-large-japanese-char-wwm を AutoTokenizer.from_pretrained() でロードしたものを指定する。
 
     Args:
         norm_text (str): 正規化されたテキスト
-        tokenizer (PreTrainedTokenizer | PreTrainedTokenizerFast): 単語分割に使うロード済みの BERT Tokenizer インスタンス
         use_jp_extra (bool, optional): False の場合、「ん」の音素を「N」ではなく「n」とする。Defaults to True.
         raise_yomi_error (bool, optional): False の場合、読めない文字が消えたような扱いとして処理される。Defaults to False.
 
@@ -66,7 +64,7 @@ def g2p(
     for i in sep_text:
         if i not in PUNCTUATIONS:
             sep_tokenized.append(
-                tokenizer.tokenize(i)
+                bert_models.load_tokenizer(Languages.JP).tokenize(i)
             )  # ここでおそらく`i`が文字単位に分割される
         else:
             sep_tokenized.append([i])

style_bert_vits2/text_processing/japanese/g2p_utils.py

@@ -1,5 +1,3 @@
-from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
-
 from style_bert_vits2.text_processing.japanese.g2p import g2p
 from style_bert_vits2.text_processing.japanese.mora_list import (
     MORA_KATA_TO_MORA_PHONEMES,
@@ -8,21 +6,19 @@ from style_bert_vits2.text_processing.japanese.mora_list import (
 from style_bert_vits2.text_processing.symbols import PUNCTUATIONS
 
 
-def g2kata_tone(norm_text: str, tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast) -> list[tuple[str, int]]:
+def g2kata_tone(norm_text: str) -> list[tuple[str, int]]:
     """
     テキストからカタカナとアクセントのペアのリストを返す。
-    推論時のみに使われるので、常に`raise_yomi_error=False`でg2pを呼ぶ。
+    推論時のみに使われるので、常に `raise_yomi_error=False` で g2p を呼ぶ。
-    tokenizer には deberta-v2-large-japanese-char-wwm を AutoTokenizer.from_pretrained() でロードしたものを指定する。
 
     Args:
         norm_text: 正規化されたテキスト。
-        tokenizer (PreTrainedTokenizer | PreTrainedTokenizerFast): 単語分割に使うロード済みの BERT Tokenizer インスタンス
 
     Returns:
         カタカナと音高のリスト。
     """
 
-    phones, tones, _ = g2p(norm_text, tokenizer, use_jp_extra=True, raise_yomi_error=False)
+    phones, tones, _ = g2p(norm_text, use_jp_extra=True, raise_yomi_error=False)
     return phone_tone2kata_tone(list(zip(phones, tones)))
 
 

text/__init__.py

@@ -1,12 +1,17 @@
+from style_bert_vits2.constants import Languages
 from style_bert_vits2.text_processing.symbols import *
 
+
 _symbol_to_id = {s: i for i, s in enumerate(SYMBOLS)}
 
 
-def cleaned_text_to_sequence(cleaned_text, tones, language):
+def cleaned_text_to_sequence(cleaned_text: str, tones: list[int], language: Languages):
-    """Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    """
+    Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+
     Args:
       text: string to convert to a sequence
+
     Returns:
       List of integers corresponding to the symbols in the text
     """
@@ -18,12 +23,19 @@ def cleaned_text_to_sequence(cleaned_text, tones, language):
     return phones, tones, lang_ids
 
 
-def get_bert(text, word2ph, language, device, assist_text=None, assist_text_weight=0.7):
+def get_bert(
-    if language == "ZH":
+    text: str,
+    word2ph,
+    language: Languages,
+    device: str,
+    assist_text: str | None = None,
+    assist_text_weight: float = 0.7,
+):
+    if language == Languages.ZH:
         from .chinese_bert import get_bert_feature
-    elif language == "EN":
+    elif language == Languages.EN:
         from .english_bert_mock import get_bert_feature
-    elif language == "JP":
+    elif language == Languages.JP:
         from .japanese_bert import get_bert_feature
     else:
         raise ValueError(f"Language {language} not supported")

text/chinese_bert.py

@@ -1,23 +1,21 @@
 import sys
 
 import torch
-from transformers import AutoModelForMaskedLM, AutoTokenizer
 
 from config import config
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.text_processing import bert_models
 
-LOCAL_PATH = "./bert/chinese-roberta-wwm-ext-large"
-
-tokenizer = AutoTokenizer.from_pretrained(LOCAL_PATH)
 
 models = dict()
 
 
 def get_bert_feature(
-    text,
+    text: str,
     word2ph,
-    device=config.bert_gen_config.device,
+    device = config.bert_gen_config.device,
-    assist_text=None,
+    assist_text: str | None = None,
-    assist_text_weight=0.7,
+    assist_text_weight: float = 0.7,
 ):
     if (
         sys.platform == "darwin"
@@ -30,8 +28,9 @@ def get_bert_feature(
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
     if device not in models.keys():
-        models[device] = AutoModelForMaskedLM.from_pretrained(LOCAL_PATH).to(device)
+        models[device] = bert_models.load_model(Languages.ZH).to(device)
     with torch.no_grad():
+        tokenizer = bert_models.load_tokenizer(Languages.ZH)
         inputs = tokenizer(text, return_tensors="pt")
         for i in inputs:
             inputs[i] = inputs[i].to(device)

text/english.py

@@ -2,16 +2,16 @@ import pickle
 import os
 import re
 from g2p_en import G2p
-from transformers import DebertaV2Tokenizer
 
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.text_processing import bert_models
 from style_bert_vits2.text_processing.symbols import PUNCTUATIONS, SYMBOLS
 
+
 current_file_path = os.path.dirname(__file__)
 CMU_DICT_PATH = os.path.join(current_file_path, "cmudict.rep")
 CACHE_PATH = os.path.join(current_file_path, "cmudict_cache.pickle")
 _g2p = G2p()
-LOCAL_PATH = "./bert/deberta-v3-large"
-tokenizer = DebertaV2Tokenizer.from_pretrained(LOCAL_PATH)
 
 arpa = {
     "AH0",
@@ -392,6 +392,7 @@ def sep_text(text):
 
 
 def text_to_words(text):
+    tokenizer = bert_models.load_tokenizer(Languages.EN)
     tokens = tokenizer.tokenize(text)
     words = []
     for idx, t in enumerate(tokens):

text/english_bert_mock.py

@@ -1,24 +1,21 @@
 import sys
 
 import torch
-from transformers import DebertaV2Model, DebertaV2Tokenizer
 
 from config import config
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.text_processing import bert_models
 
 
-LOCAL_PATH = "./bert/deberta-v3-large"
-
-tokenizer = DebertaV2Tokenizer.from_pretrained(LOCAL_PATH)
-
 models = dict()
 
 
 def get_bert_feature(
-    text,
+    text: str,
     word2ph,
-    device=config.bert_gen_config.device,
+    device = config.bert_gen_config.device,
-    assist_text=None,
+    assist_text: str | None = None,
-    assist_text_weight=0.7,
+    assist_text_weight: float = 0.7,
 ):
     if (
         sys.platform == "darwin"
@@ -31,8 +28,9 @@ def get_bert_feature(
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
     if device not in models.keys():
-        models[device] = DebertaV2Model.from_pretrained(LOCAL_PATH).to(device)
+        models[device] = bert_models.load_model(Languages.EN).to(device)
     with torch.no_grad():
+        tokenizer = bert_models.load_tokenizer(Languages.EN)
         inputs = tokenizer(text, return_tensors="pt")
         for i in inputs:
             inputs[i] = inputs[i].to(device)

text/japanese.py

@@ -1,642 +0,0 @@
-# Convert Japanese text to phonemes which is
-# compatible with Julius https://github.com/julius-speech/segmentation-kit
-import re
-import unicodedata
-
-import pyopenjtalk
-from num2words import num2words
-from transformers import AutoTokenizer
-
-from style_bert_vits2.logging import logger
-from style_bert_vits2.text_processing.japanese.mora_list import (
-    MORA_KATA_TO_MORA_PHONEMES,
-    MORA_PHONEMES_TO_MORA_KATA,
-)
-from style_bert_vits2.text_processing.japanese.user_dict import update_dict
-from style_bert_vits2.text_processing.symbols import PUNCTUATIONS
-
-# 最初にpyopenjtalkの辞書を更新
-update_dict()
-
-# 子音の集合
-COSONANTS = set(
-    [
-        cosonant
-        for cosonant, _ in MORA_KATA_TO_MORA_PHONEMES.values()
-        if cosonant is not None
-    ]
-)
-
-# 母音の集合、便宜上「ん」を含める
-VOWELS = {"a", "i", "u", "e", "o", "N"}
-
-
-class YomiError(Exception):
-    """
-    OpenJTalkで、読みが正しく取得できない箇所があるときに発生する例外。
-    基本的に「学習の前処理のテキスト処理時」には発生させ、そうでない場合は、
-    ignore_yomi_error=Trueにしておいて、この例外を発生させないようにする。
-    """
-
-    pass
-
-
-# 正規化で記号を変換するための辞書
-rep_map = {
-    "：": ",",
-    "；": ",",
-    "，": ",",
-    "。": ".",
-    "！": "!",
-    "？": "?",
-    "\n": ".",
-    "．": ".",
-    "…": "...",
-    "···": "...",
-    "・・・": "...",
-    "·": ",",
-    "・": ",",
-    "、": ",",
-    "$": ".",
-    "“": "'",
-    "”": "'",
-    '"': "'",
-    "‘": "'",
-    "’": "'",
-    "（": "'",
-    "）": "'",
-    "(": "'",
-    ")": "'",
-    "《": "'",
-    "》": "'",
-    "【": "'",
-    "】": "'",
-    "[": "'",
-    "]": "'",
-    # NFKC正規化後のハイフン・ダッシュの変種を全て通常半角ハイフン - \u002d に変換
-    "\u02d7": "\u002d",  # ˗, Modifier Letter Minus Sign
-    "\u2010": "\u002d",  # ‐, Hyphen,
-    # "\u2011": "\u002d",  # ‑, Non-Breaking Hyphen, NFKCにより\u2010に変換される
-    "\u2012": "\u002d",  # ‒, Figure Dash
-    "\u2013": "\u002d",  # –, En Dash
-    "\u2014": "\u002d",  # —, Em Dash
-    "\u2015": "\u002d",  # ―, Horizontal Bar
-    "\u2043": "\u002d",  # ⁃, Hyphen Bullet
-    "\u2212": "\u002d",  # −, Minus Sign
-    "\u23af": "\u002d",  # ⎯, Horizontal Line Extension
-    "\u23e4": "\u002d",  # ⏤, Straightness
-    "\u2500": "\u002d",  # ─, Box Drawings Light Horizontal
-    "\u2501": "\u002d",  # ━, Box Drawings Heavy Horizontal
-    "\u2e3a": "\u002d",  # ⸺, Two-Em Dash
-    "\u2e3b": "\u002d",  # ⸻, Three-Em Dash
-    # "～": "-",  # これは長音記号「ー」として扱うよう変更
-    # "~": "-",  # これも長音記号「ー」として扱うよう変更
-    "「": "'",
-    "」": "'",
-}
-
-
-def normalize_text(text):
-    """
-    日本語のテキストを正規化する。
-    結果は、ちょうど次の文字のみからなる：
-    - ひらがな
-    - カタカナ（全角長音記号「ー」が入る！）
-    - 漢字
-    - 半角アルファベット（大文字と小文字）
-    - ギリシャ文字
-    - `.` （句点`。`や`…`の一部や改行等）
-    - `,` （読点`、`や`:`等）
-    - `?` （疑問符`？`）
-    - `!` （感嘆符`！`）
-    - `'` （`「`や`」`等）
-    - `-` （`―`（ダッシュ、長音記号ではない）や`-`等）
-
-    注意点:
-    - 三点リーダー`…`は`...`に変換される（`なるほど…。` → `なるほど....`）
-    - 数字は漢字に変換される（`1,100円` → `千百円`、`52.34` → `五十二点三四`）
-    - 読点や疑問符等の位置・個数等は保持される（`??あ、、！！！` → `??あ,,!!!`）
-    """
-    res = unicodedata.normalize("NFKC", text)  # ここでアルファベットは半角になる
-    res = japanese_convert_numbers_to_words(res)  # 「100円」→「百円」等
-    # 「～」と「~」も長音記号として扱う
-    res = res.replace("~", "ー")
-    res = res.replace("～", "ー")
-
-    res = replace_punctuation(res)  # 句読点等正規化、読めない文字を削除
-
-    # 結合文字の濁点・半濁点を削除
-    # 通常の「ば」等はそのままのこされる、「あ゛」は上で「あ゙」になりここで「あ」になる
-    res = res.replace("\u3099", "")  # 結合文字の濁点を削除、る゙ → る
-    res = res.replace("\u309A", "")  # 結合文字の半濁点を削除、な゚ → な
-    return res
-
-
-def replace_punctuation(text: str) -> str:
-    """句読点等を「.」「,」「!」「?」「'」「-」に正規化し、OpenJTalkで読みが取得できるもののみ残す：
-    漢字・平仮名・カタカナ、アルファベット、ギリシャ文字
-    """
-    pattern = re.compile("|".join(re.escape(p) for p in rep_map.keys()))
-
-    # 句読点を辞書で置換
-    replaced_text = pattern.sub(lambda x: rep_map[x.group()], text)
-
-    replaced_text = re.sub(
-        # ↓ ひらがな、カタカナ、漢字
-        r"[^\u3040-\u309F\u30A0-\u30FF\u4E00-\u9FFF\u3400-\u4DBF\u3005"
-        # ↓ 半角アルファベット（大文字と小文字）
-        + r"\u0041-\u005A\u0061-\u007A"
-        # ↓ 全角アルファベット（大文字と小文字）
-        + r"\uFF21-\uFF3A\uFF41-\uFF5A"
-        # ↓ ギリシャ文字
-        + r"\u0370-\u03FF\u1F00-\u1FFF"
-        # ↓ "!", "?", "…", ",", ".", "'", "-", 但し`…`はすでに`...`に変換されている
-        + "".join(PUNCTUATIONS) + r"]+",
-        # 上述以外の文字を削除
-        "",
-        replaced_text,
-    )
-
-    return replaced_text
-
-
-_NUMBER_WITH_SEPARATOR_RX = re.compile("[0-9]{1,3}(,[0-9]{3})+")
-_CURRENCY_MAP = {"$": "ドル", "¥": "円", "£": "ポンド", "€": "ユーロ"}
-_CURRENCY_RX = re.compile(r"([$¥£€])([0-9.]*[0-9])")
-_NUMBER_RX = re.compile(r"[0-9]+(\.[0-9]+)?")
-
-
-def japanese_convert_numbers_to_words(text: str) -> str:
-    res = _NUMBER_WITH_SEPARATOR_RX.sub(lambda m: m[0].replace(",", ""), text)
-    res = _CURRENCY_RX.sub(lambda m: m[2] + _CURRENCY_MAP.get(m[1], m[1]), res)
-    res = _NUMBER_RX.sub(lambda m: num2words(m[0], lang="ja"), res)
-    return res
-
-
-def g2p(
-    norm_text: str, use_jp_extra: bool = True, raise_yomi_error: bool = False
-) -> tuple[list[str], list[int], list[int]]:
-    """
-    他で使われるメインの関数。`normalize_text()`で正規化された`norm_text`を受け取り、
-    - phones: 音素のリスト（ただし`!`や`,`や`.`等punctuationが含まれうる）
-    - tones: アクセントのリスト、0（低）と1（高）からなり、phonesと同じ長さ
-    - word2ph: 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
-    のタプルを返す。
-    ただし`phones`と`tones`の最初と終わりに`_`が入り、応じて`word2ph`の最初と最後に1が追加される。
-
-    use_jp_extra: Falseの場合、「ん」の音素を「N」ではなく「n」とする。
-    raise_yomi_error: Trueの場合、読めない文字があるときに例外を発生させる。
-    Falseの場合は読めない文字が消えたような扱いとして処理される。
-    """
-    # pyopenjtalkのフルコンテキストラベルを使ってアクセントを取り出すと、punctuationの位置が消えてしまい情報が失われてしまう：
-    # 「こんにちは、世界。」と「こんにちは！世界。」と「こんにちは！！！？？？世界……。」は全て同じになる。
-    # よって、まずpunctuation無しの音素とアクセントのリストを作り、
-    # それとは別にpyopenjtalk.run_frontend()で得られる音素リスト（こちらはpunctuationが保持される）を使い、
-    # アクセント割当をしなおすことによってpunctuationを含めた音素とアクセントのリストを作る。
-
-    # punctuationがすべて消えた、音素とアクセントのタプルのリスト（「ん」は「N」）
-    phone_tone_list_wo_punct = g2phone_tone_wo_punct(norm_text)
-
-    # sep_text: 単語単位の単語のリスト、読めない文字があったらraise_yomi_errorなら例外、そうでないなら読めない文字が消えて返ってくる
-    # sep_kata: 単語単位の単語のカタカナ読みのリスト
-    sep_text, sep_kata = text2sep_kata(norm_text, raise_yomi_error=raise_yomi_error)
-
-    # sep_phonemes: 各単語ごとの音素のリストのリスト
-    sep_phonemes = handle_long([kata2phoneme_list(i) for i in sep_kata])
-
-    # phone_w_punct: sep_phonemesを結合した、punctuationを元のまま保持した音素列
-    phone_w_punct: list[str] = []
-    for i in sep_phonemes:
-        phone_w_punct += i
-
-    # punctuation無しのアクセント情報を使って、punctuationを含めたアクセント情報を作る
-    phone_tone_list = align_tones(phone_w_punct, phone_tone_list_wo_punct)
-    # logger.debug(f"phone_tone_list:\n{phone_tone_list}")
-    # word2phは厳密な解答は不可能なので（「今日」「眼鏡」等の熟字訓が存在）、
-    # Bert-VITS2では、単語単位の分割を使って、単語の文字ごとにだいたい均等に音素を分配する
-
-    # sep_textから、各単語を1文字1文字分割して、文字のリスト（のリスト）を作る
-    sep_tokenized: list[list[str]] = []
-    for i in sep_text:
-        if i not in PUNCTUATIONS:
-            sep_tokenized.append(
-                tokenizer.tokenize(i)
-            )  # ここでおそらく`i`が文字単位に分割される
-        else:
-            sep_tokenized.append([i])
-
-    # 各単語について、音素の数と文字の数を比較して、均等っぽく分配する
-    word2ph = []
-    for token, phoneme in zip(sep_tokenized, sep_phonemes):
-        phone_len = len(phoneme)
-        word_len = len(token)
-        word2ph += distribute_phone(phone_len, word_len)
-
-    # 最初と最後に`_`記号を追加、アクセントは0（低）、word2phもそれに合わせて追加
-    phone_tone_list = [("_", 0)] + phone_tone_list + [("_", 0)]
-    word2ph = [1] + word2ph + [1]
-
-    phones = [phone for phone, _ in phone_tone_list]
-    tones = [tone for _, tone in phone_tone_list]
-
-    assert len(phones) == sum(word2ph), f"{len(phones)} != {sum(word2ph)}"
-
-    # use_jp_extraでない場合は「N」を「n」に変換
-    if not use_jp_extra:
-        phones = [phone if phone != "N" else "n" for phone in phones]
-
-    return phones, tones, word2ph
-
-
-def g2kata_tone(norm_text: str) -> list[tuple[str, int]]:
-    """
-    テキストからカタカナとアクセントのペアのリストを返す。
-    推論時のみに使われるので、常に`raise_yomi_error=False`でg2pを呼ぶ。
-    """
-    phones, tones, _ = g2p(norm_text, use_jp_extra=True, raise_yomi_error=False)
-    return phone_tone2kata_tone(list(zip(phones, tones)))
-
-
-def phone_tone2kata_tone(phone_tone: list[tuple[str, int]]) -> list[tuple[str, int]]:
-    """phone_toneをのphone部分をカタカナに変換する。ただし最初と最後の("_", 0)は無視"""
-    phone_tone = phone_tone[1:]  # 最初の("_", 0)を無視
-    phones = [phone for phone, _ in phone_tone]
-    tones = [tone for _, tone in phone_tone]
-    result: list[tuple[str, int]] = []
-    current_mora = ""
-    for phone, next_phone, tone, next_tone in zip(phones, phones[1:], tones, tones[1:]):
-        # zipの関係で最後の("_", 0)は無視されている
-        if phone in PUNCTUATIONS:
-            result.append((phone, tone))
-            continue
-        if phone in COSONANTS:  # n以外の子音の場合
-            assert current_mora == "", f"Unexpected {phone} after {current_mora}"
-            assert tone == next_tone, f"Unexpected {phone} tone {tone} != {next_tone}"
-            current_mora = phone
-        else:
-            # phoneが母音もしくは「N」
-            current_mora += phone
-            result.append((MORA_PHONEMES_TO_MORA_KATA[current_mora], tone))
-            current_mora = ""
-    return result
-
-
-def kata_tone2phone_tone(kata_tone: list[tuple[str, int]]) -> list[tuple[str, int]]:
-    """`phone_tone2kata_tone()`の逆。"""
-    result: list[tuple[str, int]] = [("_", 0)]
-    for mora, tone in kata_tone:
-        if mora in PUNCTUATIONS:
-            result.append((mora, tone))
-        else:
-            cosonant, vowel = MORA_KATA_TO_MORA_PHONEMES[mora]
-            if cosonant is None:
-                result.append((vowel, tone))
-            else:
-                result.append((cosonant, tone))
-                result.append((vowel, tone))
-    result.append(("_", 0))
-    return result
-
-
-def g2phone_tone_wo_punct(text: str) -> list[tuple[str, int]]:
-    """
-    テキストに対して、音素とアクセント（0か1）のペアのリストを返す。
-    ただし「!」「.」「?」等の非音素記号(punctuation)は全て消える（ポーズ記号も残さない）。
-    非音素記号を含める処理は`align_tones()`で行われる。
-    また「っ」は「q」に、「ん」は「N」に変換される。
-    例: "こんにちは、世界ー。。元気？！" →
-    [('k', 0), ('o', 0), ('N', 1), ('n', 1), ('i', 1), ('ch', 1), ('i', 1), ('w', 1), ('a', 1), ('s', 1), ('e', 1), ('k', 0), ('a', 0), ('i', 0), ('i', 0), ('g', 1), ('e', 1), ('N', 0), ('k', 0), ('i', 0)]
-    """
-    prosodies = pyopenjtalk_g2p_prosody(text, drop_unvoiced_vowels=True)
-    # logger.debug(f"prosodies: {prosodies}")
-    result: list[tuple[str, int]] = []
-    current_phrase: list[tuple[str, int]] = []
-    current_tone = 0
-    for i, letter in enumerate(prosodies):
-        # 特殊記号の処理
-
-        # 文頭記号、無視する
-        if letter == "^":
-            assert i == 0, "Unexpected ^"
-        # アクセント句の終わりに来る記号
-        elif letter in ("$", "?", "_", "#"):
-            # 保持しているフレーズを、アクセント数値を0-1に修正し結果に追加
-            result.extend(fix_phone_tone(current_phrase))
-            # 末尾に来る終了記号、無視（文中の疑問文は`_`になる）
-            if letter in ("$", "?"):
-                assert i == len(prosodies) - 1, f"Unexpected {letter}"
-            # あとは"_"（ポーズ）と"#"（アクセント句の境界）のみ
-            # これらは残さず、次のアクセント句に備える。
-            current_phrase = []
-            # 0を基準点にしてそこから上昇・下降する（負の場合は上の`fix_phone_tone`で直る）
-            current_tone = 0
-        # アクセント上昇記号
-        elif letter == "[":
-            current_tone = current_tone + 1
-        # アクセント下降記号
-        elif letter == "]":
-            current_tone = current_tone - 1
-        # それ以外は通常の音素
-        else:
-            if letter == "cl":  # 「っ」の処理
-                letter = "q"
-            # elif letter == "N":  # 「ん」の処理
-            #     letter = "n"
-            current_phrase.append((letter, current_tone))
-    return result
-
-
-def text2sep_kata(
-    norm_text: str, raise_yomi_error: bool = False
-) -> tuple[list[str], list[str]]:
-    """
-    `normalize_text()`で正規化済みの`norm_text`を受け取り、それを単語分割し、
-    分割された単語リストとその読み（カタカナor記号1文字）のリストのタプルを返す。
-    単語分割結果は、`g2p()`の`word2ph`で1文字あたりに割り振る音素記号の数を決めるために使う。
-    例:
-    `私はそう思う!って感じ?` →
-    ["私", "は", "そう", "思う", "!", "って", "感じ", "?"], ["ワタシ", "ワ", "ソー", "オモウ", "!", "ッテ", "カンジ", "?"]
-
-    raise_yomi_error: Trueの場合、読めない文字があるときに例外を発生させる。
-    Falseの場合は読めない文字が消えたような扱いとして処理される。
-    """
-    # parsed: OpenJTalkの解析結果
-    parsed = pyopenjtalk.run_frontend(norm_text)
-    sep_text: list[str] = []
-    sep_kata: list[str] = []
-    for parts in parsed:
-        # word: 実際の単語の文字列
-        # yomi: その読み、但し無声化サインの`’`は除去
-        word, yomi = replace_punctuation(parts["string"]), parts["pron"].replace(
-            "’", ""
-        )
-        """
-        ここで`yomi`の取りうる値は以下の通りのはず。
-        - `word`が通常単語 → 通常の読み（カタカナ）
-            （カタカナからなり、長音記号も含みうる、`アー` 等）
-        - `word`が`ー` から始まる → `ーラー` や `ーーー` など
-        - `word`が句読点や空白等 → `、`
-        - `word`がpunctuationの繰り返し → 全角にしたもの
-        基本的にpunctuationは1文字ずつ分かれるが、何故かある程度連続すると1つにまとまる。
-        他にも`word`が読めないキリル文字アラビア文字等が来ると`、`になるが、正規化でこの場合は起きないはず。
-        また元のコードでは`yomi`が空白の場合の処理があったが、これは起きないはず。
-        処理すべきは`yomi`が`、`の場合のみのはず。
-        """
-        assert yomi != "", f"Empty yomi: {word}"
-        if yomi == "、":
-            # wordは正規化されているので、`.`, `,`, `!`, `'`, `-`, `--` のいずれか
-            if not set(word).issubset(set(PUNCTUATIONS)):  # 記号繰り返しか判定
-                # ここはpyopenjtalkが読めない文字等のときに起こる
-                if raise_yomi_error:
-                    raise YomiError(f"Cannot read: {word} in:\n{norm_text}")
-                logger.warning(f"Ignoring unknown: {word} in:\n{norm_text}")
-                continue
-            # yomiは元の記号のままに変更
-            yomi = word
-        elif yomi == "？":
-            assert word == "?", f"yomi `？` comes from: {word}"
-            yomi = "?"
-        sep_text.append(word)
-        sep_kata.append(yomi)
-    return sep_text, sep_kata
-
-
-# ESPnetの実装から引用、変更点無し。「ん」は「N」なことに注意。
-# https://github.com/espnet/espnet/blob/master/espnet2/text/phoneme_tokenizer.py
-def pyopenjtalk_g2p_prosody(text: str, drop_unvoiced_vowels: bool = True) -> list[str]:
-    """Extract phoneme + prosoody symbol sequence from input full-context labels.
-
-    The algorithm is based on `Prosodic features control by symbols as input of
-    sequence-to-sequence acoustic modeling for neural TTS`_ with some r9y9's tweaks.
-
-    Args:
-        text (str): Input text.
-        drop_unvoiced_vowels (bool): whether to drop unvoiced vowels.
-
-    Returns:
-        List[str]: List of phoneme + prosody symbols.
-
-    Examples:
-        >>> from espnet2.text.phoneme_tokenizer import pyopenjtalk_g2p_prosody
-        >>> pyopenjtalk_g2p_prosody("こんにちは。")
-        ['^', 'k', 'o', '[', 'N', 'n', 'i', 'ch', 'i', 'w', 'a', '$']
-
-    .. _`Prosodic features control by symbols as input of sequence-to-sequence acoustic
-        modeling for neural TTS`: https://doi.org/10.1587/transinf.2020EDP7104
-
-    """
-    labels = pyopenjtalk.make_label(pyopenjtalk.run_frontend(text))
-    N = len(labels)
-
-    phones = []
-    for n in range(N):
-        lab_curr = labels[n]
-
-        # current phoneme
-        p3 = re.search(r"\-(.*?)\+", lab_curr).group(1)
-        # deal unvoiced vowels as normal vowels
-        if drop_unvoiced_vowels and p3 in "AEIOU":
-            p3 = p3.lower()
-
-        # deal with sil at the beginning and the end of text
-        if p3 == "sil":
-            assert n == 0 or n == N - 1
-            if n == 0:
-                phones.append("^")
-            elif n == N - 1:
-                # check question form or not
-                e3 = _numeric_feature_by_regex(r"!(\d+)_", lab_curr)
-                if e3 == 0:
-                    phones.append("$")
-                elif e3 == 1:
-                    phones.append("?")
-            continue
-        elif p3 == "pau":
-            phones.append("_")
-            continue
-        else:
-            phones.append(p3)
-
-        # accent type and position info (forward or backward)
-        a1 = _numeric_feature_by_regex(r"/A:([0-9\-]+)\+", lab_curr)
-        a2 = _numeric_feature_by_regex(r"\+(\d+)\+", lab_curr)
-        a3 = _numeric_feature_by_regex(r"\+(\d+)/", lab_curr)
-
-        # number of mora in accent phrase
-        f1 = _numeric_feature_by_regex(r"/F:(\d+)_", lab_curr)
-
-        a2_next = _numeric_feature_by_regex(r"\+(\d+)\+", labels[n + 1])
-        # accent phrase border
-        if a3 == 1 and a2_next == 1 and p3 in "aeiouAEIOUNcl":
-            phones.append("#")
-        # pitch falling
-        elif a1 == 0 and a2_next == a2 + 1 and a2 != f1:
-            phones.append("]")
-        # pitch rising
-        elif a2 == 1 and a2_next == 2:
-            phones.append("[")
-
-    return phones
-
-
-def _numeric_feature_by_regex(regex, s):
-    match = re.search(regex, s)
-    if match is None:
-        return -50
-    return int(match.group(1))
-
-
-def fix_phone_tone(phone_tone_list: list[tuple[str, int]]) -> list[tuple[str, int]]:
-    """
-    `phone_tone_list`のtone（アクセントの値）を0か1の範囲に修正する。
-    例: [(a, 0), (i, -1), (u, -1)] → [(a, 1), (i, 0), (u, 0)]
-    """
-    tone_values = set(tone for _, tone in phone_tone_list)
-    if len(tone_values) == 1:
-        assert tone_values == {0}, tone_values
-        return phone_tone_list
-    elif len(tone_values) == 2:
-        if tone_values == {0, 1}:
-            return phone_tone_list
-        elif tone_values == {-1, 0}:
-            return [
-                (letter, 0 if tone == -1 else 1) for letter, tone in phone_tone_list
-            ]
-        else:
-            raise ValueError(f"Unexpected tone values: {tone_values}")
-    else:
-        raise ValueError(f"Unexpected tone values: {tone_values}")
-
-
-def distribute_phone(n_phone: int, n_word: int) -> list[int]:
-    """
-    左から右に1ずつ振り分け、次にまた左から右に1ずつ増やし、というふうに、
-    音素の数`n_phone`を単語の数`n_word`に分配する。
-    """
-    phones_per_word = [0] * n_word
-    for _ in range(n_phone):
-        min_tasks = min(phones_per_word)
-        min_index = phones_per_word.index(min_tasks)
-        phones_per_word[min_index] += 1
-    return phones_per_word
-
-
-def handle_long(sep_phonemes: list[list[str]]) -> list[list[str]]:
-    """
-    フレーズごとに分かれた音素（長音記号がそのまま）のリストのリスト`sep_phonemes`を受け取り、
-    その長音記号を処理して、音素のリストのリストを返す。
-    基本的には直前の音素を伸ばすが、直前の音素が母音でない場合もしくは冒頭の場合は、
-    おそらく長音記号とダッシュを勘違いしていると思われるので、ダッシュに対応する音素`-`に変換する。
-    """
-    for i in range(len(sep_phonemes)):
-        if len(sep_phonemes[i]) == 0:
-            # 空白文字等でリストが空の場合
-            continue
-        if sep_phonemes[i][0] == "ー":
-            if i != 0:
-                prev_phoneme = sep_phonemes[i - 1][-1]
-                if prev_phoneme in VOWELS:
-                    # 母音と「ん」のあとの伸ばし棒なので、その母音に変換
-                    sep_phonemes[i][0] = sep_phonemes[i - 1][-1]
-                else:
-                    # 「。ーー」等おそらく予期しない長音記号
-                    # ダッシュの勘違いだと思われる
-                    sep_phonemes[i][0] = "-"
-            else:
-                # 冒頭に長音記号が来ていおり、これはダッシュの勘違いと思われる
-                sep_phonemes[i][0] = "-"
-        if "ー" in sep_phonemes[i]:
-            for j in range(len(sep_phonemes[i])):
-                if sep_phonemes[i][j] == "ー":
-                    sep_phonemes[i][j] = sep_phonemes[i][j - 1][-1]
-    return sep_phonemes
-
-
-tokenizer = AutoTokenizer.from_pretrained("./bert/deberta-v2-large-japanese-char-wwm")
-
-
-def align_tones(
-    phones_with_punct: list[str], phone_tone_list: list[tuple[str, int]]
-) -> list[tuple[str, int]]:
-    """
-    例:
-    …私は、、そう思う。
-    phones_with_punct:
-    [".", ".", ".", "w", "a", "t", "a", "sh", "i", "w", "a", ",", ",", "s", "o", "o", "o", "m", "o", "u", "."]
-    phone_tone_list:
-    [("w", 0), ("a", 0), ("t", 1), ("a", 1), ("sh", 1), ("i", 1), ("w", 1), ("a", 1), ("_", 0), ("s", 0), ("o", 0), ("o", 1), ("o", 1), ("m", 1), ("o", 1), ("u", 0))]
-    Return:
-    [(".", 0), (".", 0), (".", 0), ("w", 0), ("a", 0), ("t", 1), ("a", 1), ("sh", 1), ("i", 1), ("w", 1), ("a", 1), (",", 0), (",", 0), ("s", 0), ("o", 0), ("o", 1), ("o", 1), ("m", 1), ("o", 1), ("u", 0), (".", 0)]
-    """
-    result: list[tuple[str, int]] = []
-    tone_index = 0
-    for phone in phones_with_punct:
-        if tone_index >= len(phone_tone_list):
-            # 余ったpunctuationがある場合 → (punctuation, 0)を追加
-            result.append((phone, 0))
-        elif phone == phone_tone_list[tone_index][0]:
-            # phone_tone_listの現在の音素と一致する場合 → toneをそこから取得、(phone, tone)を追加
-            result.append((phone, phone_tone_list[tone_index][1]))
-            # 探すindexを1つ進める
-            tone_index += 1
-        elif phone in PUNCTUATIONS:
-            # phoneがpunctuationの場合 → (phone, 0)を追加
-            result.append((phone, 0))
-        else:
-            logger.debug(f"phones: {phones_with_punct}")
-            logger.debug(f"phone_tone_list: {phone_tone_list}")
-            logger.debug(f"result: {result}")
-            logger.debug(f"tone_index: {tone_index}")
-            logger.debug(f"phone: {phone}")
-            raise ValueError(f"Unexpected phone: {phone}")
-    return result
-
-
-def kata2phoneme_list(text: str) -> list[str]:
-    """
-    原則カタカナの`text`を受け取り、それをそのままいじらずに音素記号のリストに変換。
-    注意点：
-    - punctuationかその繰り返しが来た場合、punctuationたちをそのままリストにして返す。
-    - 冒頭に続く「ー」はそのまま「ー」のままにする（`handle_long()`で処理される）
-    - 文中の「ー」は前の音素記号の最後の音素記号に変換される。
-    例：
-    `ーーソーナノカーー` → ["ー", "ー", "s", "o", "o", "n", "a", "n", "o", "k", "a", "a", "a"]
-    `?` → ["?"]
-    `!?!?!?!?!` → ["!", "?", "!", "?", "!", "?", "!", "?", "!"]
-    """
-    if set(text).issubset(set(PUNCTUATIONS)):
-        return list(text)
-    # `text`がカタカナ（`ー`含む）のみからなるかどうかをチェック
-    if re.fullmatch(r"[\u30A0-\u30FF]+", text) is None:
-        raise ValueError(f"Input must be katakana only: {text}")
-    sorted_keys = sorted(MORA_KATA_TO_MORA_PHONEMES.keys(), key=len, reverse=True)
-    pattern = "|".join(map(re.escape, sorted_keys))
-
-    def mora2phonemes(mora: str) -> str:
-        cosonant, vowel = MORA_KATA_TO_MORA_PHONEMES[mora]
-        if cosonant is None:
-            return f" {vowel}"
-        return f" {cosonant} {vowel}"
-
-    spaced_phonemes = re.sub(pattern, lambda m: mora2phonemes(m.group()), text)
-
-    # 長音記号「ー」の処理
-    long_pattern = r"(\w)(ー*)"
-    long_replacement = lambda m: m.group(1) + (" " + m.group(1)) * len(m.group(2))
-    spaced_phonemes = re.sub(long_pattern, long_replacement, spaced_phonemes)
-    return spaced_phonemes.strip().split(" ")
-
-
-if __name__ == "__main__":
-    tokenizer = AutoTokenizer.from_pretrained(
-        "./bert/deberta-v2-large-japanese-char-wwm"
-    )
-    text = "こんにちは、世界。"
-    from text.japanese_bert import get_bert_feature
-
-    text = normalize_text(text)
-
-    phones, tones, word2ph = g2p(text)
-    bert = get_bert_feature(text, word2ph)
-
-    print(phones, tones, word2ph, bert.shape)

text/japanese_bert.py

@@ -1,24 +1,22 @@
 import sys
 
 import torch
-from transformers import AutoModelForMaskedLM, AutoTokenizer
 
 from config import config
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.text_processing import bert_models
 from style_bert_vits2.text_processing.japanese.g2p import text_to_sep_kata
 
-LOCAL_PATH = "./bert/deberta-v2-large-japanese-char-wwm"
-
-tokenizer = AutoTokenizer.from_pretrained(LOCAL_PATH)
 
 models = dict()
 
 
 def get_bert_feature(
-    text,
+    text: str,
     word2ph,
-    device=config.bert_gen_config.device,
+    device = config.bert_gen_config.device,
-    assist_text=None,
+    assist_text: str | None = None,
-    assist_text_weight=0.7,
+    assist_text_weight: float = 0.7,
 ):
     # 各単語が何文字かを作る`word2ph`を使う必要があるので、読めない文字は必ず無視する
     # でないと`word2ph`の結果とテキストの文字数結果が整合性が取れない
@@ -37,8 +35,9 @@ def get_bert_feature(
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
     if device not in models.keys():
-        models[device] = AutoModelForMaskedLM.from_pretrained(LOCAL_PATH).to(device)
+        models[device] = bert_models.load_model(Languages.JP).to(device)
     with torch.no_grad():
+        tokenizer = bert_models.load_tokenizer(Languages.JP)
         inputs = tokenizer(text, return_tensors="pt")
         for i in inputs:
             inputs[i] = inputs[i].to(device)