add states

app.py

@@ -179,6 +179,7 @@ def main():
     st.plotly_chart(figure, use_container_width=True)
     auth_token = os.environ.get('TOKEN') or True
     fp16 = st.checkbox("FP16", value=True)
+    st.session_state["generated_text"] = None
 
     def generate():
         text = inference(
@@ -192,8 +193,7 @@ def main():
             act_type=act_type
         )
         st.subheader("Generated text:")
-        st.write(text)
-    generate()
+        st.write(st.session_state["generated_text"])
     st.button("Generate new", on_click=generate())
 
 

generator.py

@@ -2,6 +2,7 @@ from typing import Optional, Union
 
 import torch
 import transformers
+import streamlit as st
 
 
 class Generator:
@@ -69,6 +70,7 @@ class Generator:
             progress_bar.progress((i+1)/max_length)
             if ended_sequences.all():
                 break
+            st.session_state["generated_text"] = self.tokenizer.decode(input_ids[0])
 
         return (
             [