tts/src/data/text_frontend.py

"""Shared text frontend for Yoruba TTS.

SHARED MODULE — used by BOTH data preparation (training targets) and inference
(synthesize.py, serverless handler). One source of truth guarantees the model
speaks text shaped exactly as it was trained.

v1 policy for Yoruba:
  * Unicode NFC; tone diacritics are KEPT (they are lexical/grammatical).
  * Keep sentence-level punctuation for prosody: . , ? ! ' -
  * Drop all other punctuation/symbols (quotes, brackets, …) -> space.
  * Number expansion via a verified lexicon (see numbers.py / Task 3).
"""
from __future__ import annotations

import unicodedata

import regex  # Unicode property classes \p{P}, \p{S}

from src.data.num_expand import expand_digits

# Punctuation we keep for prosody / orthography.
_KEEP = ".,?!'-"
_STRIP_RE = regex.compile(
    rf"(?:(?![{regex.escape(_KEEP)}])[\p{{P}}\p{{S}}])"
)
_WS_RE = regex.compile(r"\s+")


def normalize_for_tts(text: str, *, lexicon: dict | None = None) -> str:
    if not text:
        return ""
    t = unicodedata.normalize("NFC", text)
    t = "".join(ch if not unicodedata.category(ch).startswith("C") else " " for ch in t)
    if lexicon:
        t = expand_digits(t, lexicon)
    t = _STRIP_RE.sub(" ", t)
    t = _WS_RE.sub(" ", t).strip()
    return t