# 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 ---
@gpu_check # 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:** ■ Mask | ■ Newly Generated | ■ Previously Generated | ■ Constraint"
)
# 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