Spaces:
Sleeping
Sleeping
File size: 1,696 Bytes
5434155 0fc1701 5434155 0825778 5434155 |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
# Model name (you can swap this for another emotion model if you like)
model_name = "j-hartmann/emotion-english-distilroberta-base"
#minoosh/finetuned_bert-base-on-IEMOCAP_1
# Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
model.eval()
# Prediction function
def predict_emotion(text: str):
# Handle empty input
if not text or not text.strip():
return {"Error": "Please enter some text."}
# Tokenize
inputs = tokenizer(
text,
return_tensors="pt",
truncation=True,
padding=True,
max_length=256, # you can adjust this if needed
).to(device)
with torch.no_grad():
outputs = model(**inputs)
probs = outputs.logits.softmax(dim=-1)[0]
# Map id -> label using model config
id2label = model.config.id2label
scores = {id2label[i]: float(probs[i]) for i in range(len(probs))}
# Sort by highest probability first (optional but nice in the UI)
scores = dict(sorted(scores.items(), key=lambda x: x[1], reverse=True))
return scores
# Gradio interface
demo = gr.Interface(
fn=predict_emotion,
inputs=gr.Textbox(lines=4, label="Enter text"),
outputs=gr.Label(label="Emotion Probabilities"),
title="Emotion Classifier",
description="Enter a sentence and see the predicted emotion distribution.",
flagging_mode="never",
)
if __name__ == "__main__":
demo.launch()
|