add emotion classification

README.md

@@ -7,7 +7,7 @@ sdk: streamlit
 sdk_version: 1.10.0
 app_file: app.py
 pinned: false
-models: mrm8488/t5-base-finetuned-emotion
+models: bhadresh-savani/distilbert-base-uncased-emotion
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -8,6 +8,6 @@ name = st.text_input('Who are you?', 'Zeliboba')
 text = st.text_area('Submit your stories', '''Random symbols''')
 
 st.write('Here is your first text:', text)
-st.write(get_emotion(text))
+st.write(get_emotion(text)['label'])
 
 

emotions.py

@@ -1,16 +1,11 @@
-from transformers import AutoTokenizer, AutoModelWithLMHead
+from transformers import pipeline
 
-tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-emotion")
-
-model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-emotion")
+classifier = pipeline("text-classification",
+                      model='bhadresh-savani/distilbert-base-uncased-emotion',
+                      return_all_scores=True)
 
 
 def get_emotion(text='No text yet'):
-    input_ids = tokenizer.encode(text + '</s>', return_tensors='pt')
-
-    output = model.generate(input_ids=input_ids,
-                            max_length=2)
-
-    dec = [tokenizer.decode(ids) for ids in output]
-    label = dec[0]
-    return label
+    prediction = classifier(text)[0]
+    result = max(prediction, key=lambda x: x['score'])
+    return result