adds table

app.py

@@ -3,6 +3,7 @@ from story_gen import StoryGenerator
 import plotly.figure_factory as ff
 import plotly.express as px
 import random
+import numpy as np
 gen = StoryGenerator()
 
 st.set_page_config(page_title='Storytelling ' +
@@ -14,12 +15,13 @@ if 'count' not in st.session_state or st.session_state.count == 6:
 else:
     st.session_state.count += 1
 container_mode = st.sidebar.container()
+container_guide = st.sidebar.container()
 container_param = st.sidebar.container()
 container_button = st.sidebar.container()
 mode = container_mode.radio(
     "Select your mode",
     ('Create Statistics', 'Play Storytelling'), index=0)
-story_till_now = container_param.text_input(
+story_till_now = st.text_input(
     label='First Sentence', 
     value=random.choice([
         'Hello, I\'m a language model,',
@@ -32,7 +34,7 @@ story_till_now = container_param.text_input(
         ]))
 
 num_generation = container_param.slider(
-    label='Number of generation', min_value=1, max_value=100, value=3, step=1)
+    label='Number of generation', min_value=1, max_value=100, value=5, step=1)
 length = container_param.slider(label='Length of the generated sentence',
                                 min_value=1, max_value=100, value=10, step=1)
 if mode == 'Create Statistics':
@@ -49,19 +51,31 @@ if mode == 'Create Statistics':
     if container_button.button('Analyse'):
         gen.get_stats(story_till_now=story_till_now,
                       num_generation=num_generation, length=length, reaction_weight=reaction_weight, num_tests=num_tests)
-        if len(gen.stories) > 0:
-            for si, story in enumerate(gen.stories):
+        # if len(gen.stories) > 0:
+        #     for si, story in enumerate(gen.stories):
+        #         st.markdown(f'### Story no. {si}:', unsafe_allow_html=False)
+        #         st.markdown(story, unsafe_allow_html=False)
+        #     data=gen.stats_df[gen.stats_df.sentence_no==3]
+        #     fig = px.violin(data_frame=data, x="reaction_weight", y="num_reactions", hover_data=data.columns)
+        #     st.plotly_chart(fig, use_container_width=True)
+        #     fig2 = px.box(data_frame=data, x="reaction_weight", y="num_reactions", hover_data=data.columns)
+        #     st.plotly_chart(fig2, use_container_width=True)
+        if len(gen.data) > 0:
+            for si, story in enumerate(gen.data):
                 st.markdown(f'### Story no. {si}:', unsafe_allow_html=False)
-                st.markdown(story, unsafe_allow_html=False)
+                for i, sentence in enumerate(story):
+                    col_sentence, col_emo = st.columns([3,1])
+                    col_sentence.markdown(sentence['sentence'], unsafe_allow_html=False)
+                    col_emo.markdown(f'{sentence["emotion"]} {np.round(sentence["confidence_score"], 3)}', unsafe_allow_html=False)
+            st.table(data=gen.stats_df, )
             data=gen.stats_df[gen.stats_df.sentence_no==3]
             fig = px.violin(data_frame=data, x="reaction_weight", y="num_reactions", hover_data=data.columns)
             st.plotly_chart(fig, use_container_width=True)
             fig2 = px.box(data_frame=data, x="reaction_weight", y="num_reactions", hover_data=data.columns)
             st.plotly_chart(fig2, use_container_width=True)
     else:
-        st.markdown('### You selected statistics. Now set your parameters and click the `Analyse` button.')
+        container_guide.markdown('### You selected statistics. Now set your parameters and click the `Analyse` button.')
 elif mode == 'Play Storytelling':
-    container_mode.write('Let\'s play storytelling.')
 
     # # , placeholder="Start writing your story...")
     # story_till_now = st.text_input(
@@ -78,4 +92,6 @@ elif mode == 'Play Storytelling':
         st.text(story_till_now)
         st.markdown(f'The last sentence has the "{emotion["label"]}" **Emotion** with a confidence score of {emotion["score"]}.')
     else:
-        st.markdown('### Write the first sentence and then hit the `Run` button')
+        container_guide.markdown('### Write the first sentence and then hit the `Run` button')
+# elif mode == 'Analyse Emotions':
+#     container_mode.write('Let\'s play storytelling.')

setup.py

@@ -4,7 +4,7 @@ with open("README.md", "r") as fh:
     long_description = fh.read()
 
 setuptools.setup(
-    name="ist",  
+    name="storytelling",  
     version="0.0.1",
     author="Jitesh Gosar",
     author_email="gosar95@gmail.com",

story_gen.py

@@ -9,14 +9,14 @@ import pandas as pd
 # import nltk
 import re
 
+
 class StoryGenerator:
     def __init__(self):
         self.initialise_models()
         self.stats_df = pd.DataFrame(data=[], columns=[])
         self.stories = []
-        # nltk.download()
-    
-    
+        self.data = []
+
     def initialise_models(self):
         start = time.time()
         self.generator = pipeline('text-generation', model='gpt2')
@@ -32,11 +32,17 @@ class StoryGenerator:
         self.clear_stats()
 
     def clear_stories(self):
+        self.data = []
         self.stories = []
 
     def clear_stats(self):
         self.stats_df = pd.DataFrame(data=[], columns=[])
 
+    def get_emotion(self, text):
+        emotions = self.classifier(text)
+        emotion = max(emotions[0], key=lambda x: x['score'])
+        return emotion
+
     @staticmethod
     def get_num_token(text):
         # return len(nltk.word_tokenize(text))
@@ -60,8 +66,7 @@ class StoryGenerator:
                                                      length, num_return_sequences=1)
             story_till_now = genreate_robot_sentence[0]['generated_text']
             new_sentence = story_till_now[last_length:]
-            emotions = self.classifier(new_sentence)
-            emotion = max(emotions[0], key=lambda x: x['score'])
+            emotion = self.get_emotion(new_sentence)
             # printj.yellow(f'Sentence {i}:')
             # story_to_print = f'{printj.ColorText.cyan(story_till_now[:last_length])}{printj.ColorText.green(story_till_now[last_length:])}\n'
             # print(story_to_print)
@@ -76,6 +81,13 @@ class StoryGenerator:
         stats_dict = dict()
         num_reactions = 0
         reaction_frequency = 0
+        emotion = self.get_emotion(story_till_now)  # first line emotion
+        story_data = [{
+            'sentence': story_till_now,
+            'turn': 'first',
+            'emotion': emotion['label'],
+            'confidence_score': emotion['score'],
+        }]
         for i in range(num_generation):
             # Text generation for User
             last_length = len(story_till_now)
@@ -85,21 +97,26 @@ class StoryGenerator:
             genreate_user_sentence = self.generator(story_till_now, max_length=self.get_num_token(
                 story_till_now)+length, num_return_sequences=1)
             story_till_now = genreate_user_sentence[0]['generated_text']
-            new_sentence = story_till_now[last_length:]
+            new_sentence_user = story_till_now[last_length:]
 
             printj.red.bold_on_white(f'loop: {i}; check emotion')
             # Emotion self.classifier for User
-            emotions = self.classifier(new_sentence)
-            emotion = max(emotions[0], key=lambda x: x['score'])
-            if emotion['label'] == 'neutral':
-                show_emotion = False
+            emotion_user = self.get_emotion(new_sentence_user)
+            if emotion_user['label'] == 'neutral':
+                show_emotion_user = False
             else:
                 reaction_frequency = num_reactions/(i+1)
-                show_emotion = self.check_show_emotion(
-                    confidence_score=emotion['score'], frequency=reaction_frequency, w=reaction_weight)
-            if show_emotion:
+                show_emotion_user = self.check_show_emotion(
+                    confidence_score=emotion_user['score'], frequency=reaction_frequency, w=reaction_weight)
+            if show_emotion_user:
                 num_reactions += 1
 
+            story_data.append({
+                'sentence': new_sentence_user,
+                'turn': 'user',
+                'emotion': emotion_user['label'],
+                'confidence_score': emotion_user['score'],
+            })
             # Text generation for Robot
             last_length = len(story_till_now)
             printj.cyan(story_till_now)
@@ -109,46 +126,53 @@ class StoryGenerator:
                 story_till_now)+length, num_return_sequences=1)
             story_till_now = genreate_robot_sentence[0]['generated_text']
             new_sentence = story_till_now[last_length:]
-
-            # emotions = self.classifier(new_sentence)
-            # emotion = max(emotions[0], key=lambda x:x['score'])
-
+            emotion = self.get_emotion(new_sentence)
+
+            story_data.append({
+                'sentence': new_sentence,
+                'turn': 'robot',
+                'emotion': emotion['label'],
+                'confidence_score': emotion['score'],
+            })
+            
             stats_dict['sentence_no'] = i
-            stats_dict['show_emotion'] = show_emotion
-            stats_dict['emotion_label'] = emotion['label']
-            stats_dict['emotion_score'] = emotion['score']
+            stats_dict['sentence'] = new_sentence_user
+            stats_dict['show_emotion'] = show_emotion_user
+            stats_dict['emotion_label'] = emotion_user['label']
+            stats_dict['emotion_score'] = emotion_user['score']
             stats_dict['num_reactions'] = num_reactions
             stats_dict['reaction_frequency'] = reaction_frequency
             stats_dict['reaction_weight'] = reaction_weight
             stats_df = pd.concat(
                 [stats_df, pd.DataFrame(stats_dict, index=[f'idx_{i}'])])
-        return stats_df, story_till_now
+        return stats_df, story_till_now, story_data
 
     def get_stats(self,
                   story_till_now="Hello, I'm a language model,",
                   num_generation=4,
                   length=20, reaction_weight=-1, num_tests=2):
         use_random_w = reaction_weight == -1
-        self.stories = []
+        # self.stories = []
         try:
-            num_rows = max(self.stats_df.test_id)+1
+            num_rows = max(self.stats_df.story_id)+1
         except Exception:
-            num_rows=0
-        for test_id in range(num_tests):
+            num_rows = 0
+        for story_id in range(num_tests):
             if use_random_w:
                 # reaction_weight = np.random.random_sample()
                 reaction_weight = np.round(np.random.random_sample(), 1)
-            stats_df0, _story_till_now = self.auto_ist(
+            stats_df0, _story_till_now, story_data = self.auto_ist(
                 story_till_now=story_till_now,
                 num_generation=4,
                 length=20, reaction_weight=reaction_weight)
-            stats_df0.insert(loc=0, column='test_id', value=test_id+num_rows)
+            stats_df0.insert(loc=0, column='story_id', value=story_id+num_rows)
 
-            # stats_df0['test_id'] = test_id
+            # stats_df0['story_id'] = story_id
             self.stats_df = pd.concat([self.stats_df, stats_df0])
-            printj.yellow(f'test_id: {test_id}')
+            printj.yellow(f'story_id: {story_id}')
             printj.green(stats_df0)
             self.stories.append(_story_till_now)
+            self.data.append(story_data)
         self.stats_df = self.stats_df.reset_index(drop=True)
         print(self.stats_df)
 
@@ -159,4 +183,4 @@ class StoryGenerator:
         self.stats_df.to_excel(writer, sheet_name='IST')
 
         # Close the Pandas Excel writer and output the Excel file.
-        writer.save()
+        writer.save()