cache the models

story_gen.py

@@ -18,14 +18,23 @@ class StoryGenerator:
         self.stories = []
         self.data = []
 
-    @st.cache()
+    @staticmethod
+    @st.cache(allow_output_mutation=True)
+    def get_generator():
+        return pipeline('text-generation', model='gpt2')
+
+    @staticmethod
+    @st.cache(allow_output_mutation=True)
+    def get_classifier():
+        return pipeline("text-classification",
+                        model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)
+
     def initialise_models(self):
-        start = time.time()
-        self.generator = pipeline('text-generation', model='gpt2')
-        self.classifier = pipeline("text-classification",
-                                   model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)
-        initialising_time = time.time()-start
-        print(f'Initialising Time: {initialising_time}')
+        # start = time.time()
+        self.generator = self.get_generator()
+        self.classifier = self.get_classifier()
+        # initialising_time = time.time()-start
+        # print(f'Initialising Time: {initialising_time}')
         # set_seed(42)
         # sys.exit()