add emotional arc analysis

app.py

@@ -1,13 +1,21 @@
 import streamlit as st
 # from emoji import get_emoji
-from emotions import get_emotion
+from emotions import get_emotion, get_sentiment_arc_evaluation
+from nltk.tokenize import sent_tokenize
+nltk.download('punkt')
 
 st.title("I don't even know what this is yet")
 
-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)['label'])
+name = st.text_input('Who are you?', 'Right, who am I?')
+text = st.text_area('Submit your stories', '''Words and symbols are meant to be here''')
+emotion_result = get_emotion(text)
+st.write(f"Overall emotion of your story: {emotion_result[0]['label']}")
 
+sents = sent_tokenize(text)
+emo_arc = []
+for sent in sents:
+    emo_arc.append(emotion_result[0]['label'])
+sentiment_arc_eval = get_sentiment_arc_evaluation(emo_arc)
 
+st.write(f"Emotional arc of your story: {' - '.join(emo_arc)}")
+st.write(sentiment_arc_eval)

emotions.py

@@ -7,5 +7,30 @@ classifier = pipeline("text-classification",
 
 def get_emotion(text='No text yet'):
     prediction = classifier(text)[0]
-    result = max(prediction, key=lambda x: x['score'])
+    result = sorted(prediction, key=lambda x: x['score'])
     return result
+
+
+sentiment_map = {'anger': 'neg', 'sadness': 'neg', 'fear': 'neg',
+                 'joy': 'pos', 'love': 'pos', 'surprise': 'pos'}
+good_arcs = ['neg - pos', 'pos - neg']
+great_arcs = ['pos - neg - pos', 'neg - pos - neg']
+
+
+def get_sentiment_arc_evaluation(emotions):
+    sentiment_arc = []
+    for emo in emotions:
+        sentiment = sentiment_map[emo]
+        if sentiment_arc and sentiment_arc[-1] == sentiment:
+            continue
+        sentiment_arc.append(sentiment)
+    sentiment_arc_str = '\n'.join(sentiment_arc)
+    if sentiment_arc_str in great_arcs:
+        return 'What a great plot! Excellent! 😍'
+    elif sentiment_arc_str in good_arcs:
+        return 'Story plot seems nice! But you can do better. 😉'
+    elif len(sentiment_arc) < 2:
+        return "No judgment, but... The plot might be too simple! 🤓"
+    else:
+        return "The plot seems complicated. 🤔 But maybe I am just too stupid to understand!"
+

requirements.txt

@@ -2,4 +2,5 @@ streamlit
 transformers
 torch
 numpy
-scipy
+scipy
+nltk