from pathlib import Path

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


# Specify the path of the model
model_ckpt = Path("./distilbert-base-uncased-finetuned-emotion")

# Load the fine-tuned tokenizer and model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = AutoModelForSequenceClassification.from_pretrained(model_ckpt).to(device)

class_names = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']

# main function
def inference(text: str) -> str:
    inputs = tokenizer(text, return_tensors="pt")

    inputs = {k:v.to(device) for k,v in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1).tolist()[0]

    max_vale = max(predictions)
    idx = predictions.index(max_vale)
    return model.config.id2label[idx]
    
title = "Classify the feeling of your sentence"
description = """
<p style="text-align:center">The model has been trained to classify the feeling of the texts, between sadness, joy, love, anger, fear or surprise. Test it!</p>
"""
examples = ['Tomorrow I will celebrate my birthday!', 'I was shocked when I saw the movie']

iface = gr.Interface(fn=inference, inputs="text", outputs="text", title=title, description=description, examples=examples)
iface.launch()