Mike0021/zonos2

1"""TTS text normalization using the vendored NeMo forward-TN grammars.2 3Converts written text to its spoken form ("$5.32" -> "five dollars thirty two4cents") before tokenization. Grammar construction is expensive (~10-30s per5language) the first time; compiled WFSTs are cached as .far files under6ZONOS2_TTS_NORM_CACHE_DIR (default ~/.cache/zonos2-tts-norm) and reload in7well under a second. Use tools/build_tts_norm_fars.py to prewarm all caches.8 9Disable globally with ZONOS2_TTS_NORM=0, or per request via the10text_normalization flag.11"""12 13from __future__ import annotations14 15import os16import re17import threading18from typing import Dict, Optional19 20from zonos2.utils import init_logger21 22logger = init_logger(__name__)23 24# Server language codes -> NeMo text_normalization language packages.25SERVER_TO_NEMO_LANG: Dict[str, str] = {26    "en_us": "en",27    "en_gb": "en",28    "fr_fr": "fr",29    "de": "de",30    "es": "es",31    "it": "it",32    "pt_br": "pt",33    "ja": "ja",34    "cmn": "zh",35    "ko": "ko",36}37 38# Upstream's own tests run Korean grammars with lower_cased input; everything39# else uses cased.40LOWER_CASED_LANGS = {"ko"}41 42# zh/ja verbalizers read a cached .far but never write one (upstream quirk);43# we post-write it after first compile so later loads are fast. Note ja reads44# a "jp_" prefixed name.45_VERBALIZER_FAR_PREFIX = {"zh": "zh", "ja": "jp"}46 47# A digit directly followed by sentence punctuation confuses several upstream48# r1.2.0 grammars: pt's tagger raises FstOpError outright and de reads dates49# digit-by-digit (reproduced with upstream code and its pinned pynini50# 2.1.6.post1 — upstream bugs, not vendoring artifacts). Space the punctuation51# off before normalization; it is re-attached afterwards.52_DIGIT_PUNCT_RE = re.compile(r"(\d)([.!?,;:])(?=\s|$)")53_SPACE_PUNCT_RE = re.compile(r" +([.!?,;:])(?=\s|$)")54 55# Moses-based punct_post_process re-attaches punctuation well for these56# languages. For the European languages it also glues currency symbols to the57# following word ("5,32 € am" -> "€am"), so there we skip moses and only58# collapse the spacing we introduced ourselves.59_MOSES_POSTPROCESS_LANGS = {"en", "zh", "ja", "ko"}60 61 62def default_cache_root() -> str:63    return os.environ.get(64        "ZONOS2_TTS_NORM_CACHE_DIR",65        os.path.expanduser("~/.cache/zonos2-tts-norm"),66    )67 68 69def normalization_enabled() -> bool:70    return os.environ.get("ZONOS2_TTS_NORM", "1") != "0"71 72 73class TTSTextNormalizer:74    """Lazy per-language NeMo normalizers with .far caching.75 76    Construction and calls are serialized per language: the upstream77    Normalizer shares a mutable TokenParser and is not thread-safe.78    """79 80    def __init__(self, cache_root: str | None = None):81        self.cache_root = cache_root or default_cache_root()82        self._normalizers: Dict[str, object] = {}83        self._locks: Dict[str, threading.Lock] = {}84        self._global_lock = threading.Lock()85 86    @staticmethod87    def nemo_lang(language: str) -> Optional[str]:88        return SERVER_TO_NEMO_LANG.get(language)89 90    def supported(self, language: str) -> bool:91        return language in SERVER_TO_NEMO_LANG92 93    def _lang_lock(self, lang: str) -> threading.Lock:94        with self._global_lock:95            if lang not in self._locks:96                self._locks[lang] = threading.Lock()97            return self._locks[lang]98 99    def _build(self, lang: str):100        from zonos2.vendor.nemo_text_processing.text_normalization import (101            Normalizer,102        )103 104        input_case = "lower_cased" if lang in LOWER_CASED_LANGS else "cased"105        # One cache dir per (lang, case): upstream .far filenames collide106        # across languages (e.g. ja's tagger writes a zh_-prefixed file).107        cache_dir = os.path.join(self.cache_root, f"{lang}_{input_case}")108        os.makedirs(cache_dir, exist_ok=True)109 110        logger.info("Loading TTS text normalizer for '%s' (%s)...", lang, input_case)111        normalizer = Normalizer(112            input_case=input_case,113            lang=lang,114            cache_dir=cache_dir,115            overwrite_cache=False,116        )117 118        prefix = _VERBALIZER_FAR_PREFIX.get(lang)119        if prefix is not None:120            far_path = os.path.join(121                cache_dir, f"{prefix}_tn_True_deterministic_verbalizer.far"122            )123            if not os.path.exists(far_path):124                from zonos2.vendor.nemo_text_processing.text_normalization.en.graph_utils import (125                    generator_main,126                )127 128                generator_main(far_path, {"verbalize": normalizer.verbalizer.fst})129        return normalizer130 131    def get(self, lang: str):132        with self._lang_lock(lang):133            if lang not in self._normalizers:134                self._normalizers[lang] = self._build(lang)135            return self._normalizers[lang]136 137    def warmup(self, languages: list[str] | None = None) -> None:138        """Construct normalizers ahead of time (server codes or NeMo codes)."""139        langs = languages or sorted(set(SERVER_TO_NEMO_LANG.values()))140        for lang in langs:141            lang = SERVER_TO_NEMO_LANG.get(lang, lang)142            try:143                self.get(lang)144            except Exception:  # noqa: BLE001145                logger.exception("TTS text normalizer warmup failed for '%s'", lang)146 147    def normalize(self, text: str, language: str) -> str:148        """Normalize text for the given server language code.149 150        Returns the input unchanged for unsupported languages or on any151        normalizer error — normalization must never fail a request.152        """153        lang = SERVER_TO_NEMO_LANG.get(language)154        if lang is None or not text.strip():155            return text156        text_in = _DIGIT_PUNCT_RE.sub(r"\1 \2", text)157        use_moses = lang in _MOSES_POSTPROCESS_LANGS158        try:159            normalizer = self.get(lang)160            with self._lang_lock(lang):161                result = normalizer.normalize(text_in, punct_post_process=use_moses)162        except Exception:  # noqa: BLE001163            logger.exception(164                "TTS text normalization failed for lang=%s; using raw text", language165            )166            return text167        if isinstance(result, str):168            result = _SPACE_PUNCT_RE.sub(r"\1", result)169        if not isinstance(result, str) or not result.strip():170            return text171        logger.debug("TTS norm [%s]: %r -> %r", language, text, result)172        return result173