Update README.md

README.md

@@ -124,19 +124,33 @@ def generate(
 ):
     device = model.device
     mask_id = tokenizer.mask_token_id
-    bos_id = tokenizer.bos_token_id
     pad_id = tokenizer.pad_token_id
+    if pad_id is None:
+        pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.mask_token_id
 
-    prompt = torch.tensor(prompt, device=device).long()
-    B = 1
-    T0 = len(prompt)
-    x = prompt
+    if isinstance(prompt, torch.Tensor):
+        x = prompt.to(device).long()
+    else:
+        if isinstance(prompt[0], (list, tuple)):
+            max_len = max(len(p) for p in prompt)
+            x = torch.full((len(prompt), max_len), pad_id, device=device, dtype=torch.long)
+            for i, p in enumerate(prompt):
+                x[i, : len(p)] = torch.tensor(p, device=device)
+        else:
+            x = torch.tensor(prompt, device=device).long()
+    if x.dim() == 1:
+        x = x.unsqueeze(0)
+
+    B = x.size(0)
+    finished = torch.zeros(B, dtype=torch.bool, device=device)
 
     num_blocks = math.ceil(max_new_tokens / block_size)
     steps_per_block = math.ceil(steps / num_blocks)
     generated = 0
 
     while generated < max_new_tokens:
+        if finished.all():
+            break
         T_prefix = x.size(1)
         offset = T_prefix % block_size
         room = block_size if offset == 0 else block_size - offset
@@ -148,7 +162,6 @@ def generate(
 
         out = model(x, attention_mask=attn_pfx, position_ids=pos_pfx, use_cache=True)
         cond_past = out.past_key_values
-        prefix_logits = out.logits[:, -1:, :]
 
         if cfg_scale > 0:
             un_x = x.clone()
@@ -159,6 +172,7 @@ def generate(
             uncond_past = None
 
         block = torch.full((B, cur_len), mask_id, device=device, dtype=torch.long)
+        block[finished] = pad_id
         x = torch.cat([x, block], dim=1)
         T_total = x.size(1)
 
@@ -191,39 +205,69 @@ def generate(
                 logits, x_blk, m_blk, num_transfer[:, t], temperature, remasking
             )
             x[:, T_prefix:T_total] = x_blk_new
-
-        if (x_blk_new == tokenizer.eos_token_id).any():
-            break
+            if tokenizer.eos_token_id is not None:
+                finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1)
+            if finished.all():
+                break
 
         generated += cur_len
+        if finished.all():
+            break
 
     return x
 
 
-device = "cuda"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModelForMaskedLM.from_pretrained("dllm-collection/Qwen2.5-Coder-0.5B-Instruct-diffusion-bd3lm-v0.1", dtype=torch.bfloat16, trust_remote_code=True).to(device).eval()
 tokenizer = AutoTokenizer.from_pretrained("dllm-collection/Qwen2.5-Coder-0.5B-Instruct-diffusion-bd3lm-v0.1", trust_remote_code=True)
 
-prompt = "Lily can run 12 kilometers per hour for 4 hours. After that, she runs 6 kilometers per hour. How many kilometers can she run in 8 hours?"
-m = [
-    {"role": "system", "content": "You are a helpful AI assistant."},
-    {"role": "user", "content": prompt}
+prompts = [
+    [
+        {"role": "system", "content": "You are a helpful AI assistant."},
+        {"role": "user", "content": "Implement a BFS traversal in Python with clear inline comments."},
+    ],
+    [
+        {"role": "system", "content": "You are a helpful AI assistant."},
+        {"role": "user", "content": "Write a concise pytest that checks a Fibonacci implementation."},
+    ],
 ]
-prompt = tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
 
-input_ids = tokenizer(prompt)["input_ids"]
-input_ids = torch.tensor(input_ids).to(device).unsqueeze(0)
-text = generate(model,tokenizer, input_ids, steps=256, max_new_tokens=256, block_size=32, temperature=0.0, cfg_scale=0.0, remasking="low_confidence")
-print(tokenizer.batch_decode(text[:, input_ids.shape[1]:], skip_special_tokens=False)[0])
+encoded = [tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=True) for m in prompts]
+prompt_lens = [len(e) for e in encoded]
+max_len = max(prompt_lens)
+pad_id = tokenizer.pad_token_id
+if pad_id is None:
+    pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.mask_token_id
+input_ids = torch.full((len(encoded), max_len), pad_id, dtype=torch.long)
+for i, ids in enumerate(encoded):
+    input_ids[i, : len(ids)] = torch.tensor(ids, dtype=torch.long)
+input_ids = input_ids.to(device)
+
+max_new_tokens = 128
+text = generate(
+    model,
+    tokenizer,
+    input_ids,
+    steps=128,
+    max_new_tokens=max_new_tokens,
+    block_size=32,
+    temperature=0.0,
+    cfg_scale=0.0,
+    remasking="low_confidence",
+)
 
+new_tokens = [text[i, prompt_lens[i] : prompt_lens[i] + max_new_tokens].tolist() for i in range(len(prompt_lens))]
+for idx, decoded in enumerate(tokenizer.batch_decode(new_tokens, skip_special_tokens=False)):
+    print(f"\n[Sample {idx}]")
+    print(decoded)
 ```
 
 ## Generation Parameters
 
 | Parameter        | Description                                                                                    | Default  |
 | ---------------- | ---------------------------------------------------------------------------------------------- | -------- |
-| `max_new_tokens` | Number of tokens to generate                                                                   | 256      |
-| `steps`          | Number of diffusion denoising iterations                                                       | 256      |
+| `max_new_tokens` | Number of tokens to generate                                                                   | 128      |
+| `steps`          | Number of diffusion denoising iterations                                                       | 128      |
 | `temperature`    | Sampling temperature; set to `0.0` for deterministic generation                                | 0.0      |
 | `block_size`   | Token block size used during iterative denoising                                               | 32       |
 | `cfg_scale`      | Classifier-free guidance scale controlling instruction adherence (higher = more deterministic) | 0.0      |
@@ -236,7 +280,7 @@ Follow the Github repo's demo script [examples/a2d/bd3lm/chat.py](https://github
 ```shell
 python -u examples/a2d/bd3lm/chat.py \
   --model_name_or_path dllm-collection/Qwen2.5-Coder-0.5B-Instruct-diffusion-bd3lm-v0.1 \
-  --chat True
+  --chat_template True
 ```
 
 ## Evaluation
@@ -279,7 +323,8 @@ python -u examples/a2d/bd3lm/chat.py \
 
 To automatically evaluate Qwen2.5-Coder-0.5B-Instruct-diffusion-bd3lm-v0.1 on all benchmarks, run:
 ```shell
-bash examples/a2d/eval_bd3lm.sh \
+bash examples/a2d/bd3lm/eval.sh \
+  --model_type coder \
   --model_name_or_path dllm-collection/Qwen2.5-Coder-0.5B-Instruct-diffusion-bd3lm-v0.1
 ```