import gradio as gr
from transformers import pipeline 
import time
TASK = "text-classification"
MODEL_NAME = "Aniemore/rubert-tiny2-russian-emotion-detection"

sentiment_model = pipeline(TASK, model=MODEL_NAME)

MAX_CHARS = 2000

def runk(text):
    if text is None or not text.strip():
        return 'ошибка', None, None
    
    text = text.strip()
    if len(text) > MAX_CHARS:
        text = text[:MAX_CHARS]
    
    t0 = time.time()
    
    try:
        result = sentiment_model(text)
        latency = round((time.time() - t0) * 1000, 1)  # Исправлено вычисление задержки
        return 'okey', result, f'{latency} ms'  # Добавлен return
    except Exception as e:
        return f'Error {type(e).__name__}: {e}', None, None  # Добавлен return


with gr.Blocks() as demo:
    gr.Markdown(f'''***Задача:*** {TASK}  ***Модель:*** {MODEL_NAME}
                ''')
    inp = gr.Textbox(lines=6,
                     label='Текст сообщения',
                     placeholder='Вставьте сообщение')  
    btm = gr.Button('Обработать')  
    status = gr.Textbox(label='статус')  
    out = gr.JSON(label='результат модели') 
    latency = gr.Textbox(label='Время ответа')  
    btm.click(runk, inputs=inp, outputs=[status, out, latency]) 
    gr.Examples(
        examples=[['я люблю этот продукт, он великолепен'],
                  ['это самый худший опыт'],
                  ['ничего специфичного']],
        inputs=inp
    )
demo.launch()