Update app.py

app.py

@@ -239,7 +239,7 @@ def apply_constraints_to_state(
 @spaces.GPU # Decorator for Hugging Face Spaces GPU usage
 @torch.no_grad() # Ensure no gradients are computed during generation
 def generate_dream_response(
-    history: List[List[Optional[str]]], # Receives the latest state from _chat_history_store
+    history: List[List[Optional[str]]], # Receives the list from _chat_history_store
     gen_length: int,
     steps: int,
     constraints_text: str,
@@ -252,12 +252,16 @@ def generate_dream_response(
     ) -> List[Tuple[str, str]]:
     """ Generates text step-by-step and yields visualization states live. """
 
-    if not history or history[-1][0] is None: # Check if last user message is None or missing
-        yield history, [("Internal Error: History state invalid.", "red")], ""
+    # No history_copy needed, work directly on the input 'history' list
+    # which is a reference to the value in _chat_history_store
+
+    if not history or not history[-1][0]:
+        # Yield the original history back if there's no input
+        yield history, [("No input message found.", "red")], ""
         return
 
     # --- 1. Preparation ---
-    # History already contains the latest user message and None for the bot response
+    last_user_message = history[-1][0]
     messages_for_template = format_chat_history(history)
     parsed_constraints = parse_constraints(constraints_text)
 
@@ -266,15 +270,13 @@ def generate_dream_response(
             messages_for_template,
             return_tensors="pt",
             return_dict=True,
-            add_generation_prompt=True # Creates the '<|im_start|>assistant\n' prompt
+            add_generation_prompt=True
         )
         input_ids = inputs.input_ids.to(device)
-        # Ensure prompt_attention_mask is also on the correct device and handle missing mask
         prompt_attention_mask = inputs.attention_mask.to(device) if 'attention_mask' in inputs else torch.ones_like(input_ids)
         prompt_length = input_ids.shape[1]
     except Exception as e:
         print(f"Error applying chat template: {e}")
-        # Yield current history (with None), error message, empty text
         yield history, [("Error preparing input.", "red")], ""
         return
 
@@ -288,38 +290,33 @@ def generate_dream_response(
     initial_generation_part = torch.full((1, gen_length), MASK_ID, dtype=torch.long, device=device)
     x = torch.cat((input_ids, initial_generation_part), dim=1)
 
-    # Prepare attention mask for SDPA (float format)
-    generation_attention_mask = torch.ones((1, gen_length), dtype=prompt_attention_mask.dtype, device=device) # Match dtype
-    full_attention_mask_long = torch.cat((prompt_attention_mask, generation_attention_mask), dim=1) # Shape [B, N]
+    generation_attention_mask = torch.ones((1, gen_length), dtype=torch.long, device=device)
+    full_attention_mask_long = torch.cat((prompt_attention_mask, generation_attention_mask), dim=1)
 
-    attention_mask_for_model = full_attention_mask_long.to(model.dtype) # Convert to model's float dtype
+    attention_mask_for_model = full_attention_mask_long.to(model.dtype)
     large_neg_val = torch.finfo(model.dtype).min
     attention_mask_for_model = (1.0 - attention_mask_for_model) * large_neg_val
-    attention_mask_for_model = attention_mask_for_model.unsqueeze(1).unsqueeze(2) # Shape [B, 1, 1, N]
+    attention_mask_for_model = attention_mask_for_model.unsqueeze(1).unsqueeze(2)
 
-    # Timesteps
     timesteps = torch.linspace(1, eps, steps + 1, device=device)
-
-    # Apply initial constraints
     x = apply_constraints_to_state(x, prompt_length, total_length, parsed_constraints, current_step=-1)
 
     # --- 3. Visualization Setup ---
     previous_tokens_vis = None
     final_response_text = ""
-    # Work on a copy of the history list received as input
-    history_copy = [list(item) for item in history]
+    # history_copy removed
 
     # --- 4. Initial Yield (Masked State) ---
     initial_generated_tokens = x[0, prompt_length:].cpu()
     vis_data_initial = []
     for tok_id in initial_generated_tokens.tolist():
         display_token = MASK_TOKEN
-        color = "#444444" # Dark Gray for masks
+        color = "#444444"
         vis_data_initial.append((display_token, color))
 
     previous_tokens_vis = initial_generated_tokens
-    # Yield the initial history copy (with None placeholder), initial vis, empty text
-    yield history_copy, vis_data_initial, ""
+    # Yield the current state of the history (which has None for the bot response)
+    yield history, vis_data_initial, ""
     time.sleep(visualization_delay)
 
     # --- 5. Step-by-Step Diffusion Loop ---
@@ -334,19 +331,15 @@ def generate_dream_response(
             # --- Model Forward Pass ---
             outputs = model(
                 input_ids=x,
-                attention_mask=attention_mask_for_model, # Pass the [B, 1, 1, N] float mask
-                position_ids=None, # Let model compute default positions
+                attention_mask=attention_mask_for_model,
+                position_ids=None,
                 use_cache=False,
                 return_dict=True
             )
             logits = outputs.logits
+            logits = torch.cat([logits[:,:1], logits[:, :-1]], dim=1) # Align logits
 
-            # Align logits with the token positions they predict (logits[t] predicts token[t+1])
-            # Shift left, effectively aligning logits[t] with inputs[t]
-            logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
-
-            # Select logits for masked positions
-            mask_logits = logits[mask_index] # Shape [num_masked_tokens, V]
+            mask_logits = logits[mask_index]
             if mask_logits.numel() == 0:
                  print(f"No masked tokens found for logit selection at step {i}. Stopping.")
                  break
@@ -354,9 +347,9 @@ def generate_dream_response(
             # --- Sampling / Remasking Logic ---
             t = timesteps[i]
             s = timesteps[i + 1]
-            # Initialize the update tensor for masked positions with MASK_ID
             x_new_masked_part = torch.full_like(x[mask_index], MASK_ID, device=device, dtype=torch.long)
 
+            # [Keep sampling logic identical to previous correct version]
             if alg == 'origin':
                 p_transfer = (1.0 - s / t) if i < steps - 1 else 1.0
                 num_masked = mask_logits.shape[0]
@@ -365,13 +358,11 @@ def generate_dream_response(
 
                 if logits_to_sample.numel() > 0:
                     _, sampled_tokens = sample_tokens(logits_to_sample, temperature=temperature, top_p=top_p_val, top_k=top_k_val)
-                    # Place sampled tokens into the correct positions within the masked part update
                     x_new_masked_part[transfer_indices_relative] = sampled_tokens
 
-            else: # Confidence-based algorithms ('maskgit_plus', 'topk_margin', 'entropy')
+            else: # Confidence-based algorithms
                 use_margin = (alg == 'topk_margin')
                 use_entropy = (alg == 'entropy')
-                # Sample candidates and get confidence for all masked positions
                 confidence, x0_candidates = sample_tokens(
                     mask_logits,
                     temperature=temperature,
@@ -382,129 +373,89 @@ def generate_dream_response(
                 )
 
                 num_mask_token = mask_logits.shape[0]
-                # Calculate target number of tokens to reveal in this step
                 target_num_revealed_float = num_mask_token * (1.0 - s / t)
                 number_transfer_tokens = int(target_num_revealed_float) if i < steps - 1 else num_mask_token
 
                 if number_transfer_tokens > 0:
-                    # Determine which tokens to reveal based on confidence
-                    num_samples = min(number_transfer_tokens, num_mask_token) # Ensure k <= num_mask_token
+                    num_samples = min(number_transfer_tokens, num_mask_token)
                     if num_samples > 0:
-                        transfer_indices_relative = torch.tensor([], dtype=torch.long, device=device) # Initialize empty
-                        if alg_temp_val is None or alg_temp_val <= 0: # Use top-k confidence sorting
-                            # Sort by confidence (higher is better, except for entropy where lower is better)
+                        transfer_indices_relative = torch.tensor([], dtype=torch.long, device=device) # Initialize
+                        if alg_temp_val is None or alg_temp_val <= 0: # Top-k confidence
                             sort_metric = confidence if alg != 'entropy' else -confidence
-                            # Ensure k is not greater than the number of elements
                             k_topk = min(num_samples, sort_metric.numel())
                             if k_topk > 0:
                                 _, transfer_indices_relative = torch.topk(sort_metric, k=k_topk)
 
                         else: # Sample based on confidence temperature
-                            # Ensure confidence has elements before processing
                             if confidence.numel() > 0:
                                 conf_probs = confidence / alg_temp_val
-                                # Handle potential inf/-inf before softmax, ensure non-negative probabilities
                                 conf_probs = torch.nan_to_num(conf_probs, nan=0.0, posinf=1e9, neginf=-1e9)
-                                # Clamp to prevent large positive values causing overflow in exp
-                                conf_probs = torch.clamp(conf_probs - conf_probs.max(), min=-30) # Softmax is invariant to shift
+                                conf_probs = torch.clamp(conf_probs - conf_probs.max(), min=-30)
                                 conf_probs = F.softmax(conf_probs, dim=-1)
-                                conf_probs = torch.clamp(conf_probs, min=0.0) # Ensure non-negative
-                                conf_probs = torch.nan_to_num(conf_probs, nan=0.0) # Handle NaNs
+                                conf_probs = torch.clamp(conf_probs, min=0.0)
+                                conf_probs = torch.nan_to_num(conf_probs, nan=0.0)
 
-                                # Normalize probabilities if they don't sum to 1 (within tolerance)
                                 prob_sum = conf_probs.sum()
                                 target_sum_tensor = torch.tensor(1.0, device=device, dtype=prob_sum.dtype)
                                 if not torch.isclose(prob_sum, target_sum_tensor, atol=1e-4) and prob_sum > 0:
                                     safe_prob_sum = torch.max(prob_sum, torch.tensor(1e-12, device=device, dtype=prob_sum.dtype))
                                     conf_probs = conf_probs / safe_prob_sum
 
-                                # Check if probabilities are valid for multinomial sampling
                                 final_prob_sum_check = conf_probs.sum()
                                 if conf_probs.numel() > 0 and num_samples > 0 and torch.all(conf_probs >= 0) and torch.isclose(final_prob_sum_check, target_sum_tensor, atol=1e-4):
                                     try:
                                         transfer_indices_relative = torch.multinomial(conf_probs, num_samples=num_samples, replacement=False)
                                     except RuntimeError as e:
                                         print(f"Warning step {i}: Multinomial sampling failed ('{e}'). Falling back to top-k.")
-                                        # Fallback to top-k if multinomial fails
                                         sort_metric = confidence if alg != 'entropy' else -confidence
                                         k_multinomial_fallback = min(num_samples, sort_metric.numel())
                                         if k_multinomial_fallback > 0:
                                              _, transfer_indices_relative = torch.topk(sort_metric, k=k_multinomial_fallback)
-                                else: # Handle cases where multinomial is not possible (e.g., bad probabilities)
+                                else:
                                     # print(f"Warning step {i}: Invalid probabilities for multinomial sampling (sum={final_prob_sum_check:.4f}). Falling back to top-k.")
                                     sort_metric = confidence if alg != 'entropy' else -confidence
                                     k_multinomial_fallback = min(num_samples, sort_metric.numel())
                                     if k_multinomial_fallback > 0:
                                         _, transfer_indices_relative = torch.topk(sort_metric, k=k_multinomial_fallback)
+                            # else: # No confidence values to sample from, transfer_indices_relative remains empty
 
-                        # Apply the transfer using the selected indices, with safety checks
+                        # Apply the transfer
                         if transfer_indices_relative.numel() > 0:
-                             # Bounds check before indexing
-                             max_cand_idx = x0_candidates.shape[0] - 1
-                             max_mask_idx = x_new_masked_part.shape[0] - 1
-                             valid_indices_mask = (transfer_indices_relative >= 0) & \
-                                                  (transfer_indices_relative <= max_cand_idx) & \
-                                                  (transfer_indices_relative <= max_mask_idx)
-                             valid_transfer_indices = transfer_indices_relative[valid_indices_mask]
-
+                             valid_indices = transfer_indices_relative < x0_candidates.shape[0]
+                             valid_transfer_indices = transfer_indices_relative[valid_indices]
                              if valid_transfer_indices.numel() > 0:
-                                  x_new_masked_part[valid_transfer_indices] = x0_candidates[valid_transfer_indices].clone()
-                             # else:
-                             #    if transfer_indices_relative.numel() > 0: # Only warn if there were indices initially
-                             #         print(f"Warning step {i}: No valid transfer indices after bounds check.")
-
+                                  if valid_transfer_indices.max() < x_new_masked_part.shape[0]:
+                                       x_new_masked_part[valid_transfer_indices] = x0_candidates[valid_transfer_indices].clone()
+                                  else:
+                                       print(f"Warning step {i}: transfer_indices out of bounds for x_new_masked_part.")
+            # --- End Sampling Logic ---
 
-            # Update the global state `x` only at the masked positions
             x[mask_index] = x_new_masked_part
-
-            # --- Apply Constraints ---
-            # Constraints should be applied *after* sampling/revealing tokens for the step
             x = apply_constraints_to_state(x, prompt_length, total_length, parsed_constraints, current_step=i)
 
             # --- Yield Visualization ---
-            current_generated_tokens = x[0, prompt_length:].cpu() # Get generated part, move to CPU
+            current_generated_tokens = x[0, prompt_length:].cpu()
             vis_data = []
+            # [Keep visualization formatting logic the same]
             for j in range(gen_length):
                 current_tok_id = current_generated_tokens[j].item()
-                # Ensure previous_tokens_vis exists and index is valid
                 previous_tok_id = previous_tokens_vis[j].item() if previous_tokens_vis is not None and j < len(previous_tokens_vis) else MASK_ID
-
                 try:
-                    # Use replace='�' to handle potential bytes rendering issues in Gradio HighlightedText
                     decoded_token = tokenizer.decode([current_tok_id], skip_special_tokens=False, clean_up_tokenization_spaces=False)
                     display_token = MASK_TOKEN if current_tok_id == MASK_ID else decoded_token
-                except Exception:
-                    display_token = f"[ID:{current_tok_id}]" # Fallback
-
-                color = None
-                token_to_display = display_token
-
-                if current_tok_id == MASK_ID:
-                    color = "#444444" # Dark Gray for masks
-                elif previous_tok_id == MASK_ID: # Token was just revealed
-                    color = "#66CC66" # Light Green
-                else: # Token was already revealed
-                    color = "#6699CC" # Light Blue
-
-                # Hide special tokens (PAD/EOS) if they were already revealed (LLaDA effect)
-                # Ensure PAD_ID and EOS_ID are not None before checking
-                should_hide = False
-                if PAD_ID is not None and current_tok_id == PAD_ID: should_hide = True
-                if EOS_ID is not None and current_tok_id == EOS_ID: should_hide = True
-                # Special check: If PAD and EOS are the same, only hide if it's that ID
-                if PAD_ID == EOS_ID and PAD_ID is not None and current_tok_id == PAD_ID: should_hide = True
-
-                if should_hide and previous_tok_id == current_tok_id:
-                    token_to_display = "" # Hide by making empty
-                    color = None # No color for hidden
-
-                if token_to_display: # Avoid adding empty strings if hiding
-                    vis_data.append((token_to_display, color))
-
-            # Update previous state for the next iteration's color logic
+                except Exception: display_token = f"[ID:{current_tok_id}]"
+                color = None; token_to_display = display_token
+                if current_tok_id == MASK_ID: color = "#444444"
+                elif previous_tok_id == MASK_ID: color = "#66CC66"
+                else: color = "#6699CC"
+                should_hide = (PAD_ID is not None and current_tok_id == PAD_ID) or \
+                              (EOS_ID is not None and current_tok_id == EOS_ID)
+                if should_hide and previous_tok_id == current_tok_id: token_to_display = ""; color = None
+                if token_to_display: vis_data.append((token_to_display, color))
+            # --- End Vis Formatting ---
+
             previous_tokens_vis = current_generated_tokens
 
-            # Decode intermediate response text using the *current* state x
             intermediate_response_tokens = x[0, prompt_length:]
             intermediate_response_text = tokenizer.decode(
                 intermediate_response_tokens,
@@ -512,13 +463,10 @@ def generate_dream_response(
                 clean_up_tokenization_spaces=True
             ).strip()
 
-            # Update the *copy* of the history with the intermediate text for display purposes
-            if history_copy: # Ensure history_copy is not empty
-                 history_copy[-1][1] = intermediate_response_text # Update the None placeholder
-
-            # Yield the updated history copy, current vis, and intermediate text
-            yield history_copy, vis_data, intermediate_response_text
+            # Yield the current state of the history list (bot response still None)
+            yield history, vis_data, intermediate_response_text
             time.sleep(visualization_delay)
+        # --- End Loop ---
 
         end_time = time.time()
         print(f"Dream generation finished in {end_time - start_time:.2f} seconds.")
@@ -532,45 +480,41 @@ def generate_dream_response(
             clean_up_tokenization_spaces=True
         ).strip()
 
-        # Update the final history copy *definitively*
-        if history_copy:
-            history_copy[-1][1] = final_response_text
+        # --- CRITICAL FIX: Update history IN PLACE before final yield ---
+        if history: # Ensure history is not empty
+            history[-1][1] = final_response_text
+        # Now the list referenced by _chat_history_store is updated.
 
-        # Format the final visualization state
         final_generated_tokens = x[0, prompt_length:].cpu()
         vis_data_final = []
+        # [Keep final visualization formatting logic the same]
         for j in range(gen_length):
             current_tok_id = final_generated_tokens[j].item()
             previous_tok_id = previous_tokens_vis[j].item() if previous_tokens_vis is not None and j < len(previous_tokens_vis) else MASK_ID
             try:
                 decoded_token = tokenizer.decode([current_tok_id], skip_special_tokens=False, clean_up_tokenization_spaces=False)
                 display_token = MASK_TOKEN if current_tok_id == MASK_ID else decoded_token
-            except Exception:
-                display_token = f"[ID:{current_tok_id}]" # Fallback
-            color = None
-            token_to_display = display_token
+            except Exception: display_token = f"[ID:{current_tok_id}]"
+            color = None; token_to_display = display_token
             if current_tok_id == MASK_ID: color = "#444444"
             elif previous_tok_id == MASK_ID: color = "#66CC66"
             else: color = "#6699CC"
-
-            should_hide = False
-            if PAD_ID is not None and current_tok_id == PAD_ID: should_hide = True
-            if EOS_ID is not None and current_tok_id == EOS_ID: should_hide = True
-            if PAD_ID == EOS_ID and PAD_ID is not None and current_tok_id == PAD_ID: should_hide = True
-
-            if should_hide and previous_tok_id == current_tok_id:
-                 token_to_display = ""; color = None
+            should_hide = (PAD_ID is not None and current_tok_id == PAD_ID) or \
+                          (EOS_ID is not None and current_tok_id == EOS_ID)
+            if should_hide and previous_tok_id == current_tok_id: token_to_display = ""; color = None
             if token_to_display: vis_data_final.append((token_to_display, color))
+        # --- End Final Vis Formatting ---
 
-        # Yield the final history, final visualization, and final text
-        yield history_copy, vis_data_final, final_response_text
+        # Yield the FINAL updated history list
+        yield history, vis_data_final, final_response_text
         print("Visualization streaming complete.")
 
     except Exception as e:
-        print(f"Error during generation or processing loop: {e}")
+        print(f"Error during generation or processing: {e}")
+        import traceback
         traceback.print_exc()
-        # Yield the history as it was before the error, error vis, empty text
-        yield history_copy, [("Error during generation.", "red")], ""
+        # Yield the history state as it was when the error occurred
+        yield history, [("Error during generation.", "red")], ""
         return