cache importing the models

app.py

@@ -4,16 +4,11 @@ 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 ' +
                    u'\U0001F5BC', page_icon=u'\U0001F5BC', layout="wide")
-if 'count' not in st.session_state or st.session_state.count == 6:
-	  st.session_state.count = 0
-	  st.session_state.chat_history_ids = None
-	  st.session_state.old_response = ''
-else:
-    st.session_state.count += 1
+gen = StoryGenerator()
+
 container_mode = st.sidebar.container()
 container_guide = st.sidebar.container()
 container_param = st.sidebar.container()
@@ -22,7 +17,7 @@ mode = container_mode.radio(
     "Select your mode",
     ('Create Statistics', 'Play Storytelling'), index=0)
 story_till_now = st.text_input(
-    label='First Sentence', 
+    label='First Sentence',
     value=random.choice([
         'Hello, I\'m a language model,',
         'So I suppose you want to ask me how I did it.',
@@ -30,15 +25,15 @@ story_till_now = st.text_input(
         'My first tutor was a dragon with a terrible sense of humor.',
         'Doctors told her she could never diet again.',
         'Memory is all around us, as well as within.',
-        
-        ]))
+
+    ]))
 
 num_generation = container_param.slider(
     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':
-    
+
     num_tests = container_param.slider(
         label='Number of tests', min_value=1, max_value=1000, value=3, step=1)
     reaction_weight_mode = container_param.select_slider(
@@ -64,35 +59,41 @@ if mode == 'Create Statistics':
             for si, story in enumerate(gen.data):
                 st.markdown(f'### Story no. {si}:', unsafe_allow_html=False)
                 for i, sentence in enumerate(story):
-                    col_turn, col_sentence, col_emo = st.columns([1,8,2])
-                    col_turn.markdown(sentence['turn'], unsafe_allow_html=False)
-                    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)
+                    col_turn, col_sentence, col_emo = st.columns([1, 8, 2])
+                    col_turn.markdown(
+                        sentence['turn'], unsafe_allow_html=False)
+                    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)
+            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)
+            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:
-        container_guide.markdown('### You selected statistics. Now set your parameters and click the `Analyse` button.')
-elif mode == 'Play Storytelling':
+        container_guide.markdown(
+            '### You selected statistics. Now set your parameters and click the `Analyse` button.')
+# elif mode == 'Play Storytelling':
 
-    # # , placeholder="Start writing your story...")
-    # story_till_now = st.text_input(
-    #     label='First Sentence', value='Hello, I\'m a language model,')
+#     # # , placeholder="Start writing your story...")
+#     # story_till_now = st.text_input(
+#     #     label='First Sentence', value='Hello, I\'m a language model,')
 
-    # num_generation = st.sidebar.slider(
-    #     label='Number of generation', min_value=1, max_value=100, value=10, step=1)
-    # length = st.sidebar.slider(label='Length of the generated sentence',
-    #                            min_value=1, max_value=100, value=20, step=1)
-    if container_button.button('Run'):
-        story_till_now, emotion = gen.story(
-            story_till_now, num_generation, length)
-        st.markdown(f'### Story')
-        st.text(story_till_now)
-        st.markdown(f'The last sentence has the "{emotion["label"]}" **Emotion** with a confidence score of {emotion["score"]}.')
-    else:
-        container_guide.markdown('### Write the first sentence and then hit the `Run` button')
+#     # num_generation = st.sidebar.slider(
+#     #     label='Number of generation', min_value=1, max_value=100, value=10, step=1)
+#     # length = st.sidebar.slider(label='Length of the generated sentence',
+#     #                            min_value=1, max_value=100, value=20, step=1)
+#     if container_button.button('Run'):
+#         story_till_now, emotion = gen.story(
+#             story_till_now, num_generation, length)
+#         st.markdown(f'### Story')
+#         st.text(story_till_now)
+#         st.markdown(f'The last sentence has the "{emotion["label"]}" **Emotion** with a confidence score of {emotion["score"]}.')
+#     else:
+#         container_guide.markdown('### Write the first sentence and then hit the `Run` button')
 # elif mode == 'Analyse Emotions':
 #     container_mode.write('Let\'s play storytelling.')

story_gen.py

@@ -8,6 +8,7 @@ import numpy as np
 import pandas as pd
 # import nltk
 import re
+import streamlit as st
 
 
 class StoryGenerator:
@@ -17,6 +18,7 @@ class StoryGenerator:
         self.stories = []
         self.data = []
 
+    @st.cache()
     def initialise_models(self):
         start = time.time()
         self.generator = pipeline('text-generation', model='gpt2')
@@ -156,7 +158,7 @@ class StoryGenerator:
             stats_dict['reaction_weight'] = None
             stats_df = pd.concat(
                 [stats_df, pd.DataFrame(stats_dict, index=[f'idx_{i}'])])
-            
+
         return stats_df, story_till_now, story_data
 
     def get_stats(self,