Update app.py

app.py

@@ -41,34 +41,45 @@ def check_model_status():
         load_model()
     return model_loaded
 
-def chat(message, history):
+def chat(message, history, temperature, max_new_tokens):
     global model_pipeline
     
     # Ensure the model is loaded before proceeding
     if not check_model_status():
-        return "Model is not ready. Please try again later."
+        yield "Model is not ready. Please try again later."
+        return
     
     prompt = f"Human: {message}\n\nAssistant:"
     
-    # Generate response using the pre-loaded model
-    response = model_pipeline(
-        prompt,
-        max_new_tokens=2048,
-        temperature=0.7,
-        do_sample=True,
-        truncation=True,
-        pad_token_id=50256
-    )
+    # Stream the response
+    start_time = time.time()
+    generated_tokens = 0
     
-    try:
-        bot_text = response[0]["generated_text"]
-        bot_text = bot_text.split("Assistant:")[-1].strip()
-        if "</think>" in bot_text:
-            bot_text = bot_text.split("</think>")[-1].strip()
-    except Exception as e:
-        bot_text = f"Sorry, there was a problem generating the response: {str(e)}"
+    def generate():
+        nonlocal generated_tokens
+        for response in model_pipeline(
+            prompt,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=True,
+            truncation=True,
+            pad_token_id=50256,
+            return_full_text=False,
+            streamer=True
+        ):
+            bot_text = response[0]["generated_text"]
+            bot_text = bot_text.split("Assistant:")[-1].strip()
+            if "</think>" in bot_text:
+                bot_text = bot_text.split("</think>")[-1].strip()
+            
+            generated_tokens += len(bot_text.split())
+            elapsed_time = time.time() - start_time
+            tokens_per_second = generated_tokens / elapsed_time if elapsed_time > 0 else 0
+            
+            yield bot_text, f"Generating... Tokens/s: {tokens_per_second:.2f}"
     
-    return bot_text
+    for partial_response, status in generate():
+        yield partial_response, status
 
 def reload_model_button():
     """Reload the model manually via a button."""
@@ -82,10 +93,10 @@ def update_status_periodically(status_text):
     while True:
         time.sleep(5)  # Update every 5 seconds
         status = "Model is loaded and ready." if model_loaded else "Model is not loaded."
-        status_text.update(value=status)
+        status_text.value = status  # Update the value directly
 
 # Gradio Interface
-with gr.Blocks() as demo:
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# DeepSeek-R1 Chatbot")
     gr.Markdown("DeepSeek-R1-Distill-Qwen-1.5B 모델을 사용한 대화 테스트용 데모입니다.")
     
@@ -98,14 +109,23 @@ with gr.Blocks() as demo:
         clear_button = gr.Button("Clear")
         reload_button = gr.Button("Reload Model")
     
+    with gr.Row():
+        temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
+        max_tokens_slider = gr.Slider(minimum=32, maximum=2048, value=2048, step=32, label="Max New Tokens")
+    
     status_text = gr.Textbox(label="Model Status", value="Model not loaded yet.", interactive=False)
+    token_status = gr.Textbox(label="Token Generation Status", value="", interactive=False)
     
-    def respond(message, chat_history):
-        bot_message = chat(message, chat_history)
-        chat_history.append((message, bot_message))
-        return "", chat_history
+    def respond(message, chat_history, temperature, max_new_tokens):
+        bot_message = ""
+        status = ""
+        for partial_response, partial_status in chat(message, chat_history, temperature, max_new_tokens):
+            bot_message = partial_response
+            status = partial_status
+            token_status.update(value=status)
+            yield "", chat_history + [(message, bot_message)]
     
-    send_button.click(respond, inputs=[textbox, chatbot], outputs=[textbox, chatbot])
+    send_button.click(respond, inputs=[textbox, chatbot, temperature_slider, max_tokens_slider], outputs=[textbox, chatbot])
     clear_button.click(lambda: [], None, chatbot)
     reload_button.click(reload_model_button, None, status_text)