Fix: session state

app.py

@@ -21,9 +21,9 @@ def loadEnModel():
 class GramformerDemo:
 
     def __init__(self):
-    	if 'models_loaded' not in st.session_state:
-    	    st.session_state['models_loaded'] = False
-    		
+    	if 'gf' not in st.session_state:
+            st.session_state['gf'] = None
+
         st.set_page_config(
             page_title="Gramformer Demo",
             initial_sidebar_state="expanded",
@@ -124,10 +124,13 @@ class GramformerDemo:
             st.warning('TO BE IMPLEMENTED !!')
             st.stop()
 
-        with st.spinner('Loading model..'):
-            if not st.session_state['models_loaded']:
-            	loadEnModel()
-            gf = self.load_gf(self.model_map[model_type])
+        if not st.session_state['gf']:
+            with st.spinner('Loading model..'):
+                loadEnModel()
+                gf = self.load_gf(self.model_map[model_type])
+                st.session_state['gf'] = gf
+        else:
+            gf = st.session_state['gf']
     
         input_text = st.selectbox(
             label="Choose an example",