import argparse import html import json import re from pathlib import Path import torch from transformers import AutoModelForCausalLM, AutoTokenizer DEFAULT_MODEL_PATH = "Zigeng/DMax-Math-16B" DEFAULT_PROMPT = "Solve 37 * 48 and explain the intermediate reasoning briefly." STAGE_SIDEBAR_WIDTH = 320 STAGE_WIDTH = 720 STAGE_MIN_HEIGHT = 620 STAGE_FINAL_HEIGHT_VH = 72 STAGE_FINAL_MAX_HEIGHT = 860 TOKEN_GRID_MIN_HEIGHT = 600 FINAL_MARKDOWN_BASE_FONT_SIZE = 20 FINAL_MARKDOWN_MIN_FONT_SIZE = 11 PLAYBACK_INTERVAL_MS = 180 def parse_args(): parser = argparse.ArgumentParser( description="Visualize the full step-by-step decoding process of the diffusion LM." ) parser.add_argument("--model-path", default=DEFAULT_MODEL_PATH) parser.add_argument("--device", default="cuda:0") parser.add_argument("--prompt", default=None) parser.add_argument("--gsm8k-index", type=int, default=1) parser.add_argument("--gen-length", type=int, default=512) parser.add_argument("--block-length", type=int, default=32) parser.add_argument("--steps", type=int, default=32) parser.add_argument("--threshold", type=float, default=0.0) parser.add_argument("--output", default="dllm_demo.html") return parser.parse_args() def load_prompt(args): if args.prompt is not None: return args.prompt try: from datasets import load_dataset ds = load_dataset("openai/gsm8k", "main", split="test") return ds[args.gsm8k_index]["question"] + "\nLet's think step by step\n" except Exception: return DEFAULT_PROMPT def tokenize_prompt(tokenizer, prompt): if hasattr(tokenizer, "apply_chat_template"): return tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], add_generation_prompt=True, tokenize=True, return_tensors="pt", ) return tokenizer(prompt, return_tensors="pt").input_ids def format_token_label(tokenizer, token_id, mask_id, eos_id): if token_id == mask_id: return "[MASK]" if eos_id is not None and token_id == eos_id: return "[EOS]" text = tokenizer.decode( [token_id], skip_special_tokens=False, clean_up_tokenization_spaces=False, ) if not text: text = tokenizer.convert_ids_to_tokens(token_id) if text is None: text = str(token_id) text = ( text.replace("\n", "↵") .replace("\t", "⇥") .replace("\r", "␍") .replace(" ", "·") ) if text == "": text = "∅" if len(text) > 20: text = text[:17] + "..." return text def collect_token_labels(tokenizer, demo_trace): used_ids = set(demo_trace["prompt_token_ids"]) used_ids.update(demo_trace["generated_token_ids"]) used_ids.update(demo_trace["final_token_ids"]) for frame in demo_trace["frames"]: used_ids.update(frame.get("visible_generated_ids", [])) used_ids.update(frame.get("pre_visible_ids", [])) used_ids.update(frame.get("post_visible_ids", [])) used_ids.update(frame.get("top1_token_ids", [])) used_ids.update(frame.get("block_top1_token_ids", [])) mask_id = demo_trace["mask_id"] eos_id = demo_trace["eos_id"] token_labels = {} for token_id in sorted(used_ids): token_labels[str(token_id)] = format_token_label( tokenizer, token_id, mask_id, eos_id ) return token_labels def enrich_demo_trace_for_render(demo_trace): prompt_length = demo_trace["prompt_length"] original_frames = demo_trace["frames"] if not original_frames: return demo_trace enriched_frames = [] for frame in original_frames: generated_before_ids = frame["pre_visible_ids"][prompt_length:] generated_after_ids = frame["post_visible_ids"][prompt_length:] current_block_abs_start = max(prompt_length, frame["block_start"]) current_block_generated_start = max(0, current_block_abs_start - prompt_length) prompt_overlap = max(0, prompt_length - frame["block_start"]) visible_generated_absolute_positions = list( range(prompt_length, frame["window_end"]) ) frame["generated_before_ids"] = generated_before_ids frame["generated_after_ids"] = generated_after_ids frame["visible_generated_ids"] = generated_after_ids frame["visible_generated_absolute_positions"] = visible_generated_absolute_positions frame["current_block_absolute_start"] = current_block_abs_start frame["current_block_generated_start"] = current_block_generated_start frame["block_top1_confidence"] = frame["top1_confidence"][prompt_overlap:] frame["block_top1_token_ids"] = frame["top1_token_ids"][prompt_overlap:] frame["block_input_confidence"] = frame["input_confidence"][prompt_overlap:] frame["block_mask_index_before"] = frame["mask_index_before"][prompt_overlap:] frame["block_token_index_before"] = frame["token_index_before"][prompt_overlap:] frame["block_active_mask"] = frame["active_block_mask"][prompt_overlap:] frame["block_absolute_positions"] = list( range(current_block_abs_start, frame["block_end"]) ) frame["block_prompt_overlap"] = prompt_overlap frame["block_decoded_positions"] = [ pos - prompt_overlap for pos in frame["decoded_positions"] if pos >= prompt_overlap ] frame["decoded_absolute_positions"] = [ current_block_abs_start + pos for pos in frame["block_decoded_positions"] ] frame["state_mode"] = "after_step" enriched_frames.append(frame) first_frame = enriched_frames[0] initial_frame = { "frame_id": -1, "block_id": first_frame["block_id"], "absolute_block_id": first_frame["absolute_block_id"], "step_id": -1, "window_end": first_frame["window_end"], "block_start": first_frame["block_start"], "block_end": first_frame["block_end"], "nfe": 0, "visible_generated_ids": first_frame["generated_before_ids"], "visible_generated_absolute_positions": list( range(prompt_length, first_frame["window_end"]) ), "current_block_absolute_start": first_frame["current_block_absolute_start"], "current_block_generated_start": first_frame["current_block_generated_start"], "block_top1_confidence": [], "block_top1_token_ids": [], "block_input_confidence": first_frame["block_input_confidence"], "block_mask_index_before": first_frame["block_mask_index_before"], "block_token_index_before": first_frame["block_token_index_before"], "block_active_mask": first_frame["block_active_mask"], "block_absolute_positions": first_frame["block_absolute_positions"], "block_prompt_overlap": first_frame["block_prompt_overlap"], "block_decoded_positions": [], "decoded_absolute_positions": [], "same_as_previous": False, "all_confident": False, "converged": False, "convergence_reason": None, "state_mode": "before_first_step", } demo_trace["frames"] = [initial_frame] for idx, frame in enumerate(enriched_frames, start=1): frame["frame_id"] = idx demo_trace["frames"].append(frame) return demo_trace def markdown_to_html(text): escaped = html.escape(text.replace("\r\n", "\n").replace("\r", "\n")) code_blocks = [] def stash_code(match): code = match.group(1).strip("\n") placeholder = f"@@CODEBLOCK{len(code_blocks)}@@" code_blocks.append(f"
{code}
") return placeholder escaped = re.sub(r"```(?:[^\n`]*)\n(.*?)```", stash_code, escaped, flags=re.S) def format_inline(content): content = re.sub(r"`([^`]+)`", r"\1", content) content = re.sub(r"\*\*([^*]+)\*\*", r"\1", content) content = re.sub(r"\*([^*]+)\*", r"\1", content) return content parts = [] paragraph_lines = [] in_ul = False in_ol = False def close_lists(): nonlocal in_ul, in_ol if in_ul: parts.append("") in_ul = False if in_ol: parts.append("") in_ol = False def flush_paragraph(): nonlocal paragraph_lines if paragraph_lines: merged_parts = [] non_empty_lines = [line for line in paragraph_lines if line.strip()] for idx, line in enumerate(non_empty_lines): hard_break = line.endswith(" ") or line.endswith("\\") content = line[:-1] if line.endswith("\\") else line merged_parts.append(format_inline(content.strip())) if idx < len(non_empty_lines) - 1: merged_parts.append("
" if hard_break else " ") merged = "".join(merged_parts) if merged: parts.append(f"

{merged}

") paragraph_lines = [] for raw_line in escaped.split("\n"): line = raw_line stripped = line.strip() if not stripped: flush_paragraph() close_lists() continue heading = re.match(r"^(#{1,6})\s+(.*)$", stripped) unordered = re.match(r"^[-*+]\s+(.*)$", stripped) ordered = re.match(r"^(\d+)\.\s+(.*)$", stripped) code_block = re.fullmatch(r"@@CODEBLOCK(\d+)@@", stripped) if code_block: flush_paragraph() close_lists() parts.append(stripped) elif heading: flush_paragraph() close_lists() level = len(heading.group(1)) parts.append(f"{format_inline(heading.group(2))}") elif unordered: flush_paragraph() if in_ol: parts.append("") in_ol = False if not in_ul: parts.append("") in_ul = False if not in_ol: parts.append(f'
    ') in_ol = True parts.append(f"
  1. {format_inline(ordered.group(2))}
    2. ") else: close_lists() paragraph_lines.append(line) flush_paragraph() close_lists() rendered = "\n".join(parts) for idx, block in enumerate(code_blocks): rendered = rendered.replace(f"@@CODEBLOCK{idx}@@", block) return rendered def build_demo_payload(prompt, generated_answer, demo_trace, nfe, token_labels): generated_token_count = len(demo_trace["generated_token_ids"]) tpf = generated_token_count / nfe if nfe else 0.0 return { "title": "Diffusion Decoding Demo", "prompt": prompt, "generated_answer": generated_answer, "nfe": nfe, "tpf": tpf, "prompt_length": demo_trace["prompt_length"], "block_length": demo_trace["block_length"], "steps": demo_trace["steps"], "threshold": demo_trace["threshold"], "frames": demo_trace["frames"], "blocks": demo_trace["blocks"], "token_labels": token_labels, "mask_id": demo_trace["mask_id"], "eos_id": demo_trace["eos_id"], "generated_token_ids": demo_trace["generated_token_ids"], "generated_html": markdown_to_html(generated_answer), } def render_html(payload): payload_json = json.dumps(payload, ensure_ascii=False) html = """ Diffusion Decoding Demo

    Diffusion Language Model

    DMax Decoding Process,
    made visible.

    Inspect the full block-wise diffusion process frame by frame.
    TPF
    Generated tokens / NFE
    NFE
    Forward passes
    Generated Tokens
    Final visible answer length
    Threshold
    Mask decode cutoff
    Running

    Visible generated tokens so far

    Previously decoded blocks stay visible and the raw token sequence grows over time. The final generated answer is rendered as markdown only after the full sequence finishes decoding.

    Mask before step Existing token index Refreshed token Decoded mask Committed previous block
    """ html = html.replace("{{", "{").replace("}}", "}") html = html.replace("__PAYLOAD_JSON__", payload_json) html = html.replace("__STAGE_SIDEBAR_WIDTH__", str(STAGE_SIDEBAR_WIDTH)) html = html.replace("__STAGE_WIDTH__", str(STAGE_WIDTH)) html = html.replace("__STAGE_MIN_HEIGHT__", str(STAGE_MIN_HEIGHT)) html = html.replace("__STAGE_FINAL_HEIGHT_VH__", str(STAGE_FINAL_HEIGHT_VH)) html = html.replace("__STAGE_FINAL_MAX_HEIGHT__", str(STAGE_FINAL_MAX_HEIGHT)) html = html.replace("__TOKEN_GRID_MIN_HEIGHT__", str(TOKEN_GRID_MIN_HEIGHT)) html = html.replace( "__FINAL_MARKDOWN_BASE_FONT_SIZE__", str(FINAL_MARKDOWN_BASE_FONT_SIZE) ) html = html.replace( "__FINAL_MARKDOWN_MIN_FONT_SIZE__", str(FINAL_MARKDOWN_MIN_FONT_SIZE) ) html = html.replace("__PLAYBACK_INTERVAL_MS__", str(PLAYBACK_INTERVAL_MS)) return html def main(): args = parse_args() prompt = load_prompt(args) model = AutoModelForCausalLM.from_pretrained( args.model_path, trust_remote_code=True, device_map=args.device, ) model = model.to(torch.bfloat16) model.eval() tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True) input_ids = tokenize_prompt(tokenizer, prompt) demo_trace, nfe, generated_tokens = model.generate_uniform_demo( inputs=input_ids, gen_length=args.gen_length, block_length=args.block_length, steps=args.steps, threshold=args.threshold, ) generated_answer = tokenizer.decode( generated_tokens[0], skip_special_tokens=True, clean_up_tokenization_spaces=False, ) demo_trace = enrich_demo_trace_for_render(demo_trace) token_labels = collect_token_labels(tokenizer, demo_trace) payload = build_demo_payload( prompt=prompt, generated_answer=generated_answer, demo_trace=demo_trace, nfe=nfe, token_labels=token_labels, ) output_path = Path(args.output).expanduser().resolve() output_path.write_text(render_html(payload), encoding="utf-8") print(f"Saved demo to: {output_path}") print(f"Frames: {len(demo_trace['frames'])}") print(f"NFE: {nfe}") print(f"Generated tokens: {generated_tokens.shape[1]}") if __name__ == "__main__": main()