307 lines
12 KiB
Python
307 lines
12 KiB
Python
from __future__ import annotations
|
||
|
||
import logging
|
||
from pathlib import Path
|
||
import re
|
||
import threading
|
||
from dataclasses import dataclass
|
||
|
||
from tts_robust_normalizer_single_script import normalize_tts_text
|
||
|
||
ENGLISH_VOICES = frozenset({"Trump", "Ava", "Bella", "Adam", "Nathan"})
|
||
CUSTOM_ZH_WETEXT_CACHE_DIR = Path(__file__).resolve().parent / ".cache" / "wetext_zh_no_erhua_keep_punct"
|
||
_ZH_WETEXT_KEEP_HYPHEN = "___KEEP_HYPHEN_BEFORE_ZH_WETEXT___"
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class TextNormalizationSnapshot:
|
||
state: str
|
||
message: str
|
||
error: str | None = None
|
||
ready: bool = False
|
||
available: bool = False
|
||
|
||
@property
|
||
def failed(self) -> bool:
|
||
return self.state == "failed"
|
||
|
||
|
||
class WeTextProcessingManager:
|
||
def __init__(self) -> None:
|
||
self._lock = threading.Lock()
|
||
self._normalize_lock = threading.Lock()
|
||
self._thread: threading.Thread | None = None
|
||
self._started = False
|
||
self._state = "pending"
|
||
self._message = "Waiting for WeTextProcessing preload."
|
||
self._error: str | None = None
|
||
self._available = True
|
||
self._normalizers: dict[str, object] | None = None
|
||
|
||
def snapshot(self) -> TextNormalizationSnapshot:
|
||
with self._lock:
|
||
return TextNormalizationSnapshot(
|
||
state=self._state,
|
||
message=self._message,
|
||
error=self._error,
|
||
ready=self._state == "ready",
|
||
available=self._available,
|
||
)
|
||
|
||
def _set_state(self, *, state: str, message: str, error: str | None = None) -> None:
|
||
with self._lock:
|
||
self._state = state
|
||
self._message = message
|
||
self._error = error
|
||
|
||
def start(self) -> None:
|
||
with self._lock:
|
||
if self._started:
|
||
return
|
||
self._started = True
|
||
self._thread = threading.Thread(target=self._run, name="wetext-preload", daemon=True)
|
||
self._thread.start()
|
||
|
||
def ensure_ready(self) -> TextNormalizationSnapshot:
|
||
with self._lock:
|
||
if not self._started:
|
||
self._started = True
|
||
self._thread = threading.Thread(target=self._run, name="wetext-preload", daemon=True)
|
||
self._thread.start()
|
||
thread = self._thread
|
||
if thread is not None and thread.is_alive():
|
||
thread.join()
|
||
return self.snapshot()
|
||
|
||
def close(self) -> None:
|
||
return
|
||
|
||
def _run(self) -> None:
|
||
if not self._available:
|
||
self._set_state(
|
||
state="failed",
|
||
message="WeTextProcessing unavailable.",
|
||
error="installed WeTextProcessing modules are unavailable",
|
||
)
|
||
return
|
||
try:
|
||
self._set_state(state="running", message="Loading WeTextProcessing graphs.", error=None)
|
||
self._ensure_normalizers_loaded()
|
||
self._set_state(state="ready", message="WeTextProcessing ready. languages=zh,en", error=None)
|
||
except Exception as exc:
|
||
logging.exception("WeTextProcessing preload failed")
|
||
self._set_state(state="failed", message="WeTextProcessing preload failed.", error=str(exc))
|
||
|
||
def _ensure_normalizers_loaded(self) -> dict[str, object]:
|
||
with self._lock:
|
||
if self._normalizers is not None:
|
||
return self._normalizers
|
||
|
||
from tn.chinese.normalizer import Normalizer as ZhNormalizer
|
||
from tn.english.normalizer import Normalizer as EnNormalizer
|
||
|
||
logging.getLogger().setLevel(logging.INFO)
|
||
self._normalizers = {
|
||
"zh": ZhNormalizer(
|
||
cache_dir=str(CUSTOM_ZH_WETEXT_CACHE_DIR),
|
||
overwrite_cache=False,
|
||
remove_interjections=False,
|
||
remove_erhua=False,
|
||
full_to_half=False,
|
||
),
|
||
"en": EnNormalizer(overwrite_cache=False),
|
||
}
|
||
return self._normalizers
|
||
|
||
def normalize(self, *, text: str, prompt_text: str, language: str) -> tuple[str, str]:
|
||
snapshot = self.ensure_ready()
|
||
if not snapshot.ready:
|
||
raise RuntimeError(snapshot.error or snapshot.message)
|
||
|
||
with self._normalize_lock:
|
||
normalizers = self._ensure_normalizers_loaded()
|
||
if language not in normalizers:
|
||
raise ValueError(f"Unsupported text normalization language: {language}")
|
||
normalizer = normalizers[language]
|
||
normalized_text = normalizer.normalize(text) if text else ""
|
||
normalized_prompt_text = normalizer.normalize(prompt_text) if prompt_text else ""
|
||
return normalized_text, normalized_prompt_text
|
||
|
||
|
||
def resolve_text_normalization_language(*, text: str, voice: str) -> str:
|
||
if re.search(r"[\u3400-\u9fff]", text):
|
||
return "zh"
|
||
if re.search(r"[A-Za-z]", text):
|
||
return "en"
|
||
if voice in ENGLISH_VOICES:
|
||
return "en"
|
||
return "zh"
|
||
|
||
|
||
def _rewrite_hyphens_before_zh_wetext(text: str) -> str:
|
||
"""Avoid Chinese WeText reading non-numeric hyphens as '减'."""
|
||
rewritten = str(text or "")
|
||
if "-" not in rewritten:
|
||
return rewritten
|
||
|
||
# Preserve start-of-text negatives like `-2`.
|
||
rewritten = re.sub(
|
||
r"(^\s*)-\s*(?=\d)",
|
||
rf"\1{_ZH_WETEXT_KEEP_HYPHEN}",
|
||
rewritten,
|
||
)
|
||
|
||
# Preserve negatives after common delimiters like `x=-2` or `(-2)`.
|
||
rewritten = re.sub(
|
||
r"([=:+*/,(，:：；;（【\[{])\s*-\s*(?=\d)",
|
||
rf"\1{_ZH_WETEXT_KEEP_HYPHEN}",
|
||
rewritten,
|
||
)
|
||
|
||
# Preserve Chinese-context negatives like `为-2` / `计算-2`.
|
||
rewritten = re.sub(
|
||
r"([\u3400-\u9fff])\s*-\s*(?=\d)",
|
||
rf"\1{_ZH_WETEXT_KEEP_HYPHEN}",
|
||
rewritten,
|
||
)
|
||
|
||
# Preserve numeric ranges/dates like `10-3` / `2024-05-01`.
|
||
rewritten = re.sub(
|
||
r"(\d)\s*-\s*(?=\d)",
|
||
rf"\1{_ZH_WETEXT_KEEP_HYPHEN}",
|
||
rewritten,
|
||
)
|
||
|
||
# Chinese compound phrases sound more natural with a pause boundary.
|
||
rewritten = re.sub(
|
||
r"([\u3400-\u9fff])\s*-\s*(?=[\u3400-\u9fff])",
|
||
r"\1，",
|
||
rewritten,
|
||
)
|
||
|
||
# Remaining token-internal hyphens are flattened to spaces for zh WeText.
|
||
rewritten = re.sub(
|
||
r"([^\s-])\s*-\s*(?=[^\s-])",
|
||
r"\1 ",
|
||
rewritten,
|
||
)
|
||
|
||
rewritten = re.sub(r" {2,}", " ", rewritten).strip()
|
||
return rewritten.replace(_ZH_WETEXT_KEEP_HYPHEN, "-")
|
||
|
||
|
||
def prepare_tts_request_texts(
|
||
*,
|
||
text: str,
|
||
prompt_text: str = "",
|
||
voice: str = "",
|
||
enable_wetext: bool,
|
||
enable_normalize_tts_text: bool = True,
|
||
text_normalizer_manager: WeTextProcessingManager | None,
|
||
) -> dict[str, object]:
|
||
raw_text = str(text or "")
|
||
raw_prompt_text = str(prompt_text or "")
|
||
|
||
normalization_stages: list[str] = []
|
||
normalization_language = ""
|
||
intermediate_text = raw_text
|
||
intermediate_prompt_text = raw_prompt_text
|
||
|
||
if enable_normalize_tts_text and enable_wetext:
|
||
pre_robust_text = normalize_tts_text(raw_text)
|
||
pre_robust_prompt_text = normalize_tts_text(raw_prompt_text) if raw_prompt_text else ""
|
||
if pre_robust_text != raw_text:
|
||
logging.info(
|
||
"normalized text chars_before=%d chars_after=%d stage=robust_pre",
|
||
len(raw_text),
|
||
len(pre_robust_text),
|
||
)
|
||
if raw_prompt_text and pre_robust_prompt_text != raw_prompt_text:
|
||
logging.info(
|
||
"normalized prompt_text chars_before=%d chars_after=%d stage=robust_pre",
|
||
len(raw_prompt_text),
|
||
len(pre_robust_prompt_text),
|
||
)
|
||
intermediate_text = pre_robust_text
|
||
intermediate_prompt_text = pre_robust_prompt_text
|
||
normalization_stages.append("robust_pre")
|
||
|
||
if enable_wetext:
|
||
if text_normalizer_manager is None:
|
||
raise RuntimeError("WeTextProcessing manager is unavailable.")
|
||
wetext_input_text = intermediate_text
|
||
wetext_input_prompt_text = intermediate_prompt_text
|
||
normalization_language = resolve_text_normalization_language(text=wetext_input_text, voice=voice)
|
||
if normalization_language == "zh":
|
||
rewritten_wetext_input_text = _rewrite_hyphens_before_zh_wetext(wetext_input_text)
|
||
rewritten_wetext_input_prompt_text = _rewrite_hyphens_before_zh_wetext(wetext_input_prompt_text)
|
||
if rewritten_wetext_input_text != wetext_input_text:
|
||
logging.info(
|
||
"rewrote zh wetext text hyphens chars_before=%d chars_after=%d stage=zh_wetext_hyphen_guard",
|
||
len(wetext_input_text),
|
||
len(rewritten_wetext_input_text),
|
||
)
|
||
if wetext_input_prompt_text and rewritten_wetext_input_prompt_text != wetext_input_prompt_text:
|
||
logging.info(
|
||
"rewrote zh wetext prompt_text hyphens chars_before=%d chars_after=%d stage=zh_wetext_hyphen_guard",
|
||
len(wetext_input_prompt_text),
|
||
len(rewritten_wetext_input_prompt_text),
|
||
)
|
||
wetext_input_text = rewritten_wetext_input_text
|
||
wetext_input_prompt_text = rewritten_wetext_input_prompt_text
|
||
intermediate_text, intermediate_prompt_text = text_normalizer_manager.normalize(
|
||
text=wetext_input_text,
|
||
prompt_text=wetext_input_prompt_text,
|
||
language=normalization_language,
|
||
)
|
||
if intermediate_text != wetext_input_text:
|
||
logging.info(
|
||
"normalized text chars_before=%d chars_after=%d stage=wetext language=%s",
|
||
len(wetext_input_text),
|
||
len(intermediate_text),
|
||
normalization_language,
|
||
)
|
||
if wetext_input_prompt_text and intermediate_prompt_text != wetext_input_prompt_text:
|
||
logging.info(
|
||
"normalized prompt_text chars_before=%d chars_after=%d stage=wetext language=%s",
|
||
len(wetext_input_prompt_text),
|
||
len(intermediate_prompt_text),
|
||
normalization_language,
|
||
)
|
||
normalization_stages.append(f"wetext:{normalization_language}" if normalization_language else "wetext")
|
||
|
||
final_text = intermediate_text
|
||
final_prompt_text = intermediate_prompt_text
|
||
if enable_normalize_tts_text:
|
||
final_text = normalize_tts_text(intermediate_text)
|
||
final_prompt_text = normalize_tts_text(intermediate_prompt_text) if intermediate_prompt_text else ""
|
||
robust_stage_name = "robust_post" if enable_wetext else "robust"
|
||
|
||
if final_text != intermediate_text:
|
||
logging.info(
|
||
"normalized text chars_before=%d chars_after=%d stage=%s",
|
||
len(intermediate_text),
|
||
len(final_text),
|
||
robust_stage_name,
|
||
)
|
||
if intermediate_prompt_text and final_prompt_text != intermediate_prompt_text:
|
||
logging.info(
|
||
"normalized prompt_text chars_before=%d chars_after=%d stage=%s",
|
||
len(intermediate_prompt_text),
|
||
len(final_prompt_text),
|
||
robust_stage_name,
|
||
)
|
||
normalization_stages.append(robust_stage_name)
|
||
|
||
return {
|
||
"text": final_text,
|
||
"prompt_text": final_prompt_text,
|
||
"normalized_text": final_text,
|
||
"normalized_prompt_text": final_prompt_text,
|
||
"normalization_method": "+".join(normalization_stages) if normalization_stages else "none",
|
||
"text_normalization_language": normalization_language,
|
||
"text_normalization_enabled": bool(enable_wetext or enable_normalize_tts_text),
|
||
"wetext_processing_enabled": bool(enable_wetext),
|
||
"normalize_tts_text_enabled": bool(enable_normalize_tts_text),
|
||
}
|