import gradio as gr
from runorm import RUNorm

models = {
    "RUNorm-small": "small",
    "RUNorm-medium": "medium",
    "RUNorm-big": "big"
}

normalizers = {}

for model_name, model_size in models.items():
    normalizer = RUNorm()
    normalizer.load(model_size=model_size, device="cpu")
    normalizers[model_name] = normalizer

def normalize_text(input_text, model_name):
    normalizer = normalizers[model_name]
    normalized_text = normalizer.norm(input_text)
    return normalized_text

iface = gr.Interface(
    fn=normalize_text,
    inputs=[
        gr.components.Textbox(lines=5, placeholder="Введите текст для нормализации"),
        gr.components.Radio(list(models.keys()), label="Модель", value="RUNorm-medium")
    ],
    outputs="text",
    title="RUNorm",
    description="Демо RUNorm - модели для нормализации русского текста перед синтезом речи. Исходный код доступен на [GitHub](https://github.com/Den4ikAI/runorm).",
)

iface.launch()