joaopdrm's picture
Update app.py
184cc66
import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
class Emotionclass:
def __init__(self, model: str):
self.model = AutoModelForSequenceClassification.from_pretrained(model)
self.tokenizer = AutoTokenizer.from_pretrained(model)
self.pipeline = pipeline(
"text-classification",
model=self.model,
tokenizer=self.tokenizer,
return_all_scores=True,
)
def predict(self, input: str):
output = self.pipeline(input)[0]
result = {
"sad": output[0]["score"],
"joy": output[1]["score"],
"love": output[2]["score"],
"anger": output[3]["score"],
"fear": output[4]["score"],
"surprise": output[5]["score"],
}
return result
def main():
model = Emotionclass("bhadresh-savani/bert-base-uncased-emotion")
iface = gr.Interface(
fn=model.predict,
inputs=gr.inputs.Textbox(
lines=3,
placeholder="type here",
label="Input",
),
outputs="label",
title="Sentiment Classification",
)
iface.launch()
if __name__ == "__main__":
main()