import gradio as gr
from transformers import pipeline
import torch

MODEL = "AnasAlokla/multilingual_go_emotions"
classifier = pipeline(
    "text-classification",
    model=MODEL,
    tokenizer=MODEL,
    return_all_scores=True,
    device=0 if torch.cuda.is_available() else -1
)

def detect_emotions(text):
    if not text or not text.strip():
        return {"top_emotion": "none", "top_score": 0.0, "all_emotions": {}}
    
    out = classifier(text)[0]  # list of dicts: each with label & score
    
    # Find the emotion with highest score
    top_emotion = max(out, key=lambda x: x['score'])
    
    return {
        "top_emotion": top_emotion['label'],
        "top_score": round(top_emotion['score'], 3),
        "all_emotions": {e['label']: round(e['score'], 3) for e in out}
    }

demo = gr.Interface(
    fn=detect_emotions,
    inputs=gr.Textbox(lines=2, placeholder="Введите English или Russian..."),
    outputs="json",
    title="Emotion Detector",
    description="Detects emotions using multilingual_go_emotions - shows top emotion"
)

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