app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load the model from Hugging Face Model Hub
model_name = "SamLowe/roberta-base-go_emotions"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Define emotion labels used by the model
emotion_labels = ["admiration", "amusement", "anger", "annoyance", "approval",
                  "caring", "confusion", "curiosity", "desire", "disappointment",
                  "disapproval", "disgust", "embarrassment", "excitement",
                  "fear", "gratitude", "grief", "joy", "love", "nervousness",
                  "optimism", "pride", "realization", "relief", "remorse",
                  "sadness", "surprise", "neutral"]

def predict_emotion(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = logits.argmax().item()

    predicted_emotion = emotion_labels[predicted_class]
    confidence = logits.softmax(dim=1).squeeze()[predicted_class].item()
    return predicted_emotion, confidence  # Return the predicted emotion and confidence

def get_confidence_color(confidence):
    # Define a color scale based on confidence
    if confidence >= 0.8:
        return "green"
    elif confidence >= 0.5:
        return "orange"
    else:
        return "red"

iface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Textbox(),
    outputs=[
        gr.Textbox(),
        "text",
        gr.Textbox("Confidence Score:", default=""),
    ],
    live=True,
    title="Emotion Prediction",
    description="Enter a sentence for emotion prediction.",
)

# Customize the interface appearance
iface.style(
    confidence_score=get_confidence_color
)

iface.launch()