Update app.py

app.py

@@ -531,8 +531,8 @@ def create_chatbot_demo():
             "[[Blog](https://hkunlp.github.io/blog/2025/dream/)]" # Note: Link might be hypothetical
         )
 
-        # STATE MANAGEMENT
-        _chat_history_store = gr.State([]) # Hidden state to store actual history list
+        # STATE MANAGEMENT: Use a hidden state variable to store the actual list
+        _chat_history_store = gr.State([])
 
         # UI COMPONENTS
         with gr.Row():
@@ -565,7 +565,7 @@ def create_chatbot_demo():
                     label="Denoising Process Visualization",
                     combine_adjacent=False,
                     show_legend=True,
-                    interactive=False,
+                    interactive=False
                 )
                 response_text_display = gr.Textbox(
                     label="Generated Response",
@@ -582,9 +582,8 @@ def create_chatbot_demo():
                 temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.05, label="Temperature (0 = greedy)")
                 alg_temp = gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.05, label="Remasking Temp (Confidence Algs)")
              with gr.Row():
-                # Adjusted label for clarity on disabling top_p
-                top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, label="Top-P (>0 & <1 to enable)")
-                top_k = gr.Slider(minimum=0, maximum=200, value=0, step=5, label="Top-K (>0 to enable)")
+                top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, label="Top-P (0 disables)")
+                top_k = gr.Slider(minimum=0, maximum=200, value=0, step=5, label="Top-K (0 disables)")
              with gr.Row():
                  remasking_strategy = gr.Radio(choices=['origin', 'maskgit_plus', 'topk_margin', 'entropy'], value='entropy', label="Remasking Strategy (Algorithm)")
              with gr.Row():
@@ -596,78 +595,72 @@ def create_chatbot_demo():
         # --- Event Handlers ---
 
         def add_user_message_to_history(message: str, history_store: List[List[Optional[str]]]):
-            """Adds user message TO STATE, clears input, prepares for bot response."""
+            """Adds user message, clears input, prepares for bot response."""
             if not message.strip():
                 gr.Warning("Please enter a message.")
-                # Return unchanged state, but clear inputs/outputs for next step
-                # Outputs: _chat_history_store, user_input, output_vis, response_text_display
-                return history_store, message, [], "" # Return original message to keep it in input if invalid
-
-            # Add user message with placeholder for bot response TO THE STATE
-            history_store.append([message.strip(), None]) # Ensure message is stripped
-            # Return updated history store, clear input box, clear vis, clear response text
-            # Outputs: _chat_history_store, user_input, output_vis, response_text_display
-            return history_store, "", [], "" # Clear user_input only on success
+                # Return unchanged history store, history for UI, empty input, empty vis, empty response
+                return history_store, history_store, "", [], ""
+            # Append user message with placeholder for bot response
+            history_store.append([message, None])
+            # Return updated history store for the state, history for the UI, clear other fields
+            return history_store, history_store, "", [], ""
 
         def clear_conversation():
-            """Clears the chat history state and UI elements."""
-            # Outputs: _chat_history_store, chatbot_ui, user_input, output_vis, response_text_display
-            return [], [], "", [], "" # Clear everything
-
+            """Clears the chat history, visualization, and response text."""
+            # Return empty values for all relevant outputs
+            return [], [], "", [], "" # History store, chatbot UI, input, vis, response text
 
         # --- Connect UI elements ---
 
-        # Inputs for the generation function
+        # Define the inputs for the generation function once
         generation_inputs = [
             _chat_history_store, gen_length, steps, constraints_input,
             temperature, top_p, top_k, remasking_strategy, alg_temp,
             visualization_delay
         ]
-        # Outputs for the generation function (yields history, vis_data, text)
-        generation_outputs = [chatbot_ui, output_vis, response_text_display]
-
-        # Outputs for add_user_message_to_history
-        add_message_outputs = [
-            _chat_history_store, # Update state
-            user_input,          # Clear input (or return original if invalid)
-            output_vis,          # Clear visualization
-            response_text_display # Clear response text
-        ]
+        # Define the outputs for the generation function's yield values
+        # Yields: history_list, vis_data, response_text
+        # CRITICAL FIX: Map the first yielded value (history_list) back to _chat_history_store
+        # as well as chatbot_ui.
+        generation_outputs = [_chat_history_store, chatbot_ui, output_vis, response_text_display]
 
         # Handle Textbox Submission (Enter key)
         submit_listener = user_input.submit(
             fn=add_user_message_to_history,
             inputs=[user_input, _chat_history_store],
-            outputs=add_message_outputs, # Step 1: Update state, clear inputs/vis/response
-            queue=True # Ensure intermediate steps are processed
-        ).then(
+            # Outputs of step 1: update store, UI, clear inputs/outputs
+            outputs=[_chat_history_store, chatbot_ui, user_input, output_vis, response_text_display]
+        )
+        # Chain the bot response generation after the user message is added
+        submit_listener.then(
             fn=generate_dream_response,
-            inputs=generation_inputs, # Takes the updated state
-            outputs=generation_outputs, # Step 2: Generate response and stream history/vis/text to UI
-            show_progress="hidden", # Hide default progress as we have live vis
-            queue=True # Ensure generation runs in the queue
+            inputs=generation_inputs, # Passes the updated _chat_history_store
+            # Outputs of step 2: **UPDATE STORE**, UI, vis, response text
+            outputs=generation_outputs, # Maps the final yielded history back to the store
+            show_progress="hidden"
         )
 
         # Handle Send Button Click
         click_listener = send_btn.click(
             fn=add_user_message_to_history,
             inputs=[user_input, _chat_history_store],
-            outputs=add_message_outputs, # Step 1: Update state, clear inputs/vis/response
-            queue=True # Ensure intermediate steps are processed
-        ).then(
+            outputs=[_chat_history_store, chatbot_ui, user_input, output_vis, response_text_display]
+        )
+        # Chain the bot response generation after the user message is added
+        click_listener.then(
             fn=generate_dream_response,
-            inputs=generation_inputs, # Takes the updated state
-            outputs=generation_outputs, # Step 2: Generate response and stream history/vis/text to UI
-            show_progress="hidden", # Hide default progress as we have live vis
-            queue=True # Ensure generation runs in the queue
+            inputs=generation_inputs,
+            # Outputs of step 2: **UPDATE STORE**, UI, vis, response text
+            outputs=generation_outputs, # Map final history back to store here too
+            show_progress="hidden"
         )
 
         # Clear Button Action
         clear_btn.click(
             clear_conversation,
             inputs=[],
-            outputs=[_chat_history_store, chatbot_ui, user_input, output_vis, response_text_display],
-            queue=False # Clearing can be immediate
+            # Ensure clear updates the state variable too
+            outputs=[_chat_history_store, chatbot_ui, user_input, output_vis, response_text_display]
         )
 
     return demo