chat template

app.py

@@ -12,7 +12,7 @@ import torch
 from huggingface_hub import login
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-from utils_chatbot import organize_messages, stream2display_text
+from utils_chatbot import organize_messages_from_messages, stream2display_text
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -44,6 +44,7 @@ def gpu_generate_stream(inputs, history, temperature, top_p):
     )
     model_inputs = tokenizer([prompt_text], return_tensors="pt").to("cuda")
 
+    history.append({"role": "assistant", "content": ""})
     yield "", history
 
     streamer = TextIteratorStreamer(
@@ -77,17 +78,17 @@ def gpu_generate_stream(inputs, history, temperature, top_p):
         elapsed = time.time() - start_time
         token_per_sec = gen_tk_count / elapsed if elapsed > 0 else 0
         display_text = stream2display_text(stream_text, token_per_sec)
-        history[-1] = (history[-1][0], display_text)
+        history[-1]["content"] = display_text
         yield "", history
 
     thread.join()
-    history[-1] = (history[-1][0], stream_text.replace("<|im_end|>", ""))
+    history[-1]["content"] = stream_text.replace("<|im_end|>", "")
     yield "", history
 
 
 def gen_response_stream(message, history, temperature, top_p):
-    chat_msg_ls = organize_messages(message, history)
-    history.append((message, ""))
+    chat_msg_ls = organize_messages_from_messages(message, history)
+    history.append({"role": "user", "content": message})
     yield from gpu_generate_stream(
         chat_msg_ls, history,
         temperature=temperature,
@@ -99,32 +100,7 @@ def create_app():
     assets_path = Path.cwd().absolute() / "assets"
     gr.set_static_paths(paths=[assets_path])
 
-    theme = gr.themes.Soft(
-        primary_hue="blue",
-        secondary_hue="gray",
-        neutral_hue="slate",
-        font=[gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"],
-    )
-
-    with gr.Blocks(
-        theme=theme,
-        css="""
-        .logo-container {
-            text-align: center;
-            margin: 0.5rem 0 1rem 0;
-        }
-        .logo-container img {
-            height: 96px;
-            width: auto;
-            max-width: 200px;
-            display: inline-block;
-        }
-        .input-box {
-            border: 1px solid #2f63b8;
-            border-radius: 8px;
-        }
-        """,
-    ) as demo:
+    with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column(scale=1):
                 gr.HTML(
@@ -165,6 +141,32 @@ def create_app():
     return demo
 
 
+THEME = gr.themes.Soft(
+    primary_hue="blue",
+    secondary_hue="gray",
+    neutral_hue="slate",
+    font=[gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"],
+)
+
+CSS = """
+.logo-container {
+    text-align: center;
+    margin: 0.5rem 0 1rem 0;
+}
+.logo-container img {
+    height: 96px;
+    width: auto;
+    max-width: 200px;
+    display: inline-block;
+}
+.input-box {
+    border: 1px solid #2f63b8;
+    border-radius: 8px;
+}
+"""
+
+
+demo = create_app()
+
 if __name__ == "__main__":
-    demo = create_app()
-    demo.launch()
+    demo.launch(theme=THEME, css=CSS)

utils_chatbot.py

@@ -1,13 +1,12 @@
-def organize_messages(message, history):
-    msg_ls = [dict(
-        role="system",
-        content="You are a helpful assistant.",
-    )]
-    for user, assistant in history:
-        msg_ls.append(dict(role="user", content=user))
-        if assistant:
-            msg_ls.append(dict(role="assistant", content=assistant))
-    msg_ls.append(dict(role="user", content=message))
+def organize_messages_from_messages(message, history):
+    """Build chat messages from Gradio 6.x messages-format history."""
+    msg_ls = [{"role": "system", "content": "You are a helpful assistant."}]
+    for msg in history:
+        role = msg.get("role", "")
+        content = msg.get("content", "")
+        if role in ("user", "assistant") and content:
+            msg_ls.append({"role": role, "content": content})
+    msg_ls.append({"role": "user", "content": message})
     return msg_ls