import gradio as gr
from transformers import pipeline
import pandas as pd
import time
TASK = "fill-mask"
MODEL_OPTIONS = [
    "ai-forever/ruBert-large",
    "google-bert/bert-base-multilingual-cased",
    "FacebookAI/xlm-roberta-large"
]



MAX_CHARS = 500

#функция обработки текста и параметра top_k
def run_mask(text, top_k, model_name):

    fill_mask = pipeline(TASK, model=model_name)  #fill-mask для выбранной модел

    if text is None or not text.strip(): # проверка на пустой ввод
        return "ошибка", None, None

    
    if "[MASK]" not in text: # проверка наличия маски [MASK] во входном тексте
        return "ошибка. Добавьте [MASK] в текст", None, "-"

  
    text = text.strip()   # удаление пробелов в начале и конце

    if len(text) > MAX_CHARS:
        text = text[:MAX_CHARS]

   
    processed_text = text.replace("[MASK]", fill_mask.tokenizer.mask_token)  # замена [MASK] на токен модели (чтобы работало, если у новой выбранной модели будет другая маска)

   
    t0 = time.time()  # фиксация времени начала работы 

    try:
        # запуск модели 
        preds = fill_mask(processed_text, top_k=int(top_k))

        
        latency = round((time.time() - t0) * 1000, 1) # текущее время (сек) - время начала модели (сек) и перевод в миллимекундды с округлением


        rows = []# список для хранения результатов

        for p in preds:
         
            rows.append({
                "Вариант": p["token_str"],
                "Вероятность": round(p["score"], 3),
                "Предложение": p["sequence"]
            })

        # список в таблицу pandas
        df = pd.DataFrame(rows)

        # возврат статуса, таблицы результатов и времени ответа
        return "okey", df, f"{latency} ms"

    # обработка возможных ошибок 
    except Exception as e:
        return f"Error {type(e).__name__}: {e}", None, None


# функция очистки всех полей интерфейса
def clear_all():
    return "", "", None, ""


# создание интерфейса gradio на основе блоков
with gr.Blocks() as demo:

    # текстовое описание 
    gr.Markdown(
        f"""
        ### Заполнение маски (Fill Mask)
        **Задача:** {TASK}  
        """
    )

    # текстовое поле для ввода предложения с маской
    inp = gr.Textbox(
        lines=4,
        label="Введите текст с [MASK]",
        placeholder="На ужин я приготовил [MASK] с овощами."
    )

    # слайдер для выбора количества вариантов предсказания
    top_k = gr.Slider(
        minimum=1,
        maximum=10,
        value=5,
        step=1,
        label="Количество вариантов (top_k)"
    )


    # dropdown для выбора модели
    model_select = gr.Dropdown(
        choices=MODEL_OPTIONS,
        value=MODEL_OPTIONS[0],
        label="Выберите модель"
    )

    #кнопки
    with gr.Row():
        
        btn_run = gr.Button("Обработать")# кнопка запуска обработки
        btn_clear = gr.Button("Очистить") # кнопка очистки полей

    #  поля вывода
    status = gr.Textbox(label="Статус")
    out = gr.DataFrame(label="Результат модели")
    latency = gr.Textbox(label="Время ответа")



    # кнопка запускает функцию обработки текста и выводит её результаты в соответствующие поля интерфейса
    btn_run.click(
        run_mask,
        inputs=[inp, top_k,  model_select],
        outputs=[status, out, latency]
    )


    # кнопка запускает функцию очистки (clear_all)
    btn_clear.click(
        clear_all,
        outputs=[inp, status, out, latency]
    )

    #  примеры 
    gr.Examples(
        examples=[
            ["Терпение и труд всё [MASK]."],
            ["В шкафу лежал старый [MASK], который никто не использовал."],
            ["Этот фильм оказался очень [MASK]."],
            ["Лучше один раз увидеть, чем сто раз [MASK]."],
            ["Сквозь шум города прорывался тихий [MASK] надежды."]
        
        ],
        inputs=inp
    )


demo.launch()