Create app.py

app.py

@@ -0,0 +1,411 @@
+import torch
+import numpy as np
+import gradio as gr
+import spaces
+from transformers import AutoTokenizer, AutoModel
+import time
+import re
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"Using device: {device}")
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained('mlx-community/LLaDA-8B-Instruct-mlx-4bit', trust_remote_code=True)
+model = AutoModel.from_pretrained('mlx-community/LLaDA-8B-Instruct-mlx-4bit', trust_remote_code=True, 
+                                  torch_dtype=torch.bfloat16).to(device)
+
+# Constants
+MASK_TOKEN = "[MASK]"
+MASK_ID = 126336  # The token ID of [MASK] in LLaDA
+
+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:
+        if ':' not in part:
+            continue
+        pos_str, word = part.split(':', 1)
+        try:
+            pos = int(pos_str.strip())
+            word = word.strip()
+            if word and pos >= 0:
+                constraints[pos] = word
+        except ValueError:
+            continue
+    
+    return constraints
+
+def format_chat_history(history):
+    """
+    Format chat history for the LLaDA model
+    
+    Args:
+        history: List of [user_message, assistant_message] pairs
+        
+    Returns:
+        Formatted conversation for the model
+    """
+    messages = []
+    for user_msg, assistant_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:  # Skip if None (for the latest user message)
+            messages.append({"role": "assistant", "content": assistant_msg})
+    
+    return messages
+
+@spaces.GPU
+def generate_response_with_visualization(model, tokenizer, device, messages, gen_length=64, steps=32, constraints=None):
+    """
+    Generate text with LLaDA model with visualization of the denoising process
+    
+    Args:
+        messages: List of message dictionaries with 'role' and 'content'
+        
+    Returns:
+        List of visualization states showing the progression and final text
+    """
+    
+    # Process constraints
+    if constraints is None:
+        constraints = {}
+        
+    # Convert any string constraints to token IDs
+    processed_constraints = {}
+    for pos, word in constraints.items():
+        tokens = tokenizer.encode(" " + word, add_special_tokens=False)
+        for i, token_id in enumerate(tokens):
+            processed_constraints[pos + i] = token_id
+    
+    # Prepare the prompt using chat template
+    chat_input = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+    input_ids = tokenizer(chat_input)['input_ids']
+    input_ids = torch.tensor(input_ids).to(device).unsqueeze(0)
+    
+    # For generation
+    prompt_length = input_ids.shape[1]
+    
+    # Initialize the sequence with masks for the response part
+    x = torch.full((1, prompt_length + gen_length), MASK_ID, dtype=torch.long).to(device)
+    x[:, :prompt_length] = input_ids.clone()
+    
+    # Initialize visualization states for just the response part
+    visualization_states = []
+    
+    # Add initial state (all masked) - only for the response part
+    initial_state = [(MASK_TOKEN, "#444444") for _ in range(gen_length)]
+    visualization_states.append(initial_state)
+    
+    # Apply constraints to the initial state
+    for pos, token_id in processed_constraints.items():
+        absolute_pos = prompt_length + pos
+        if absolute_pos < x.shape[1]:
+            x[:, absolute_pos] = token_id
+    
+    # Calculate timesteps
+    timesteps = torch.linspace(1.0, 0.0, steps + 1)[:-1]
+    
+    # Keep track of already revealed tokens
+    revealed_tokens = torch.zeros(1, gen_length, dtype=torch.bool).to(device)
+    
+    for step, t in enumerate(timesteps):
+        # Current t to next t
+        s = t - 1.0 / steps if step < steps - 1 else 0
+        
+        # Get all mask positions in the current sequence
+        mask_indices = (x == MASK_ID)
+        
+        # Skip if no masks
+        if not mask_indices.any():
+            break
+            
+        # Get logits from the model
+        logits = model(x).logits
+        
+        # Get the top predictions
+        x0 = torch.argmax(logits, dim=-1)
+        
+        # Get probabilities for visualization
+        probs = torch.softmax(logits, dim=-1)
+        top_probs = torch.max(probs, dim=-1)[0]
+        
+        # Apply the predictions where we have masks
+        x_old = x.clone()
+        x = torch.where(mask_indices, x0, x)
+        
+        # Calculate how many tokens should remain masked at next step
+        total_len = gen_length
+        current_t_value = float(t)
+        next_t_value = float(s)
+        
+        # Linear schedule: t=1 → all masked, t=0 → none masked
+        current_masks_expected = int(current_t_value * total_len)
+        next_masks_expected = int(next_t_value * total_len)
+        
+        # How many to unmask in this step
+        tokens_to_unmask = current_masks_expected - next_masks_expected
+        
+        if tokens_to_unmask > 0 and mask_indices.any():
+            # Get confidence scores for currently masked tokens
+            confidence_scores = top_probs[mask_indices]
+            
+            # Sort confidence scores
+            sorted_indices = torch.argsort(confidence_scores, descending=True)
+            
+            # Select which tokens to keep masked (the lowest confidence ones)
+            indices_to_remask = sorted_indices[tokens_to_unmask:]
+            
+            # Get the actual indices in the sequence
+            mask_positions = torch.where(mask_indices)[1]
+            positions_to_remask = mask_positions[indices_to_remask]
+            
+            # Remask these positions
+            x[:, positions_to_remask] = MASK_ID
+        
+        # Ensure constraints are maintained
+        for pos, token_id in processed_constraints.items():
+            absolute_pos = prompt_length + pos
+            if absolute_pos < x.shape[1]:
+                x[:, absolute_pos] = token_id
+        
+        # Create visualization state ONLY for the response part
+        current_state = []
+        
+        # Update which tokens are newly revealed in this step
+        for i in range(gen_length):
+            pos = prompt_length + i  # Absolute position in the sequence
+            
+            if x[0, pos] == MASK_ID:
+                # Still masked
+                current_state.append((MASK_TOKEN, "#444444"))  # Dark gray for masks
+                
+            elif x_old[0, pos] == MASK_ID:
+                # Newly revealed in this step
+                token = tokenizer.decode([x[0, pos].item()], skip_special_tokens=True)
+                confidence = float(top_probs[0, pos].cpu())
+                
+                # Color based on confidence: red (low) to green (high)
+                if confidence < 0.3:
+                    color = "#FF6666"  # Light red
+                elif confidence < 0.7:
+                    color = "#FFAA33"  # Orange
+                else:
+                    color = "#66CC66"  # Light green
+                    
+                current_state.append((token, color))
+                revealed_tokens[0, i] = True
+                
+            else:
+                # Previously revealed
+                token = tokenizer.decode([x[0, pos].item()], skip_special_tokens=True)
+                current_state.append((token, "#6699CC"))  # Light blue
+        
+        visualization_states.append(current_state)
+    
+    # Extract final text (just the assistant's response)
+    response_tokens = x[0, prompt_length:]
+    response_text = tokenizer.decode(response_tokens, skip_special_tokens=True)
+    
+    # Clean the response text
+    final_text = tokenizer.decode(response_tokens, 
+                                 skip_special_tokens=True,
+                                 clean_up_tokenization_spaces=True)
+    
+    return visualization_states, final_text
+
+css = '''
+.category-legend{display:none}
+button{height: 60px}
+'''
+def create_chatbot_demo():
+    with gr.Blocks(css=css) as demo:
+        gr.Markdown("# LLaDA - Large Language Diffusion Model demo")
+        gr.Markdown("[model](https://huggingface.co/GSAI-ML/LLaDA-8B-Instruct), [project page](https://ml-gsai.github.io/LLaDA-demo/)")
+        
+        # STATE MANAGEMENT - IMPORTANT
+        # We use a dedicated state to track the full conversation history
+        chat_history = gr.State([])
+        
+        # UI COMPONENTS
+        # Chatbot for displaying messages
+        with gr.Row():
+            with gr.Column(scale=3):
+                chatbot_ui = gr.Chatbot(label="Conversation", height=500)
+                
+                # Message input
+                with gr.Group():
+                    with gr.Row():
+                        user_input = gr.Textbox(
+                            label="Your Message", 
+                            placeholder="Type your message here...",
+                            show_label=False
+                        )
+                        send_btn = gr.Button("Send")
+                
+                constraints_input = gr.Textbox(
+                    label="Word Constraints", 
+                    info="This model allows for placing specific words at specific positions using 'position:word' format. Example: 1st word once, 6th word 'upon' and 11th word 'time', would be: '0:Once, 5:upon, 10:time",
+                    placeholder="0:Once, 5:upon, 10:time",
+                    value=""
+                )
+            with gr.Column(scale=2):
+                output_vis = gr.HighlightedText(
+                    label="Denoising Process Visualization",
+                    combine_adjacent=False,
+                    show_legend=True,
+                )
+        # Visualization and response components
+        with gr.Accordion("Generation Settings", open=False):
+            with gr.Row():
+                gen_length = gr.Slider(
+                    minimum=16, maximum=128, value=64, step=8,
+                    label="Generation Length"
+                )
+                steps = gr.Slider(
+                    minimum=8, maximum=64, value=32, step=4,
+                    label="Denoising Steps"
+                )
+            
+            
+            visualization_delay = gr.Slider(
+                minimum=0.0, maximum=1.0, value=0.1, step=0.1, visible=False,
+                label="Visualization Delay (seconds)"
+            )
+        
+        # Current response text box
+        current_response = gr.Textbox(
+            label="Current Response",
+            placeholder="The assistant's response will appear here...",
+            lines=3,
+            visible=False
+        )
+        
+        # Clear button
+        clear_btn = gr.Button("Clear Conversation")
+        
+        # HELPER FUNCTIONS
+        def add_message(history, message, response):
+            """Add a message pair to the history and return the updated history"""
+            history = history.copy()
+            history.append([message, response])
+            return history
+            
+        def user_message_submitted(message, history, gen_length, steps, constraints, delay):
+            """Process a submitted user message"""
+            # Skip empty messages
+            if not message.strip():
+                # Return current state unchanged
+                history_for_display = history.copy()
+                return history, history_for_display, "", [], ""
+                
+            # Add user message to history
+            history = add_message(history, message, None)
+            
+            # Format for display - temporarily show user message with empty response
+            history_for_display = history.copy()
+            
+            # Clear the input
+            message_out = ""
+            
+            # Return immediately to update UI with user message
+            return history, history_for_display, message_out, [], ""
+            
+        def bot_response(history, gen_length, steps, constraints, delay):
+            """Generate bot response for the latest message"""
+            if not history:
+                return history, [], ""
+                
+            # Get the last user message
+            last_user_message = history[-1][0]
+            
+            try:
+                # Format all messages except the last one (which has no response yet)
+                messages = format_chat_history(history[:-1])
+                
+                # Add the last user message
+                messages.append({"role": "user", "content": last_user_message})
+                
+                # Parse constraints
+                parsed_constraints = parse_constraints(constraints)
+                
+                # Generate response with visualization
+                vis_states, response_text = generate_response_with_visualization(
+                    model, tokenizer, device, 
+                    messages, 
+                    gen_length=gen_length, 
+                    steps=steps,
+                    constraints=parsed_constraints
+                )
+                
+                # Update history with the assistant's response
+                history[-1][1] = response_text
+                
+                # Return the initial state immediately
+                yield history, vis_states[0], response_text
+                
+                # Then animate through visualization states
+                for state in vis_states[1:]:
+                    time.sleep(delay)
+                    yield history, state, response_text
+                    
+            except Exception as e:
+                error_msg = f"Error: {str(e)}"
+                print(error_msg)
+                
+                # Show error in visualization
+                error_vis = [(error_msg, "red")]
+                
+                # Don't update history with error
+                yield history, error_vis, error_msg
+        
+        def clear_conversation():
+            """Clear the conversation history"""
+            return [], [], "", []
+        
+        # EVENT HANDLERS
+        
+        # Clear button handler
+        clear_btn.click(
+            fn=clear_conversation,
+            inputs=[],
+            outputs=[chat_history, chatbot_ui, current_response, output_vis]
+        )
+        
+        # User message submission flow (2-step process)
+        # Step 1: Add user message to history and update UI
+        msg_submit = user_input.submit(
+            fn=user_message_submitted,
+            inputs=[user_input, chat_history, gen_length, steps, constraints_input, visualization_delay],
+            outputs=[chat_history, chatbot_ui, user_input, output_vis, current_response]
+        )
+        
+        # Also connect the send button
+        send_click = send_btn.click(
+            fn=user_message_submitted,
+            inputs=[user_input, chat_history, gen_length, steps, constraints_input, visualization_delay],
+            outputs=[chat_history, chatbot_ui, user_input, output_vis, current_response]
+        )
+        
+        # Step 2: Generate bot response
+        # This happens after the user message is displayed
+        msg_submit.then(
+            fn=bot_response,
+            inputs=[chat_history, gen_length, steps, constraints_input, visualization_delay],
+            outputs=[chatbot_ui, output_vis, current_response]
+        )
+        
+        send_click.then(
+            fn=bot_response,
+            inputs=[chat_history, gen_length, steps, constraints_input, visualization_delay],
+            outputs=[chatbot_ui, output_vis, current_response]
+        )
+        
+    return demo
+
+# Launch the demo
+if __name__ == "__main__":
+    demo = create_chatbot_demo()
+    demo.queue().launch(share=True)