multiturn chat

app.py

@@ -1,8 +1,7 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import time
-import random
 from threading import Thread
 import sys
 import os
@@ -45,19 +44,19 @@ def chat_predict(message, history, max_length, temperature, top_p, repetition_pe
         yield "⚠️ Please wait for the model to finish loading..."
         return
     
-    if not message.strip():
-        yield "⚠️ Please enter a prompt."
-        return
-    
     try:
-        # Build conversation history
+        # Prepare messages for chat template
         messages = []
         if system_prompt:
             messages.append({"role": "system", "content": system_prompt})
             
-        for user_msg, assistant_msg in history:
-            messages.append({"role": "user", "content": user_msg})
-            messages.append({"role": "assistant", "content": assistant_msg})
+        # history is a list of dicts: [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
+        for msg in history:
+            # Clean up history messages (remove stats if they were appended)
+            content = msg["content"]
+            if "\n\n---\n*Generated" in content:
+                content = content.split("\n\n---\n*Generated")[0]
+            messages.append({"role": msg["role"], "content": content})
             
         messages.append({"role": "user", "content": message})
         
@@ -89,13 +88,25 @@ def chat_predict(message, history, max_length, temperature, top_p, repetition_pe
         generated_text = ""
         start_time = time.time()
         token_count = 0
+        last_update_time = start_time
+        current_stats = ""
         
         for new_text in streamer:
             generated_text += new_text
             token_count += 1
-            yield generated_text
+            
+            # Update stats every 0.2 seconds
+            current_time = time.time()
+            if current_time - last_update_time > 0.2:
+                elapsed = current_time - start_time
+                if elapsed > 0:
+                    tps = token_count / elapsed
+                    current_stats = f"\n\n---\n*Generating... ({tps:.1f} t/s)*"
+                last_update_time = current_time
+                
+            yield generated_text + current_stats
 
-        # Append stats after generation is complete
+        # Final stats
         elapsed_time = time.time() - start_time
         if elapsed_time > 0:
             tps = token_count / elapsed_time
@@ -121,7 +132,7 @@ custom_theme = gr.themes.Soft(
 )
 
 # Build the Gradio interface
-with gr.Blocks(theme=custom_theme) as demo:
+with gr.Blocks(theme=custom_theme, fill_height=True) as demo:
     gr.Markdown(
         """
         # 🤖 Smol LLM Chat
@@ -133,6 +144,8 @@ with gr.Blocks(theme=custom_theme) as demo:
     # Chat Interface
     chat_interface = gr.ChatInterface(
         fn=chat_predict,
+        type="messages",
+        fill_height=True,
         additional_inputs=[
             gr.Slider(
                 minimum=50,
@@ -175,8 +188,9 @@ with gr.Blocks(theme=custom_theme) as demo:
 load_status = load_model()
 print(f"Startup load status: {load_status}")
 
-# Launch the application
-demo.launch(
-    share=False,
-    show_error=True
-)
+if __name__ == "__main__":
+    # Launch the application
+    demo.launch(
+        share=False,
+        show_error=True
+    )