Update app.py

app.py

@@ -1,20 +1,159 @@
-# Replace the existing dream_generate_response_with_visualization function
-# in the previous code block with this updated version.
+import torch
+# import numpy as np # Not strictly needed anymore
+import gradio as gr
+import spaces
+from transformers import AutoTokenizer, AutoModel
+import time
+import re # Keep for parsing constraints
+
+# Use try-except for space deployment vs local
+try:
+    # Used for spaces deployment with GPU
+    gpu_check = spaces.GPU
+    print("Running in Gradio Spaces with GPU environment.")
+except AttributeError:
+    # Fallback for local execution or environments without spaces.GPU
+    print("Running in local environment or without spaces.GPU.")
+    # Define a dummy decorator if spaces.GPU is not available
+    def gpu_check(func):
+        return func
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"Using device: {device}")
+
+# --- Load DREAM Model and Tokenizer ---
+# Ensure sufficient VRAM, Dream 7B needs ~16GB+ VRAM in bfloat16
+model_path = "Dream-org/Dream-v0-Instruct-7B"
+print(f"Loading model: {model_path}...")
+try:
+    model = AutoModel.from_pretrained(
+        model_path,
+        torch_dtype=torch.bfloat16, # Use bfloat16 for efficiency
+        trust_remote_code=True,
+        # device_map='auto' # Consider if running into OOM errors, might split across GPUs/CPU
+    ).to(device).eval()
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path,
+        trust_remote_code=True
+    )
+    print("Model and tokenizer loaded successfully.")
+except Exception as e:
+    print(f"Error loading model or tokenizer: {e}")
+    print("Please ensure you have enough GPU memory and the model files are accessible.")
+    # Exit or raise if loading fails
+    raise e
+
+
+# --- Constants for DREAM ---
+# Find the mask token and ID from the DREAM tokenizer
+if tokenizer.mask_token is None:
+    print("Warning: Mask token not found in tokenizer. Attempting to add '[MASK]'.")
+    # This might require retraining or fine-tuning if the model didn't see this token
+    num_added = tokenizer.add_special_tokens({'mask_token': '[MASK]'})
+    if num_added > 0:
+        print(f"Added '{tokenizer.mask_token}' to tokenizer.")
+        # Resize model embeddings if vocab changed
+        model.resize_token_embeddings(len(tokenizer))
+        print("Resized model token embeddings.")
+    else:
+        # Fallback or error if adding failed or mask token still None
+        # It's possible a different token serves this purpose in DREAM's training
+        print("Error: Could not set a mask token. Visualization might be inaccurate.")
+        # You might need to identify which token ID DREAM uses internally for masking tasks
+        # For now, we'll proceed but this is a potential issue.
+        MASK_TOKEN = "<?>" # Placeholder symbol
+        MASK_ID = -1 # Invalid ID indicates issue
+    if tokenizer.mask_token is None:
+         raise ValueError("Could not set a mask token for the tokenizer.")
+
+MASK_TOKEN = tokenizer.mask_token
+MASK_ID = tokenizer.mask_token_id
+print(f"Using MASK_TOKEN='{MASK_TOKEN}' with ID={MASK_ID}")
+
+# Identify other special tokens to potentially hide/show
+eos_token_id = tokenizer.eos_token_id
+pad_token_id = tokenizer.pad_token_id
+special_token_ids_set = {MASK_ID} # Start with Mask ID
+if eos_token_id is not None:
+    special_token_ids_set.add(eos_token_id)
+    print(f"EOS token ID: {eos_token_id} ({tokenizer.decode([eos_token_id])})")
+if pad_token_id is not None:
+    special_token_ids_set.add(pad_token_id)
+    print(f"PAD token ID: {pad_token_id} ({tokenizer.decode([pad_token_id])})")
+# Add other common special tokens if needed (e.g., BOS, UNK)
+if tokenizer.bos_token_id is not None:
+    special_token_ids_set.add(tokenizer.bos_token_id)
+    print(f"BOS token ID: {tokenizer.bos_token_id} ({tokenizer.decode([tokenizer.bos_token_id])})")
+if tokenizer.unk_token_id is not None:
+    special_token_ids_set.add(tokenizer.unk_token_id)
+    print(f"UNK token ID: {tokenizer.unk_token_id} ({tokenizer.decode([tokenizer.unk_token_id])})")
+
+print(f"Identified special token IDs: {special_token_ids_set}")
+
+# --- Helper Functions (Constraint Parsing, History Formatting) ---
+
+def parse_constraints(constraints_text):
+    """Parse constraints in format: 'position:word, position:word, ...'"""
+    constraints = {}
+    if not constraints_text:
+        return constraints
+
+    parts = constraints_text.split(',')
+    for part in parts:
+        part = part.strip() # Trim whitespace
+        if ':' not in part:
+            continue
+        try:
+            pos_str, word = part.split(':', 1)
+            pos = int(pos_str.strip())
+            word = word.strip()
+            # Allow empty words if needed? Forcing empty seems odd. Let's require a word.
+            if word and pos >= 0:
+                constraints[pos] = word
+        except ValueError:
+            print(f"Warning: Could not parse constraint part: '{part}'")
+            continue
+
+    return constraints
+
+def format_chat_history(history):
+    """
+    Format chat history for the DREAM model (standard messages format)
+
+    Args:
+        history: List of [user_message, assistant_message] pairs
 
-@gpu_check
+    Returns:
+        Formatted conversation for the model (list of dictionaries)
+    """
+    messages = []
+    # Check if a system prompt is appropriate for Dream-Instruct
+    # From demo_completion.py example, it seems it uses system prompt via template
+    # messages.append({"role": "system", "content": "You are a helpful assistant."})
+    for user_msg, assistant_msg in history:
+        if user_msg is not None: # Handle potential None message if clearing failed
+             messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:  # Skip if None (for the latest user message awaiting response)
+            messages.append({"role": "assistant", "content": assistant_msg})
+
+    return messages
+
+# --- Core Generation Logic for DREAM with Visualization ---
+
+@gpu_check # Use the potentially dummy decorator
+@torch.no_grad() # Disable gradient calculations for inference
 def dream_generate_response_with_visualization(
     messages,
-    gen_length=64,
-    steps=64,
+    gen_length=128,
+    steps=128, # Default based on DREAM examples
     constraints=None,
-    temperature=0.6,
-    top_p=0.95,
-    alg="entropy",
-    alg_temp=0.0,
+    temperature=0.6, # Default based on DREAM examples
+    top_p=0.95, # Default based on DREAM examples
+    alg="entropy", # Default based on DREAM examples
+    alg_temp=0.1, # Default based on DREAM examples
 ):
     """
-    Generate text with DREAM model with visualization using the generation hook,
-    including special token handling (show once, then hide).
+    Generate text with DREAM model with visualization using the generation hook.
 
     Args:
         messages: List of message dictionaries with 'role' and 'content'
@@ -29,239 +168,606 @@ def dream_generate_response_with_visualization(
     Returns:
         Tuple: (List of visualization states, final generated text string)
     """
-    print("--- Starting DREAM Generation ---")
-    print(f"Parameters: gen_length={gen_length}, steps={steps}, temperature={temperature}, top_p={top_p}, alg='{alg}', alg_temp={alg_temp}")
+    print("\n--- Starting DREAM Generation ---")
+    print(f"Params: len={gen_length}, steps={steps}, temp={temperature}, top_p={top_p}, alg='{alg}', alg_temp={alg_temp}")
     print(f"Constraints: {constraints}")
 
     # --- Input Preparation ---
     if constraints is None:
         constraints = {}
 
+    # Convert word constraints to token IDs (handle multi-token words)
     processed_constraints = {}
+    constraint_token_lengths = {} # Store length for multi-token constraints
     print("Processing constraints:")
     for pos, word in constraints.items():
+        # Prepend space for potentially better tokenization consistency
+        # (though apply_chat_template should handle spacing)
         tokens = tokenizer.encode(" " + word, add_special_tokens=False)
         if not tokens:
             print(f"  Warning: Could not tokenize constraint word '{word}' at position {pos}. Skipping.")
             continue
-        print(f"  Pos {pos}, Word '{word}' -> Tokens {tokens}")
+        print(f"  Pos {pos}, Word '{word}' -> Tokens {tokens} ({tokenizer.convert_ids_to_tokens(tokens)})")
+        constraint_token_lengths[pos] = len(tokens)
         for i, token_id in enumerate(tokens):
-            if pos + i not in processed_constraints:
-                 processed_constraints[pos + i] = token_id
+            target_pos = pos + i
+            if target_pos in processed_constraints:
+                 print(f"  Warning: Overlapping constraint token at position {target_pos}. Keeping first constraint's token ({processed_constraints[target_pos]}).")
             else:
-                 print(f"  Warning: Overlapping constraint at position {pos+i}. Keeping first.")
+                 processed_constraints[target_pos] = token_id
 
+    # Prepare the prompt using chat template
     try:
         inputs = tokenizer.apply_chat_template(
             messages,
             return_tensors="pt",
             return_dict=True,
-            add_generation_prompt=True
+            add_generation_prompt=True # Crucial for Dream-Instruct
         )
         input_ids = inputs.input_ids.to(device=device)
+        # Use the attention mask generated by the template
         attention_mask = inputs.attention_mask.to(device=device)
         prompt_length = input_ids.shape[1]
         print(f"Input prompt length: {prompt_length}")
-        # print(f"Input IDs: {input_ids}") # Verbose
+        # print(f"Input IDs: {input_ids}")
+        # print(f"Attention Mask: {attention_mask}") # Verify mask covers prompt
     except Exception as e:
         print(f"Error applying chat template: {e}")
         return [([("Error applying chat template.", "red")],)], f"Error: {e}"
 
-    if prompt_length + gen_length > 2048:
-         print(f"Warning: Requested length ({prompt_length + gen_length}) exceeds model max length (2048). Truncating gen_length.")
-         gen_length = 2048 - prompt_length
+    # Check context length (DREAM uses 2048 default)
+    model_max_length = getattr(model.config, 'max_position_embeddings', 2048)
+    if prompt_length + gen_length > model_max_length:
+         print(f"Warning: Requested length ({prompt_length + gen_length}) exceeds model max length ({model_max_length}). Truncating gen_length.")
+         gen_length = model_max_length - prompt_length
          if gen_length <= 0:
              print("Error: Prompt is already too long.")
              return [([("Prompt too long.", "red")],)], "Error: Prompt too long."
 
     # --- State for Visualization Hook ---
     visualization_states = []
-    last_x = None
-
-    # Define special token IDs to hide after first reveal
-    # Using a set for efficient lookup. Add others if needed (e.g., pad_token_id).
-    special_token_ids_to_hide = {
-        tokenizer.eos_token_id,
-        tokenizer.pad_token_id,
-        # tokenizer.bos_token_id # Usually not generated mid-sequence
-    }
-    # Filter out None values if some special tokens aren't defined
-    special_token_ids_to_hide = {tid for tid in special_token_ids_to_hide if tid is not None}
-    print(f"Special token IDs to hide visually after reveal: {special_token_ids_to_hide}")
+    last_x = None # Store the full sequence (prompt + generation) from the previous step
 
-
-    initial_x_part = torch.full((1, gen_length), MASK_ID, dtype=torch.long, device=device)
+    # Initial state: Prompt + all masks for generation part
+    initial_gen_part = torch.full((1, gen_length), MASK_ID, dtype=torch.long, device=device)
+    # Apply initial constraints to the masked part *before* the first visualization state
     for pos, token_id in processed_constraints.items():
-        absolute_pos = pos
+        absolute_pos = pos # Position relative to start of generation
         if 0 <= absolute_pos < gen_length:
-            initial_x_part[0, absolute_pos] = token_id
+            initial_gen_part[0, absolute_pos] = token_id
 
+    # Create the first visualization state (only the generation part)
     initial_state_vis = []
     for i in range(gen_length):
-        token_id = initial_x_part[0, i].item()
+        token_id = initial_gen_part[0, i].item()
         if token_id == MASK_ID:
             initial_state_vis.append((MASK_TOKEN, "#444444")) # Mask color
         else:
-            token_str = tokenizer.decode([token_id], skip_special_tokens=True)
+            # This must be a constraint applied initially
+            # Decode without skipping special to see raw constraint if needed
+            token_str = tokenizer.decode([token_id], skip_special_tokens=False).strip()
             initial_state_vis.append((token_str if token_str else "?", "#800080")) # Constraint color (purple)
     visualization_states.append(initial_state_vis)
 
     # --- Define the Hook Function ---
     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}") # Verbose
+        # print(f"Hook step {step}") # Keep console less noisy
+
+        current_x = x.clone() # Full sequence (prompt + generation) at this step
 
-        current_x = x.clone()
+        # 1. Apply Constraints to the current sequence
         constrained_x = current_x.clone()
-        prompt_len = current_x.shape[1] - gen_length
-        if prompt_len < 0:
-            print("Warning: prompt_len negative in hook, skipping constraints/vis.")
-            return current_x
+        current_prompt_len = current_x.shape[1] - gen_length # Recalculate prompt length based on current x
+        if current_prompt_len < 0:
+            print(f"Warning: prompt_len {current_prompt_len} negative in hook step {step}, skipping constraints/vis.")
+            return current_x # Return unmodified if something is wrong
 
-        # 1. Apply Constraints
         for pos, token_id in processed_constraints.items():
-            absolute_pos = prompt_len + pos
-            if prompt_len <= absolute_pos < current_x.shape[1]:
+            # pos is relative to the start of the *generation* part
+            absolute_pos = current_prompt_len + pos
+            # Ensure position is within the bounds of the *current* sequence 'x'
+            if current_prompt_len <= absolute_pos < current_x.shape[1]:
                 if constrained_x[0, absolute_pos] != token_id:
                     constrained_x[0, absolute_pos] = token_id
+                    # print(f"  Constraint enforced at pos {pos} ({absolute_pos}) -> {token_id}")
 
-        # 2. Generate Visualization State for *this* step
+
+        # 2. Generate Visualization State for *this* step (generation part only)
         current_state_vis = []
-        gen_part_current = current_x[0, prompt_len:]
-        gen_part_last = last_x[0, prompt_len:] if last_x is not None else None
+        # Compare current_x (before explicit constraint application in *this* hook call)
+        # with last_x (state from *previous* hook call / initial state)
+        gen_part_current = current_x[0, current_prompt_len:]
+        # Ensure last_x exists and has the same shape for comparison
+        gen_part_last = last_x[0, current_prompt_len:] if (last_x is not None and last_x.shape == current_x.shape) else None
 
         for i in range(gen_length):
-            current_token_id = gen_part_current[i].item()
+             # Ensure index i is valid for the current generation part
+            if i >= gen_part_current.shape[0]:
+                 print(f"Warning: Index {i} out of bounds for gen_part_current (shape {gen_part_current.shape}) in step {step}.")
+                 continue # Skip if index is invalid
 
-             # Basic check for safety, though unlikely needed with prompt_len check
-            if current_token_id is None:
-                current_state_vis.append((MASK_TOKEN, "#444444"))
-                continue
+            current_token_id = gen_part_current[i].item()
+             # Handle case where last_x was None or had different shape
+            last_token_id = gen_part_last[i].item() if gen_part_last is not None and i < gen_part_last.shape[0] else MASK_ID # Assume mask initially
 
             is_constrained = i in processed_constraints
-            is_special = current_token_id in special_token_ids_to_hide
-
-            # Decode carefully: don't skip specials initially for display text
-            raw_token_str = tokenizer.decode([current_token_id], skip_special_tokens=False).strip()
-            # Use MASK_TOKEN string for MASK_ID, otherwise use decoded string or '?'
-            display_token = MASK_TOKEN if current_token_id == MASK_ID else (raw_token_str if raw_token_str else "?")
+            is_special = current_token_id in special_token_ids_set
+            is_mask = current_token_id == MASK_ID
+            was_mask = last_token_id == MASK_ID or last_x is None # Treat first step as coming from mask
 
-            # Determine the state based on current and previous token
-            previous_token_id = gen_part_last[i].item() if gen_part_last is not None else None
+            display_token = ""
+            color = ""
 
-            # Determine color and potentially modify display_token for hiding
-            if current_token_id == MASK_ID:
-                color = "#444444" # Dark Gray
+            # Determine display token and color based on state transitions
+            if is_mask:
                 display_token = MASK_TOKEN
-            elif is_constrained and processed_constraints.get(i) == current_token_id:
-                 color = "#800080" # Purple - keep constraint visible
-                 # Even if special, show the constraint for clarity
-            elif previous_token_id == MASK_ID or previous_token_id is None:
-                # --- Newly revealed in this step ---
-                if is_special:
-                    # Newly revealed special token: Show it this time
-                    color = "#FF8C00" # Dark Orange (distinct color for first reveal)
-                    # display_token is already the raw special token string
-                else:
-                    # Newly revealed regular token
-                    color = "#66CC66" # Light Green
-                    # display_token is already the regular token string
+                color = "#444444" # Dark Gray
+            elif is_constrained and processed_constraints[i] == current_token_id:
+                # Always show the constrained token, color purple
+                # Decide whether to show raw special tokens when constrained
+                raw_decode = tokenizer.decode([current_token_id], skip_special_tokens=False).strip()
+                display_token = raw_decode if raw_decode else "?"
+                color = "#800080" # Purple
             elif is_special:
-                # --- Was revealed previously AND is special: Hide it now ---
-                color = "#FFFFFF" # White background / Transparent conceptually
-                display_token = "" # Make it disappear visually
-                # Alternative: Subtle placeholder
-                # display_token = "."
-                # color = "#EEEEEE"
+                if was_mask:
+                    # Newly revealed special token: Show its representation once
+                    display_token = tokenizer.decode([current_token_id], skip_special_tokens=False).strip() # Show raw special
+                    color = "#FF8C00" # DarkOrange
+                else:
+                    # Already revealed special token: Hide it by showing a space
+                    display_token = " " # Effectively hides it
+                    color = "#6699CC" # Use 'Old' color (Light Blue) but content is hidden space
+            elif was_mask:
+                # Newly revealed normal token
+                display_token = tokenizer.decode([current_token_id], skip_special_tokens=True).strip()
+                color = "#66CC66" # Light Green
             else:
-                # --- Previously revealed regular token ---
+                # Previously revealed normal token
+                display_token = tokenizer.decode([current_token_id], skip_special_tokens=True).strip()
                 color = "#6699CC" # Light Blue
-                # display_token is already the regular token string
+
+            # Fallback for empty decodes of non-special, non-mask tokens
+            if not display_token and not is_mask and not (is_special and not was_mask):
+                 display_token = "?" # Use question mark for unexpected empty decodes
 
             current_state_vis.append((display_token, color))
 
         visualization_states.append(current_state_vis)
 
         # 3. Update last_x for the *next* step's comparison
+        # Store the state *after* applying constraints for accurate comparison next time
         last_x = constrained_x.clone()
 
-        # 4. Return the sequence with constraints applied
-        return constrained_x
+        # 4. Return the sequence with constraints applied for the model's next step
+        return constrained_x # Return the sequence with constraints enforced
 
 
     # --- Run DREAM Generation ---
     try:
         print("Calling model.diffusion_generate...")
-        initial_full_x = torch.cat([input_ids, initial_x_part], dim=1)
-        last_x = initial_full_x.clone() # Initialize last_x for the first hook call
+        # Make sure last_x is initialized correctly before the first hook call
+        # It should represent the state *before* the first diffusion step.
+        initial_full_x = torch.cat([input_ids, initial_gen_part], dim=1)
+        last_x = initial_full_x.clone() # Initialize last_x with prompt + initial masked/constrained gen part
 
         output = model.diffusion_generate(
-            input_ids,
-            attention_mask=attention_mask,
+            input_ids=input_ids,
+            attention_mask=attention_mask, # Pass the correct attention mask
             max_new_tokens=gen_length,
-            output_history=False,
+            output_history=False, # We build history in the hook
             return_dict_in_generate=True,
             steps=steps,
             temperature=temperature,
             top_p=top_p,
             alg=alg,
+            # alg_temp is only relevant for confidence-based algs (not 'origin')
             alg_temp=alg_temp if alg != "origin" else 0.0,
             generation_tokens_hook_func=generation_tokens_hook_func
+            # Ensure generation doesn't run past eos_token if not desired
+            # eos_token_id=eos_token_id, # This might stop generation early
+            # pad_token_id=tokenizer.eos_token_id # Often pad is same as eos for LLMs
         )
         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]
+        # Handle potential length mismatch if generation stopped early
+        actual_gen_len = final_sequence.shape[0] - prompt_length
         response_token_ids = final_sequence[prompt_length:]
 
-        # Decode final text *skipping* special tokens for the chatbot display
+        # Decode the final response, skipping special tokens like EOS/PAD
         final_text = tokenizer.decode(
             response_token_ids,
             skip_special_tokens=True,
-            clean_up_tokenization_spaces=True
+            clean_up_tokenization_spaces=True # Recommended for cleaner output
         ).strip()
-        print(f"Final generated text (cleaned): {final_text}")
+        print(f"Final generated text: '{final_text}'")
 
-        # Add the very final state to visualization if needed (safeguard)
-        # This uses the same logic as the hook for consistency
-        if len(visualization_states) <= steps:
-             print("Adding final visualization state manually (safeguard).")
+        # Add the very final state to visualization if the hook didn't capture it
+        # (Mainly a safeguard, hook should run 'steps' times or until completion)
+        if len(visualization_states) <= steps: # Hook might run 'steps' times
              final_state_vis = []
-             final_gen_part = final_sequence[prompt_length:]
-             # Need the state *before* this final one to know what was 'new'
-             gen_part_last_final = last_x[0, prompt_len:] if last_x is not None else None
+             final_gen_part = response_token_ids # Use the extracted response tokens
 
-             for i in range(gen_length):
-                 current_token_id = final_gen_part[i].item()
+             for i in range(len(final_gen_part)): # Iterate over actual generated tokens
+                 token_id = final_gen_part[i].item()
                  is_constrained = i in processed_constraints
-                 is_special = current_token_id in special_token_ids_to_hide
-                 raw_token_str = tokenizer.decode([current_token_id], skip_special_tokens=False).strip()
-                 display_token = MASK_TOKEN if current_token_id == MASK_ID else (raw_token_str if raw_token_str else "?")
-                 previous_token_id = gen_part_last_final[i].item() if gen_part_last_final is not None else None
-
-                 if current_token_id == MASK_ID:
-                     color = "#444444"
-                     display_token = MASK_TOKEN
-                 elif is_constrained and processed_constraints.get(i) == current_token_id:
-                     color = "#800080"
-                 elif previous_token_id == MASK_ID or previous_token_id is None: # Newly revealed
-                     color = "#FF8C00" if is_special else "#66CC66"
-                 elif is_special: # Previously revealed special
-                     color = "#FFFFFF"
-                     display_token = ""
-                 else: # Previously revealed regular
-                     color = "#6699CC"
-
+                 is_special = token_id in special_token_ids_set
+                 is_mask = token_id == MASK_ID # Should not happen in final output
+
+                 display_token = ""
+                 color = ""
+
+                 if is_mask: color = "#444444"; display_token = MASK_TOKEN
+                 elif is_constrained and processed_constraints.get(i) == token_id:
+                      raw_decode = tokenizer.decode([token_id], skip_special_tokens=False).strip()
+                      display_token = raw_decode if raw_decode else "?"; color = "#800080" # Purple
+                 elif is_special:
+                      # Hide special tokens in the *final* display state for cleaner look
+                      display_token = " "; color = "#6699CC" # Hide as 'Old' blue
+                 else:
+                      display_token = tokenizer.decode([token_id], skip_special_tokens=True).strip()
+                      color = "#6699CC" # Final state uses 'Old' blue
+
+                 if not display_token: display_token = "?" # Fallback
                  final_state_vis.append((display_token, color))
-             visualization_states.append(final_state_vis)
+
+             # Pad the final state visualization if actual gen len < requested gen_length
+             # This shouldn't be necessary if HighlightedText handles shorter lists
+             # while len(final_state_vis) < gen_length:
+             #     final_state_vis.append((" ", "#FFFFFF")) # Add empty space
+
+             if final_state_vis: # Only append if we generated something
+                 visualization_states.append(final_state_vis)
 
 
     except Exception as e:
-        print(f"Error during generation: {e}")
+        print(f"\n--- Error during generation ---")
         import traceback
         traceback.print_exc()
-        error_msg = f"Error during generation: {str(e)}"
+        # Add error message to visualization
+        error_msg = f"Generation Error: Check Logs"
+        # Append error to visualization states if possible
         visualization_states.append([("Error", "red")])
         final_text = f"Generation failed: {e}"
 
-    print("--- DREAM Generation Finished ---")
-    return visualization_states, final_text
+    print("--- DREAM Generation Finished ---\n")
+    # Ensure we always return a list (even if empty) and a string
+    if not isinstance(visualization_states, list): visualization_states = []
+    if not isinstance(final_text, str): final_text = str(final_text)
+
+    return visualization_states, final_text
+
+
+# --- Gradio UI Setup ---
+
+css = '''
+.category-legend{display:none}
+/* Increase overall base font size */
+body, .gradio-container { font-size: 105%; }
+/* Make buttons slightly larger */
+/* button { min-height: 40px; } */
+.small_btn {
+    min-width: 60px; /* Adjust as needed */
+    max-width: 100px;
+    height: 42px;    /* Adjust height */
+    flex-grow: 0 !important; /* Prevent button from growing */
+    margin-left: 5px !important; /* Add some space */
+    font-size: 100%; /* Match button font size */
+    padding: 0 10px !important; /* Adjust padding */
+}
+.chat-input-row {
+    display: flex;
+    align-items: center; /* Vertically align items */
+    margin-top: 10px; /* Add space above input row */
+}
+/* Ensure Textbox takes up space */
+.chat-input-row .gr-textbox {
+    flex-grow: 1;
+    margin-right: 5px;
+}
+/* Chatbot styling */
+.gr-chatbot .message {
+    font-size: 100%; /* Ensure chat message font size is reasonable */
+    padding: 10px !important;
+    border-radius: 8px !important;
+}
+.gr-chatbot .message.user { background-color: #E0F7FA !important; align-self: flex-end; } /* Light cyan for user */
+.gr-chatbot .message.bot { background-color: #F1F8E9 !important; align-self: flex-start; } /* Light green for bot */
+/* HighlightedText styling */
+.gr-highlightedtext span {
+    padding: 1px 2px; /* Minimal padding */
+    margin: 0 1px; /* Minimal margin */
+    border-radius: 3px;
+    font-family: monospace; /* Use monospace font for better alignment */
+    font-size: 95%; /* Slightly smaller font for dense vis */
+    line-height: 1.4; /* Adjust line spacing */
+}
+.gr-highlightedtext {
+    padding: 10px;
+    border: 1px solid #E0E0E0;
+    border-radius: 5px;
+    background-color: #FAFAFA; /* Light background for the container */
+}
+/* Legend Styling */
+.legend {
+    font-size: 90%;
+    margin-top: 5px;
+    color: #555;
+}
+.legend span {
+    display: inline-block; /* Keep legend items inline */
+    margin-right: 10px;
+    white-space: nowrap; /* Prevent wrapping */
+}
+.legend span::before { /* Style the color square */
+    content: '■';
+    display: inline-block;
+    margin-right: 4px;
+    font-size: 120%; /* Make square slightly larger */
+    vertical-align: middle; /* Align square with text */
+}
+'''
+def create_chatbot_demo():
+    with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+        gr.Markdown("## Dream 7B - Diffusion Language Model Demo")
+        gr.Markdown("Interact with the Dream 7B instruction-tuned model and watch the diffusion process unfold step-by-step. "
+                    "You can optionally constrain specific words at certain positions.")
+        with gr.Row():
+            gr.Markdown("[Model Card](https://huggingface.co/Dream-org/Dream-v0-Instruct-7B)", scale=1)
+            gr.Markdown("[Blog Post](https://hkunlp.github.io/blog/2025/dream/)", scale=1)
+
+        # STATE MANAGEMENT
+        chat_history = gr.State([]) # Stores conversation [[user, bot], ...]
+
+        # UI LAYOUT
+        with gr.Row():
+            # Left Column: Chat Interface
+            with gr.Column(scale=3):
+                chatbot_ui = gr.Chatbot(
+                    label="Conversation",
+                    height=550,
+                    bubble_full_width=False,
+                    show_copy_button=True,
+                    render=False # Rendered explicitly later for streaming
+                 )
+                chatbot_ui.render() # Manually render after setting parameters
+
+                # Message input Row
+                with gr.Row(elem_classes="chat-input-row"):
+                        user_input = gr.Textbox(
+                            label="Your Message",
+                            placeholder="Type your message and press Enter, or click Send...",
+                            scale=4, # Give textbox more space relative to button
+                            container=False,
+                            show_label=False
+                        )
+                        send_btn = gr.Button("Send", scale=1, elem_classes="small_btn", variant="primary")
+
+                constraints_input = gr.Textbox(
+                    label="Word Constraints (Optional)",
+                    info="Force 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",
+                    lines=1
+                )
+
+            # Right Column: Visualization and Settings
+            with gr.Column(scale=2):
+                gr.Markdown("### Denoising Process Visualization")
+                output_vis = gr.HighlightedText(
+                    label="Generation Steps",
+                    show_label=False, # Label provided by Markdown above
+                    combine_adjacent=False,
+                    show_legend=False, # Using custom HTML legend below
+                    # color_map is not directly used due to show_legend=False, but useful for reference
+                    color_map={
+                        "Mask": "#444444",
+                        "New": "#66CC66",
+                        "Old": "#6699CC",
+                        "Constraint": "#800080",
+                        "Special (New)": "#FF8C00",
+                        "Error": "red"
+                    }
+                )
+                # Custom HTML Legend
+                gr.HTML(
+                    """
+                    <div class='legend'>
+                        <span style="color:#444444;">■ Mask</span> |
+                        <span style='color:#66CC66;'>■ New</span> |
+                        <span style='color:#FF8C00;'>■ Special (New)</span> |
+                        <span style='color:#6699CC;'>■ Old</span> |
+                        <span style='color:#800080;'>■ Constraint</span>
+                    </div>
+                    """,
+                    elem_id="legend-html"
+                )
+
+                # Generation Settings Accordion
+                with gr.Accordion("Generation Settings", open=False):
+                    gen_length = gr.Slider(
+                        minimum=16, maximum=512, value=128, step=16,
+                        label="Max New Tokens", info="Max response length."
+                    )
+                    steps = gr.Slider(
+                        minimum=8, maximum=512, value=128, step=8,
+                        label="Diffusion Steps", info="More steps = finer generation (potentially slower)."
+                    )
+                    temperature = gr.Slider(
+                        minimum=0.0, maximum=1.5, value=0.6, step=0.05,
+                        label="Temperature", info="Controls randomness. Lower=more deterministic."
+                    )
+                    top_p = gr.Slider(
+                        minimum=0.0, maximum=1.0, value=0.95, step=0.05,
+                        label="Top-P (Nucleus)", info="Filters vocabulary probabilistically. Lower=less diverse."
+                    )
+                    # Map UI choices to DREAM's alg parameters
+                    remasking_strategy = gr.Radio(
+                        choices=[
+                            ("Random", "origin"), # User friendly name -> actual param
+                            ("Entropy", "entropy"),
+                            ("MaskGit+", "maskgit_plus"),
+                            ("TopK Margin", "topk_margin"),
+                        ],
+                        value="entropy", # Default
+                        label="Generation Order Strategy (alg)",
+                        info="How the model decides which tokens to generate first."
+                    )
+                    alg_temp = gr.Slider(
+                        minimum=0.0, maximum=1.0, value=0.1, step=0.05,
+                        label="Order Randomness (alg_temp)" ,
+                        info="Adds randomness to confidence-based strategies (Entropy, MaskGit+, TopK). Ignored for Random."
+                    )
+                    visualization_delay = gr.Slider(
+                        minimum=0.0, maximum=0.5, value=0.05, step=0.01,
+                        label="Visualization Delay (sec)", info="Pause between steps in visualization."
+                    )
+
+        # Clear button Row
+        with gr.Row():
+            clear_btn = gr.Button("Clear Conversation", variant="stop", icon="🗑️")
+
+
+        # --- Event Handlers ---
+
+        # Helper to add message to history state
+        def add_message_to_history(history_state, user_message, bot_message):
+            # history_state is the raw list from gr.State
+            history_state.append([user_message, bot_message])
+            return history_state
+
+        # Function when user submits message (Enter or Send button)
+        def handle_user_message(message, history_state):
+            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
+                 # history_state, chatbot_ui, user_input, output_vis
+                 return history_state, history_state, "", [] # No change to chatbot UI yet
+
+            # Add user message to history state (with None for bot response initially)
+            updated_history_state = add_message_to_history(history_state, message, None)
+
+            # Prepare updated history for display in Chatbot UI
+            # We only display the user message now, bot response comes later
+            chatbot_display = updated_history_state.copy()
+
+            # Clear the input textbox and visualization
+            return updated_history_state, chatbot_display, "", []
+
+        # Function to generate bot response (triggered after user message is handled)
+        # Uses yield for streaming visualization updates
+        def generate_bot_response(
+            history_state, # The current state [[user, None], ...]
+            gen_length_val, steps_val, constraints_text, delay_val,
+            temperature_val, top_p_val, alg_val, alg_temp_val
+            ):
+            print("\n--- Streaming Bot Response ---")
+            if not history_state or history_state[-1][1] is not None:
+                print("History empty or last message already has response. Skipping generation.")
+                # Yield current state if called unnecessarily
+                yield history_state, [] # Chatbot UI, Visualization
+                return
+
+            # Get the conversation history in the format the model expects
+            messages_for_model = format_chat_history(history_state) # Includes the latest user query
+
+            # Parse constraints from the textbox
+            parsed_constraints = parse_constraints(constraints_text)
+
+            # Generate response with visualization (this function handles the core logic)
+            vis_states, response_text = dream_generate_response_with_visualization(
+                messages_for_model,
+                gen_length=gen_length_val,
+                steps=steps_val,
+                constraints=parsed_constraints,
+                temperature=temperature_val,
+                top_p=top_p_val,
+                alg=alg_val,
+                alg_temp=alg_temp_val
+            )
+
+            # Update the history state with the final bot response (critical!)
+            history_state[-1][1] = response_text.strip()
+
+            # Stream the updates
+            if vis_states:
+                 # Yield the initial visualization state first
+                 yield history_state, vis_states[0] # Update chatbot UI (implicitly via history), update visualization
+
+                 # Then animate through the rest of the visualization states
+                 for state in vis_states[1:]:
+                     time.sleep(delay_val)
+                     yield history_state, state # Update chatbot UI, update visualization
+            else:
+                 # Handle case where generation failed or produced no visualization
+                 print("Warning: No visualization states generated.")
+                 yield history_state, [("No visualization generated.", "orange")] # Update chatbot UI, show warning in vis
+
+            print("--- Streaming Complete ---")
+
+
+        # Function to clear everything
+        def clear_conversation_state():
+            print("Clearing conversation.")
+            # Reset state and UI components
+            return [], [], "", [] # chat_history (State), chatbot_ui, user_input, output_vis
+
+        # --- Wire UI elements to functions ---
+
+        # Define shared inputs for generation to avoid repetition
+        generation_inputs = [
+            chat_history, gen_length, steps, constraints_input, visualization_delay,
+            temperature, top_p, remasking_strategy, alg_temp
+        ]
+        # Define shared outputs for streaming
+        streaming_outputs = [chatbot_ui, output_vis]
+
+        # Typing in Textbox and pressing Enter
+        user_input.submit(
+            fn=handle_user_message,
+            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 user input immediately
+        ).then(
+            fn=generate_bot_response,
+            inputs=generation_inputs,
+            outputs=streaming_outputs, # Stream updates to chatbot and visualization
+            #api_name="generate_stream" # Optional: Name for API endpoint
+        )
+
+        # Clicking the Send button
+        send_btn.click(
+            fn=handle_user_message,
+            inputs=[user_input, chat_history],
+            outputs=[chat_history, chatbot_ui, user_input, output_vis],
+            queue=False
+        ).then(
+            fn=generate_bot_response,
+            inputs=generation_inputs,
+            outputs=streaming_outputs,
+           # api_name="generate_stream_click" # Optional
+        )
+
+        # Clicking the Clear button
+        clear_btn.click(
+            fn=clear_conversation_state,
+            inputs=[],
+            outputs=[chat_history, chatbot_ui, user_input, output_vis],
+            queue=False # Clearing should be instant
+        )
+
+    return demo
+
+# --- Launch the Gradio App ---
+if __name__ == "__main__":
+    print("Creating Gradio demo...")
+    gradio_demo = create_chatbot_demo()
+    print("Launching Gradio demo...")
+    # Use queue() for handling concurrent users and potentially long generation times
+    # share=True generates a public link (useful for Colab/Spaces)
+    # debug=True provides helpful Gradio logs in the console
+    gradio_demo.queue().launch(share=True, debug=False) # Set debug=True for more verbose logs if needed