UI fixes/workaround

app.py

@@ -1,6 +1,6 @@
 import os
 from functools import lru_cache
-from typing import Any, Iterator
+from typing import Any
 
 import gradio as gr
 from huggingface_hub import hf_hub_download
@@ -128,7 +128,7 @@ def build_messages(
     return messages
 
 
-def respond(
+def generate_reply(
     message: str,
     history: list[Any],
     system_message: str,
@@ -136,7 +136,16 @@ def respond(
     temperature: float,
     top_p: float,
     repeat_penalty: float,
-) -> Iterator[str]:
+) -> Any:
+    if not message.strip():
+        yield history, ""
+        return
+
+    visible_history = list(history or [])
+    visible_history.append({"role": "user", "content": message})
+    visible_history.append({"role": "assistant", "content": ""})
+    yield visible_history, ""
+
     llm = get_llm()
     messages = build_messages(message, history, system_message or DEFAULT_SYSTEM_PROMPT)
     output = ""
@@ -155,54 +164,80 @@ def respond(
         token = delta.get("content") or ""
         if token:
             output += token
-            yield output
+            visible_history[-1] = {"role": "assistant", "content": output}
+            yield visible_history, ""
+
 
+def clear_chat() -> tuple[list[dict[str, str]], str]:
+    return [], ""
+
+
+with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        "## Dreadzone\nCPU-only local chat with NVIDIA Nemotron 3 Nano 4B GGUF."
+    )
 
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    title="Dreadzone",
-    description="CPU-only local chat with NVIDIA Nemotron 3 Nano 4B GGUF.",
-    chatbot=gr.Chatbot(
+    chat = gr.Chatbot(
         type="messages",
         height=420,
         show_copy_button=True,
         container=True,
-    ),
-    additional_inputs=[
-        gr.Textbox(
+        label="Chat",
+    )
+
+    with gr.Row():
+        message = gr.Textbox(
+            placeholder="Type a message...",
+            show_label=False,
+            scale=8,
+            container=False,
+        )
+        submit = gr.Button("Send", variant="primary", scale=1)
+
+    with gr.Accordion("Generation settings", open=False):
+        system_message = gr.Textbox(
             value=DEFAULT_SYSTEM_PROMPT,
             label="System message",
             lines=3,
-        ),
-        gr.Slider(
+        )
+        max_tokens = gr.Slider(
             minimum=32,
             maximum=1024,
             value=384,
             step=16,
             label="Max new tokens",
-        ),
-        gr.Slider(
+        )
+        temperature = gr.Slider(
             minimum=0.0,
             maximum=1.5,
             value=0.7,
             step=0.05,
             label="Temperature",
-        ),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"),
-        gr.Slider(
+        )
+        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p")
+        repeat_penalty = gr.Slider(
             minimum=1.0,
             maximum=1.3,
             value=1.1,
             step=0.01,
             label="Repeat penalty",
-        ),
-    ],
-    additional_inputs_accordion=gr.Accordion("Generation settings", open=False),
-    css=CSS,
-    theme=gr.themes.Soft(),
-)
+        )
+
+    clear = gr.Button("Clear")
+
+    inputs = [
+        message,
+        chat,
+        system_message,
+        max_tokens,
+        temperature,
+        top_p,
+        repeat_penalty,
+    ]
+    submit.click(generate_reply, inputs=inputs, outputs=[chat, message])
+    message.submit(generate_reply, inputs=inputs, outputs=[chat, message])
+    clear.click(clear_chat, outputs=[chat, message])
 
 
 if __name__ == "__main__":
-    chatbot.queue(max_size=8).launch(ssr_mode=False)
+    demo.queue(max_size=8).launch(ssr_mode=False)