import gradio as gr
from transformers import pipeline

# Define model names
models = {
    "ModernBERT Base (Go-Emotions)": "cirimus/modernbert-base-go-emotions",
    "ModernBERT Large (Go-Emotions)": "cirimus/modernbert-large-go-emotions"
}

# Function to load the selected model and classify text
def classify_text(model_name, text):
    classifier = pipeline("text-classification", model=models[model_name], top_k=None)
    predictions = classifier(text)
    return {pred["label"]: pred["score"] for pred in predictions[0]}

# Create the Gradio interface
interface = gr.Interface(
    fn=classify_text,
    inputs=[
        gr.Dropdown(
            list(models.keys()),
            label="Select Model",
            value="ModernBERT Base (Go-Emotions)"
        ),
        gr.Textbox(
            lines=2,
            placeholder="Enter text to analyze emotions...",
            value="I am thrilled to be a part of this amazing journey!"
        )
    ],
    outputs=gr.Label(num_top_classes=5),
    title="🎭 ModernBERT Emotion Classifier",
    description="Select a model and enter a sentence to see its associated emotions and confidence scores.",
)

# Launch the app
interface.launch()