omnivoice/utils/text.py

#!/usr/bin/env python3
# Copyright    2026  Xiaomi Corp.        (authors:  Han Zhu)
#
# See ../../LICENSE for clarification regarding multiple authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Text processing utilities for TTS inference.

Provides:
- ``chunk_text_punctuation()``: Splits long text into model-friendly chunks at
  sentence boundaries, with abbreviation-aware punctuation splitting.
- ``add_punctuation()``: Appends missing end punctuation (Chinese or English).
- ``normalize_text()``: Optional text normalization (numbers, dates, currency,
  etc.) into their spoken form, while preserving inline control syntax.
"""

import logging
import re
from typing import Callable, List, Optional

logger = logging.getLogger(__name__)


SPLIT_PUNCTUATION = set(".,;:!?。，；：！？")
CLOSING_MARKS = set("\"'“”‘’）]》>」】")

END_PUNCTUATION = {
    ";",
    ":",
    ",",
    ".",
    "!",
    "?",
    "…",
    ")",
    "]",
    "}",
    '"',
    "'",
    "“",
    "”",
    "‘",
    "’",
    "；",
    "：",
    "，",
    "。",
    "！",
    "？",
    "、",
    "……",
    "）",
    "】",
}


ABBREVIATIONS = {
    "Mr.",
    "Mrs.",
    "Ms.",
    "Dr.",
    "Prof.",
    "Sr.",
    "Jr.",
    "Rev.",
    "Fr.",
    "Hon.",
    "Pres.",
    "Gov.",
    "Capt.",
    "Gen.",
    "Sen.",
    "Rep.",
    "Col.",
    "Maj.",
    "Lt.",
    "Cmdr.",
    "Sgt.",
    "Cpl.",
    "Co.",
    "Corp.",
    "Inc.",
    "Ltd.",
    "Est.",
    "Dept.",
    "St.",
    "Ave.",
    "Blvd.",
    "Rd.",
    "Mt.",
    "Ft.",
    "No.",
    "Jan.",
    "Feb.",
    "Mar.",
    "Apr.",
    "Aug.",
    "Sep.",
    "Sept.",
    "Oct.",
    "Nov.",
    "Dec.",
    "i.e.",
    "e.g.",
    "vs.",
    "Vs.",
    "Etc.",
    "approx.",
    "fig.",
    "def.",
}


def chunk_text_punctuation(
    text: str,
    chunk_len: int,
    min_chunk_len: Optional[int] = None,
) -> List[str]:
    """
    Splits the input tokens list into chunks according to punctuations,
    avoiding splits on common abbreviations (e.g., Mr., No.).
    """

    # 1. Split the tokens according to punctuations.
    sentences = []
    current_sentence = []

    tokens_list = list(text)

    for token in tokens_list:
        # If the first token of current sentence is punctuation,
        # append it to the end of the previous sentence.
        if (
            len(current_sentence) == 0
            and len(sentences) != 0
            and (token in SPLIT_PUNCTUATION or token in CLOSING_MARKS)
        ):
            sentences[-1].append(token)
        # Otherwise, append the current token to the current sentence.
        else:
            current_sentence.append(token)

            # Split the sentence in positions of punctuations.
            if token in SPLIT_PUNCTUATION:
                is_abbreviation = False

                if token == ".":
                    temp_str = "".join(current_sentence).strip()
                    if temp_str:
                        last_word = temp_str.split()[-1]
                        if last_word in ABBREVIATIONS:
                            is_abbreviation = True

                if not is_abbreviation:
                    sentences.append(current_sentence)
                    current_sentence = []
    # Assume the last few tokens are also a sentence
    if len(current_sentence) != 0:
        sentences.append(current_sentence)

    # 2. Merge short sentences.
    merged_chunks = []
    current_chunk = []
    for sentence in sentences:
        if len(current_chunk) + len(sentence) <= chunk_len:
            current_chunk.extend(sentence)
        else:
            if len(current_chunk) > 0:
                merged_chunks.append(current_chunk)
            current_chunk = sentence

    if len(current_chunk) > 0:
        merged_chunks.append(current_chunk)

    # 4. Post-process: Check for undersized chunks and merge them
    #  with the previous chunk or next chunk (if it's the first chunk).
    if min_chunk_len is not None:
        first_chunk_short_flag = (
            len(merged_chunks) > 0 and len(merged_chunks[0]) < min_chunk_len
        )
        final_chunks = []
        for i, chunk in enumerate(merged_chunks):
            if i == 1 and first_chunk_short_flag:
                final_chunks[-1].extend(chunk)
            else:
                if len(chunk) >= min_chunk_len:
                    final_chunks.append(chunk)
                else:
                    if len(final_chunks) == 0:
                        final_chunks.append(chunk)
                    else:
                        final_chunks[-1].extend(chunk)
    else:
        final_chunks = merged_chunks

    chunk_strings = [
        "".join(chunk).strip() for chunk in final_chunks if "".join(chunk).strip()
    ]
    return chunk_strings


def add_punctuation(text: str):
    """Add punctuation if there is not in the end of text"""
    text = text.strip()

    if not text:
        return text

    if text[-1] not in END_PUNCTUATION:
        is_chinese = any("\u4e00" <= char <= "\u9fff" for char in text)

        text += "。" if is_chinese else "."

    return text


# ---------------------------------------------------------------------------
# Optional text normalization (opt-in via ``generate(normalize_text=True)``)
# ---------------------------------------------------------------------------
#
# Arabic numerals, dates, currency, etc. are converted into their spoken form
# so the model reads them correctly (e.g. "2345" -> "twenty three forty five",
# "199" -> the Chinese reading). Chinese/English go through WeTextProcessing;
# any other language falls back to ``num2words`` for bare integers when
# available.
#
# The OmniVoice inline control syntax must survive normalization:
#   * bracketed non-verbal tags, e.g. ``[laughter]``, ``[sigh]``;
#   * bracketed CMU pronunciation overrides, e.g. ``[B EY1 S]`` -- the stress
#     digit would otherwise be read as a number;
#   * Chinese pinyin tone markers (uppercase pinyin + tone digit) -- likewise.
# Protected spans are held out and re-inserted verbatim around normalization.

# Any ``[...]`` span covers both non-verbal tags and CMU pronunciation.
_BRACKET_TAG_RE = re.compile(r"\[[^\[\]]*\]")
# Uppercase pinyin followed by a tone digit 1-5 (Chinese pronunciation control).
_PINYIN_TONE_RE = re.compile(r"[A-Z]+[1-5]")
_CJK_RE = re.compile(r"[\u4e00-\u9fff]")

_TN_INSTALL_MSG = (
    "Text normalization (normalize_text=True) requires WeTextProcessing, which "
    "is not installed.\n"
    "  pip install WeTextProcessing         # or:  pip install 'omnivoice[tn]'\n"
    "WeTextProcessing depends on pynini, which has no prebuilt wheel for macOS "
    "arm64 (Apple Silicon). On macOS, install pynini from conda-forge first:\n"
    "  conda install -c conda-forge pynini\n"
    "then:  pip install WeTextProcessing"
)

# Normalizer construction builds FSTs and is comparatively slow, so instances
# are cached per language for the lifetime of the process.
_ZH_NORMALIZER = None
_EN_NORMALIZER = None


def _get_zh_normalizer():
    global _ZH_NORMALIZER
    if _ZH_NORMALIZER is None:
        try:
            from tn.chinese.normalizer import Normalizer
        except ImportError as e:  # pragma: no cover - depends on optional extra
            raise ImportError(_TN_INSTALL_MSG) from e
        # Conservative flags: normalize numbers/symbols only. Keep interjections
        # and erhua (they are spoken), keep the user's original characters, and
        # do not delete or rewrite anything beyond numeric/symbolic tokens.
        _ZH_NORMALIZER = Normalizer(
            remove_interjections=False,
            remove_erhua=False,
            traditional_to_simple=False,
            remove_puncts=False,
            full_to_half=False,
        )
    return _ZH_NORMALIZER


def _get_en_normalizer():
    global _EN_NORMALIZER
    if _EN_NORMALIZER is None:
        try:
            from tn.english.normalizer import Normalizer
        except ImportError as e:  # pragma: no cover - depends on optional extra
            raise ImportError(_TN_INSTALL_MSG) from e
        _EN_NORMALIZER = Normalizer()
    return _EN_NORMALIZER


def _resolve_lang_code(language: Optional[str], text: str) -> str:
    """Map a language name/code to ``"zh"``/``"en"``/other code.

    When ``language`` is ``None`` (or unrecognized), fall back to detecting
    Chinese vs. English by the presence of CJK characters.
    """
    if language is not None:
        code = language.strip().lower()
        if code and code != "none":
            if code in ("zh", "en"):
                return code
            try:
                from omnivoice.utils.lang_map import LANG_IDS, LANG_NAME_TO_ID

                if code in LANG_IDS:
                    return code
                if code in LANG_NAME_TO_ID:
                    return LANG_NAME_TO_ID[code]
            except Exception:  # pragma: no cover - lang_map should be importable
                pass
            return code  # assume it is already a language id, e.g. "ja", "de"
    return "zh" if _CJK_RE.search(text) else "en"


def _num2words_segment(text: str, lang: str) -> str:
    """Best-effort integer-to-words fallback for non zh/en languages."""
    try:
        from num2words import num2words
    except ImportError:
        return text  # fallback is best-effort; silently skip when unavailable

    def _repl(match):
        try:
            return num2words(int(match.group()), lang=lang)
        except Exception:
            return match.group()  # unsupported language / value: leave as-is

    return re.sub(r"\d+", _repl, text)


def _normalize_segment(fn: Callable[[str], str], segment: str) -> str:
    """Normalize one non-protected segment, never raising on bad input.

    Leading/trailing whitespace is preserved explicitly because the underlying
    normalizers strip it, which would otherwise glue words to an adjacent
    protected span (e.g. ``the [B EY1 S] guitar`` -> ``the[B EY1 S]guitar``).
    """
    if not segment.strip():
        return segment
    lead = segment[: len(segment) - len(segment.lstrip())]
    trail = segment[len(segment.rstrip()) :]
    try:
        core = fn(segment.strip())
    except Exception as e:  # pragma: no cover - defensive
        logger.warning(
            "Text normalization failed on a segment (%s); keeping it unchanged.",
            type(e).__name__,
        )
        return segment
    return lead + core + trail


def _apply_with_protection(
    text: str, fn: Callable[[str], str], protect_pinyin: bool
) -> str:
    """Run ``fn`` on ``text`` while holding out protected control spans."""
    spans = [m.span() for m in _BRACKET_TAG_RE.finditer(text)]
    if protect_pinyin:
        spans += [m.span() for m in _PINYIN_TONE_RE.finditer(text)]
    if not spans:
        return _normalize_segment(fn, text)

    # Merge overlapping/adjacent protected spans, then normalize the gaps.
    spans.sort()
    merged: List[List[int]] = []
    for start, end in spans:
        if merged and start <= merged[-1][1]:
            merged[-1][1] = max(merged[-1][1], end)
        else:
            merged.append([start, end])

    out: List[str] = []
    last = 0
    for start, end in merged:
        if start > last:
            out.append(_normalize_segment(fn, text[last:start]))
        out.append(text[start:end])  # protected span, verbatim
        last = end
    if last < len(text):
        out.append(_normalize_segment(fn, text[last:]))
    return "".join(out)


def normalize_text(text: str, language: Optional[str] = None) -> str:
    """Normalize numbers, dates, currency, etc. into their spoken form.

    Chinese is routed to WeTextProcessing's ``ZhNormalizer`` and English to its
    ``EnNormalizer`` (configured to only rewrite numeric/symbolic tokens). Any
    other language falls back to ``num2words`` for bare integers when it is
    installed, otherwise the text is returned unchanged.

    Inline OmniVoice control syntax is preserved: bracketed non-verbal tags
    (``[laughter]``) and CMU pronunciation overrides (``[B EY1 S]``) are passed
    through untouched, and Chinese pinyin tone markers (uppercase pinyin +
    tone digit) are protected so the tone digit is not read as a number.

    Args:
        text: Input text.
        language: Language code (``"en"``/``"zh"``) or full name (``"English"``).
            ``None`` auto-detects Chinese vs. English by script.

    Returns:
        The normalized text.

    Raises:
        ImportError: For Chinese/English when the optional ``omnivoice[tn]``
            dependency (WeTextProcessing) is not installed.
    """
    if not text or not text.strip():
        return text

    code = _resolve_lang_code(language, text)
    if code == "zh":
        normalizer = _get_zh_normalizer()
        return _apply_with_protection(text, normalizer.normalize, protect_pinyin=True)
    if code == "en":
        normalizer = _get_en_normalizer()
        return _apply_with_protection(text, normalizer.normalize, protect_pinyin=False)
    # Other languages: best-effort integer conversion via num2words.
    return _apply_with_protection(
        text, lambda s: _num2words_segment(s, code), protect_pinyin=False
    )