import gradio as gr
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Load model from the Hub
model_name = "enansari/emotion_roberta_weighted" 
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

emotion_labels = ["sadness", "joy", "love", "anger", "fear", "surprise"]

def classify_emotion(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_id = torch.argmax(logits, dim=-1).item()
    return emotion_labels[predicted_id]

iface = gr.Interface(
    fn=classify_emotion,
    inputs=gr.Textbox(lines=2, placeholder="Type a sentence...", label="Input Text"),
    outputs=gr.Textbox(label="Predicted Emotion"),
    title="Sentiment Analysis with RoBERTa",
    description="Detect emotions (joy, sadness, anger, etc.) in text.",
    examples=[
        ["I am extremely happy today!"],
        ["This creates a lot of frustration."],
        ["I feel so lonely and cold."]
    ]
)

iface.launch()