speech-uk/text-normalization

1import gradio as gr2import os3import ctranslate24from transformers import M2M100Tokenizer5from huggingface_hub import snapshot_download6import time7 8# Model settings9MODEL_PATH = os.getenv("MODEL_PATH", "skypro1111/m2m100-ukr-verbalization-ct2")10TOKENIZER_PATH = os.getenv("TOKENIZER_PATH", "skypro1111/m2m100-ukr-verbalization")11 12# Initialize translator and tokenizer globally13def init_model(model_path=MODEL_PATH, tokenizer_path=TOKENIZER_PATH):14    print("\nInitializing CTranslate2 model and tokenizer...")15    16    # Download the model from HuggingFace Hub17    local_model_path = snapshot_download(18        repo_id=model_path,19        allow_patterns=["*.bin", "*.json", "tokenizer.json", "vocab.json"],20    )21    22    # Initialize translator with optimizations for CPU23    translator = ctranslate2.Translator(24        local_model_path,25        device="cpu",26        compute_type="int8",  # Use INT8 quantization for CPU27        intra_threads=4,      # Adjust based on CPU cores available28    )29    30    # Load tokenizer31    tokenizer = M2M100Tokenizer.from_pretrained(tokenizer_path)32    tokenizer.src_lang = "uk"33    34    return translator, tokenizer35 36translator, tokenizer = init_model()37 38def process_text(text: str):39    """Process a single text input using the CTranslate2 model."""40    start_time = time.time()41    42    # Tokenize input43    source = tokenizer.convert_ids_to_tokens(tokenizer.encode(text))44    target_prefix = [tokenizer.lang_code_to_token["uk"]]45    46    # Run inference47    results = translator.translate_batch(48        [source],49        target_prefix=[target_prefix],50        beam_size=1,51        num_hypotheses=1,52        use_vmap=True,53    )54    55    # Get target tokens and decode56    target = results[0].hypotheses[0][1:]  # Remove language token57    output = tokenizer.decode(tokenizer.convert_tokens_to_ids(target))58    59    inference_time = time.time() - start_time60    return output, inference_time61 62def inference(text):63    """Gradio interface function."""64    output, inference_time = process_text(text)65    return output, f"{inference_time:.3f} секунд"66 67# Define examples from inference_ct2.py68examples = [69    ["Моя бабуся народилася 07.11.1919, у важкий післявоєнний час."],70    ["Зустріч призначена на 15:30 12.05.2025 у конференц-залі №3."],71    ["Телефонуйте нам за номером +380 (44) 123-45-67 або 0800 500 123."],72    ["Температура повітря сьогодні становить +25°C, а тиск 750 мм.рт.ст."],73    ["ТОВ «Мрія» було засновано 28/06/2022 з початковим капіталом 50 тис. грн."],74    ["Швидкість вітру 15 м/с, видимість 10 км, вологість 65%."],75    ["Потяг №743 Київ-Львів відправляється о 08:45 з платформи №2."],76    ["Ціна на пальне зросла на 2,5 грн/л і становить 54,99 грн."],77    ["Площа квартири 75,5 м², висота стелі 2,75 м."],78    ["Відстань між містами становить 450 км або 280 миль."],79    ["Станом на 7:00 15 квітня 2025 року температура повітря становить +18°C, вологість 60%."]80]81 82# Define Gradio interface83with gr.Blocks() as interface:84    gr.Markdown("# Вербалізація українського тексту (M2M100-CT2)")85    gr.Markdown("""Модель для перетворення чисел, дат, одиниць вимірювання та інших символьних позначень 86    у їх текстовий запис українською мовою. Використовує оптимізовану CTranslate2 версію для швидкого інференсу.""")87    88    with gr.Row():89        with gr.Column():90            input_text = gr.Textbox(91                label="Введіть текст для вербалізації",92                placeholder="Наприклад: Температура повітря сьогодні становить +25°C"93            )94            submit_btn = gr.Button("Вербалізувати")95            96            gr.Examples(97                examples=examples,98                inputs=input_text,99            )100            101        with gr.Column():102            output_text = gr.Textbox(label="Вербалізований текст")103            inference_time = gr.Textbox(label="Час обробки")104    105    gr.Markdown("""106    ### Можливості моделі:107    - Дати та час108    - Телефонні номери109    - Одиниці вимірювання110    - Грошові суми111    - Числові вирази112    113    ### Технічні деталі:114    - Базова модель: facebook/m2m100_418M115    - Оптимізація: CTranslate2 з INT8 квантизацією116    """)117    118    submit_btn.click(119        fn=inference,120        inputs=input_text,121        outputs=[output_text, inference_time],122    )123 124# Launch the interface125interface.launch(share=False)