# required imports
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# loading the model from the repository
REPO_ID = "adith-ds/deberta-classifier-v1"
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model = AutoModelForSequenceClassification.from_pretrained(REPO_ID)

# denoting the labels
LABEL_COLUMNS = ['anger', 'fear', 'joy', 'sadness', 'surprise']

#helper function to make a prediction using the model
def predict(text):
    input = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**input)
        logits = outputs.logits
    probs = torch.sigmoid(logits)[0]
    return {label: float(prob) for label, prob in zip(LABEL_COLUMNS, probs)}

# gradio configuration stuff
demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Enter your text here"),
    outputs=gr.Label(num_top_classes=5, label="Emotions detected"),
    title="Emotion Classifier (using finetuned DeBERTa)",
    description="Classify your text into the following emotions: Anger, Fear, Joy, Sadness, Surprise",
    examples=[
        ["Yikes! What a scare."],
        ["I hated this movie, it was so boring."],
        ["Oh wow! The effects were incredible. "]
    ]
)

# running the application
demo.launch()