Update app.py

app.py

@@ -11,10 +11,10 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 phi4_model = AutoModelForCausalLM.from_pretrained(phi4_model_path, device_map="auto", torch_dtype="auto")
 phi4_tokenizer = AutoTokenizer.from_pretrained(phi4_model_path)
 
-@spaces.GPU(duration=120)
-def generate_response(user_message, max_tokens, temperature, top_k, top_p, repetition_penalty, history_state):
+@spaces.GPU(duration=60)
+def generate_response(user_message, max_tokens, temperature, top_k, top_p, repetition_penalty, history):
     if not user_message.strip():
-        return history_state, history_state
+        return history, history
 
     model = phi4_model
     tokenizer = phi4_tokenizer
@@ -32,12 +32,17 @@ Use LaTeX for any formulas or values (e.g., $\\text{BMI} = \\frac{\\text{weight
 
 Now, analyze the following case:"""
 
+    # Build conversation history in the format the model expects
     prompt = f"{start_tag}system{sep_tag}{system_message}{end_tag}"
-    for message in history_state:
-        if message["role"] == "user":
-            prompt += f"{start_tag}user{sep_tag}{message['content']}{end_tag}"
-        elif message["role"] == "assistant" and message["content"]:
-            prompt += f"{start_tag}assistant{sep_tag}{message['content']}{end_tag}"
+    
+    # Convert chat history format from the Gradio Chatbot format to prompt format
+    for user_msg, bot_msg in history:
+        if user_msg:
+            prompt += f"{start_tag}user{sep_tag}{user_msg}{end_tag}"
+        if bot_msg:
+            prompt += f"{start_tag}assistant{sep_tag}{bot_msg}{end_tag}"
+    
+    # Add the current user message
     prompt += f"{start_tag}user{sep_tag}{user_message}{end_tag}{start_tag}assistant{sep_tag}"
 
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
@@ -58,19 +63,19 @@ Now, analyze the following case:"""
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    assistant_response = ""
-    new_history = history_state + [
-        {"role": "user", "content": user_message},
-        {"role": "assistant", "content": ""}
-    ]
+    # Create a new history with the current user message
+    new_history = history.copy() + [[user_message, ""]]
     
+    # Collect the generated response
+    assistant_response = ""
     for new_token in streamer:
         cleaned_token = new_token.replace("<|im_start|>", "").replace("<|im_sep|>", "").replace("<|im_end|>", "")
         assistant_response += cleaned_token
-        new_history[-1]["content"] = assistant_response.strip()
-        yield new_history, new_history
-
-    yield new_history, new_history
+        # Update the last message in history with the current response
+        new_history[-1][1] = assistant_response.strip()
+        
+    # Return the updated history
+    return new_history, new_history
 
 
 example_messages = {
@@ -120,7 +125,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     </script>
     """)
 
-    history_state = gr.State([])
+    chatbot = gr.Chatbot(label="Chat", render_markdown=True, show_copy_button=True)
+    history = gr.State([])
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -133,7 +139,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 repetition_penalty_slider = gr.Slider(1.0, 2.0, value=1.0, label="Repetition Penalty")
 
         with gr.Column(scale=4):
-            chatbot = gr.Chatbot(label="Chat", render_markdown=True, show_copy_button=True)
             with gr.Row():
                 user_input = gr.Textbox(label="Describe symptoms or ask a medical question", placeholder="Type your message here...", scale=3)
                 submit_button = gr.Button("Send", variant="primary", scale=1)
@@ -145,26 +150,19 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 example3 = gr.Button("Abdominal pain")
                 example4 = gr.Button("BMI calculation")
 
-    # Replace the streaming approach with a standard click event
-    def process_input(message, max_tokens, temperature, top_k, top_p, repetition_penalty, history):
-        # Call the generator function and get the final result
-        generator = generate_response(message, max_tokens, temperature, top_k, top_p, repetition_penalty, history)
-        result = None
-        for res in generator:
-            result = res
-        return result
-
+    # Use click instead of stream
     submit_button.click(
-        fn=process_input,
-        inputs=[user_input, max_tokens_slider, temperature_slider, top_k_slider, top_p_slider, repetition_penalty_slider, history_state],
-        outputs=[chatbot, history_state]
+        fn=generate_response,
+        inputs=[user_input, max_tokens_slider, temperature_slider, top_k_slider, top_p_slider, 
+                repetition_penalty_slider, history],
+        outputs=[chatbot, history]
     ).then(
         fn=lambda: gr.update(value=""),
         inputs=None,
         outputs=user_input
     )
 
-    clear_button.click(fn=lambda: ([], []), inputs=None, outputs=[chatbot, history_state])
+    clear_button.click(fn=lambda: ([], []), inputs=None, outputs=[chatbot, history])
 
     example1.click(lambda: gr.update(value=example_messages["Headache case"]), None, user_input)
     example2.click(lambda: gr.update(value=example_messages["Chest pain"]), None, user_input)