Fix stress using model

app.py

@@ -2,7 +2,7 @@ import tempfile
 import gradio as gr
 from datetime import datetime
 from enum import Enum
-from ukrainian_tts.tts import TTS
+from ukrainian_tts.tts import TTS, Stress, Voices
 from torch.cuda import is_available
 
 class StressOption(Enum):
@@ -32,18 +32,21 @@ def tts(text: str, voice: str, stress: str):
     print("Voice", voice)
     print("Stress:", stress)
     print("Time:", datetime.utcnow())
-    autostress_with_model = (
-        True if stress == StressOption.AutomaticStressWithModel.value else False
-    )
+
     voice_mapping = {
-        VoiceOption.Olena.value: "olena",
-        VoiceOption.Mykyta.value: "mykyta",
-        VoiceOption.Lada.value: "lada",
-        VoiceOption.Dmytro.value: "dmytro",
-        VoiceOption.Olga.value: "olga",
+        VoiceOption.Olena.value: Voices.Olena.value,
+        VoiceOption.Mykyta.value: Voices.Mykyta.value,
+        VoiceOption.Lada.value: Voices.Lada.value,
+        VoiceOption.Dmytro.value: Voices.Dmytro.value,
+        VoiceOption.Olga.value: Voices.Olga.value,
+    }
+    stress_mapping = {
+        StressOption.AutomaticStress.value: Stress.Dictionary.value,
+        StressOption.AutomaticStressWithModel.value: Stress.Model.value
     }
-    speaker_name = voice_mapping[voice]
 
+    speaker_name = voice_mapping[voice]
+    stress_selected = stress_mapping[stress]
     text_limit = 7200
     text = (
         text if len(text) < text_limit else text[0:text_limit]
@@ -51,7 +54,7 @@ def tts(text: str, voice: str, stress: str):
     
 
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
-        _, text = ukr_tts.tts(text, speaker_name, autostress_with_model, fp)
+        _, text = ukr_tts.tts(text, speaker_name, stress_selected, fp)
         return fp.name, text
 
 

ukrainian_tts/formatter.py

@@ -1,7 +1,6 @@
 import num2words
 import re
-from .stress import sentence_to_stress, stress_dict
-from .stress_with_model import stress_with_model
+from .stress import sentence_to_stress, stress_dict, stress_with_model
 
 
 def preprocess_text(text, use_autostress_model=False):

ukrainian_tts/stress.py

@@ -1,6 +1,6 @@
 from typing import List
-import numpy as np
 from ukrainian_word_stress import Stressifier, StressSymbol
+import ukrainian_accentor as accentor
 
 stressify = Stressifier(stress_symbol=StressSymbol.CombiningAcuteAccent)
 
@@ -10,10 +10,7 @@ special = "'"
 alphabet = vowels + consonants + special
 
 
-def stress_dict(sentence: str):
-    stressed = stressify(sentence.replace("+", "")).replace(
-        StressSymbol.CombiningAcuteAccent, "+"
-    )
+def _shift_stress(stressed):
     new_stressed = ""
     start = 0
     last = 0
@@ -33,6 +30,19 @@ def stress_dict(sentence: str):
     return new_stressed
 
 
+def stress_with_model(text: str):
+    text = text.lower()
+    result = accentor.process(text, mode="plus")
+    return result
+
+
+def stress_dict(sentence: str):
+    stressed = stressify(sentence.replace("+", "")).replace(
+        StressSymbol.CombiningAcuteAccent, "+"
+    )
+    return _shift_stress(stressed)
+
+
 def sentence_to_stress(sentence: str, stress_function=stress_dict) -> str:
     # save custom stress positions
     all_stresses = []
@@ -68,6 +78,7 @@ def sentence_to_stress(sentence: str, stress_function=stress_dict) -> str:
 
 
 if __name__ == "__main__":
+    # TODO: move it to unit tests
     sentence = "Кам'янець-Подільський - місто в Хмельницькій області України, центр Кам'янець-Подільської міської об'єднаної територіальної громади і Кам'янець-Подільського району."
     print(sentence_to_stress(sentence))
     sentence = "Привіт, як тебе звати?"
@@ -84,3 +95,20 @@ if __name__ == "__main__":
     print(sentence_to_stress(sentence))
     sentence = "Н тльк в крн тк мж бт."
     print(sentence_to_stress(sentence))
+
+    sentence = "Кам'янець-Подільський - місто в Хмельницькій області України, центр Кам'янець-Подільської міської об'єднаної територіальної громади і Кам'янець-Подільського району."
+    print(stress_with_model(sentence))
+    sentence = "Привіт, як тебе звати?"
+    print(stress_with_model(sentence))
+    sentence = "АННА - український панк-рок гурт"
+    print(stress_with_model(sentence))
+    sentence = "Не тільки в Україні таке може бути."
+    print(stress_with_model(sentence))
+    sentence = "Не тільки в +Укра+їні т+аке може бути."
+    print(stress_with_model(sentence))
+    sentence = "два + два"
+    print(stress_with_model(sentence))
+    sentence = "Н тльк в крн тк мж бт."
+    print(stress_with_model(sentence))
+    sentence = "Н тльк в крн тк мж бт."
+    print(stress_with_model(sentence))

ukrainian_tts/stress_with_model.py

@@ -1,26 +0,0 @@
-import ukrainian_accentor as accentor
-
-# run
-def stress_with_model(text: str):
-    text = text.lower()
-    result = accentor.process(text, mode="plus")
-    return result
-
-
-if __name__ == "__main__":
-    sentence = "Кам'янець-Подільський - місто в Хмельницькій області України, центр Кам'янець-Подільської міської об'єднаної територіальної громади і Кам'янець-Подільського району."
-    print(stress_with_model(sentence))
-    sentence = "Привіт, як тебе звати?"
-    print(stress_with_model(sentence))
-    sentence = "АННА - український панк-рок гурт"
-    print(stress_with_model(sentence))
-    sentence = "Не тільки в Україні таке може бути."
-    print(stress_with_model(sentence))
-    sentence = "Не тільки в +Укра+їні т+аке може бути."
-    print(stress_with_model(sentence))
-    sentence = "два + два"
-    print(stress_with_model(sentence))
-    sentence = "Н тльк в крн тк мж бт."
-    print(stress_with_model(sentence))
-    sentence = "Н тльк в крн тк мж бт."
-    print(stress_with_model(sentence))

ukrainian_tts/tts.py

@@ -15,7 +15,7 @@ class Voices(Enum):
     Olga = "olga"
 
 
-class StressOption(Enum):
+class Stress(Enum):
     """Options how to stress sentence.
     - `dictionary` - performs lookup in dictionary, taking into account grammatical case of a word and its' neighbors
     - `model` - stress using transformer model"""
@@ -40,17 +40,21 @@ class TTS:
         Run a Text-to-Speech engine and output to `output_fp` BytesIO-like object.
         - `text` - your model input text.
         - `voice` - one of predefined voices from `Voices` enum.
-        - `stress` - stress method options, predefined in `StressOption` enum.
+        - `stress` - stress method options, predefined in `Stress` enum.
         - `output_fp` - file-like object output. Stores in RAM by default.
         """
-        autostress_with_model = (
-            True if stress == StressOption.Model.value else False
-        )
 
+        if stress not in [option.value for option in Stress]:
+            raise ValueError(f"Invalid value for stress option selected! Please use one of the following values: {', '.join([option.value for option in Stress])}.")
+
+        if stress == Stress.Model.value:
+            stress = True
+        else:
+            stress = False
         if voice not in [option.value for option in Voices]:
             raise ValueError(f"Invalid value for voice selected! Please use one of the following values: {', '.join([option.value for option in Voices])}.")
 
-        text = preprocess_text(text, autostress_with_model)
+        text = preprocess_text(text, stress)
 
         with no_grad():
             wavs = self.synthesizer.tts(text, speaker_name=voice)