Deploy BYOD-Llama-3.1-8B full-precision demo

app.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 import json
 import os
 import sys
+import time
 from pathlib import Path
 
 import gradio as gr
@@ -35,9 +36,11 @@ def _duration(*args) -> int:
     """Reserve enough GPU time for the requested number of denoising steps."""
     try:
         steps = int(args[3])
+        pause_per_step = float(args[8])
     except (IndexError, TypeError, ValueError):
         steps = 64
-    return min(300, max(30, steps * 2))
+        pause_per_step = 0.0
+    return min(300, max(30, round(steps * (2.0 + pause_per_step))))
 
 
 @spaces.GPU(size="large", duration=_duration)
@@ -50,12 +53,13 @@ def generate(
     temperature: float,
     top_k: int,
     seed: int,
-    show_trajectory: bool,
+    pause_per_step: float,
+    delay_eos_eot: bool,
 ):
     """Stream iterative masked-diffusion generation from the fixed model."""
     question = question.strip() or "What do you know about Amsterdam?"
     block_length = min(int(block_length), int(max_new_tokens))
-    latest = ("", "Starting…", "")
+    first_step = True
     for text, status, trajectory_html in denoise_stream(
         SESSION,
         question=question,
@@ -70,15 +74,21 @@ def generate(
         confidence_guided=True,
         proportional_unmask=False,
         early_stopping=False,
-        confidence_eos_eot_inf=True,
+        # This delays retention of predicted endings; it does not alter their
+        # sampling probability. Keep it optional so answers can end naturally.
+        confidence_eos_eot_inf=bool(delay_eos_eot),
         freeze_retained_tokens=True,
         repetition_penalty=1.0,
         eos_eot_prediction_penalty=1.0,
-        include_pre_remask_prediction=show_trajectory,
+        include_pre_remask_prediction=False,
         block_length=block_length,
     ):
-        latest = (text, status, trajectory_html if show_trajectory else "")
-        yield latest
+        if not first_step and float(pause_per_step) > 0:
+            # The sleep occurs inside this one decorated generator invocation,
+            # so ZeroGPU remains allocated for the entire denoising run.
+            time.sleep(float(pause_per_step))
+        first_step = False
+        yield text, status, trajectory_html
 
 
 with gr.Blocks(title=f"{DISPLAY_NAME} · masked diffusion") as demo:
@@ -104,14 +114,26 @@ with gr.Blocks(title=f"{DISPLAY_NAME} · masked diffusion") as demo:
                 value="You are a helpful assistant.",
                 lines=2,
             )
-            max_new_tokens = gr.Slider(16, 512, value=128, step=16, label="New tokens")
+            max_new_tokens = gr.Slider(16, 512, value=64, step=16, label="New tokens")
             num_steps = gr.Slider(1, 512, value=64, step=1, label="Denoising steps")
-            block_length = gr.Slider(16, 512, value=128, step=16, label="Block length")
+            block_length = gr.Slider(16, 512, value=64, step=16, label="Block length")
             temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature")
             top_k = gr.Slider(1, 100, value=3, step=1, label="Top-k")
             seed = gr.Number(value=1234, precision=0, label="Seed")
-            show_trajectory = gr.Checkbox(value=False, label="Show inference trajectory")
-    trajectory = gr.HTML(label="Token trajectory", visible=True)
+            pause_per_step = gr.Slider(
+                0.0,
+                1.0,
+                value=0.0,
+                step=0.05,
+                label="Pause between denoising steps (seconds)",
+                info="Slows the visualization while keeping one GPU allocation for the full run.",
+            )
+            delay_eos_eot = gr.Checkbox(
+                value=False,
+                label="Delay EOS/EOT retention (longer answers)",
+                info="Preferentially re-masks predicted endings; it does not lower their prediction probability.",
+            )
+    trajectory = gr.HTML(label="Live confidence-colored denoising")
     gr.Markdown(
         "The first request may take longer while the base model and adapter are loaded. "
         f"[Model card](https://huggingface.co/{MODEL_REPO_ID}) · "
@@ -127,7 +149,8 @@ with gr.Blocks(title=f"{DISPLAY_NAME} · masked diffusion") as demo:
         temperature,
         top_k,
         seed,
-        show_trajectory,
+        pause_per_step,
+        delay_eos_eot,
     ]
     run.click(generate, inputs=inputs, outputs=[output, status, trajectory])
     question.submit(generate, inputs=inputs, outputs=[output, status, trajectory])

src/diffusion_lm/inference.py

@@ -402,7 +402,12 @@ def load_llada_session(repo_id: str = "GSAI-ML/LLaDA-8B-Instruct", device_name:
     return session
 
 
-def _sample(logits: torch.Tensor, temperature: float, top_k: int, generator: torch.Generator | None) -> tuple[torch.Tensor, torch.Tensor]:
+def _sample(
+    logits: torch.Tensor,
+    temperature: float,
+    top_k: int,
+    generator: torch.Generator | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Top-k sample token IDs and return their normalized sampling confidence."""
     logits = logits / max(temperature, 1e-5)
     vocab_size = logits.shape[-1]
@@ -412,7 +417,7 @@ def _sample(logits: torch.Tensor, temperature: float, top_k: int, generator: tor
     picked_local = torch.multinomial(probabilities, 1, generator=generator)
     picked = indices.gather(-1, picked_local).squeeze(-1)
     confidence = probabilities.gather(-1, picked_local).squeeze(-1)
-    return picked, confidence
+    return picked, confidence, indices, probabilities
 
 
 def _llada_gumbel_noise(logits: torch.Tensor, temperature: float) -> torch.Tensor:
@@ -800,22 +805,31 @@ def render_denoising_step(tokens: list[int], confidences: list[float], answer_st
         output_token_count += 1
         token_text = escape(tokenizer.decode([token], skip_special_tokens=False)).replace("\n", "↵ ")
         if token == mask_token_id:
-            style, token_text = "background:#d1d5db;color:#111827;border-radius:3px;padding:1px 4px", "MASK"
-        elif frozen and offset in frozen:
-            style = "color:#1d4ed8;font-weight:700"
-        elif retained and offset in retained:
-            style = "color:#7c3aed;font-weight:700"
+            style, token_text = (
+                "display:inline-block;background:#d1d5db;color:#4b5563;"
+                "border:1px solid #9ca3af;border-radius:4px;padding:0 4px;"
+                "font-size:.78em;line-height:1.45;margin:0 1px;vertical-align:baseline",
+                "mask",
+            )
+            title = f"position {offset} · masked"
         else:
             confidence = max(0.0, min(1.0, float(confidences[offset]))) if offset < len(confidences) else 0.0
             hue = int(confidence * 120)
             style = f"color:hsl({hue},90%,30%);font-weight:{'600' if confidence > .8 else '400'}"
-        pieces.append(f"<span style='{style}' title='position {offset}'>{token_text}</span>")
+            title = f"position {offset} · sampling probability {confidence:.1%}"
+            if frozen and offset in frozen:
+                style += ";border-bottom:2px solid #2563eb"
+                title += " · retained and locked"
+            elif retained and offset in retained:
+                style += ";border-bottom:2px solid #7c3aed"
+                title += " · retained and editable"
+        pieces.append(f"<span style='{style}' title='{title}'>{token_text}</span>")
     pct = int(100 * step / max(total_steps, 1))
     return (f"<div style='font-family:system-ui;padding:14px;border:1px solid #d1d5db;border-radius:9px;background:#fafafa'>"
             f"<div style='font-weight:700;color:#2563eb;margin-bottom:7px'>Denoising step {step}/{total_steps} · {output_token_count} output tokens</div>"
             f"<div style='background:#e5e7eb;border-radius:4px;height:7px;margin-bottom:10px'><div style='background:#2563eb;width:{pct}%;height:7px;border-radius:4px'></div></div>"
             f"<div style='line-height:2;font-size:15px;white-space:pre-wrap'>{''.join(pieces)}</div>"
-            f"<div style='font-size:11px;color:#6b7280;margin-top:8px'>Green hues indicate confidence; gray tokens are MASK; purple tokens are retained but editable; blue tokens are retained and locked.</div></div>")
+            f"<div style='font-size:11px;color:#6b7280;margin-top:8px'>Red-to-green text indicates lower-to-higher sampling probability; gray chips are masks; purple underlines indicate retained/editable tokens; blue underlines indicate retained/locked tokens. Hover for details.</div></div>")
 
 
 def decode_denoising_state(
@@ -916,7 +930,12 @@ def denoise_stream(session: InferenceSession, question: str, system_prompt: str,
                 session.tokenizer.eos_token_id,
                 eot_token_id,
             )
-            sampled, confidence = _sample(logits, float(temperature), int(top_k), None)
+            sampled, confidence, candidate_ids, candidate_probabilities = _sample(
+                logits,
+                float(temperature),
+                int(top_k),
+                None,
+            )
         retention_confidence = confidence
         if confidence_eos_eot_inf:
             special_prediction = sampled == session.tokenizer.eos_token_id
@@ -929,6 +948,14 @@ def denoise_stream(session: InferenceSession, question: str, system_prompt: str,
         if freeze_retained_tokens:
             for offset, token in frozen.items():
                 ids[answer_start + offset] = token
+        display_confidence = confidence.clone()
+        if freeze_retained_tokens:
+            # The sampled token may be overwritten by a previously locked one.
+            # Report the probability of the token actually displayed, rather
+            # than the confidence of the discarded sample.
+            for offset, token in frozen.items():
+                matches = candidate_ids[offset] == int(token)
+                display_confidence[offset] = candidate_probabilities[offset][matches].sum()
         predicted_text = decode_denoising_state(
             ids[answer_start:],
             session.tokenizer,
@@ -1032,7 +1059,7 @@ def denoise_stream(session: InferenceSession, question: str, system_prompt: str,
             status += " · block stopped early (same prediction for 3 iterations)"
         html = render_denoising_step(
             ids,
-            confidence.tolist(),
+            display_confidence.tolist(),
             answer_start,
             session.tokenizer,
             session.mask_token_id,