Spaces:
Runtime error
Runtime error
import gradio as gr | |
from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
from torch.nn.functional import softmax | |
class TextEmotionDetector: | |
def __init__(self, model_path): | |
self.model = AutoModelForSequenceClassification.from_pretrained(model_path) | |
self.tokenizer = AutoTokenizer.from_pretrained(model_path) | |
self.label_key = {0: 'surprise', 1: 'anger', 2: 'love', 3: 'joy', 4: 'fear', 5: 'sadness'} | |
def inference(self, text): | |
embedding = self.tokenizer(text, return_tensors="pt") | |
logits = self.model(input_ids=embedding['input_ids'], attention_mask=embedding['attention_mask'], | |
token_type_ids=embedding['token_type_ids'])['logits'] | |
probabilities = softmax(logits) | |
predication = { self.label_key[i]: float(prob) for i, prob in enumerate(probabilities[0]) } | |
return predication | |
detector = TextEmotionDetector("emotion-nlp-model") | |
text = gr.inputs.Textbox(lines=1, placeholder="Please enter text") | |
label = gr.outputs.Label() | |
interface = gr.Interface(fn=detector.inference, inputs=text, outputs=label, title="Text Emotion Classifier") | |
interface.launch(inline=False) |