Add dictionary stress mode and tagged number expansion

README.md

@@ -11,9 +11,9 @@ tags:
 
 # TeraTTSv2 ONNX
 
-TeraTTSv2 text-to-speech with ONNX Runtime, selectable teacher/distilled
-diffusion samplers, precomputed voice styles, CPU/CUDA execution, and streamed
-audio chunks.
+TeraTTSv2 is a self-contained ONNX Runtime text-to-speech release with
+selectable diffusion samplers, six voice styles, Russian stress marking, and
+streamed audio output.
 
 ## Installation
 
@@ -47,10 +47,14 @@ tts.save_wav("teratts.wav", waveform)
 `waveform` is a mono `float32` NumPy array at 44,100 Hz. `save_wav` writes
 standard signed-16-bit PCM WAV without an extra audio package.
 
-## Voices and samplers
+## Controls
 
-Available `voice` values are `en_female`, `en_male`, `ru_female`, `ru_male`,
-`mita`, and `masha_ls`.
+| Control | Values | Effect |
+| --- | --- | --- |
+| `voice` | `en_female`, `en_male`, `ru_female`, `ru_male`, `mita`, `masha_ls` | Selects a bundled precomputed voice style. |
+| `duration_scale` | Positive float, default `1` | Higher values produce slower, longer speech. |
+| `diffusion_model` | `distilled` (default), `teacher` | Distilled is faster; teacher supports adjustable CFG. |
+| `ruaccent_mode` | `full` (default), `dictionary` | Full uses RUAccent neural ONNX graphs plus dictionaries; dictionary mode loads dictionaries only. |
 
 The default `diffusion_model="distilled"` is the fast eight-step sampler. To
 use the teacher sampler, choose it while loading:
@@ -68,19 +72,38 @@ teacher_tts = AutoModel.from_pretrained(
 `guidance` can be adjusted when generating with the teacher sampler. The
 distilled sampler has CFG 3 baked into its graph.
 
-Russian `<ru>…</ru>` text is automatically stress-marked by the bundled
-RUAccent-derived ONNX runtime before inference. It makes no secondary model
-download and does not require the `ruaccent` package. Manual `+` markers
-remain unchanged:
+## Language tags, numbers, and Russian stress
+
+Wrap text in `<en>…</en>` or `<ru>…</ru>`. Numbers inside those tags are
+expanded to words in the matching language before synthesis:
 
 ```python
 waveform = tts.generate_speech(
-    "<ru>На двери висит замок.</ru>",
+    "<ru>У меня 21 яблоко.</ru> <en>I have 42 apples.</en>",
     voice="ru_female",
     duration_scale=1,
 )
 ```
 
+Russian text is automatically stress-marked by the bundled RUAccent-derived
+runtime. Manual `+` markers remain authoritative. For a lower-memory,
+deterministic dictionary-only path, choose the mode while loading:
+
+```python
+dictionary_tts = AutoModel.from_pretrained(
+    "TeraSpace/TeraTTSv2",
+    trust_remote_code=True,
+    provider="CPUExecutionProvider",
+    threads=6,
+    ruaccent_mode="dictionary",
+)
+```
+
+Dictionary mode does not load RUAccent neural ONNX graphs. It marks known
+words and applies deterministic `ё` replacements, while unknown words and
+ambiguous homographs are left unchanged. Set `russian_stress=False` to disable
+automatic Russian stress processing entirely.
+
 ## Stream audio
 
 ```python

modeling_teratts.py

@@ -43,6 +43,7 @@ class TeraTTSModel(PreTrainedModel):
         russian_stress: bool = True,
         ruaccent_model_size: str = "turbo3.1",
         ruaccent_device: str = "CPU",
+        ruaccent_mode: str = "full",
         **kwargs: object,
     ) -> "TeraTTSModel":
         """Download/load the release and initialize reusable ONNX sessions.
@@ -94,6 +95,7 @@ class TeraTTSModel(PreTrainedModel):
             russian_stress=russian_stress,
             ruaccent_model_size=ruaccent_model_size,
             ruaccent_device=ruaccent_device,
+            ruaccent_mode=ruaccent_mode,
         )
         return instance
 

requirements.txt

@@ -3,4 +3,5 @@ onnxruntime>=1.16
 torch>=2.1
 transformers>=4.40
 huggingface_hub>=0.23
+num2words>=0.5.13
 sounddevice>=0.5

teratts.py

@@ -29,6 +29,8 @@ DEFAULT_STREAM_CHUNK_FRAMES = 16
 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)
+TAGGED_NUMBER = re.compile(r"(?<![\w.])[-−]?\d+(?:[.,]\d+)?(?![\w.])")
 
 
 def prepare_raw_text(raw_text: str) -> tuple[str, str]:
@@ -41,11 +43,12 @@ def load_ruaccent(
     model_size: str = "turbo3.1",
     device: str = "CPU",
     workdir: Path | None = None,
+    mode: str = "full",
 ) -> object:
     """Load the bundled RUAccent-derived ONNX models without downloading."""
     if workdir is None:
         raise ValueError("load_ruaccent requires the release's ruaccent asset directory")
-    return RUAccent(workdir, model_size=model_size, device=device)
+    return RUAccent(workdir, model_size=model_size, device=device, mode=mode)
 
 
 def add_russian_stress(text: str, accentizer: object | None) -> str:
@@ -65,6 +68,39 @@ def add_russian_stress(text: str, accentizer: object | None) -> str:
     return RUSSIAN_TAG.sub(accent, text)
 
 
+def expand_tagged_numbers(text: str) -> str:
+    """Spell out numeric literals inside ``<ru>`` and ``<en>`` text spans.
+
+    Language tags are intentionally required: this avoids guessing a language
+    for bare text or for identifiers such as versions and file names.
+    """
+    spans = list(LANGUAGE_TAG.finditer(text))
+    if not any(TAGGED_NUMBER.search(match.group(2)) for match in spans):
+        return text
+    try:
+        from num2words import num2words
+    except ImportError as error:
+        raise RuntimeError(
+            "number expansion requires num2words; install the model requirements"
+        ) from error
+
+    def expand_span(match: re.Match[str]) -> str:
+        language, content = match.groups()
+
+        def expand_number(number: re.Match[str]) -> str:
+            literal = number.group(0).replace("−", "-")
+            value: int | float
+            if "." in literal or "," in literal:
+                value = float(literal.replace(",", "."))
+            else:
+                value = int(literal)
+            return str(num2words(value, lang=language))
+
+        return f"<{language}>{TAGGED_NUMBER.sub(expand_number, content)}</{language}>"
+
+    return LANGUAGE_TAG.sub(expand_span, text)
+
+
 class UnicodeIndexer:
     def __init__(self, indexer_path: Path):
         self.table = json.loads(indexer_path.read_text())
@@ -207,13 +243,16 @@ def load_model(
     ruaccent_model_size: str = "turbo3.1",
     ruaccent_device: str = "CPU",
     ruaccent_workdir: Path | None = None,
+    ruaccent_mode: str = "full",
 ) -> LoadedTTS:
     """Load reusable ONNX sessions; call once before generating many utterances.
 
     ``threads`` controls CPU intra-op parallelism. ``None`` chooses a
     physical-core-scale default; CUDA ignores this value. Russian ``<ru>``
     spans receive automatic ``+`` stress markers when ``russian_stress`` is
-    enabled; manually supplied markers are preserved.
+    enabled; manually supplied markers are preserved. ``ruaccent_mode`` is
+    ``"full"`` (neural ONNX models plus dictionaries) or ``"dictionary"``
+    (dictionaries only, with no accentuation-model ONNX sessions).
     """
     if model not in {"teacher", "distilled"}:
         raise ValueError("model must be 'teacher' or 'distilled'")
@@ -240,6 +279,7 @@ def load_model(
                 model_size=ruaccent_model_size,
                 device=ruaccent_device,
                 workdir=ruaccent_workdir or release / "ruaccent",
+                mode=ruaccent_mode,
             )
             if russian_stress
             else None
@@ -320,7 +360,10 @@ def _generate_latent(
     if style_ttl.shape != (1, 50, 256) or style_dp.shape != (1, 8, 16):
         raise ValueError("style assets have unexpected shapes")
 
-    model_text, duration_text = prepare_raw_text(add_russian_stress(text, loaded.accentizer))
+    expanded_text = expand_tagged_numbers(text)
+    model_text, duration_text = prepare_raw_text(
+        add_russian_stress(expanded_text, loaded.accentizer)
+    )
     text_ids, text_mask = loaded.indexer.batch(model_text)
     duration_ids, duration_mask = loaded.indexer.batch(duration_text)
     text_emb = loaded.text_encoder.run(

teratts_ruaccent.py

@@ -168,10 +168,20 @@ class RUAccent:
     _tokens = re.compile(r"\w*(?:\+\w+)*|[^\w\s]+")
     _sentence = re.compile(r"[^.!?…]+[.!?…]*[\"»“]*")
 
-    def __init__(self, root: Path, *, model_size: str = "turbo3.1", device: str = "CPU") -> None:
+    def __init__(
+        self,
+        root: Path,
+        *,
+        model_size: str = "turbo3.1",
+        device: str = "CPU",
+        mode: str = "full",
+    ) -> None:
         root = root.resolve()
         if model_size != "turbo3.1":
             raise ValueError("the bundled RUAccent model is turbo3.1")
+        if mode not in {"full", "dictionary"}:
+            raise ValueError("RUAccent mode must be 'full' or 'dictionary'")
+        self.mode = mode
         dictionary = root / "dictionary"
         self.accents = json.load(gzip.open(dictionary / "accents.json.gz"))
         self.omographs = json.load(gzip.open(dictionary / "omographs.json.gz"))
@@ -179,10 +189,11 @@ class RUAccent:
         self.yo_words = json.load(gzip.open(dictionary / "yo_words.json.gz"))
         self.yo_homographs = json.load(gzip.open(dictionary / "yo_homographs.json.gz"))
         self.accents.update({"о": "+о", "О": "+О"})
-        self.accent_model = _CharAccentModel(root / "nn" / "nn_accent", device)
-        self.omograph_model = _OmographModel(root / "nn" / "nn_omograph" / model_size, device)
-        self.stress_usage = _TokenClassifier(root / "nn" / "nn_stress_usage_predictor", device)
-        self.yo_classifier = _TokenClassifier(root / "nn" / "nn_yo_homograph_resolver", device)
+        if mode == "full":
+            self.accent_model = _CharAccentModel(root / "nn" / "nn_accent", device)
+            self.omograph_model = _OmographModel(root / "nn" / "nn_omograph" / model_size, device)
+            self.stress_usage = _TokenClassifier(root / "nn" / "nn_stress_usage_predictor", device)
+            self.yo_classifier = _TokenClassifier(root / "nn" / "nn_yo_homograph_resolver", device)
 
     @staticmethod
     def _remaining(sentence: str, matches: list[re.Match[str]]) -> tuple[list[str], list[str]]:
@@ -253,8 +264,30 @@ class RUAccent:
             words[index] = "".join(target)
         return words
 
+    def _dictionary_word(self, match: re.Match[str]) -> str:
+        """Apply deterministic ``ё`` and accent dictionary entries only."""
+        word = match.group(0)
+        normalized = _fix_capital(word, self.yo_words.get(word.lower(), word))
+        accented = self.accents.get(normalized.lower(), normalized.lower())
+        if accented == normalized.lower():
+            return normalized
+        target = list(normalized)
+        inserted = 0
+        for marker in re.finditer(r"\+", accented):
+            target.insert(marker.start() + inserted, "+")
+            inserted += 1
+        return "".join(target)
+
+    def _process_dictionary(self, text: str) -> str:
+        # Dictionary mode deliberately avoids loading or calling every neural
+        # RUAccent ONNX graph. Unknown words and unresolved homographs remain
+        # untouched rather than receiving a neural prediction.
+        return re.sub(r"[A-Za-zА-Яа-яЁё]+", self._dictionary_word, text)
+
     def process_all(self, text: str) -> str:
         text = self._normalize.sub("", text)
+        if self.mode == "dictionary":
+            return self._process_dictionary(text)
         output: list[str] = []
         # Keep sentence delimiters attached to the sentence, matching the
         # original RUAccent intent without its optional razdel dependency.