File size: 1,116 Bytes
70ddccd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import gradio as gr
# Load model and tokenizer
model_name = "nateraw/bert-base-uncased-emotion"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
# Emotion labels
labels = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
# Prediction function
def predict_emotion(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=1)
pred_class = torch.argmax(probs).item()
emotion = labels[pred_class]
return f"{emotion} ({probs[0][pred_class].item()*100:.2f}% confidence)"
# Gradio Interface
interface = gr.Interface(
fn=predict_emotion,
inputs=gr.Textbox(lines=2, placeholder="Type something here..."),
outputs="text",
title="BERT-based Emotion Detection",
description="A web app that uses a fine-tuned BERT model to detect emotions from text."
)
interface.launch()
|