import gradio as gr
from transformers import pipeline 

sentimrnt_model = pipeline("text-classification", model="Aniemore/rubert-tiny2-russian-emotion-detection")

LABELS = {
    'neutral': 'нейтральный',
    'happiness': 'счастье', 
    'sadness': "грустный", 
    'enthusiasm': 'энтузиазм', 
    'fear': 'страх', 
    'anger': 'злость', 
    'disgust': 'отвращение'
}

def analyze_message(text):
    text = text.strip()
    
    if text == '':
        return 'Введите текст сообщения'
    try:
        result = sentimrnt_model(text[:512])[0]
        english_label = result['label'] 
        score = result['score']

        russian_label = LABELS.get(english_label, english_label)
    except Exception:
        russian_label = 'неизвестно'
        score = 0.0
    
    return russian_label, score

demo = gr.Interface(
    fn = analyze_message,
    inputs = gr.Textbox(lines = 3,
                        label = 'Текст сообщения',
                        placeholder = 'Вставьте сообщение'),
    outputs = [gr.Label(label = 'тональность'),
               gr.Textbox(label = 'оценка'),
              ],
    title = 'AI инспектор'
)

if __name__=="__main__":
    demo.launch()