import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces import torch import torch.nn.functional as F from transformers import AutoTokenizer, AutoConfig import gradio as gr import sys import time # The HF model repo has configuration_sdar.py but NOT modeling_sdar.py. # We provide our patched modeling_sdar.py locally (removes flash_attn hard dep). sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from configuration_sdar import SDARConfig from modeling_sdar import SDARForCausalLM MODEL_ID = "SJTU-DENG-Lab/MBD-Math-SDAR-8B-Chat-b32" BLOCK_SIZE = 32 MASK_TOKEN_ID = 151669 EOS_TOKEN_ID = 151643 MAX_NEW_BLOCKS = 8 # 8 blocks × 32 tokens = 256 tokens max NUM_DIFFUSION_STEPS = 32 # denoising steps per block print("Loading tokenizer...") tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) tokenizer.padding_side = "left" print("Loading config...") config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True) # Force SDPA attention config._attn_implementation = "sdpa" config.attn_implementation = "sdpa" print("Loading model...") model = SDARForCausalLM.from_pretrained( MODEL_ID, config=config, torch_dtype=torch.bfloat16, ).to("cuda") model.eval() print("Model loaded successfully!") def block_diffusion_generate( input_ids: torch.Tensor, attention_mask: torch.Tensor, max_new_blocks: int = 8, num_diffusion_steps: int = 32, temperature: float = 0.0, ) -> torch.Tensor: """ Block diffusion generation: generate new blocks of tokens one at a time. Each block is initialized with MASK tokens and iteratively denoised. """ batch_size, seq_len = input_ids.shape # Pad sequence to block boundary remainder = seq_len % BLOCK_SIZE if remainder > 0: pad_len = BLOCK_SIZE - remainder pad_ids = torch.full((batch_size, pad_len), tokenizer.pad_token_id, dtype=input_ids.dtype, device=input_ids.device) input_ids = torch.cat([input_ids, pad_ids], dim=1) pad_mask = torch.zeros((batch_size, pad_len), dtype=attention_mask.dtype, device=attention_mask.device) attention_mask = torch.cat([attention_mask, pad_mask], dim=1) seq_len = input_ids.shape[1] generated_ids = input_ids gen_attention_mask = attention_mask for block_idx in range(max_new_blocks): # Create a new block of MASK tokens new_block = torch.full((batch_size, BLOCK_SIZE), MASK_TOKEN_ID, dtype=generated_ids.dtype, device=generated_ids.device) new_mask = torch.ones((batch_size, BLOCK_SIZE), dtype=gen_attention_mask.dtype, device=gen_attention_mask.device) # Append the masked block current_ids = torch.cat([generated_ids, new_block], dim=1) current_mask = torch.cat([gen_attention_mask, new_mask], dim=1) # Iteratively denoise the masked block for step in range(num_diffusion_steps): with torch.no_grad(): outputs = model( input_ids=current_ids, attention_mask=current_mask, use_cache=False, ) logits = outputs.logits # [batch, seq, vocab] # Only look at logits for masked positions mask_positions = (current_ids[0] == MASK_TOKEN_ID).nonzero(as_tuple=True)[0] if len(mask_positions) == 0: break # Get logits at masked positions masked_logits = logits[0, mask_positions, :] # [num_masked, vocab] if temperature == 0.0: # Greedy: pick the argmax new_tokens = masked_logits.argmax(dim=-1) else: # Sample probs = F.softmax(masked_logits / temperature, dim=-1) new_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1) # Determine how many tokens to reveal this step # In block diffusion, we reveal a fraction of masked tokens each step num_to_reveal = max(1, len(mask_positions) // (num_diffusion_steps - step)) if step < num_diffusion_steps - 1: # Reveal only some tokens reveal_indices = mask_positions[:num_to_reveal] current_ids[0, reveal_indices] = new_tokens[:num_to_reveal] else: # Final step: reveal all remaining current_ids[0, mask_positions] = new_tokens # Check if the new block is all EOS (generation complete) new_block_tokens = current_ids[0, seq_len:seq_len + BLOCK_SIZE] generated_ids = current_ids[:, :seq_len + BLOCK_SIZE] gen_attention_mask = current_mask[:, :seq_len + BLOCK_SIZE] seq_len = generated_ids.shape[1] # If the first token of the new block is EOS, stop if new_block_tokens[0].item() == EOS_TOKEN_ID: break return generated_ids, gen_attention_mask def format_response(text: str) -> str: """Clean up the model response.""" # Remove special tokens for token in ["<|MASK|>", "<|im_start|>", "<|im_end|>", "<|endoftext|>"]: text = text.replace(token, "") return text.strip() @spaces.GPU(duration=120) def chat( message: str, history: list, system_prompt: str, temperature: float, max_blocks: int, num_steps: int, ): """Chat function for Gradio ChatInterface.""" messages = [] if system_prompt: messages.append({"role": "system", "content": system_prompt}) for h in history: if h["role"] == "user": messages.append({"role": "user", "content": h["content"]}) elif h["role"] == "assistant": messages.append({"role": "assistant", "content": h["content"]}) messages.append({"role": "user", "content": message}) # Apply chat template input_ids = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", ).to("cuda") attention_mask = torch.ones_like(input_ids) start_time = time.time() # Generate using block diffusion output_ids, _ = block_diffusion_generate( input_ids=input_ids, attention_mask=attention_mask, max_new_blocks=int(max_blocks), num_diffusion_steps=int(num_steps), temperature=temperature, ) elapsed = time.time() - start_time # Decode only the new tokens new_tokens = output_ids[0, input_ids.shape[1]:] response = tokenizer.decode(new_tokens, skip_special_tokens=True) response = format_response(response) if not response: response = "(model produced no output)" return response EXAMPLES = [ ["What is 2 + 2?"], ["Solve the equation 3x - 6 = 12 for x."], ["Explain the Pythagorean theorem."], ["What is the derivative of x^2 + 3x?"], ["If a train travels 60 mph for 2.5 hours, how far does it go?"], ] demo = gr.ChatInterface( fn=chat, type="messages", title="Multi-Block Diffusion Language Models", description=( "Chat with **MBD-Math-SDAR-8B-Chat-b32**, a Multi-Block Diffusion Language Model " "from SJTU-DENG-Lab. This model generates text via iterative block diffusion " "(non-autoregressive) instead of standard token-by-token generation. " "Trained on math reasoning data.\n\n" "📄 Paper: [Multi-Block Diffusion Language Models](https://huggingface.co/papers/2606.29215)\n" "🐙 GitHub: [SJTU-DENG-Lab/mbd-lms](https://github.com/SJTU-DENG-Lab/mbd-lms)\n" "🤗 Model: [SJTU-DENG-Lab/MBD-Math-SDAR-8B-Chat-b32](https://huggingface.co/SJTU-DENG-Lab/MBD-Math-SDAR-8B-Chat-b32)" ), additional_inputs=[ gr.Textbox(value="You are a helpful math assistant.", label="System Prompt"), gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Temperature (0 = greedy)"), gr.Slider(1, 16, value=8, step=1, label="Max New Blocks (×32 tokens)"), gr.Slider(8, 64, value=32, step=4, label="Diffusion Steps per Block"), ], examples=EXAMPLES, cache_examples=False, ) if __name__ == "__main__": demo.launch()