minor

app.py

@@ -11,7 +11,7 @@ import uuid
 import json
 
 # Constants
-SYSTEM_PROMPT = """You are SmallThinker-3B, a helpful AI assistant. You try to follow instructions as much as possible while being accurate and brief."""
+SYSTEM_PROMPT = """You are a helpful assistant."""
 device = "cuda" if torch.cuda.is_available() else "cpu"
 TITLE = "<h1><center>SmallThinker-3B Chat</center></h1>"
 MODEL_PATH = "PowerInfer/SmallThinker-3B-Preview"
@@ -53,7 +53,7 @@ button {
 # Load model and tokenizer
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_PATH,
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.float16,
 ).to(device)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
 
@@ -129,7 +129,13 @@ def stream_chat(
     #     with logs_file.open("a") as f:
     #         f.write(json.dumps({"input": input_text.replace(SYSTEM_PROMPT, ""), "output": buffer.replace(SYSTEM_PROMPT, ""), "model": "SmallThinker-3B"}))
     #         f.write("\n")
-
+def stop_generation():
+    stop_event.set()
+    return {
+        stop_btn: gr.Button.update(interactive=False),
+        submit: gr.Button.update(interactive=True),
+        textbox: gr.Textbox.update(interactive=True)
+    }
 def clear_input():
     return ""
 
@@ -140,7 +146,12 @@ def add_message(message: str, history: list):
 
 def clear_session() -> Tuple[str, List]:
     return '', []
-
+def on_submit(textbox, chatbot, *args):
+    return {
+        textbox: gr.Textbox.update(value="", interactive=False),
+        submit: gr.Button.update(interactive=False),
+        stop_btn: gr.Button.update(interactive=True),
+    }
 def main():
     with gr.Blocks(css=CSS, theme="soft") as demo:
         gr.HTML(TITLE)
@@ -167,39 +178,63 @@ def main():
                 with gr.Row():
                     clear_history = gr.Button("🧹 Clear History")
                     submit = gr.Button("🚀 Send")
+                    stop_btn = gr.Button("🛑 Stop", interactive=False)
                 
                 # Chain of events for submit button
                 submit_event = submit.click(
-                    fn=add_message,
+                    fn=on_submit,
                     inputs=[textbox, chatbot],
-                    outputs=chatbot,
+                    outputs=[textbox, submit, stop_btn],
                     queue=False
                 ).then(
-                    fn=clear_input,
-                    outputs=textbox,
+                    fn=add_message,
+                    inputs=[textbox, chatbot],
+                    outputs=chatbot,
                     queue=False
                 ).then(
                     fn=stream_chat,
                     inputs=[textbox, chatbot, temperature, max_new_tokens, top_p, top_k, repetition_penalty],
                     outputs=chatbot,
                     show_progress=True
+                ).then(
+                    fn=lambda: {
+                        textbox: gr.Textbox.update(interactive=True),
+                        submit: gr.Button.update(interactive=True),
+                        stop_btn: gr.Button.update(interactive=False)
+                    },
+                    outputs=[textbox, submit, stop_btn]
                 )
 
                 # Chain of events for enter key
                 enter_event = textbox.submit(
-                    fn=add_message,
+                    fn=on_submit,
                     inputs=[textbox, chatbot],
-                    outputs=chatbot,
+                    outputs=[textbox, submit, stop_btn],
                     queue=False
                 ).then(
-                    fn=clear_input,
-                    outputs=textbox,
+                    fn=add_message,
+                    inputs=[textbox, chatbot],
+                    outputs=chatbot,
                     queue=False
                 ).then(
                     fn=stream_chat,
                     inputs=[textbox, chatbot, temperature, max_new_tokens, top_p, top_k, repetition_penalty],
                     outputs=chatbot,
                     show_progress=True
+                ).then(
+                    fn=lambda: {
+                        textbox: gr.Textbox.update(interactive=True),
+                        submit: gr.Button.update(interactive=True),
+                        stop_btn: gr.Button.update(interactive=False)
+                    },
+                    outputs=[textbox, submit, stop_btn]
+                )
+                
+                # Stop button event
+                stop_btn.click(
+                    fn=stop_generation,
+                    outputs=[stop_btn, submit, textbox],
+                    queue=False
                 )
                 
                 clear_history.click(fn=clear_session,