| import gradio as gr | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline | |
| class EmotionClassifier: | |
| def __init__(self, model_name: str): | |
| self.model = AutoModelForSequenceClassification.from_pretrained(model_name) | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| 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 = { | |
| "Sadness 😭": pred[0]["score"], | |
| "Joy 😂": pred[1]["score"], | |
| "Love 😍": pred[2]["score"], | |
| "Anger 😠": pred[3]["score"], | |
| "Fear 😨": pred[4]["score"], | |
| "Surprise 😲": pred[5]["score"], | |
| } | |
| return result | |
| def main(): | |
| model = EmotionClassifier("bhadresh-savani/bert-base-uncased-emotion") | |
| 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=[ | |
| "I get so down when I'm alone", | |
| "I believe that today everything will work out", | |
| "It was so dark there I was afraid to go", | |
| "I loved the gift you gave me", | |
| "I was very surprised by your presentation.", | |
| ], | |
| ) | |
| iface.launch() | |
| if __name__ == "__main__": | |
| main() | |