Spaces:
Runtime error
Runtime error
| # llada_app.py -> dream_app.py | |
| import torch | |
| import numpy as np | |
| import gradio as gr | |
| import spaces | |
| # import torch.nn.functional as F # Not needed for DREAM's basic visualization | |
| 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 --- | |
| model_path = "Dream-org/Dream-v0-Instruct-7B" | |
| print(f"Loading model: {model_path}") | |
| model = AutoModel.from_pretrained(model_path, torch_dtype=torch.bfloat16, trust_remote_code=True).to(device).eval() | |
| 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.") | |
| 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.") | |
| MASK_TOKEN = tokenizer.mask_token | |
| MASK_ID = tokenizer.mask_token_id | |
| print(f"Using MASK_TOKEN='{MASK_TOKEN}' with ID={MASK_ID}") | |
| # --- 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, but usually we want 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 | |
| Returns: | |
| Formatted conversation for the model (list of dictionaries) | |
| """ | |
| messages = [] | |
| # Add system prompt if desired (check DREAM examples/recommendations) | |
| # messages.append({"role": "system", "content": "You are a helpful assistant."}) # Optional | |
| for user_msg, assistant_msg in history: | |
| if user_msg: # 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 --- | |
| # Use the potentially dummy decorator | |
| def dream_generate_response_with_visualization( | |
| messages, | |
| gen_length=64, | |
| steps=64, # Default based on DREAM examples | |
| constraints=None, | |
| 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.0, # Default based on DREAM examples | |
| ): | |
| """ | |
| Generate text with DREAM model with visualization using the generation hook. | |
| Args: | |
| messages: List of message dictionaries with 'role' and 'content' | |
| gen_length: Length of text to generate (max_new_tokens) | |
| steps: Number of diffusion steps | |
| constraints: Dictionary mapping positions (relative to response start) to words | |
| temperature: Sampling temperature | |
| top_p: Nucleus sampling p | |
| alg: Remasking algorithm ('origin', 'maskgit_plus', 'topk_margin', 'entropy') | |
| alg_temp: Temperature for confidence-based algorithms | |
| 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(f"Constraints: {constraints}") | |
| # --- Input Preparation --- | |
| if constraints is None: | |
| constraints = {} | |
| # Convert word constraints to token IDs (handle multi-token words) | |
| processed_constraints = {} | |
| print("Processing constraints:") | |
| for pos, word in constraints.items(): | |
| # Prepend space for consistent tokenization, similar to LLaDA example | |
| 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}") | |
| for i, token_id in enumerate(tokens): | |
| # Ensure we don't overwrite parts of multi-token constraints accidentally | |
| if pos + i not in processed_constraints: | |
| processed_constraints[pos + i] = token_id | |
| else: | |
| 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, | |
| return_tensors="pt", | |
| return_dict=True, | |
| add_generation_prompt=True # Crucial for instruction-tuned models like Dream-Instruct | |
| ) | |
| input_ids = inputs.input_ids.to(device=device) | |
| 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}") | |
| 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}" | |
| if prompt_length + gen_length > 2048: # Check context length (DREAM uses 2048) | |
| print(f"Warning: Requested length ({prompt_length + gen_length}) exceeds model max length (2048). Truncating gen_length.") | |
| 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." | |
| # --- State for Visualization Hook --- | |
| visualization_states = [] | |
| last_x = None # Store the sequence from the previous step | |
| # Initial state: Prompt + all masks | |
| initial_x_part = torch.full((1, gen_length), MASK_ID, dtype=torch.long, device=device) | |
| # Apply initial constraints to the masked part *before* showing the first state | |
| 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) | |
| # --- 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}") | |
| current_x = x.clone() # Work on a copy for comparison | |
| # 1. Apply Constraints *before* generating visualization | |
| # 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: | |
| 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]: | |
| 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 | |
| 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 | |
| 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 "?" | |
| # Check if this position is constrained | |
| is_constrained = i in processed_constraints | |
| if current_token_id == MASK_ID: | |
| color = "#444444" # Dark Gray for masks | |
| 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 | |
| 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 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 | |
| # --- Run DREAM Generation --- | |
| try: | |
| 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. | |
| initial_full_x = torch.cat([input_ids, initial_x_part], dim=1) | |
| last_x = initial_full_x.clone() | |
| output = model.diffusion_generate( | |
| input_ids, | |
| attention_mask=attention_mask, | |
| max_new_tokens=gen_length, | |
| 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=alg_temp if alg != "origin" else 0.0, # alg_temp only for confidence algs | |
| generation_tokens_hook_func=generation_tokens_hook_func | |
| ) | |
| 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 | |
| final_text = tokenizer.decode( | |
| response_token_ids, | |
| skip_special_tokens=True, | |
| clean_up_tokenization_spaces=True # Recommended for cleaner output | |
| ).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) | |
| except Exception as e: | |
| print(f"Error during generation: {e}") | |
| import traceback | |
| traceback.print_exc() | |
| # Add error message to visualization | |
| error_msg = f"Error during generation: {str(e)}" | |
| visualization_states.append([("Error", "red")]) | |
| final_text = f"Generation failed: {e}" | |
| print("--- DREAM Generation Finished ---") | |
| return visualization_states, final_text | |
| # --- Gradio UI Setup --- | |
| css = ''' | |
| .category-legend{display:none} | |
| /* button{height: 60px} */ /* Optional: Adjust button height */ | |
| .small_btn { | |
| max-width: 100px; /* Adjust as needed */ | |
| height: 40px; /* Adjust as needed */ | |
| flex-grow: 0; /* Prevent button from growing */ | |
| margin-left: 5px; /* Add some space */ | |
| } | |
| .chat-input-row { | |
| display: flex; | |
| align-items: center; /* Vertically align items */ | |
| } | |
| .chat-input-row > * { | |
| margin-right: 5px; /* Space between textbox and button */ | |
| } | |
| .chat-input-row > *:last-child { | |
| margin-right: 0; | |
| } | |
| ''' | |
| def create_chatbot_demo(): | |
| with gr.Blocks(css=css) as demo: | |
| gr.Markdown("# Dream 7B - Diffusion Language Model Demo") | |
| gr.Markdown("A demonstration of the Dream 7B diffusion-based language model. Watch the text generate step-by-step.") | |
| gr.Markdown("[Model Card](https://huggingface.co/Dream-org/Dream-v0-Instruct-7B) - [Blog Post](https://hkunlp.github.io/blog/2025/dream/)") | |
| # STATE MANAGEMENT | |
| chat_history = gr.State([]) | |
| # UI COMPONENTS | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| chatbot_ui = gr.Chatbot( | |
| label="Conversation", | |
| height=500, | |
| bubble_full_width=False # Improves layout for shorter messages | |
| ) | |
| # Message input Row | |
| with gr.Row(elem_classes="chat-input-row"): | |
| 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 | |
| show_label=False | |
| ) | |
| send_btn = gr.Button("Send", scale=1, elem_classes="small_btn") | |
| constraints_input = gr.Textbox( | |
| 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 | |
| ) | |
| 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>" | |
| ) | |
| # 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 | |
| label="Max New Tokens" | |
| ) | |
| steps = gr.Slider( | |
| minimum=8, maximum=512, value=128, step=8, # Increased max steps | |
| 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 | |
| label="Temperature" | |
| ) | |
| top_p = gr.Slider( | |
| minimum=0.0, maximum=1.0, value=0.95, step=0.05, | |
| 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 | |
| ("Entropy", "entropy"), | |
| ("MaskGit+", "maskgit_plus"), | |
| ("TopK Margin", "topk_margin"), | |
| ], | |
| value="entropy", # Default | |
| label="Generation Order Strategy (alg)" | |
| ) | |
| 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 non-Random strategies. Ignored for Random." | |
| ) | |
| with gr.Row(): | |
| visualization_delay = gr.Slider( | |
| minimum=0.0, maximum=0.5, value=0.05, step=0.01, | |
| label="Visualization Delay (seconds)" | |
| ) | |
| # 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 | |
| def add_message_to_history(history, message, response): | |
| history = history.copy() # Modify copy | |
| history.append([message, response]) | |
| return history | |
| # Function when user submits message (Enter or Send button) | |
| def user_message_submitted(message, history): | |
| 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 | |
| return history, history_for_display, message_out, vis_clear | |
| # Function to generate bot response (triggered after user message is processed) | |
| def bot_response_generator( | |
| history, gen_length, steps, constraints_text, delay, | |
| temperature, top_p, alg, alg_temp | |
| ): | |
| 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." | |
| 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 | |
| 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, | |
| steps=steps, | |
| constraints=parsed_constraints, | |
| temperature=temperature, | |
| top_p=top_p, | |
| alg=alg, | |
| 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")] | |
| # 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 | |
| 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 | |
| # Function to clear everything | |
| def clear_conversation(): | |
| print("Clearing conversation.") | |
| return [], [], "", [] # chat_history, chatbot_ui, user_input, output_vis | |
| # --- Wire UI elements to functions --- | |
| # Typing in Textbox and pressing Enter | |
| 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 | |
| ) | |
| # Clicking the Send button | |
| 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=[ | |
| chat_history, gen_length, steps, constraints_input, visualization_delay, | |
| temperature, top_p, remasking_strategy, alg_temp | |
| ], | |
| outputs=[chatbot_ui, output_vis] | |
| ) | |
| # Clicking the Clear button | |
| clear_btn.click( | |
| fn=clear_conversation, | |
| inputs=[], | |
| outputs=[chat_history, chatbot_ui, user_input, output_vis], | |
| queue=False | |
| ) | |
| return demo | |
| # --- Launch the Gradio App --- | |
| if __name__ == "__main__": | |
| print("Creating Gradio demo...") | |
| demo = create_chatbot_demo() | |
| print("Launching Gradio demo...") | |
| # Use queue for potentially long generation times | |
| # share=True generates a public link (useful for Colab/Spaces) | |
| demo.queue().launch(share=True, debug=True) # Add debug=True for more logs |