Add ONNX Runtime implementation

teratts.py

@@ -8,6 +8,7 @@ import math
 import os
 import re
 import unicodedata
+import warnings
 import wave
 from collections.abc import Iterator
 from dataclasses import dataclass
@@ -30,7 +31,10 @@ SPEED = 1.05
 SEED = 1234
 RUSSIAN_TAG = re.compile(r"<ru>(.*?)</ru>", flags=re.DOTALL)
 LANGUAGE_TAG = re.compile(r"<(ru|en)>(.*?)</\1>", flags=re.DOTALL)
+LANGUAGE_TAG_TOKEN = re.compile(r"<(/?)([a-z]{2})>")
 TAGGED_NUMBER = re.compile(r"(?<![\w.])[-−]?\d+(?:[.,]\d+)?(?![\w.])")
+PUNCTUATION_NEEDS_SPACE = re.compile(r"[,.!?;:…](?=[^\s<])")
+NUMBER_NEEDS_SPACE = re.compile(r"(?<=\d)(?=[A-Za-zА-Яа-яЁё])")
 
 
 def prepare_raw_text(raw_text: str) -> tuple[str, str]:
@@ -38,6 +42,96 @@ def prepare_raw_text(raw_text: str) -> tuple[str, str]:
     return model_text, model_text.replace("+", "")
 
 
+def _add_punctuation_spaces(text: str) -> str:
+    """Separate punctuation without splitting decimal literals or closing tags."""
+    def space_after(match: re.Match[str]) -> str:
+        punctuation = match.group(0)
+        index = match.start()
+        previous = text[index - 1] if index else ""
+        following = text[index + 1] if index + 1 < len(text) else ""
+        if punctuation in ".," and previous.isdigit() and following.isdigit():
+            return punctuation
+        return punctuation + " "
+
+    return PUNCTUATION_NEEDS_SPACE.sub(space_after, text)
+
+
+def validate_language_tags(text: str) -> None:
+    """Require balanced ``<ru>`` / ``<en>`` spans for all public synthesis."""
+    tokens = list(LANGUAGE_TAG_TOKEN.finditer(text))
+    if not tokens or not LANGUAGE_TAG.search(text):
+        raise ValueError(
+            "text must contain a language tag: wrap text in <ru>...</ru> or <en>...</en>"
+        )
+    stack: list[str] = []
+    for token in tokens:
+        closing, language = token.groups()
+        if language not in {"ru", "en"}:
+            raise ValueError(f"unsupported language tag <{language}>; use <ru> or <en>")
+        if not closing:
+            stack.append(language)
+        elif not stack or stack.pop() != language:
+            raise ValueError("language tags must be balanced: use <ru>...</ru> or <en>...</en>")
+    if stack:
+        raise ValueError("language tags must be balanced: use <ru>...</ru> or <en>...</en>")
+    # Angle brackets that did not form a valid tag would be accepted by the
+    # character vocabulary but are not meaningful model input.
+    if "<" in LANGUAGE_TAG_TOKEN.sub("", text) or ">" in LANGUAGE_TAG_TOKEN.sub("", text):
+        raise ValueError("invalid language tags; use only <ru>...</ru> or <en>...</en>")
+
+
+def _skip_unsupported_characters(
+    text: str,
+    indexer: "UnicodeIndexer",
+    *,
+    preserve_digits: bool = False,
+) -> str:
+    """Return supported text and issue one clear warning for skipped characters."""
+    kept: list[str] = []
+    skipped: list[str] = []
+    for character in text:
+        # The released table was trained on NFKD text.  Keep the human-readable
+        # NFC spelling here (especially ``й`` and ``ё``) as long as all of its
+        # decomposed codepoints exist in the table.  RUAccent must receive this
+        # spelling: passing ``и`` + COMBINING BREVE makes its text cleaner drop
+        # the breve and turn ``й`` into ``и``.
+        encoded = unicodedata.normalize("NFKD", character)
+        supported = bool(encoded) and all(
+            (indexer.table[ord(item)] if ord(item) < 65_536 else -1) >= 0
+            for item in encoded
+        )
+        if not supported and not (preserve_digits and character.isdigit()):
+            skipped.append(character)
+        else:
+            kept.append(character)
+    if skipped:
+        labels = ", ".join(
+            f"{character!r} (U+{ord(character):04X})" for character in sorted(set(skipped))
+        )
+        warnings.warn(
+            f"skipped unsupported characters not present in the TeraTTS vocabulary: {labels}",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+    return "".join(kept)
+
+
+def normalize_input_text(raw_text: str, indexer: "UnicodeIndexer") -> str:
+    """Normalize spacing and skip unsupported vocabulary characters with a warning."""
+    if not isinstance(raw_text, str) or not raw_text.strip():
+        raise ValueError("text must not be empty; use <ru>...</ru> or <en>...</en>")
+    # Retain composed characters through RUAccent.  ``prepare_raw_text``
+    # performs the required NFKD conversion immediately before ONNX encoding.
+    text = unicodedata.normalize("NFC", raw_text)
+    text = _add_punctuation_spaces(text)
+    text = NUMBER_NEEDS_SPACE.sub(" ", text)
+    # Digits are retained only long enough for tagged ``num2words`` expansion;
+    # any remaining unsupported digits are skipped after that expansion.
+    text = _skip_unsupported_characters(text, indexer, preserve_digits=True)
+    validate_language_tags(text)
+    return text
+
+
 def load_ruaccent(
     *,
     model_size: str = "turbo3.1",
@@ -103,7 +197,10 @@ def expand_tagged_numbers(text: str) -> str:
 
 def normalize_text(loaded: "LoadedTTS", text: str) -> str:
     """Return the exact text tensorized by the text encoder for an utterance."""
-    expanded_text = expand_tagged_numbers(text)
+    normalized_input = normalize_input_text(text, loaded.indexer)
+    expanded_text = _skip_unsupported_characters(
+        expand_tagged_numbers(normalized_input), loaded.indexer
+    )
     model_text, _ = prepare_raw_text(add_russian_stress(expanded_text, loaded.accentizer))
     return model_text