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

model_checkpoint = "luckyp71/bert_base_uncased_emotion_classification"

# device agnostic code
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model, tokenizer = None, None


def load_model_tokenizer(checkpoint):
    global model, tokenizer
    model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    model.to(device)


# load model when space starts
load_model_tokenizer(model_checkpoint)


@spaces.GPU
def prediction(text):
    encoded_text = tokenizer(text, return_tensors="pt").to(device)
    with torch.inference_mode():
        output = model(**encoded_text)
        logits = output.logits
        pred_ids = torch.argmax(logits, dim=1).item()

    return model.config.id2label[pred_ids].upper()


demo = gr.Interface(
    fn=prediction,
    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence..."),
    outputs=gr.Label(label="Predicted Emotion"),
    title="Emotion Classifier",
    description="Enter a sentence to predict the emotion using BERT fine-tuned on emotion text data."
)

demo.launch()