Update app.py

app.py

@@ -1,4 +1,4 @@
-# llada_app.py -> dream_app.py
+# dream_app.py (Updated)
 
 import torch
 import numpy as np
@@ -32,21 +32,32 @@ tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 print("Model and tokenizer loaded.")
 
 # --- Constants for DREAM ---
-# Find the mask token and ID from the DREAM tokenizer
 if tokenizer.mask_token is None:
-    # Handle cases where a mask token might not be explicitly set
-    # You might need to choose a suitable placeholder or investigate further
-    # For now, let's try adding one if it's missing and check its id
-    # This is speculative and might depend on the specific tokenizer setup
-    print("Warning: Mask token not found in tokenizer. Attempting to add.")
+    print("Warning: Mask token not found in tokenizer. Attempting to add '[MASK]'.")
     tokenizer.add_special_tokens({'mask_token': '[MASK]'})
     model.resize_token_embeddings(len(tokenizer)) # Important if vocab size changed
-    if tokenizer.mask_token is None:
-         raise ValueError("Could not set a mask token for the tokenizer.")
+    if tokenizer.mask_token is None or tokenizer.mask_token_id is None:
+         raise ValueError("Could not set or find ID for a mask token for the tokenizer.")
 
 MASK_TOKEN = tokenizer.mask_token
 MASK_ID = tokenizer.mask_token_id
+EOS_TOKEN = tokenizer.eos_token # Get EOS token string
+EOS_ID = tokenizer.eos_token_id # Get EOS token ID
+# Add other special tokens if needed for visualization
+SPECIAL_TOKENS_MAP = {
+    tokenizer.eos_token_id: "[EOS]",
+    tokenizer.bos_token_id: "[BOS]",
+    tokenizer.pad_token_id: "[PAD]",
+    tokenizer.unk_token_id: "[UNK]",
+    MASK_ID: MASK_TOKEN # Map mask ID back to its string representation
+}
+# Add None key to handle cases where token IDs might be None (shouldn't happen with tensors)
+SPECIAL_TOKENS_MAP[None] = "[NONE]"
+
+
 print(f"Using MASK_TOKEN='{MASK_TOKEN}' with ID={MASK_ID}")
+print(f"Using EOS_TOKEN='{EOS_TOKEN}' with ID={EOS_ID}")
+
 # --- Helper Functions (Constraint Parsing, History Formatting) ---
 
 def parse_constraints(constraints_text):
@@ -136,6 +147,7 @@ def dream_generate_response_with_visualization(
     print("Processing constraints:")
     for pos, word in constraints.items():
         # Prepend space for consistent tokenization, similar to LLaDA example
+        # Important: use add_special_tokens=False for constraints
         tokens = tokenizer.encode(" " + word, add_special_tokens=False)
         if not tokens:
             print(f"  Warning: Could not tokenize constraint word '{word}' at position {pos}. Skipping.")
@@ -149,7 +161,6 @@ def dream_generate_response_with_visualization(
                  print(f"  Warning: Overlapping constraint at position {pos+i}. Keeping first.")
 
     # Prepare the prompt using chat template
-    # Note: DREAM examples use add_generation_prompt=True
     try:
         inputs = tokenizer.apply_chat_template(
             messages,
@@ -161,17 +172,10 @@ def dream_generate_response_with_visualization(
         attention_mask = inputs.attention_mask.to(device=device) # Get attention mask
         prompt_length = input_ids.shape[1]
         print(f"Input prompt length: {prompt_length}")
-        print(f"Input IDs: {input_ids}")
+        # print(f"Input IDs: {input_ids}") # Keep commented unless debugging
     except Exception as e:
         print(f"Error applying chat template: {e}")
-        # Provide a fallback or raise the error
-        # Fallback: Simple concatenation (less ideal for instruction models)
-        # chat_input = "".join([f"{msg['role']}: {msg['content']}\n" for msg in messages]) + "assistant:"
-        # input_ids = tokenizer(chat_input, return_tensors="pt").input_ids.to(device)
-        # attention_mask = torch.ones_like(input_ids)
-        # prompt_length = input_ids.shape[1]
-        # print(f"Warning: Using basic concatenation due to template error. Prompt length: {prompt_length}")
-        return [([("Error applying chat template.", "red")],)], f"Error: {e}"
+        return [([("Error applying chat template.", "Error")],)], f"Error: {e}"
 
 
     if prompt_length + gen_length > 2048: # Check context length (DREAM uses 2048)
@@ -179,7 +183,7 @@ def dream_generate_response_with_visualization(
          gen_length = 2048 - prompt_length
          if gen_length <= 0:
              print("Error: Prompt is already too long.")
-             return [([("Prompt too long.", "red")],)], "Error: Prompt too long."
+             return [([("Prompt too long.", "Error")],)], "Error: Prompt too long."
 
 
     # --- State for Visualization Hook ---
@@ -192,74 +196,80 @@ def dream_generate_response_with_visualization(
     for pos, token_id in processed_constraints.items():
         absolute_pos = pos # Position relative to start of generation
         if 0 <= absolute_pos < gen_length:
-            initial_x_part[0, absolute_pos] = token_id
-
-    initial_state_vis = []
-    for i in range(gen_length):
-        token_id = initial_x_part[0, i].item()
-        if token_id == MASK_ID:
-            initial_state_vis.append((MASK_TOKEN, "#444444")) # Mask color
-        else:
-            # This must be a constraint applied initially
-            token_str = tokenizer.decode([token_id], skip_special_tokens=True)
-            initial_state_vis.append((token_str if token_str else "?", "#800080")) # Constraint color (purple)
-    visualization_states.append(initial_state_vis)
+             # Check if the constraint token itself is special
+             if token_id in SPECIAL_TOKENS_MAP:
+                 print(f"  Note: Constraint at pos {pos} is a special token: {SPECIAL_TOKENS_MAP[token_id]}")
+             initial_x_part[0, absolute_pos] = token_id
+
 
     # --- Define the Hook Function ---
+    # This function will be called at each diffusion step
     def generation_tokens_hook_func(step, x, logits):
         nonlocal last_x, visualization_states # Allow modification of outer scope variables
-        print(f"Hook called for step {step}")
+        # print(f"Hook called for step {step}") # Keep commented unless debugging
 
-        current_x = x.clone() # Work on a copy for comparison
+        current_x = x.clone() # Work on a copy for comparison/modification
 
-        # 1. Apply Constraints *before* generating visualization
+        # 1. Apply Constraints *before* generating visualization for this step
         # Constraints are relative to the start of the *generated* part
         constrained_x = current_x.clone()
-        prompt_len = current_x.shape[1] - gen_length # Recalculate just in case
-        if prompt_len < 0:
+        current_prompt_len = current_x.shape[1] - gen_length # Recalculate actual prompt length
+        if current_prompt_len < 0:
             print("Warning: prompt_len negative in hook, skipping constraints/vis.")
             return current_x # Return unmodified if something is wrong
 
-        constraints_applied_this_step = False
         for pos, token_id in processed_constraints.items():
-            absolute_pos = prompt_len + pos
-            if prompt_len <= absolute_pos < current_x.shape[1]:
+            absolute_pos = current_prompt_len + pos
+            if current_prompt_len <= absolute_pos < current_x.shape[1]:
+                # Apply constraint if the current token doesn't match
                 if constrained_x[0, absolute_pos] != token_id:
                     constrained_x[0, absolute_pos] = token_id
-                    constraints_applied_this_step = True
                     # print(f"  Constraint applied at pos {pos} ({absolute_pos}) -> token {token_id}")
 
 
         # 2. Generate Visualization State for *this* step
+        # Compare current_x (output of diffusion for this step, before constraints applied *in this call*)
+        # with last_x (state from *previous* hook call / initial state, *after* constraints were applied then)
         current_state_vis = []
-        # Compare current_x (before explicit constraint application in *this* hook call)
-        # with last_x (state from *previous* hook call / initial state)
-        # Generate based on the state *before* reapplying constraints here,
-        # but *after* the model's diffusion step determined current_x.
-        gen_part_current = current_x[0, prompt_len:]
-        gen_part_last = last_x[0, prompt_len:] if last_x is not None else None
+        gen_part_current = current_x[0, current_prompt_len:]
+        gen_part_last = last_x[0, current_prompt_len:] if last_x is not None else None
 
         for i in range(gen_length):
             current_token_id = gen_part_current[i].item()
-            token_str = tokenizer.decode([current_token_id], skip_special_tokens=True).strip()
-            # Use a placeholder if decoding results in empty string
-            display_token = token_str if token_str else MASK_TOKEN if current_token_id == MASK_ID else "?"
+            last_token_id = gen_part_last[i].item() if gen_part_last is not None else MASK_ID # Assume mask initially
 
-            # Check if this position is constrained
+            # Determine display string - Handle special tokens explicitly
+            if current_token_id in SPECIAL_TOKENS_MAP:
+                display_token = SPECIAL_TOKENS_MAP[current_token_id]
+            else:
+                # Decode non-special tokens, skipping special tokens in the *output string*
+                # and stripping whitespace
+                display_token = tokenizer.decode([current_token_id],
+                                                 skip_special_tokens=True,
+                                                 clean_up_tokenization_spaces=True).strip()
+                # If decoding results in empty string for a non-special token, use a space perhaps
+                if not display_token:
+                    display_token = " " # Use a single space as placeholder
+
+
+            # Determine category (label) for color mapping
+            category = "Old" # Default assume it was revealed before
             is_constrained = i in processed_constraints
 
             if current_token_id == MASK_ID:
-                color = "#444444" # Dark Gray for masks
+                category = "Mask"
             elif is_constrained and processed_constraints[i] == current_token_id:
-                 color = "#800080" # Purple for correctly constrained tokens
-            elif gen_part_last is None or gen_part_last[i].item() == MASK_ID:
-                # Newly revealed (was mask in previous step or initial state)
-                color = "#66CC66" # Light Green
-            else:
-                # Previously revealed and not masked
-                color = "#6699CC" # Light Blue
+                # Check if it was *just* constrained or already was correct
+                # We mark as 'Constraint' if it matches the required token, regardless of when it appeared
+                 category = "Constraint"
+            elif last_token_id == MASK_ID and current_token_id != MASK_ID:
+                # It was a mask before, now it's not -> Newly revealed
+                # (Unless it's a constraint, handled above)
+                 category = "New"
+            # else: category remains "Old"
+
 
-            current_state_vis.append((display_token, color))
+            current_state_vis.append((display_token, category))
 
         visualization_states.append(current_state_vis)
 
@@ -268,7 +278,6 @@ def dream_generate_response_with_visualization(
         last_x = constrained_x.clone()
 
         # 4. Return the sequence with constraints applied for the model's next step
-        # print(f"Hook returning constrained_x: {constrained_x[:, prompt_len:]}")
         return constrained_x # Return the sequence with constraints enforced
 
 
@@ -277,8 +286,30 @@ def dream_generate_response_with_visualization(
         print("Calling model.diffusion_generate...")
         # Make sure last_x is initialized correctly before the first hook call
         # It should represent the state *before* the first diffusion step.
+        # Create the initial full sequence (prompt + initial masked/constrained part)
         initial_full_x = torch.cat([input_ids, initial_x_part], dim=1)
-        last_x = initial_full_x.clone()
+        last_x = initial_full_x.clone() # Initialize last_x with the state before step 0
+
+        # Add the very first visualization state (prompt + initial masks/constraints)
+        # This state corresponds to the `last_x` *before* the first hook call.
+        initial_state_vis = []
+        initial_gen_part = initial_full_x[0, prompt_length:]
+        for i in range(gen_length):
+             token_id = initial_gen_part[i].item()
+             category = "Mask"
+             display_token = MASK_TOKEN
+             if token_id != MASK_ID:
+                 # This must be an initial constraint
+                 category = "Constraint"
+                 if token_id in SPECIAL_TOKENS_MAP:
+                     display_token = SPECIAL_TOKENS_MAP[token_id]
+                 else:
+                     display_token = tokenizer.decode([token_id], skip_special_tokens=True).strip()
+                     if not display_token: display_token = " " # Placeholder
+
+             initial_state_vis.append((display_token, category))
+        visualization_states.append(initial_state_vis)
+
 
         output = model.diffusion_generate(
             input_ids,
@@ -296,63 +327,55 @@ def dream_generate_response_with_visualization(
         print("model.diffusion_generate finished.")
 
         # Extract final generated sequence (response part only)
-        # The hook ensures the returned sequence has constraints applied
         final_sequence = output.sequences[0]
         response_token_ids = final_sequence[prompt_length:]
 
-        # Decode the final response
+        # Decode the final response, skipping special tokens for the final output text
         final_text = tokenizer.decode(
             response_token_ids,
             skip_special_tokens=True,
-            clean_up_tokenization_spaces=True # Recommended for cleaner output
+            clean_up_tokenization_spaces=True
         ).strip()
         print(f"Final generated text: {final_text}")
 
-        # Add the very final state to visualization if the hook didn't capture it
-        # (Should be captured, but as a safeguard)
-        if len(visualization_states) <= steps: # Hook might run 'steps' times
-             final_state_vis = []
-             final_gen_part = final_sequence[prompt_length:]
-             for i in range(gen_length):
-                 token_id = final_gen_part[i].item()
-                 token_str = tokenizer.decode([token_id], skip_special_tokens=True).strip()
-                 display_token = token_str if token_str else MASK_TOKEN if token_id == MASK_ID else "?"
-                 is_constrained = i in processed_constraints
-
-                 if token_id == MASK_ID: color = "#444444"
-                 elif is_constrained and processed_constraints[i] == token_id: color = "#800080"
-                 else: color = "#6699CC" # Default to blue for final state tokens
-                 final_state_vis.append((display_token, color))
-             visualization_states.append(final_state_vis)
+        # The hook should have added the last state, no need for safeguard typically
 
 
     except Exception as e:
         print(f"Error during generation: {e}")
         import traceback
         traceback.print_exc()
-        # Add error message to visualization
+        # Add error message to visualization using the "Error" category
         error_msg = f"Error during generation: {str(e)}"
-        visualization_states.append([("Error", "red")])
+        visualization_states.append([("Error", "Error")]) # Use 'Error' category
         final_text = f"Generation failed: {e}"
 
     print("--- DREAM Generation Finished ---")
+    # Return states list (already built by hook) and final text
     return visualization_states, final_text
 
 
 # --- Gradio UI Setup ---
 
 css = '''
-.category-legend{display:none}
-/* button{height: 60px} */ /* Optional: Adjust button height */
+/* Hide the default legend */
+.gradio-container .output-markdown table { display: none !important; }
+
 .small_btn {
     max-width: 100px; /* Adjust as needed */
+    min-width: 60px;  /* Ensure button doesn't collapse */
     height: 40px;    /* Adjust as needed */
-    flex-grow: 0;  /* Prevent button from growing */
-    margin-left: 5px; /* Add some space */
+    flex-grow: 0 !important;  /* Prevent button from growing */
+    margin-left: 5px !important; /* Add some space */
+    margin-top: auto; /* Align button bottom with textbox */
+    margin-bottom: auto; /* Align button bottom with textbox */
+    line-height: 1; /* Adjust line height if text vertical align is off */
+    padding: 0 10px; /* Adjust padding */
 }
 .chat-input-row {
     display: flex;
     align-items: center; /* Vertically align items */
+    margin-bottom: 10px; /* Add space below input row */
 }
 .chat-input-row > * {
     margin-right: 5px; /* Space between textbox and button */
@@ -360,7 +383,43 @@ css = '''
 .chat-input-row > *:last-child {
     margin-right: 0;
 }
+/* Style HighlightedText elements */
+.token-hl span {
+    padding: 2px 1px; /* Minimal padding */
+    margin: 0 1px; /* Minimal margin */
+    border-radius: 3px;
+    display: inline-block; /* Ensure background covers token */
+    line-height: 1.2; /* Adjust for better vertical spacing */
+}
+/* Custom legend styling */
+.custom-legend span {
+    display: inline-block;
+    margin-right: 15px;
+    font-size: 0.9em;
+}
+.custom-legend span::before {
+    content: "■";
+    margin-right: 4px;
+    font-size: 1.1em; /* Make square slightly larger */
+    vertical-align: middle; /* Align square with text */
+}
 '''
+# Define color map mapping CATEGORY names to colors
+color_map = {
+    "Mask": "#A0A0A0",      # Darker Gray for masks
+    "New": "#77DD77",       # Light Green for new tokens
+    "Old": "#AEC6CF",       # Light Blue/Gray for old tokens
+    "Constraint": "#C3A0E0", # Purple for constraints
+    "Error": "#FF6961"      # Light Red for errors
+}
+
+# Create the custom legend HTML string
+legend_html = "<div class='custom-legend'>"
+for category, color in color_map.items():
+    legend_html += f"<span style='color:{color};'>{category}</span>"
+legend_html += "</div>"
+
+
 def create_chatbot_demo():
     with gr.Blocks(css=css) as demo:
         gr.Markdown("# Dream 7B - Diffusion Language Model Demo")
@@ -376,7 +435,7 @@ def create_chatbot_demo():
                 chatbot_ui = gr.Chatbot(
                     label="Conversation",
                     height=500,
-                    bubble_full_width=False # Improves layout for shorter messages
+                    bubble_full_width=False
                  )
 
                 # Message input Row
@@ -384,8 +443,8 @@ def create_chatbot_demo():
                         user_input = gr.Textbox(
                             label="Your Message",
                             placeholder="Type your message here and press Enter...",
-                            scale=4, # Give textbox more space
-                            container=False, # Remove container background/padding
+                            scale=4,
+                            container=False,
                             show_label=False
                         )
                         send_btn = gr.Button("Send", scale=1, elem_classes="small_btn")
@@ -394,41 +453,34 @@ def create_chatbot_demo():
                     label="Word Constraints (Optional)",
                     info="Force specific words at positions (0-indexed from response start). Format: 'pos:word, pos:word'. Example: '0:Once, 5:upon, 10:time'",
                     placeholder="e.g., 0:Hello, 6:world",
-                    value="" # Default empty
+                    value=""
                 )
             with gr.Column(scale=2):
                 output_vis = gr.HighlightedText(
                     label="Denoising Process Visualization",
-                    combine_adjacent=False,
-                    show_legend=False, # Keep legend off as requested
-                    # Color map for legend (though hidden)
-                    # color_map={
-                    #     "Mask": "#444444",
-                    #     "New": "#66CC66",
-                    #     "Old": "#6699CC",
-                    #     "Constraint": "#800080",
-                    #     "Error": "red"
-                    # }
-                )
-                gr.Markdown(
-                    "**Color Legend:** <span style='color:#444444'>■ Mask</span> | <span style='color:#66CC66'>■ Newly Generated</span> | <span style='color:#6699CC'>■ Previously Generated</span> | <span style='color:#800080'>■ Constraint</span>"
+                    combine_adjacent=False, # Keep tokens separate
+                    show_legend=False,      # Hide default legend table
+                    color_map=color_map,    # Provide the color map
+                    elem_classes="token-hl" # Add class for token styling
                 )
+                # Use Markdown to display the custom legend
+                gr.Markdown(legend_html)
 
 
         # Advanced generation settings
         with gr.Accordion("Generation Settings", open=False):
             with gr.Row():
                 gen_length = gr.Slider(
-                    minimum=16, maximum=512, value=128, step=8, # Increased max length
+                    minimum=16, maximum=512, value=128, step=8,
                     label="Max New Tokens"
                 )
                 steps = gr.Slider(
-                    minimum=8, maximum=512, value=128, step=8, # Increased max steps
+                    minimum=8, maximum=512, value=128, step=8,
                     label="Diffusion Steps"
                 )
             with gr.Row():
                 temperature = gr.Slider(
-                    minimum=0.0, maximum=1.5, value=0.6, step=0.05, # Wider range for temp
+                    minimum=0.0, maximum=1.5, value=0.6, step=0.05,
                     label="Temperature"
                 )
                 top_p = gr.Slider(
@@ -436,15 +488,14 @@ def create_chatbot_demo():
                     label="Top-P (Nucleus Sampling)"
                 )
             with gr.Row():
-                 # Map UI choices to DREAM's alg parameters
                 remasking_strategy = gr.Radio(
                     choices=[
-                        ("Random", "origin"), # User friendly name -> actual param
+                        ("Random", "origin"),
                         ("Entropy", "entropy"),
                         ("MaskGit+", "maskgit_plus"),
                         ("TopK Margin", "topk_margin"),
                     ],
-                    value="entropy", # Default
+                    value="entropy",
                     label="Generation Order Strategy (alg)"
                 )
                 alg_temp = gr.Slider(
@@ -462,9 +513,6 @@ def create_chatbot_demo():
         # Clear button
         clear_btn = gr.Button("Clear Conversation")
 
-        # Hidden textbox to potentially store intermediate response (might not be needed)
-        # current_response = gr.Textbox(visible=False)
-
         # --- Event Handlers ---
 
         # Helper to add message to history state
@@ -478,22 +526,12 @@ def create_chatbot_demo():
             print(f"User submitted: '{message}'")
             if not message or not message.strip():
                  print("Empty message submitted, doing nothing.")
-                 # Return unchanged state if message is empty
-                 # Need to return values for all outputs of the .submit/.click
                  return history, history, "", [] # history, chatbot_ui, user_input, output_vis
 
-            # Add user message to history (with None for bot response initially)
             history = add_message_to_history(history, message, None)
-
-            # Prepare updated history for display in Chatbot UI
             history_for_display = history.copy()
-
-            # Clear the input textbox
             message_out = ""
-             # Clear the visualization
-            vis_clear = []
-
-            # Return updated history state, chatbot display, cleared input, cleared visualization
+            vis_clear = [] # Clear visualization when new message submitted
             return history, history_for_display, message_out, vis_clear
 
         # Function to generate bot response (triggered after user message is processed)
@@ -504,18 +542,13 @@ def create_chatbot_demo():
             print("--- Generating Bot Response ---")
             if not history or history[-1][1] is not None:
                 print("History empty or last message already has response. Skipping generation.")
-                # Yield current state if called unnecessarily
-                yield history, [], "No response generated."
+                yield history, [], "No response generated." # Yield current state if called unnecessarily
                 return
 
-            # Get the conversation history in the format the model expects
-            messages = format_chat_history(history) # Includes the latest user query
-
-            # Parse constraints from the textbox
+            messages = format_chat_history(history)
             parsed_constraints = parse_constraints(constraints_text)
 
             try:
-                # Generate response with visualization
                 vis_states, response_text = dream_generate_response_with_visualization(
                     messages,
                     gen_length=gen_length,
@@ -527,31 +560,30 @@ def create_chatbot_demo():
                     alg_temp=alg_temp
                 )
 
-                # Update the history state with the final bot response
-                history[-1][1] = response_text.strip()
-
-                # Yield the initial visualization state immediately
-                if vis_states:
-                    yield history, vis_states[0] # Update chatbot, update visualization
-                else:
-                    # Handle case where generation failed before first state
-                    yield history, [("Generation failed.", "red")]
+                # Update the history state only ONCE with the final bot response
+                final_history = history.copy() # Create copy to modify
+                final_history[-1][1] = response_text.strip() # Update the last element
 
-                # Then animate through the rest of the visualization states
-                for state in vis_states[1:]:
-                    time.sleep(delay)
-                    yield history, state # Update chatbot (implicitly via history), update visualization
+                # Yield visualization states one by one
+                # Important: Yield the *original* history for all intermediate steps,
+                # only yield the final_history with the *last* visualization state.
+                num_states = len(vis_states)
+                for i, state in enumerate(vis_states):
+                    current_chatbot_state = history if i < num_states - 1 else final_history
+                    yield current_chatbot_state, state
+                    if delay > 0 and i < num_states - 1: # Don't sleep after last state
+                        time.sleep(delay)
 
             except Exception as e:
                 print(f"Error in bot_response_generator: {e}")
                 import traceback
                 traceback.print_exc()
                 error_msg = f"Error: {str(e)}"
-                 # Show error in visualization
-                error_vis = [(error_msg, "red")]
-                 # Update history with error message? Optional.
-                # history[-1][1] = error_msg
-                yield history, error_vis
+                error_vis = [(error_msg, "Error")] # Use Error category
+                # Update history with error message? Optional.
+                final_history_error = history.copy()
+                final_history_error[-1][1] = error_msg # Add error to chatbot too
+                yield final_history_error, error_vis
 
         # Function to clear everything
         def clear_conversation():
@@ -561,34 +593,39 @@ def create_chatbot_demo():
         # --- Wire UI elements to functions ---
 
         # Typing in Textbox and pressing Enter
-        user_input.submit(
+        submit_event = user_input.submit(
             fn=user_message_submitted,
             inputs=[user_input, chat_history],
-            outputs=[chat_history, chatbot_ui, user_input, output_vis], # Update history state, chatbot display, clear input, clear vis
-            queue=False # Process immediately
-        ).then(
-            fn=bot_response_generator,
-            inputs=[
-                chat_history, gen_length, steps, constraints_input, visualization_delay,
-                temperature, top_p, remasking_strategy, alg_temp
-            ],
-            outputs=[chatbot_ui, output_vis] # Update chatbot display (with new response), update visualization
-            # Note: history state is updated implicitly by bot_response_generator modifying its input
+            outputs=[chat_history, chatbot_ui, user_input, output_vis],
+            queue=False # Show user message immediately
         )
 
         # Clicking the Send button
-        send_btn.click(
+        click_event = send_btn.click(
             fn=user_message_submitted,
             inputs=[user_input, chat_history],
             outputs=[chat_history, chatbot_ui, user_input, output_vis],
             queue=False
-        ).then(
-            fn=bot_response_generator,
-            inputs=[
+        )
+
+        # Chain the generation after user message is processed (for both submit and click)
+        # Use .then() to trigger the generator
+        generation_inputs = [
                 chat_history, gen_length, steps, constraints_input, visualization_delay,
                 temperature, top_p, remasking_strategy, alg_temp
-            ],
-            outputs=[chatbot_ui, output_vis]
+            ]
+        generation_outputs = [chatbot_ui, output_vis]
+
+        submit_event.then(
+            fn=bot_response_generator,
+            inputs=generation_inputs,
+            outputs=generation_outputs
+        )
+
+        click_event.then(
+            fn=bot_response_generator,
+            inputs=generation_inputs,
+            outputs=generation_outputs
         )
 
         # Clicking the Clear button