Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch.nn.functional as F | |
| import torch | |
| from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast | |
| def translate(text): | |
| model_name = 'sbenel/emotion-distilbert' | |
| tokenizer = DistilBertTokenizerFast.from_pretrained(model_name) | |
| model= DistilBertForSequenceClassification.from_pretrained(model_name) | |
| input = tokenizer(text, return_tensors="pt") | |
| labels = torch.tensor([1]).unsqueeze(0) # Batch size 1 | |
| output = model(**input, labels=labels) | |
| logits = output.logits | |
| prediction = F.softmax(logits, dim=1) | |
| y_pred = torch.argmax(prediction).numpy() | |
| class_names = ['sad','joy','love','anger','fear','surprise'] | |
| return class_names[y_pred] | |
| # output = model.generate(input["input_ids"], max_length=40, num_beams=4, early_stopping=True) | |
| title = "Text Emotion Classification" | |
| inputs = gr.inputs.Textbox(lines=1, label="Text") | |
| outputs = [gr.outputs.Textbox(label="Emotions")] | |
| description = "Here use the [emotion-distilbert](https://huggingface.co/sbenel/emotion-distilbert) that was trained with [emotion dataset](https://huggingface.co/datasets/emotion)." | |
| iface = gr.Interface(fn=translate, inputs=inputs, outputs=outputs, theme="grass", title=title, description=description) | |
| iface.launch(enable_queue=True) |