Update app.py

app.py

@@ -48,8 +48,11 @@ def respond(message, max_tokens, temperature, top_p):
     thread = threading.Thread(target=run)
     thread.start()
 
+    response = ""
     for new_text in streamer:
-        yield new_text
+        response += new_text
+        # prepend model name on its own line
+        yield f"**{current_model_name}**\n\n{response}"
 
 # User input handler
 def user(message, chat_history):
@@ -57,18 +60,12 @@ def user(message, chat_history):
     return "", chat_history
 
 # Bot response handler
-def bot(chat_history, max_tokens, temperature, top_p):
-    # Insert model name bubble and placeholder for output
-    chat_history.append([None, f"**{current_model_name}**"])
-    chat_history.append([None, ""])
-    # Render model name immediately
-    yield chat_history
-
-    # Stream generation into the last bubble
-    message = chat_history[-3][0]  # Original user message
-    for chunk in respond(message, max_tokens, temperature, top_p):
-        chat_history[-1][1] += chunk
-        yield chat_history
+def bot(chatbot, max_tokens, temperature, top_p):
+    message = chatbot[-1][0]
+    response_generator = respond(message, max_tokens, temperature, top_p)
+    for response in response_generator:
+        chatbot[-1][1] = response
+        yield chatbot
 
 # Model selector handler
 def update_model(model_name):
@@ -76,7 +73,7 @@ def update_model(model_name):
     return []
 
 # Gradio UI
-with gr.Blocks(css=".gr-chatbot .message { margin: 2px 0 !important; }", title=$1) as demo:
+with gr.Blocks(title="LeCarnet - Chat Interface") as demo:
     with gr.Row():
         gr.HTML("""
         <div style="text-align: center; width: 100%;">
@@ -85,7 +82,6 @@ with gr.Blocks(css=".gr-chatbot .message { margin: 2px 0 !important; }", title=$
         """ )
 
     with gr.Row():
-        # Options column
         with gr.Column(scale=1, min_width=150):
             model_selector = gr.Dropdown(
                 choices=["LeCarnet-3M", "LeCarnet-8M", "LeCarnet-21M"],
@@ -97,7 +93,6 @@ with gr.Blocks(css=".gr-chatbot .message { margin: 2px 0 !important; }", title=$
             top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p Sampling")
             clear_button = gr.Button("Clear Chat")
 
-        # Chat column
         with gr.Column(scale=4):
             chatbot = gr.Chatbot(
                 bubble_full_width=False,
@@ -117,25 +112,11 @@ with gr.Blocks(css=".gr-chatbot .message { margin: 2px 0 !important; }", title=$
                 label="Example Prompts"
             )
 
-    # Event handlers
     model_selector.change(fn=update_model, inputs=[model_selector], outputs=[])
-
-    msg_input.submit(
-        fn=user,
-        inputs=[msg_input, chatbot],
-        outputs=[msg_input, chatbot],
-        queue=False
-    ).then(
-        fn=bot,
-        inputs=[chatbot, max_tokens, temperature, top_p],
-        outputs=[chatbot]
-    )
-    clear_button.click(
-        fn=lambda: [],
-        inputs=None,
-        outputs=chatbot,
-        queue=False
+    msg_input.submit(fn=user, inputs=[msg_input, chatbot], outputs=[msg_input, chatbot], queue=False).then(
+        fn=bot, inputs=[chatbot, max_tokens, temperature, top_p], outputs=[chatbot]
     )
+    clear_button.click(fn=lambda: None, inputs=None, outputs=chatbot, queue=False)
 
 if __name__ == "__main__":
     demo.queue(default_concurrency_limit=10, max_size=10).launch(ssr_mode=False, max_threads=10)