macapa's picture
Update app.py
a83ebf6 verified
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
model_name = "macapa/emotion-classifier"
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
labels = {0: 'sadness',
1: 'joy',
2: 'love',
3: 'anger',
4: 'fear',
5: 'surprise'}
def predict(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
inputs = inputs.to(device)
outputs = model(**inputs)
predictions = torch.argmax(outputs.logits, dim=1)
label = labels[predictions.item()]
return label
# Create the Gradio interface
iface = gr.Interface(
fn=predict,
inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
outputs="textbox",
title="Emotion Classification",
description="Enter some text and the model will predict the emotion",
)
# Launch the interface
iface.launch()