upload demo

app.py

@@ -0,0 +1,648 @@
+# dream_app.py
+import torch
+import numpy as np
+import gradio as gr
+#import spaces # Ensure spaces is installed if needed for GPU decorator
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModel, AutoConfig
+import time
+import re
+from typing import List, Dict, Tuple, Optional, Any # Added Any
+import torch.distributions as dists # Added import
+import traceback # For better error printing
+
+# --- START: Copied Helper functions from generation_utils.py ---
+# [Keep the copied functions: top_p_logits, top_k_logits, sample_tokens]
+def top_p_logits(logits, top_p=None):
+    """ Applies top-p filtering to logits. """
+    if top_p is None or top_p >= 1.0:
+        return logits
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+    sorted_indices_to_remove = cumulative_probs > top_p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = 0
+    mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)
+    mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove)
+    logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min)
+    return logits
+
+def top_k_logits(logits, top_k=None):
+    """ Applies top-k filtering to logits. """
+    if top_k is None or top_k <= 0:
+        return logits
+    top_k = min(top_k, logits.size(-1))
+    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+    logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)
+    return logits
+
+def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False, use_ori_logits = True):
+    """ Samples tokens based on logits and calculates confidence. """
+    original_dtype = logits.dtype
+    logits = logits.to(torch.float32)
+    if use_ori_logits:
+        ori_logits = logits.clone()
+    if temperature > 0:
+        safe_temp = max(temperature, 1e-6)
+        logits = logits / safe_temp
+    if top_p is not None and 0.0 < top_p < 1.0:
+        logits = top_p_logits(logits, top_p)
+    if top_k is not None and top_k > 0:
+        logits = top_k_logits(logits, top_k)
+
+    is_all_neg_inf = torch.all(logits <= torch.finfo(logits.dtype).min, dim=-1, keepdim=True)
+    if torch.any(is_all_neg_inf):
+        uniform_logits = torch.zeros_like(logits)
+        logits = torch.where(is_all_neg_inf, uniform_logits, logits)
+
+    probs = torch.softmax(logits, dim=-1)
+    probs = torch.clamp(probs, min=0.0)
+    prob_sum = probs.sum(dim=-1, keepdim=True)
+    safe_prob_sum = torch.max(prob_sum, torch.tensor(1e-12, device=probs.device, dtype=probs.dtype))
+    probs = probs / safe_prob_sum
+    probs = torch.nan_to_num(probs, nan=0.0)
+
+    if temperature > 0:
+        x0 = dists.Categorical(probs=probs).sample()
+        if use_ori_logits:
+            confidence = torch.gather(torch.softmax(ori_logits, dim=-1), -1, x0.unsqueeze(-1)).squeeze(-1)
+        else:
+            confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
+    else:
+        confidence, x0 = probs.max(dim=-1)
+
+    if margin_confidence:
+        if use_ori_logits:
+            sorted_probs, _ = torch.sort(torch.softmax(ori_logits, dim=-1), dim=-1, descending=True)
+        else:
+            sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)
+        top1_probs = sorted_probs[..., 0]
+        top2_probs = sorted_probs[..., 1] if sorted_probs.shape[-1] > 1 else top1_probs
+        confidence = top1_probs - top2_probs
+    elif neg_entropy: # Use elif to avoid calculating entropy if margin_confidence was True
+        epsilon = 1e-10
+        log_probs = torch.log(probs + epsilon)
+        confidence = torch.sum(probs * log_probs, dim=-1) # Negative entropy
+    # Else: confidence is just the probability of the sampled token if temperature > 0, or max prob otherwise
+
+    confidence = torch.nan_to_num(confidence, nan=0.0)
+    return confidence, x0
+# --- END: Copied Helper functions ---
+
+
+# --- Model Loading and Constants ---
+# Load model configuration to get special token IDs
+config = AutoConfig.from_pretrained("Dream-org/Dream-v0-Instruct-7B", trust_remote_code=True)
+model_path = "Dream-org/Dream-Coder-v0-Instruct-7B"
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"Using device: {device}")
+
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+print("Loading model...")
+model = AutoModel.from_pretrained(
+    model_path,
+    torch_dtype=torch.bfloat16 if device == 'cuda' else torch.float32,
+    trust_remote_code=True,
+    attn_implementation="sdpa" # Explicitly request SDPA
+)
+model = model.to(device).eval()
+print("Model loaded.")
+
+MASK_TOKEN = tokenizer.mask_token
+MASK_ID = tokenizer.mask_token_id
+PAD_ID = tokenizer.pad_token_id
+EOS_ID = tokenizer.eos_token_id
+
+if MASK_ID is None:
+    raise ValueError("Cannot determine MASK_ID. Check model's tokenizer configuration.")
+
+SPECIAL_TOKEN_IDS = {PAD_ID, EOS_ID, MASK_ID}
+try:
+    IM_START_ID = tokenizer.convert_tokens_to_ids("<|im_start|>")
+    IM_END_ID = tokenizer.convert_tokens_to_ids("<|im_end|>")
+    SPECIAL_TOKEN_IDS.add(IM_START_ID)
+    SPECIAL_TOKEN_IDS.add(IM_END_ID)
+except KeyError:
+    print("Warning: <|im_start|> or <|im_end|> not found in tokenizer vocab.")
+    IM_START_ID = None
+    IM_END_ID = None
+
+
+# --- Helper Functions ---
+def parse_constraints(constraints_text: str) -> Dict[int, List[int]]:
+    """ Parses word constraints. """
+    constraints = {}
+    if not constraints_text: return constraints
+    parts = constraints_text.split(',')
+    for part in parts:
+        part = part.strip()
+        if ':' not in part: continue
+        pos_str, word = part.split(':', 1)
+        try:
+            pos = int(pos_str.strip())
+            word = word.strip()
+            token_ids = []
+            if word:
+                text_to_encode = (" " + word) if (pos > 0 and not word.startswith(" ")) else word
+                token_ids = tokenizer.encode(text_to_encode, add_special_tokens=False)
+            if token_ids and pos >= 0: constraints[pos] = token_ids
+            elif not token_ids and word: print(f"Warning: Could not tokenize constraint word '{word}'")
+        except ValueError: print(f"Warning: Invalid position '{pos_str}' in constraint part '{part}'")
+        except Exception as e: print(f"Warning: Error processing constraint '{part}': {e}")
+    return constraints
+
+# Removed format_chat_history as the state will now be in the correct format
+
+def apply_constraints_to_state(
+    x: torch.Tensor,
+    prompt_length: int,
+    total_length: int,
+    parsed_constraints: Dict[int, List[int]],
+    current_step: Optional[int] = None
+) -> torch.Tensor:
+    """ Applies constraints directly to the state tensor `x`. """
+    modified_x = x.clone()
+    for rel_pos, word_token_ids in parsed_constraints.items():
+        abs_start_pos = prompt_length + rel_pos
+        abs_end_pos = abs_start_pos + len(word_token_ids)
+        if abs_start_pos < total_length and abs_end_pos <= total_length:
+            try:
+                constraint_tensor = torch.tensor(word_token_ids, dtype=torch.long, device=modified_x.device)
+                modified_x[0, abs_start_pos:abs_end_pos] = constraint_tensor
+            except IndexError: print(f"Warning (Step {current_step}): Constraint at {rel_pos} ('{tokenizer.decode(word_token_ids)}') goes out of bounds.")
+            except Exception as e: print(f"Warning (Step {current_step}): Failed to apply constraint at {rel_pos}: {e}")
+    return modified_x
+
+
+# --- Core Generation Logic with Live Visualization ---
+
+#@spaces.GPU
+@torch.no_grad()
+def generate_dream_response(
+    history_dict_list: List[Dict[str, str]], # Now expects list of dicts
+    gen_length: int,
+    steps: int,
+    constraints_text: str,
+    temperature: float,
+    top_p: Optional[float],
+    top_k: Optional[int],
+    alg: str,
+    alg_temp: Optional[float],
+    pad_penalty: Optional[float],
+    visualization_delay: float
+    ) -> List[Tuple[str, str]]:
+    """ Generates text step-by-step and yields visualization states live. """
+
+    if not history_dict_list or history_dict_list[-1]['role'] != 'user':
+        # Handle cases where history is empty or doesn't end with user message
+        # This check might be redundant if add_user_message handles it, but good for safety.
+        yield history_dict_list, [("No user message found.", "red")], ""
+        return
+
+    # --- 1. Preparation ---
+    parsed_constraints = parse_constraints(constraints_text)
+
+    # Prepare history for the model template (don't include the empty assistant msg yet)
+    history_for_template = history_dict_list # Already in list-of-dicts format
+
+    try:
+        inputs = tokenizer.apply_chat_template(
+            history_for_template, # Pass the list of dicts directly
+            return_tensors="pt",
+            return_dict=True,
+            add_generation_prompt=True # Crucial: Adds the '<|im_start|>assistant\n' turn
+        )
+        input_ids = inputs.input_ids.to(device)
+        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}")
+        traceback.print_exc()
+        yield history_dict_list, [("Error preparing input.", "red")], ""
+        return
+
+    eps = 1e-3
+    top_p_val = top_p if top_p is not None and 0.0 < top_p < 1.0 else None
+    top_k_val = top_k if top_k is not None and top_k > 0 else None
+    alg_temp_val = alg_temp if alg in ['maskgit_plus', 'topk_margin', 'entropy'] and alg_temp is not None and alg_temp > 0 else None
+
+    # --- 2. Initialize Generation State ---
+    total_length = prompt_length + gen_length
+    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)
+
+    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)
+    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) # [B, 1, 1, N]
+
+    timesteps = torch.linspace(1, eps, steps + 1, device=device)
+    x = apply_constraints_to_state(x, prompt_length, total_length, parsed_constraints, current_step=-1)
+
+    # --- 3. Visualization & History Setup ---
+    previous_tokens_vis = None
+    final_response_text = ""
+    # The history_dict_list is the state we update and yield for the chatbot UI
+    # Add the empty assistant message placeholder *to the history state* now
+    history_dict_list.append({"role": "assistant", "content": ""})
+
+    # --- 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"
+        vis_data_initial.append((display_token, color))
+
+    previous_tokens_vis = initial_generated_tokens
+    # Yield the history (which now includes the empty assistant turn)
+    yield history_dict_list, vis_data_initial, ""
+    time.sleep(visualization_delay)
+
+    # --- 5. Step-by-Step Diffusion Loop ---
+    eps = 1e-3
+    try:
+        start_time = time.time()
+        for i in range(steps):
+            mask_index = (x == MASK_ID)
+            if not mask_index.any():
+                 print(f"No mask tokens left at step {i}. Stopping early.")
+                 break
+
+            outputs = model(
+                input_ids=x,
+                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)
+
+            mask_logits = logits[mask_index]
+            if mask_logits.numel() == 0:
+                 print(f"No masked tokens found for logit selection at step {i}. Stopping.")
+                 break
+
+            t = timesteps[i]
+            s = timesteps[i + 1]
+            mask_logits[:, PAD_ID] += pad_penalty * torch.log(1 - t + eps)
+            x_new_masked_part = torch.full_like(x[mask_index], MASK_ID, device=device, dtype=torch.long)
+
+            # [Keep sampling logic the same - 'origin' and confidence-based]
+            if alg == 'origin':
+                p_transfer = (1.0 - s / t) if i < steps - 1 else 1.0
+                num_masked = mask_logits.shape[0]
+                transfer_indices_relative = torch.rand(num_masked, device=device) < p_transfer
+                logits_to_sample = mask_logits[transfer_indices_relative]
+                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)
+                    if transfer_indices_relative.sum() == sampled_tokens.numel(): # Basic check
+                        x_new_masked_part[transfer_indices_relative] = sampled_tokens
+                    else: print(f"Warning step {i} (origin): Mismatch transfer indices and sampled tokens.")
+
+
+            else: # Confidence-based
+                use_margin = (alg == 'topk_margin')
+                use_entropy = (alg == 'entropy')
+                confidence, x0_candidates = sample_tokens(mask_logits, temperature=temperature, top_p=top_p_val, top_k=top_k_val, margin_confidence=use_margin, neg_entropy=use_entropy)
+
+                num_mask_token = mask_logits.shape[0]
+                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:
+                    num_samples = min(number_transfer_tokens, num_mask_token)
+                    if num_samples > 0:
+                        transfer_indices_relative = torch.tensor([], dtype=torch.long, device=device) # Init empty
+                        if alg_temp_val is None or alg_temp_val <= 0: # Top-k
+                            sort_metric = confidence
+                            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 temp
+                            if confidence.numel() > 0:
+                                conf_probs = confidence / alg_temp_val
+                                conf_probs = torch.nan_to_num(conf_probs, nan=0.0, posinf=1e9, neginf=-1e9)
+                                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)
+                                conf_probs = torch.nan_to_num(conf_probs, nan=0.0)
+                                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
+                                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.")
+                                        sort_metric = 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: # Fallback if probs invalid for multinomial
+                                    # print(f"Warning step {i}: Invalid probabilities for multinomial sampling (sum={final_prob_sum_check:.4f}). Falling back to top-k.")
+                                    sort_metric = 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)
+
+                        # Apply transfer
+                        if transfer_indices_relative.numel() > 0:
+                             if x0_candidates.numel() > 0 and transfer_indices_relative.max() < x0_candidates.shape[0]:
+                                 if transfer_indices_relative.max() < x_new_masked_part.shape[0]:
+                                      x_new_masked_part[transfer_indices_relative] = x0_candidates[transfer_indices_relative].clone()
+                                 else: print(f"Warning step {i}: transfer_indices out of bounds for x_new_masked_part.")
+                             else: print(f"Warning step {i}: transfer_indices out of bounds for x0_candidates or x0_candidates empty.")
+
+
+            x[mask_index] = x_new_masked_part
+            x = apply_constraints_to_state(x, prompt_length, total_length, parsed_constraints, current_step=i)
+
+            # --- Yield Visualization & Update History ---
+            current_generated_tokens = x[0, prompt_length:].cpu()
+            vis_data = []
+            # [Visualization formatting logic remains the same]
+            for j in range(gen_length):
+                current_tok_id = current_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}]"
+                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))
+
+            previous_tokens_vis = current_generated_tokens
+
+            intermediate_response_tokens = x[0, prompt_length:]
+            intermediate_response_text = tokenizer.decode(
+                intermediate_response_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=True
+            ).strip()
+
+            # --- Update the *last* message in history_dict_list ---
+            history_dict_list[-1]['content'] = intermediate_response_text
+
+            # Yield the updated history list (for chatbot UI), vis data, and response text
+            yield history_dict_list, vis_data, intermediate_response_text
+            time.sleep(visualization_delay)
+
+        end_time = time.time()
+        print(f"Dream generation finished in {end_time - start_time:.2f} seconds.")
+
+        # --- 6. Final Processing & Yield ---
+        final_sequence = x[0]
+        response_tokens = final_sequence[prompt_length:]
+        final_response_text = tokenizer.decode(
+            response_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=True
+        ).strip()
+
+        # Ensure the final text is in the history object before the last yield
+        history_dict_list[-1]['content'] = final_response_text
+
+        final_generated_tokens = x[0, prompt_length:].cpu()
+        vis_data_final = []
+        # [Final visualization formatting logic remains 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}]"
+            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_final.append((token_to_display, color))
+
+        yield history_dict_list, vis_data_final, final_response_text
+        print("Visualization streaming complete.")
+
+    except Exception as e:
+        print(f"Error during generation or processing: {e}")
+        traceback.print_exc()
+        # Attempt to add error message to history if possible
+        if history_dict_list and history_dict_list[-1]['role'] == 'assistant':
+             history_dict_list[-1]['content'] = f"Error: {e}"
+        yield history_dict_list, [("Error during generation.", "red")], f"Error: {e}" # Also show error in text box
+        return
+
+input_examples = ["""You will be given a question (problem specification) and will generate a correct Python program that matches the specification and passes all tests. You will NOT return anything except for the program.
+
+Question:
+There are N cells in a row, numbered 1 to N.
+For each 1 \\leq i < N, cells i and i+1 are adjacent.
+Initially, cell i is painted with color i.
+You are given Q queries. Process them in order. Each query is of one of the following two types.
+
+- 1 x c: Repaint the following to color c: all reachable cells reachable from cell x by repeatedly moving to an adjacent cell painted in the same color as the current cell.
+- 2 c: Print the number of cells painted with color c.
+
+Input
+
+The input is given from Standard Input in the following format:
+N Q
+\\mathrm{query}_1
+\\vdots
+\\mathrm{query}_Q
+
+Each query is given in one of the following two formats:
+1 x c
+
+2 c
+
+Output
+
+Let q be the number of queries of the second type. Print q lines.
+The i-th line should contain the answer to the i-th such query.
+
+Constraints
+
+
+- 1 \\leq N \\leq 5 \times 10^5
+- 1 \\leq Q \\leq 2 \times 10^5
+- In queries of the first type, 1 \\leq x \\leq N.
+- In queries of the first and second types, 1 \\leq c \\leq N.
+- There is at least one query of the second type.
+- All input values are integers.
+
+Sample Input 1
+
+5 6
+1 5 4
+1 4 2
+2 2
+1 3 2
+1 2 3
+2 3
+
+Sample Output 1
+
+3
+4
+
+The queries recolor the cells as shown in the figure.
+
+Read the inputs from stdin solve the problem and write the answer to stdout (do not directly test on the sample inputs). Enclose your code within delimiters as follows. Ensure that when the python program runs, it reads the inputs, runs the algorithm and writes output to STDOUT.
+
+```python
+# YOUR CODE HERE
+```
+""","""Please provide a self-contained Python script that solves the following problem in a markdown code block:
+Calculates the average of the sums of absolute differences between each pair of consecutive numbers for all permutations of a given list. Each permutation is shuffled before calculating the differences. Args: - numbers (list): A list of numbers. Default is numbers from 1 to 10.
+The function should output with:
+    float: The average of the sums of absolute differences for each shuffled permutation of the list.
+You should write self-contained code starting with:
+```
+import itertools
+from random import shuffle
+def task_func(numbers=list(range(1, 3))):
+```""","You are given a Python function and an assertion containing an input to the function. Complete the assertion with a literal (no unsimplified expressions, no function calls) containing the output when executing the provided code on the given input, even if the function is incorrect or incomplete. Do NOT output any extra information. Execute the program step by step before arriving at an answer, and provide the full assertion with the correct output in [ANSWER] and [/ANSWER] tags, following the examples.\n\n[PYTHON]\ndef f(s):\n    s = s + s\n    return \"b\" + s + \"a\"\nassert f(\"hi\") == ??\n[/PYTHON]\n[THOUGHT]\nLet's execute the code step by step:\n\n1. The function f is defined, which takes a single argument s.\n2. The function is called with the argument \"hi\", so within the function, s is initially \"hi\".\n3. Inside the function, s is concatenated with itself, so s becomes \"hihi\".\n4. The function then returns a new string that starts with \"b\", followed by the value of s (which is now \"hihi\"), and ends with \"a\".\n5. The return value of the function is therefore \"bhihia\".\n[/THOUGHT]\n[ANSWER]\nassert f(\"hi\") == \"bhihia\"\n[/ANSWER]\n\n[PYTHON]\ndef f(nums):\n    output = []\n    for n in nums:\n        output.append((nums.count(n), n))\n    output.sort(reverse=True)\n    return output\nassert f([1, 1, 3, 1, 3, 1]) == ??\n[/PYTHON]\n[THOUGHT]\n", "Write a quick sort algorithm."]
+
+labels = [ 'Sketch-First Generation (from LiveCodeBench)', 'Left-to-Right Generation (from BigCodeBench)',' Interleaved Reasoning Generation (from CRUXEval)', ' Quicksort algorithm']
+
+# --- Gradio UI ---
+css = '''
+.category-legend{display:none}
+'''
+def create_chatbot_demo():
+    with gr.Blocks(css=css) as demo:
+        gr.Markdown("# Dream-Coder-7B-Instruct ")
+        gr.Markdown(
+            "[[Model Card](https://huggingface.co/Dream-org/Dream-Coder-v0-Instruct-7B)] "
+            "[[Blog](https://hkunlp.github.io/blog/2025/dream-coder/)]"
+            "[[Github](https://github.com/DreamLM/Dream-Coder)]"
+        )
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                chatbot_ui = gr.Chatbot(
+                    label="Conversation",
+                    height=500,
+                    show_copy_button=True,
+                    bubble_full_width=False,
+                    value=[], # Initialize empty
+                    type="messages" # Crucial: Use the messages format
+                )
+                with gr.Group():
+                    with gr.Row():
+                        user_input = gr.Textbox(
+                            label="Your Message", placeholder="Type your message here...",
+                            scale=7, autofocus=True, show_label=False, container=False
+                        )
+                        send_btn = gr.Button("Send", scale=1, variant="primary")
+                constraints_input = gr.Textbox(
+                    label="Word Constraints (Optional)",
+                    info="Format: 'pos:word, pos:word,...'. Example: '0:Once, 5:upon'",
+                    placeholder="0:Hello, 10:world", value="",
+                    visible=False
+                )
+                clear_btn = gr.Button("Clear Conversation")
+                examples = gr.Examples(
+                    examples = input_examples,
+                    example_labels = labels,
+                    inputs = user_input
+                )
+                
+
+            with gr.Column(scale=2):
+                output_vis = gr.HighlightedText(
+                    label="Denoising Process Visualization", combine_adjacent=False,
+                    show_legend=True, interactive=False
+                )
+                response_text_display = gr.Textbox(
+                    label="Current/Final Response", interactive=False, lines=5, visible=False
+                )
+
+        # [Keep Accordion with Generation Settings the same]
+        with gr.Accordion("Generation Settings", open=False):
+             with gr.Row():
+                gen_length = gr.Slider(minimum=16, maximum=1024, value=512, step=8, label="Max New Tokens")
+                steps = gr.Slider(minimum=8, maximum=512, value=512, step=8, label="Diffusion Steps")
+             with gr.Row():
+                temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.05, label="Temperature (0 = greedy)")
+                alg_temp = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.05, label="Algorithm Temperature")
+             with gr.Row():
+                top_p = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.05, label="Top-P (0 disables)")
+                top_k = gr.Slider(minimum=0, maximum=200, value=0, step=5, label="Top-K (0 disables)")
+             with gr.Row():
+                pad_penalty = gr.Slider(minimum=0, maximum=5,value=3, step=0.5, label="Pad Penalty")
+                remasking_strategy = gr.Radio(choices=['origin', 'maskgit_plus', 'topk_margin', 'entropy'], value='entropy', label="Generation Algorithm")
+             with gr.Row():
+                visualization_delay = gr.Slider(minimum=0.0, maximum=0.5, value=0.0, step=0.01, label="Visualization Delay (s)")
+
+        
+
+        # --- Event Handlers ---
+
+        # User function: Appends user message to the history (list of dicts)
+        def add_user_message(message: str, history: List[Dict[str, str]]):
+            if not message.strip():
+                gr.Warning("Please enter a message.")
+                return history, "" # Return unchanged history, empty input
+            history.append({"role": "user", "content": message})
+            # Return updated history for chatbot UI, and clear input box
+            return history, ""
+
+        # Bot function (now the generator)
+        # Inputs: Chatbot history (list of dicts), generation params
+        # Outputs: Chatbot history (updated list of dicts), visualization, response text
+        generation_inputs = [
+            chatbot_ui, # Pass chatbot state directly (list of dicts)
+            gen_length, steps, constraints_input,
+            temperature, top_p, top_k, remasking_strategy, alg_temp, pad_penalty,
+            visualization_delay
+        ]
+        generation_outputs = [chatbot_ui, output_vis, response_text_display]
+
+        # --- Connect UI elements ---
+
+        # Textbox Submission (Enter key)
+        submit_listener = user_input.submit(
+            fn=add_user_message,
+            inputs=[user_input, chatbot_ui],
+            outputs=[chatbot_ui, user_input] # Update chatbot UI and clear input
+        ).then(
+            fn=generate_dream_response,
+            inputs=generation_inputs,
+            outputs=generation_outputs,
+            show_progress="hidden" # Hide default progress bar
+        )
+
+        # Send Button Click
+        click_listener = send_btn.click(
+            fn=add_user_message,
+            inputs=[user_input, chatbot_ui],
+            outputs=[chatbot_ui, user_input] # Update chatbot UI and clear input
+        ).then(
+            fn=generate_dream_response,
+            inputs=generation_inputs,
+            outputs=generation_outputs,
+            show_progress="hidden"
+        )
+
+        # Clear Button Action
+        clear_btn.click(
+            lambda: ([], [], ""), # Function to return empty values
+            inputs=[],
+            outputs=[chatbot_ui, output_vis, response_text_display], # Clear chatbot, vis, text
+            queue=False # No need to queue clearing usually
+        )
+
+
+
+    return demo
+
+# --- Launch ---
+if __name__ == "__main__":
+    demo = create_chatbot_demo()
+    demo.queue().launch(debug=True, share=True)

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+accelerate
+sentencepiece