text_normalization_pipeline.py

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),
    }