app.py

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']


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]

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