Update app.py

app.py

@@ -30,7 +30,6 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.bfloat16,
 )
 model.eval()
-
 @spaces.GPU(duration=90)
 def generate(
     message: str,
@@ -61,89 +60,87 @@ def generate(
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
-    outputs = []
+    full_message = message
     for text in streamer:
-        outputs.append(text)
-        yield message + "".join(outputs)
-
-with gr.Blocks(css="style.css", fill_height=True) as demo:
-    gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            input_box = gr.Textbox(
-                label="Enter your text",
-                placeholder="Type your message here...",
-                lines=5
-            )
-            output_box = gr.Textbox(
-                label="Completed text",
-                lines=10,
-                interactive=False
-            )
-        with gr.Column(scale=1):
-            max_new_tokens = gr.Slider(
-                label="Max new tokens",
-                minimum=1,
-                maximum=MAX_MAX_NEW_TOKENS,
-                step=1,
-                value=DEFAULT_MAX_NEW_TOKENS,
-            )
-            temperature = gr.Slider(
-                label="Temperature",
-                minimum=0.1,
-                maximum=4.0,
-                step=0.1,
-                value=0.6,
-            )
-            top_p = gr.Slider(
-                label="Top-p (nucleus sampling)",
-                minimum=0.05,
-                maximum=1.0,
-                step=0.05,
-                value=0.9,
-            )
-            top_k = gr.Slider(
-                label="Top-k",
-                minimum=1,
-                maximum=1000,
-                step=1,
-                value=50,
-            )
-            repetition_penalty = gr.Slider(
-                label="Repetition penalty",
-                minimum=1.0,
-                maximum=2.0,
-                step=0.05,
-                value=1.2,
+        full_message += text
+        yield full_message
+        with gr.Blocks(css="style.css", fill_height=True) as demo:
+            gr.Markdown(DESCRIPTION)
+            gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
+
+            with gr.Row():
+                with gr.Column(scale=4):
+                    text_box = gr.Textbox(
+                        label="Enter your text",
+                        placeholder="Type your message here...",
+                        lines=10
+                    )
+                with gr.Column(scale=1):
+                    max_new_tokens = gr.Slider(
+                        label="Max new tokens",
+                        minimum=1,
+                        maximum=MAX_MAX_NEW_TOKENS,
+                        step=1,
+                        value=DEFAULT_MAX_NEW_TOKENS,
+                    )
+                    temperature = gr.Slider(
+                        label="Temperature",
+                        minimum=0.1,
+                        maximum=4.0,
+                        step=0.1,
+                        value=0.6,
+                    )
+                    top_p = gr.Slider(
+                        label="Top-p (nucleus sampling)",
+                        minimum=0.05,
+                        maximum=1.0,
+                        step=0.05,
+                        value=0.9,
+                    )
+                    top_k = gr.Slider(
+                        label="Top-k",
+                        minimum=1,
+                        maximum=1000,
+                        step=1,
+                        value=50,
+                    )
+                    repetition_penalty = gr.Slider(
+                        label="Repetition penalty",
+                        minimum=1.0,
+                        maximum=2.0,
+                        step=0.05,
+                        value=1.2,
+                    )
+
+            with gr.Row():
+                complete_btn = gr.Button("Complete")
+                stop_btn = gr.Button("Stop Generation")
+
+            stop_click = stop_btn.click(fn=None, cancels=[complete_btn.click])
+
+            complete_btn.click(
+                fn=generate,
+                inputs=[
+                    text_box,
+                    max_new_tokens,
+                    temperature,
+                    top_p,
+                    top_k,
+                    repetition_penalty
+                ],
+                outputs=text_box
             )
 
-    complete_btn = gr.Button("Complete")
-
-    complete_btn.click(
-        fn=generate,
-        inputs=[
-            input_box,
-            max_new_tokens,
-            temperature,
-            top_p,
-            top_k,
-            repetition_penalty
-        ],
-        outputs=output_box
-    )
-
-    gr.Examples(
-        examples=[
-            "Hello there! How are you doing?",
-            "Can you explain briefly to me what is the Python programming language?",
-            "Explain the plot of Cinderella in a sentence.",
-            "How many hours does it take a man to eat a Helicopter?",
-            "Write a 100-word article on 'Benefits of Open-Source in AI research'",
-        ],
-        inputs=input_box
-    )
+            gr.Examples(
+                examples=[
+                    "Hello there! How are you doing?",
+                    "Can you explain briefly to me what is the Python programming language?",
+                    "Explain the plot of Cinderella in a sentence.",
+                    "How many hours does it take a man to eat a Helicopter?",
+                    "Write a 100-word article on 'Benefits of Open-Source in AI research'",
+                ],
+                inputs=text_box
+            )
 
 if __name__ == "__main__":
     demo.queue(max_size=20).launch()