app.py

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()