MND_Tweet_Space / app.py
barbieheimer's picture
Update app.py
7fc03a0
raw history blame
No virus
1.48 kB
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()