Add verbalizer

app.py

@@ -1,13 +1,125 @@
 import glob
 import os
 import gradio as gr
-from infer import inference
-
+from infer import inference, split_to_parts
+import onnxruntime
+from transformers import AutoTokenizer
+from huggingface_hub import hf_hub_download
+import numpy as np
 
 prompts_dir = 'voices'
 prompts_list = sorted(glob.glob(os.path.join(prompts_dir, '*.wav')))
 prompts_list = ['.'.join(p.split('/')[-1].split('.')[:-1]) for p in prompts_list]
 
+verbalizer_model_name = "skypro1111/mbart-large-50-verbalization"
+
+def cache_model_from_hf(repo_id, model_dir="./"):
+    """Download ONNX models from HuggingFace Hub."""
+    files = ["onnx/encoder_model.onnx", "onnx/decoder_model.onnx", "onnx/decoder_model.onnx_data"]
+    
+    for file in files:
+        hf_hub_download(
+            repo_id=repo_id,
+            filename=file,
+            local_dir=model_dir,
+        )
+
+
+def create_onnx_session(model_path, use_gpu=True):
+    """Create an ONNX inference session."""
+    session_options = onnxruntime.SessionOptions()
+    session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+    session_options.enable_mem_pattern = True
+    session_options.enable_mem_reuse = True
+    session_options.intra_op_num_threads = 8
+    session_options.log_severity_level = 1
+    
+    cuda_provider_options = {
+        'device_id': 0,
+        'arena_extend_strategy': 'kSameAsRequested',
+        'gpu_mem_limit': 0,  # 0 means no limit
+        'cudnn_conv_algo_search': 'DEFAULT',
+        'do_copy_in_default_stream': True,
+    }
+    
+    if use_gpu and 'CUDAExecutionProvider' in onnxruntime.get_available_providers():
+        providers = [('CUDAExecutionProvider', cuda_provider_options)]
+    else:
+        providers = ['CPUExecutionProvider']
+    
+    session = onnxruntime.InferenceSession(
+        model_path,
+        providers=providers,
+        sess_options=session_options
+    )
+    
+    return session
+
+def init_verbalizer():
+    cache_model_from_hf(verbalizer_model_name)
+
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(verbalizer_model_name)
+    tokenizer.src_lang = "uk_UA"
+    tokenizer.tgt_lang = "uk_UA"
+
+    print("Creating ONNX sessions...")
+    encoder_session = create_onnx_session("onnx/encoder_model.onnx")
+    decoder_session = create_onnx_session("onnx/decoder_model.onnx")
+    return tokenizer, encoder_session, decoder_session
+
+tokenizer, encoder_session, decoder_session = init_verbalizer()
+
+
+def generate_text(text):
+    """Generate text for a single input."""
+    # Prepare input
+    inputs = tokenizer(text, return_tensors="np", padding=True, truncation=True, max_length=512)
+    input_ids = inputs["input_ids"].astype(np.int64)
+    attention_mask = inputs["attention_mask"].astype(np.int64)
+    
+    # Run encoder
+    encoder_outputs = encoder_session.run(
+        output_names=["last_hidden_state"],
+        input_feed={
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        }
+    )[0]
+    
+    # Initialize decoder input
+    decoder_input_ids = np.array([[tokenizer.pad_token_id]], dtype=np.int64)
+    
+    # Generate sequence
+    for _ in range(512):
+        # Run decoder
+        decoder_outputs = decoder_session.run(
+            output_names=["logits"],
+            input_feed={
+                "input_ids": decoder_input_ids,
+                "encoder_hidden_states": encoder_outputs,
+                "encoder_attention_mask": attention_mask,
+            }
+        )[0]
+        
+        # Get next token
+        next_token = decoder_outputs[:, -1:].argmax(axis=-1)
+        decoder_input_ids = np.concatenate([decoder_input_ids, next_token], axis=-1)
+        
+        # Check if sequence is complete
+        if tokenizer.eos_token_id in decoder_input_ids[0]:
+            break
+    
+    # Decode sequence
+    output_text = tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)
+    return output_text
+
+def verbalize(text):
+    parts = split_to_parts(text)
+    verbalized = ''
+    for part in parts:
+        verbalized += generate_text(part)
+    return verbalized
 
 description = f'''
 <h1 style="text-align:center;">StyleTTS2 ukrainian demo</h1><br>
@@ -20,7 +132,9 @@ description = f'''
 examples = [
     ["Решта окупантів звернула на Вокзальну — центральну вулицю Бучі. Тільки уявіть їхній настрій, коли перед ними відкрилася ця пасторальна картина! Невеличкі котеджі й просторіші будинки шикуються обабіч, перед ними вивищуються голі липи та електро-стовпи, тягнуться газони й жовто-чорні бордюри. Доглянуті сади визирають із-поза зелених парканів, гавкотять собаки, співають птахи… На дверях будинку номер тридцять шість досі висить різдвяний вінок.", 1.0],
     ["Одна дівчинка стала королевою Франції. Звали її Анна, і була вона донькою Ярослава Му+дрого, великого київського князя. Він опі+кувався літературою та культурою в Київській Русі+, а тоді переважно про таке не дбали – більше воювали і споруджували фортеці.", 1.0],
-    ["Одна дівчинка народилася і виросла в Америці, та коли стала дорослою, зрозуміла, що дуже любить українські вірші й найбільше хоче робити вистави про Україну. Звали її Вірляна. Дід Вірляни був український мовознавець і педагог Кость Кисілевський, котрий навчався в Лейпцизькому та Віденському університетах і, після Другої світової війни виїхавши до США, започаткував систему шкіл українознавства по всій Америці. Тож Вірляна зростала в українському середовищі, а окрім того – в середовищі вихідців з інших країн.", 1.0]
+    ["Одна дівчинка народилася і виросла в Америці, та коли стала дорослою, зрозуміла, що дуже любить українські вірші й найбільше хоче робити вистави про Україну. Звали її Вірляна. Дід Вірляни був український мовознавець і педагог Кость Кисілевський, котрий навчався в Лейпцизькому та Віденському університетах і, після Другої світової війни виїхавши до США, започаткував систему шкіл українознавства по всій Америці. Тож Вірляна зростала в українському середовищі, а окрім того – в середовищі вихідців з інших країн.", 1.0],
+    ["За інформацією від Державної служби з надзвичайних ситуацій станом на 7 ранку 15 липня.", 1.0],
+    ["Очікується, що цей застосунок буде запущено 22.08.2025.", 1.0],
 ]
 
 def synthesize_multi(text, voice_audio, speed, progress=gr.Progress()):
@@ -55,8 +169,10 @@ with gr.Blocks() as single:
     with gr.Row():
         with gr.Column(scale=1):
             input_text = gr.Text(label='Text:', lines=5, max_lines=10)
+            verbalize_button = gr.Button("Вербалізувати(beta)")
             speed = gr.Slider(label='Швидкість:', maximum=1.3, minimum=0.7, value=1.0)
-            synthesise_button = gr.Button("Синтезувати")
+            verbalize_button.click(verbalize, inputs=[input_text], outputs=[input_text])
+            
         with gr.Column(scale=1):
             output_audio = gr.Audio(
                     label="Audio:",
@@ -64,6 +180,7 @@ with gr.Blocks() as single:
                     streaming=False,
                     type="numpy",
                 )
+            synthesise_button = gr.Button("Синтезувати")
             
             synthesise_button.click(synthesize_single, inputs=[input_text, speed], outputs=[output_audio])
     
@@ -75,8 +192,10 @@ with gr.Blocks() as multy:
     with gr.Row():
         with gr.Column(scale=1):
             input_text = gr.Text(label='Text:', lines=5, max_lines=10)
+            verbalize_button = gr.Button("Вербалізувати(beta)")
             speed = gr.Slider(label='Швидкість:', maximum=1.3, minimum=0.7, value=1.0)
             speaker = gr.Dropdown(label="Голос:", choices=prompts_list, value=prompts_list[0])
+            verbalize_button.click(verbalize, inputs=[input_text], outputs=[input_text])
 
         with gr.Column(scale=1):
             output_audio = gr.Audio(

requirements.txt

@@ -16,4 +16,6 @@ transformers
 git+https://github.com/patriotyk/ukrainian-word-stress.git
 git+https://github.com/patriotyk/ipa-uk.git
 spaces
-numpy<2
+numpy<2
+huggingface_hub
+onnxruntime