readded state

app.py

@@ -19,9 +19,12 @@ generation_config = GenerationConfig.from_pretrained('gpt2-medium')
 generation_config.max_new_tokens = response_length
 generation_config.pad_token_id = generation_config.eos_token_id
 
-def generate_response(story_so_far, new_prompt):
 
-    truncated_story = story_so_far[:1024 - response_length]
+
+
+def generate_response(outputs, new_prompt):
+
+    story_so_far = "\n".join(outputs[:int(1024 / response_length + 1)])
 
     set_seed(random.randint(0, 4000000000))
     inputs = tokenizer.encode(story_so_far + '\n' + new_prompt if story_so_far else new_prompt,
@@ -29,9 +32,21 @@ def generate_response(story_so_far, new_prompt):
                               max_length=1024 - response_length)
 
     output = model.generate(inputs, do_sample=True, generation_config=generation_config)
-    response = clean_paragraph(tokenizer.batch_decode(output)[0][((len(truncated_story) + 1) if truncated_story else 0):])
 
-    return (story_so_far + '\n' if story_so_far else '') + response
+    response = clean_paragraph(tokenizer.batch_decode(output)[0][((len(story_so_far) + 1) if story_so_far else 0):])
+    outputs.append(response)
+    return {
+        user_outputs: outputs,
+        story: ((story_so_far + '\n' if story_so_far else '') + response).replace('\n', '\n\n')
+    }
+
+def undo(outputs):
+
+    outputs = outputs[:-1]
+    return {
+        user_outputs: outputs,
+        story: "\n".join(outputs)
+    }
 
 def clean_paragraph(entry):
     paragraphs = entry.split('\n')
@@ -44,9 +59,10 @@ def clean_paragraph(entry):
     return capitalize_first_char("\n".join(paragraphs))
 
 def reset():
-    # global outputs
-    # outputs = []
-    return None
+    return {
+        user_outputs: None,
+        story: None
+    }
 
 def capitalize_first_char(entry):
     for i in range(len(entry)):
@@ -58,17 +74,19 @@ with gr.Blocks() as demo:
     story = gr.Textbox(interactive=False, lines=20)
     story.style(show_copy_button=True)
 
+    user_outputs = gr.State()
+
     prompt = gr.Textbox(placeholder="Continue the story here!", lines=3, max_lines=3)
 
     with gr.Row():
         gen_button = gr.Button('Generate')
-        # undo_button = gr.Button("Undo")
+        undo_button = gr.Button("Undo")
         res_button = gr.Button("Reset")
 
-    prompt.submit(generate_response, [story, prompt], story, scroll_to_output=True)
-    gen_button.click(generate_response, [story, prompt], story, scroll_to_output=True)
-    # undo_button.click(undo, [], story, scroll_to_output=True)
-    res_button.click(reset, [], story, scroll_to_output=True)
+    prompt.submit(generate_response, [user_outputs, prompt], [user_outputs, story], scroll_to_output=True)
+    gen_button.click(generate_response, [user_outputs, prompt], [user_outputs, story], scroll_to_output=True)
+    undo_button.click(undo, user_outputs, [user_outputs, story], scroll_to_output=True)
+    res_button.click(reset, [], [user_outputs, story], scroll_to_output=True)
 
 demo.launch(inbrowser=True)