Spaces:
Runtime error
Runtime error
import gradio as gr | |
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline | |
class EmotionClassifier: | |
def __init__(self): | |
self.model = AutoModelForSequenceClassification.from_pretrained("barbieheimer/MND_TweetEvalBert_model") | |
self.tokenizer = AutoTokenizer.from_pretrained("barbieheimer/MND_TweetEvalBert_model") | |
self.pipeline = pipeline( | |
"text-classification", | |
model=self.model, | |
tokenizer=self.tokenizer, | |
return_all_scores=True, | |
) | |
def predict(self, input_text: str): | |
pred = self.pipeline(input_text)[0] | |
result = { | |
"Anger π ": pred[0]["score"], | |
"Joy π": pred[1]["score"], | |
"Surprise π²": pred[2]["score"], | |
"Sadness π": pred[3]["score"], | |
} | |
return result | |
def main(): | |
model = EmotionClassifier() | |
iface = gr.Interface( | |
fn=model.predict, | |
inputs=gr.inputs.Textbox( | |
lines=3, | |
placeholder="Type a phrase that has some emotion", | |
label="Input Text", | |
), | |
outputs="label", | |
title="Emotion Classification", | |
examples=[ | |
["The movie was a bummer."], | |
["I cannot wait to watch all these movies!"], | |
["The ending of the movie really irks me, gives me the ick fr."], | |
["The protagonist seems to have a lot of hope...."] | |
], | |
) | |
iface.launch() | |
if __name__ == "__main__": | |
main() | |