ahadjawaid
Initial Commit
1ea8a0d
raw
history blame contribute delete
No virus
1.2 kB
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)