Spaces:

Ruurd
/

llada-8b-instruct

Running on Zero

Add optional expanded-corpus bigram and exact trigram guidance

b31a17a verified about 4 hours ago

23.5 kB

"""Shared full-precision ZeroGPU demo for BYOD adapters and native LLaDA."""
from __future__ import annotations

import json
from functools import lru_cache
import os
import secrets
import sys
import time
from pathlib import Path
from urllib.parse import urlencode

import gradio as gr
import spaces

sys.path.insert(0, str(Path(__file__).parent / "src"))

from diffusion_lm.bigram import BigramScorer, BigramViterbiSelector, RepetitionPenalizedBigramSelector
from diffusion_lm.bigram_evaluation import ProfiledBigramSelector
from diffusion_lm.trigram import TrigramScorer
from diffusion_lm.inference import denoise_stream, load_hub_adapter_session, load_llada_session, _ending_token_ids


SPACE = json.loads((Path(__file__).parent / "space_model.json").read_text())
BIGRAM_ENABLED = bool(SPACE.get("bigram_enabled", False))
SHOW_BIGRAM_EVALUATION_CONTROLS = bool(SPACE.get("show_bigram_evaluation_controls", False))
BIGRAM_PATH = os.getenv("BIGRAM_MODEL_PATH", str(Path(__file__).parent / "bigram.pt"))
# Bigram defaults: edit these here to change the hosted configuration.
# Scoring: "Conditional bigram" or "PMI bigram".
# Retention: "Diffusion probability" or "Viterbi path margin".
BIGRAM_SCORING_METHOD = "PMI bigram"
BIGRAM_REMASKING_CONFIDENCE = "Diffusion probability"
PMI_MIN_COUNT = 1
PMI_CLIP = 2.0
BIGRAM_ENDING_DELAY_STRENGTH = 10.0
# Absolute iterations per block: full on iteration 1, zero from iteration 6.
REPETITION_FADE_ITERATIONS = 6
EXPANDED_BIGRAM_PATH = str(Path(__file__).parent / "bigram_expanded.pt")
TRIGRAM_PATH = str(Path(__file__).parent / "trigram_expanded.pt")
TRIGRAM_RELATIVE_WEIGHT = 0.5
TRIGRAM_MIN_COUNT = 2
TRIGRAM_CLIP = 2.0
if SPACE.get("ngram_scorer_repo"):
    from huggingface_hub import hf_hub_download
    # Download at startup, outside timed GPU inference; pin the artifact commit.
    EXPANDED_BIGRAM_PATH = hf_hub_download(SPACE["ngram_scorer_repo"], "bigram_expanded.pt", revision=SPACE["ngram_scorer_revision"])
    TRIGRAM_PATH = hf_hub_download(SPACE["ngram_scorer_repo"], "trigram_expanded.pt", revision=SPACE["ngram_scorer_revision"])
# Preserve the candidate pool and balance used in the matched Viterbi tests.
DEFAULT_TEMPERATURE = (1.0 if SPACE.get("model_backend") == "llada" else 2.0) if BIGRAM_ENABLED else 0.7
DEFAULT_TOP_K = 16 if BIGRAM_ENABLED else 3

MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", SPACE["model_repo_id"])
DISPLAY_NAME = SPACE["display_name"]
SHOW_REPOSITORY_LINKS = bool(SPACE.get("show_repository_links", True))
SOURCE_URL = SPACE.get("source_url")
VISITOR_COUNTER_ID = SPACE.get("visitor_counter_id")


def _visitor_badge(counter_id: str | None) -> str:
    """Return a persistent, client-loaded visit counter for this Space."""
    if not counter_id:
        return ""
    query = urlencode(
        {
            "page_id": counter_id,
            "left_text": "visitors",
            "left_color": "#4b5563",
            "right_color": "#6d28d9",
        }
    )
    return f"![Visitor count](https://visitor-badge.laobi.icu/badge?{query})"

# ZeroGPU recommends constructing and placing the root module on CUDA at module
# scope. The models are loaded without quantization.
print(f"Loading {MODEL_REPO_ID} in full precision...")
if SPACE.get("model_backend", "byod") == "llada":
    SESSION = load_llada_session(
        MODEL_REPO_ID, device_name="cuda", revision=SPACE["model_revision"], preflight=False,
    )
else:
    SESSION = load_hub_adapter_session(
        MODEL_REPO_ID,
        device_name="cuda",
        quantization="none",
        # A forward pass is only valid after @spaces.GPU allocates hardware.
        preflight=False,
    )
print(f"Loaded {DISPLAY_NAME} ({SESSION.compute_dtype}, unquantized).")


@lru_cache(maxsize=2)
def _bigram_scorer(path):
    # weights_only loading and tokenizer validation happen once per checkpoint.
    return BigramScorer.load(path, tokenizer=SESSION.tokenizer)


@lru_cache(maxsize=1)
def _trigram_scorer(path, bigram_path):
    return TrigramScorer.load(path, bigram=_bigram_scorer(bigram_path))


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[7])
        remasking_strategy = str(args[9])
        ar_check_interval = max(1, int(args[11]))
    except (IndexError, TypeError, ValueError):
        steps = 64
        pause_per_step = 0.0
        remasking_strategy = "Confidence-guided"
        ar_check_interval = 1
    seconds_per_step = 2.0 + (2.0 / ar_check_interval if remasking_strategy == "Autoregressive-confidence" else 0.0)
    return min(300, max(30, round(steps * (seconds_per_step + pause_per_step))))


@spaces.GPU(size="large", duration=_duration)
def generate(
    question: str,
    system_prompt: str,
    max_new_tokens: int,
    num_steps: int,
    block_length: int,
    temperature: float,
    top_k: int,
    pause_per_step: float,
    trajectory_color_mode: str,
    remasking_strategy: str,
    revisable_tokens: bool,
    ar_check_interval: int,
    delay_eos_eot: bool,
    early_stopping: bool,
    bigram_model_path: str = BIGRAM_PATH,
    bigram_scoring_method: str = BIGRAM_SCORING_METHOD,
    pmi_min_count: int = PMI_MIN_COUNT,
    pmi_clip: float = PMI_CLIP,
    generation_seed: int = -1,
    measure_decoding: bool = False,
    bigram_eos_delay_penalty: float = BIGRAM_ENDING_DELAY_STRENGTH,
    bigram_remasking_confidence: str = BIGRAM_REMASKING_CONFIDENCE,
    repetition_penalty_strength: float = 1.0,
    adaptive_parallel_decoding: bool = False,
    sequence_sampling: bool = False,
    count_all_candidate_bigrams: bool = False,
    bigram_weight: float = 1.0,
    repetition_risk_penalty: bool = False,
    decay_repetition_penalty: bool = False,
    apd_aggressiveness: float = 0.0,
    ngram_guidance: str = "Original bigram",
):
    """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))
    seed = secrets.randbelow(2**63 - 1) if int(generation_seed) < 0 else int(generation_seed)
    selector = None
    if remasking_strategy in {"Bigram Viterbi", "PMI Viterbi + Repetition Penalty"}:
        if not BIGRAM_ENABLED:
            raise gr.Error("Bigram Viterbi is not enabled for this Space.")
        try:
            trigram_options = {}
            if ngram_guidance not in {"Original bigram", "Expanded bigram", "Expanded bigram + trigram"}:
                raise ValueError("Unknown n-gram guidance option")
            if ngram_guidance != "Original bigram":
                if not getattr(SESSION, "llada", False):
                    raise ValueError("Expanded n-gram checkpoints require the LLaDA tokenizer")
                bigram_model_path = EXPANDED_BIGRAM_PATH
            if ngram_guidance == "Expanded bigram + trigram":
                trigram_options = dict(
                    trigram_scorer=_trigram_scorer(TRIGRAM_PATH, bigram_model_path).to(SESSION.device),
                    trigram_weight=TRIGRAM_RELATIVE_WEIGHT,
                    trigram_min_count=TRIGRAM_MIN_COUNT, trigram_clip=TRIGRAM_CLIP,
                )
            selector = BigramViterbiSelector(
                _bigram_scorer(bigram_model_path).to(SESSION.device),
                float(bigram_weight) if getattr(SESSION, "llada", False) else 0.5,
                scoring_method=("pmi" if remasking_strategy == "PMI Viterbi + Repetition Penalty" else
                    {"Conditional bigram": "conditional", "PMI bigram": "pmi"}.get(bigram_scoring_method, "")),
                min_count=int(pmi_min_count), pmi_clip=float(pmi_clip),
                sequence_sampling=bool(sequence_sampling),
                independent_temperature=bool(getattr(SESSION, "llada", False)),
                remasking_confidence={"Diffusion probability": "diffusion", "Viterbi path margin": "viterbi_margin"}.get(bigram_remasking_confidence, ""),
                **trigram_options,
            )
            if remasking_strategy == "PMI Viterbi + Repetition Penalty":
                selector = RepetitionPenalizedBigramSelector(
                    selector, float(repetition_penalty_strength),
                    excluded_token_ids=_ending_token_ids(SESSION.tokenizer),
                    decay_steps=(REPETITION_FADE_ITERATIONS if decay_repetition_penalty and getattr(SESSION, "llada", False) else 0),
                    count_mode=("risk" if repetition_risk_penalty else
                        "candidate_occurrences" if count_all_candidate_bigrams else "expected"),
                )
            if measure_decoding:
                selector = ProfiledBigramSelector(selector)
        except (FileNotFoundError, ValueError) as exc:
            raise gr.Error(str(exc)) from exc
    first_step = True
    inference_started = time.perf_counter() if measure_decoding else None
    forward_passes = 0
    for text, status, trajectory_html in denoise_stream(
        SESSION,
        question=question,
        system_prompt=system_prompt,
        max_new_tokens=int(max_new_tokens),
        num_steps=int(num_steps),
        noise_level=1.0,
        temperature=float(temperature),
        top_k=int(top_k),
        seed=int(seed),
        permanent_unmask=not bool(revisable_tokens),
        confidence_guided=remasking_strategy in {"Confidence-guided", "Bigram Viterbi", "PMI Viterbi + Repetition Penalty"},
        proportional_unmask=False,
        early_stopping=bool(early_stopping),
        # Bigram mode also fades an ending-prediction penalty to counter its
        # deterministic tendency to choose the same short ending repeatedly.
        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=False,
        block_length=block_length,
        trajectory_color_mode=trajectory_color_mode,
        autoregressive_guided=remasking_strategy == "Autoregressive-confidence",
        revisable_tokens=bool(revisable_tokens),
        autoregressive_check_interval=int(ar_check_interval),
        bigram_selector=selector,
        bigram_eos_delay_penalty=float(bigram_eos_delay_penalty),
        adaptive_parallel_decoding=bool(adaptive_parallel_decoding),
        apd_aggressiveness=float(apd_aggressiveness),
    ):
        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
        forward_passes += 1
        if measure_decoding:
            status += "\n\nInference metrics: " + json.dumps(dict(
                denoising_forward_passes=forward_passes, elapsed_ms=(time.perf_counter() - inference_started) * 1000,
                adaptive_parallel_decoding=bool(adaptive_parallel_decoding),
            ))
        if measure_decoding and selector is not None:
            status += "\n\nDecoding metrics: " + json.dumps(selector.metrics())
        yield status, trajectory_html


def show_loading():
    """Immediately acknowledge a request while ZeroGPU prepares generation."""
    return (
        "⏳ Requesting a GPU and loading the model…",
        "<div style='font-family:system-ui;padding:14px;border:1px solid #d1d5db;"
        "border-radius:9px;background:#fafafa;color:#4b5563'>"
        "Preparing generation… The first request may take a little longer while "
        "the model is loaded.</div>",
    )


with gr.Blocks(title=f"{DISPLAY_NAME} · masked diffusion") as demo:
    gr.Markdown(
        f"# {DISPLAY_NAME}\n"
        "A full-precision masked-diffusion model converted from an "
        "autoregressive language model with LoRA. This demo uses the exact "
        "published training checkpoint and does **not** use 4-bit quantization."
    )
    with gr.Row():
        with gr.Column(scale=3):
            question = gr.Textbox(
                label="Prompt",
                value="What do you know about Amsterdam?",
                lines=4,
            )
            run = gr.Button("Generate", variant="primary")
            status = gr.Markdown("Ready")
            trajectory = gr.HTML(label="Live denoising")
        with gr.Column(scale=2):
            system_prompt = gr.Textbox(
                label="System prompt",
                value="You are a helpful assistant.",
                lines=2,
            )
            max_new_tokens = gr.Slider(1, 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(1, 512, value=64, step=1, label="Block length")
            temperature = gr.Slider(0.0, 2.0, value=DEFAULT_TEMPERATURE, step=0.05, label="Temperature",
                info=("For sequence sampling: higher values increase randomness; 0 selects the best path. Bigram weight controls PMI influence. Maximum-path Viterbi ignores temperature."
                    if SESSION.llada else "In Viterbi, lower values favor diffusion predictions over bigrams; 0 ignores bigram scores. Viterbi remains deterministic.") if BIGRAM_ENABLED else None)
            bigram_weight = gr.Slider(0.0, 5.0, value=1.0, step=0.05, label="Bigram weight",
                visible=SESSION.llada and BIGRAM_ENABLED,
                info="PMI coefficient: 0 disables PMI; 1 gives equal diffusion and PMI score coefficients. Independent of temperature.")
            top_k = gr.Slider(1, 100, value=DEFAULT_TOP_K, step=1, label="Top-k",
                info="Number of candidate tokens per position for sampling or Viterbi. At 1, only the model’s highest-scoring token is available." if BIGRAM_ENABLED else None)
            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.",
            )
            trajectory_color_mode = gr.Radio(
                choices=["No coloring", "Prediction probability", "Prediction iteration"],
                value="Prediction iteration",
                label="Token coloring",
                info="Hover over any token to see its position, prediction iteration, and probability.",
            )
            remasking_strategy = gr.Radio(
                choices=["Confidence-guided"] + ([] if SESSION.llada else ["Autoregressive-confidence"]) + ["Random"] + (["Bigram Viterbi", "PMI Viterbi + Repetition Penalty"] if BIGRAM_ENABLED else []),
                value="Bigram Viterbi" if BIGRAM_ENABLED else "Confidence-guided",
                label="Remasking strategy",
            )
            ngram_guidance = gr.Dropdown(
                choices=["Original bigram", "Expanded bigram", "Expanded bigram + trigram"],
                value="Original bigram", label="N-gram guidance", visible=SESSION.llada and BIGRAM_ENABLED,
                info="For bigram decoding strategies. Compare the original corpus with expanded English/code/math counts and optional trigram corrections. Trigrams cost more at large Top-k; start at 16.")
            repetition_penalty_strength = gr.Slider(
                0.0, 10.0, value=1.0, step=0.1, label="Repetition penalty strength (γ)",
                visible=False, info="Penalize repeated bigram transitions; EOS/EOT are exempt. At 0, disables the repetition penalty.",
            )
            sequence_sampling = gr.Checkbox(value=False, visible=SESSION.llada and BIGRAM_ENABLED,
                label="Sample bigram-guided sequences",
                info="Sample complete sequences using diffusion, PMI and any repetition scores. No score perturbations. Applies to bigram decoding options.")
            count_all_candidate_bigrams = gr.Checkbox(value=False, visible=SESSION.llada and BIGRAM_ENABLED,
                label="Count every candidate bigram occurrence",
                info="For repetition penalty: count each earlier adjacent candidate-list occurrence as 1, regardless of probability. EOS/EOT are exempt.")
            repetition_risk_penalty = gr.Checkbox(value=False, visible=SESSION.llada and BIGRAM_ENABLED,
                label="Repetition-risk penalty",
                info="For the repetition strategy: estimate risk from joint diffusion + PMI probabilities. Overrides candidate-count mode; EOS/EOT are exempt.")
            decay_repetition_penalty = gr.Checkbox(value=True, visible=SESSION.llada and BIGRAM_ENABLED,
                label="Fade repetition penalty over early iterations",
                info=f"For repetition decoding: logistic fade, full on iteration 1 and zero from iteration {REPETITION_FADE_ITERATIONS}. Resets per block. Turn off for constant strength.")
            if SHOW_BIGRAM_EVALUATION_CONTROLS and BIGRAM_ENABLED:
                with gr.Group():
                    bigram_eos_delay_penalty = gr.Number(value=BIGRAM_ENDING_DELAY_STRENGTH, minimum=1,
                        label="Bigram ending delay strength",
                        info="With Delay EOS/EOT enabled, discourages ending predictions early and fades to 1 at the final block step. Set 1 for retention-only delay.")
                    bigram_model_path = gr.Textbox(value=BIGRAM_PATH, label="Bigram checkpoint path")
                    with gr.Accordion("Reproducible evaluation", open=False):
                        generation_seed = gr.Number(value=-1, minimum=-1, maximum=2**31-1, precision=0, label="Seed (-1 for random)")
                        measure_decoding = gr.Checkbox(value=False, label="Measure decoding time and changes from argmax")
            revisable_tokens = gr.Checkbox(
                value=False,
                label="Allow revising retained tokens",
                visible=not SESSION.llada,
                info="When enabled, previously visible tokens can be re-masked on later iterations.",
                interactive=False,
            )
            ar_check_interval = gr.Slider(
                1,
                16,
                value=4,
                step=1,
                label="AR check interval",
                visible=not SESSION.llada,
                info="Use the original causal model every N denoising steps; diffusion confidence is used between checks.",
                interactive=False,
            )
            delay_eos_eot = gr.Checkbox(
                value=True,
                label="Delay EOS/EOT retention (longer answers)",
                info="Retains ending tokens last to give the answer room to develop.",
            )
            early_stopping = gr.Checkbox(
                value=True,
                label="Early stopping",
                info="Converge a block after two matching answers; stop all remaining blocks when EOS/EOT is finalized with a complete preceding answer.",
            )
            adaptive_parallel_decoding = gr.Checkbox(
                value=False, label="Adaptive Parallel Decoding (APD)",
                info="Commit confident tokens together, with at least ceil(block length / block iterations) per step. Finish blocks as soon as all positions are resolved.",
            )
            apd_aggressiveness = gr.Slider(0.0, 1.0, value=0.0, step=0.05,
                label="APD aggressiveness", visible=SESSION.llada, interactive=False,
                info="0 preserves APD's current thresholds. Higher values commit lower-confidence tokens sooner; 1 commits all positions in one pass per block. May reduce quality.")
    footer = "The first request may take longer while the model is loaded."
    if SHOW_REPOSITORY_LINKS:
        footer += f" [Model card](https://huggingface.co/{MODEL_REPO_ID})"
        if SOURCE_URL:
            footer += f" · [Source code]({SOURCE_URL})"
    visitor_badge = _visitor_badge(VISITOR_COUNTER_ID)
    if visitor_badge:
        footer += f"\n\n{visitor_badge}"
    gr.Markdown(footer)

    # Hidden inputs preserve programmatic evaluation overrides. The interface
    # uses the editable defaults above without exposing expert controls.
    bigram_scoring_method = gr.Textbox(value=BIGRAM_SCORING_METHOD, visible=False)
    bigram_remasking_confidence = gr.Textbox(value=BIGRAM_REMASKING_CONFIDENCE, visible=False)
    pmi_min_count = gr.Number(value=PMI_MIN_COUNT, precision=0, visible=False)
    pmi_clip = gr.Number(value=PMI_CLIP, visible=False)
    if not (SHOW_BIGRAM_EVALUATION_CONTROLS and BIGRAM_ENABLED):
        bigram_eos_delay_penalty = gr.Number(value=BIGRAM_ENDING_DELAY_STRENGTH, visible=False)
        bigram_model_path = gr.Textbox(value=BIGRAM_PATH, visible=False)
        generation_seed = gr.Number(value=-1, precision=0, visible=False)
        measure_decoding = gr.Checkbox(value=False, visible=False)

    inputs = [
        question,
        system_prompt,
        max_new_tokens,
        num_steps,
        block_length,
        temperature,
        top_k,
        pause_per_step,
        trajectory_color_mode,
        remasking_strategy,
        revisable_tokens,
        ar_check_interval,
        delay_eos_eot,
        early_stopping,
        bigram_model_path,
        bigram_scoring_method,
        pmi_min_count,
        pmi_clip,
        generation_seed,
        measure_decoding,
        bigram_eos_delay_penalty,
        bigram_remasking_confidence,
        repetition_penalty_strength,
        adaptive_parallel_decoding,
        sequence_sampling,
        count_all_candidate_bigrams,
        bigram_weight,
        repetition_risk_penalty,
        decay_repetition_penalty,
        apd_aggressiveness,
        ngram_guidance,
    ]
    def update_ar_controls(strategy, apd=False):
        """Expose AR-only controls and clear revision when another strategy is selected."""
        enabled = strategy == "Autoregressive-confidence" and not apd
        return gr.update(interactive=enabled, value=False), gr.update(interactive=enabled)

    remasking_strategy.change(
        update_ar_controls,
        inputs=[remasking_strategy, adaptive_parallel_decoding],
        outputs=[revisable_tokens, ar_check_interval],
        queue=False,
    )
    adaptive_parallel_decoding.change(
        update_ar_controls, inputs=[remasking_strategy, adaptive_parallel_decoding],
        outputs=[revisable_tokens, ar_check_interval], queue=False,
    )
    def update_apd_aggressiveness(enabled):
        return gr.update(interactive=bool(enabled))

    adaptive_parallel_decoding.change(
        update_apd_aggressiveness, inputs=[adaptive_parallel_decoding],
        outputs=[apd_aggressiveness], queue=False,
    )
    def update_repetition_control(strategy):
        return gr.update(visible=strategy == "PMI Viterbi + Repetition Penalty")

    remasking_strategy.change(
        update_repetition_control, inputs=[remasking_strategy],
        outputs=[repetition_penalty_strength], queue=False,
    )
    run_event = run.click(show_loading, outputs=[status, trajectory], queue=False)
    run_event.then(generate, inputs=inputs, outputs=[status, trajectory])
    submit_event = question.submit(show_loading, outputs=[status, trajectory], queue=False)
    submit_event.then(generate, inputs=inputs, outputs=[status, trajectory])

demo.queue(default_concurrency_limit=1).launch(ssr_mode=False)

	"""Shared full-precision ZeroGPU demo for BYOD adapters and native LLaDA."""
	from __future__ import annotations

	import json
	from functools import lru_cache
	import os
	import secrets
	import sys
	import time
	from pathlib import Path
	from urllib.parse import urlencode

	import gradio as gr
	import spaces

	sys.path.insert(0, str(Path(__file__).parent / "src"))

	from diffusion_lm.bigram import BigramScorer, BigramViterbiSelector, RepetitionPenalizedBigramSelector
	from diffusion_lm.bigram_evaluation import ProfiledBigramSelector
	from diffusion_lm.trigram import TrigramScorer
	from diffusion_lm.inference import denoise_stream, load_hub_adapter_session, load_llada_session, _ending_token_ids


	SPACE = json.loads((Path(__file__).parent / "space_model.json").read_text())
	BIGRAM_ENABLED = bool(SPACE.get("bigram_enabled", False))
	SHOW_BIGRAM_EVALUATION_CONTROLS = bool(SPACE.get("show_bigram_evaluation_controls", False))
	BIGRAM_PATH = os.getenv("BIGRAM_MODEL_PATH", str(Path(__file__).parent / "bigram.pt"))
	# Bigram defaults: edit these here to change the hosted configuration.
	# Scoring: "Conditional bigram" or "PMI bigram".
	# Retention: "Diffusion probability" or "Viterbi path margin".
	BIGRAM_SCORING_METHOD = "PMI bigram"
	BIGRAM_REMASKING_CONFIDENCE = "Diffusion probability"
	PMI_MIN_COUNT = 1
	PMI_CLIP = 2.0
	BIGRAM_ENDING_DELAY_STRENGTH = 10.0
	# Absolute iterations per block: full on iteration 1, zero from iteration 6.
	REPETITION_FADE_ITERATIONS = 6
	EXPANDED_BIGRAM_PATH = str(Path(__file__).parent / "bigram_expanded.pt")
	TRIGRAM_PATH = str(Path(__file__).parent / "trigram_expanded.pt")
	TRIGRAM_RELATIVE_WEIGHT = 0.5
	TRIGRAM_MIN_COUNT = 2
	TRIGRAM_CLIP = 2.0
	if SPACE.get("ngram_scorer_repo"):
	from huggingface_hub import hf_hub_download
	# Download at startup, outside timed GPU inference; pin the artifact commit.
	EXPANDED_BIGRAM_PATH = hf_hub_download(SPACE["ngram_scorer_repo"], "bigram_expanded.pt", revision=SPACE["ngram_scorer_revision"])
	TRIGRAM_PATH = hf_hub_download(SPACE["ngram_scorer_repo"], "trigram_expanded.pt", revision=SPACE["ngram_scorer_revision"])
	# Preserve the candidate pool and balance used in the matched Viterbi tests.
	DEFAULT_TEMPERATURE = (1.0 if SPACE.get("model_backend") == "llada" else 2.0) if BIGRAM_ENABLED else 0.7
	DEFAULT_TOP_K = 16 if BIGRAM_ENABLED else 3

	MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", SPACE["model_repo_id"])
	DISPLAY_NAME = SPACE["display_name"]
	SHOW_REPOSITORY_LINKS = bool(SPACE.get("show_repository_links", True))
	SOURCE_URL = SPACE.get("source_url")
	VISITOR_COUNTER_ID = SPACE.get("visitor_counter_id")


	def _visitor_badge(counter_id: str \| None) -> str:
	"""Return a persistent, client-loaded visit counter for this Space."""
	if not counter_id:
	return ""
	query = urlencode(
	{
	"page_id": counter_id,
	"left_text": "visitors",
	"left_color": "#4b5563",
	"right_color": "#6d28d9",
	}
	)
	return f"![Visitor count](https://visitor-badge.laobi.icu/badge?{query})"

	# ZeroGPU recommends constructing and placing the root module on CUDA at module
	# scope. The models are loaded without quantization.
	print(f"Loading {MODEL_REPO_ID} in full precision...")
	if SPACE.get("model_backend", "byod") == "llada":
	SESSION = load_llada_session(
	MODEL_REPO_ID, device_name="cuda", revision=SPACE["model_revision"], preflight=False,
	)
	else:
	SESSION = load_hub_adapter_session(
	MODEL_REPO_ID,
	device_name="cuda",
	quantization="none",
	# A forward pass is only valid after @spaces.GPU allocates hardware.
	preflight=False,
	)
	print(f"Loaded {DISPLAY_NAME} ({SESSION.compute_dtype}, unquantized).")


	@lru_cache(maxsize=2)
	def _bigram_scorer(path):
	# weights_only loading and tokenizer validation happen once per checkpoint.
	return BigramScorer.load(path, tokenizer=SESSION.tokenizer)


	@lru_cache(maxsize=1)
	def _trigram_scorer(path, bigram_path):
	return TrigramScorer.load(path, bigram=_bigram_scorer(bigram_path))


	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[7])
	remasking_strategy = str(args[9])
	ar_check_interval = max(1, int(args[11]))
	except (IndexError, TypeError, ValueError):
	steps = 64
	pause_per_step = 0.0
	remasking_strategy = "Confidence-guided"
	ar_check_interval = 1
	seconds_per_step = 2.0 + (2.0 / ar_check_interval if remasking_strategy == "Autoregressive-confidence" else 0.0)
	return min(300, max(30, round(steps * (seconds_per_step + pause_per_step))))


	@spaces.GPU(size="large", duration=_duration)
	def generate(
	question: str,
	system_prompt: str,
	max_new_tokens: int,
	num_steps: int,
	block_length: int,
	temperature: float,
	top_k: int,
	pause_per_step: float,
	trajectory_color_mode: str,
	remasking_strategy: str,
	revisable_tokens: bool,
	ar_check_interval: int,
	delay_eos_eot: bool,
	early_stopping: bool,
	bigram_model_path: str = BIGRAM_PATH,
	bigram_scoring_method: str = BIGRAM_SCORING_METHOD,
	pmi_min_count: int = PMI_MIN_COUNT,
	pmi_clip: float = PMI_CLIP,
	generation_seed: int = -1,
	measure_decoding: bool = False,
	bigram_eos_delay_penalty: float = BIGRAM_ENDING_DELAY_STRENGTH,
	bigram_remasking_confidence: str = BIGRAM_REMASKING_CONFIDENCE,
	repetition_penalty_strength: float = 1.0,
	adaptive_parallel_decoding: bool = False,
	sequence_sampling: bool = False,
	count_all_candidate_bigrams: bool = False,
	bigram_weight: float = 1.0,
	repetition_risk_penalty: bool = False,
	decay_repetition_penalty: bool = False,
	apd_aggressiveness: float = 0.0,
	ngram_guidance: str = "Original bigram",
	):
	"""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))
	seed = secrets.randbelow(2**63 - 1) if int(generation_seed) < 0 else int(generation_seed)
	selector = None
	if remasking_strategy in {"Bigram Viterbi", "PMI Viterbi + Repetition Penalty"}:
	if not BIGRAM_ENABLED:
	raise gr.Error("Bigram Viterbi is not enabled for this Space.")
	try:
	trigram_options = {}
	if ngram_guidance not in {"Original bigram", "Expanded bigram", "Expanded bigram + trigram"}:
	raise ValueError("Unknown n-gram guidance option")
	if ngram_guidance != "Original bigram":
	if not getattr(SESSION, "llada", False):
	raise ValueError("Expanded n-gram checkpoints require the LLaDA tokenizer")
	bigram_model_path = EXPANDED_BIGRAM_PATH
	if ngram_guidance == "Expanded bigram + trigram":
	trigram_options = dict(
	trigram_scorer=_trigram_scorer(TRIGRAM_PATH, bigram_model_path).to(SESSION.device),
	trigram_weight=TRIGRAM_RELATIVE_WEIGHT,
	trigram_min_count=TRIGRAM_MIN_COUNT, trigram_clip=TRIGRAM_CLIP,
	)
	selector = BigramViterbiSelector(
	_bigram_scorer(bigram_model_path).to(SESSION.device),
	float(bigram_weight) if getattr(SESSION, "llada", False) else 0.5,
	scoring_method=("pmi" if remasking_strategy == "PMI Viterbi + Repetition Penalty" else
	{"Conditional bigram": "conditional", "PMI bigram": "pmi"}.get(bigram_scoring_method, "")),
	min_count=int(pmi_min_count), pmi_clip=float(pmi_clip),
	sequence_sampling=bool(sequence_sampling),
	independent_temperature=bool(getattr(SESSION, "llada", False)),
	remasking_confidence={"Diffusion probability": "diffusion", "Viterbi path margin": "viterbi_margin"}.get(bigram_remasking_confidence, ""),
	**trigram_options,
	)
	if remasking_strategy == "PMI Viterbi + Repetition Penalty":
	selector = RepetitionPenalizedBigramSelector(
	selector, float(repetition_penalty_strength),
	excluded_token_ids=_ending_token_ids(SESSION.tokenizer),
	decay_steps=(REPETITION_FADE_ITERATIONS if decay_repetition_penalty and getattr(SESSION, "llada", False) else 0),
	count_mode=("risk" if repetition_risk_penalty else
	"candidate_occurrences" if count_all_candidate_bigrams else "expected"),
	)
	if measure_decoding:
	selector = ProfiledBigramSelector(selector)
	except (FileNotFoundError, ValueError) as exc:
	raise gr.Error(str(exc)) from exc
	first_step = True
	inference_started = time.perf_counter() if measure_decoding else None
	forward_passes = 0
	for text, status, trajectory_html in denoise_stream(
	SESSION,
	question=question,
	system_prompt=system_prompt,
	max_new_tokens=int(max_new_tokens),
	num_steps=int(num_steps),
	noise_level=1.0,
	temperature=float(temperature),
	top_k=int(top_k),
	seed=int(seed),
	permanent_unmask=not bool(revisable_tokens),
	confidence_guided=remasking_strategy in {"Confidence-guided", "Bigram Viterbi", "PMI Viterbi + Repetition Penalty"},
	proportional_unmask=False,
	early_stopping=bool(early_stopping),
	# Bigram mode also fades an ending-prediction penalty to counter its
	# deterministic tendency to choose the same short ending repeatedly.
	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=False,
	block_length=block_length,
	trajectory_color_mode=trajectory_color_mode,
	autoregressive_guided=remasking_strategy == "Autoregressive-confidence",
	revisable_tokens=bool(revisable_tokens),
	autoregressive_check_interval=int(ar_check_interval),
	bigram_selector=selector,
	bigram_eos_delay_penalty=float(bigram_eos_delay_penalty),
	adaptive_parallel_decoding=bool(adaptive_parallel_decoding),
	apd_aggressiveness=float(apd_aggressiveness),
	):
	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
	forward_passes += 1
	if measure_decoding:
	status += "\n\nInference metrics: " + json.dumps(dict(
	denoising_forward_passes=forward_passes, elapsed_ms=(time.perf_counter() - inference_started) * 1000,
	adaptive_parallel_decoding=bool(adaptive_parallel_decoding),
	))
	if measure_decoding and selector is not None:
	status += "\n\nDecoding metrics: " + json.dumps(selector.metrics())
	yield status, trajectory_html


	def show_loading():
	"""Immediately acknowledge a request while ZeroGPU prepares generation."""
	return (
	"⏳ Requesting a GPU and loading the model…",
	"<div style='font-family:system-ui;padding:14px;border:1px solid #d1d5db;"
	"border-radius:9px;background:#fafafa;color:#4b5563'>"
	"Preparing generation… The first request may take a little longer while "
	"the model is loaded.</div>",
	)


	with gr.Blocks(title=f"{DISPLAY_NAME} · masked diffusion") as demo:
	gr.Markdown(
	f"# {DISPLAY_NAME}\n"
	"A full-precision masked-diffusion model converted from an "
	"autoregressive language model with LoRA. This demo uses the exact "
	"published training checkpoint and does not use 4-bit quantization."
	)
	with gr.Row():
	with gr.Column(scale=3):
	question = gr.Textbox(
	label="Prompt",
	value="What do you know about Amsterdam?",
	lines=4,
	)
	run = gr.Button("Generate", variant="primary")
	status = gr.Markdown("Ready")
	trajectory = gr.HTML(label="Live denoising")
	with gr.Column(scale=2):
	system_prompt = gr.Textbox(
	label="System prompt",
	value="You are a helpful assistant.",
	lines=2,
	)
	max_new_tokens = gr.Slider(1, 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(1, 512, value=64, step=1, label="Block length")
	temperature = gr.Slider(0.0, 2.0, value=DEFAULT_TEMPERATURE, step=0.05, label="Temperature",
	info=("For sequence sampling: higher values increase randomness; 0 selects the best path. Bigram weight controls PMI influence. Maximum-path Viterbi ignores temperature."
	if SESSION.llada else "In Viterbi, lower values favor diffusion predictions over bigrams; 0 ignores bigram scores. Viterbi remains deterministic.") if BIGRAM_ENABLED else None)
	bigram_weight = gr.Slider(0.0, 5.0, value=1.0, step=0.05, label="Bigram weight",
	visible=SESSION.llada and BIGRAM_ENABLED,
	info="PMI coefficient: 0 disables PMI; 1 gives equal diffusion and PMI score coefficients. Independent of temperature.")
	top_k = gr.Slider(1, 100, value=DEFAULT_TOP_K, step=1, label="Top-k",
	info="Number of candidate tokens per position for sampling or Viterbi. At 1, only the model’s highest-scoring token is available." if BIGRAM_ENABLED else None)
	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.",
	)
	trajectory_color_mode = gr.Radio(
	choices=["No coloring", "Prediction probability", "Prediction iteration"],
	value="Prediction iteration",
	label="Token coloring",
	info="Hover over any token to see its position, prediction iteration, and probability.",
	)
	remasking_strategy = gr.Radio(
	choices=["Confidence-guided"] + ([] if SESSION.llada else ["Autoregressive-confidence"]) + ["Random"] + (["Bigram Viterbi", "PMI Viterbi + Repetition Penalty"] if BIGRAM_ENABLED else []),
	value="Bigram Viterbi" if BIGRAM_ENABLED else "Confidence-guided",
	label="Remasking strategy",
	)
	ngram_guidance = gr.Dropdown(
	choices=["Original bigram", "Expanded bigram", "Expanded bigram + trigram"],
	value="Original bigram", label="N-gram guidance", visible=SESSION.llada and BIGRAM_ENABLED,
	info="For bigram decoding strategies. Compare the original corpus with expanded English/code/math counts and optional trigram corrections. Trigrams cost more at large Top-k; start at 16.")
	repetition_penalty_strength = gr.Slider(
	0.0, 10.0, value=1.0, step=0.1, label="Repetition penalty strength (γ)",
	visible=False, info="Penalize repeated bigram transitions; EOS/EOT are exempt. At 0, disables the repetition penalty.",
	)
	sequence_sampling = gr.Checkbox(value=False, visible=SESSION.llada and BIGRAM_ENABLED,
	label="Sample bigram-guided sequences",
	info="Sample complete sequences using diffusion, PMI and any repetition scores. No score perturbations. Applies to bigram decoding options.")
	count_all_candidate_bigrams = gr.Checkbox(value=False, visible=SESSION.llada and BIGRAM_ENABLED,
	label="Count every candidate bigram occurrence",
	info="For repetition penalty: count each earlier adjacent candidate-list occurrence as 1, regardless of probability. EOS/EOT are exempt.")
	repetition_risk_penalty = gr.Checkbox(value=False, visible=SESSION.llada and BIGRAM_ENABLED,
	label="Repetition-risk penalty",
	info="For the repetition strategy: estimate risk from joint diffusion + PMI probabilities. Overrides candidate-count mode; EOS/EOT are exempt.")
	decay_repetition_penalty = gr.Checkbox(value=True, visible=SESSION.llada and BIGRAM_ENABLED,
	label="Fade repetition penalty over early iterations",
	info=f"For repetition decoding: logistic fade, full on iteration 1 and zero from iteration {REPETITION_FADE_ITERATIONS}. Resets per block. Turn off for constant strength.")
	if SHOW_BIGRAM_EVALUATION_CONTROLS and BIGRAM_ENABLED:
	with gr.Group():
	bigram_eos_delay_penalty = gr.Number(value=BIGRAM_ENDING_DELAY_STRENGTH, minimum=1,
	label="Bigram ending delay strength",
	info="With Delay EOS/EOT enabled, discourages ending predictions early and fades to 1 at the final block step. Set 1 for retention-only delay.")
	bigram_model_path = gr.Textbox(value=BIGRAM_PATH, label="Bigram checkpoint path")
	with gr.Accordion("Reproducible evaluation", open=False):
	generation_seed = gr.Number(value=-1, minimum=-1, maximum=2**31-1, precision=0, label="Seed (-1 for random)")
	measure_decoding = gr.Checkbox(value=False, label="Measure decoding time and changes from argmax")
	revisable_tokens = gr.Checkbox(
	value=False,
	label="Allow revising retained tokens",
	visible=not SESSION.llada,
	info="When enabled, previously visible tokens can be re-masked on later iterations.",
	interactive=False,
	)
	ar_check_interval = gr.Slider(
	1,
	16,
	value=4,
	step=1,
	label="AR check interval",
	visible=not SESSION.llada,
	info="Use the original causal model every N denoising steps; diffusion confidence is used between checks.",
	interactive=False,
	)
	delay_eos_eot = gr.Checkbox(
	value=True,
	label="Delay EOS/EOT retention (longer answers)",
	info="Retains ending tokens last to give the answer room to develop.",
	)
	early_stopping = gr.Checkbox(
	value=True,
	label="Early stopping",
	info="Converge a block after two matching answers; stop all remaining blocks when EOS/EOT is finalized with a complete preceding answer.",
	)
	adaptive_parallel_decoding = gr.Checkbox(
	value=False, label="Adaptive Parallel Decoding (APD)",
	info="Commit confident tokens together, with at least ceil(block length / block iterations) per step. Finish blocks as soon as all positions are resolved.",
	)
	apd_aggressiveness = gr.Slider(0.0, 1.0, value=0.0, step=0.05,
	label="APD aggressiveness", visible=SESSION.llada, interactive=False,
	info="0 preserves APD's current thresholds. Higher values commit lower-confidence tokens sooner; 1 commits all positions in one pass per block. May reduce quality.")
	footer = "The first request may take longer while the model is loaded."
	if SHOW_REPOSITORY_LINKS:
	footer += f" [Model card](https://huggingface.co/{MODEL_REPO_ID})"
	if SOURCE_URL:
	footer += f" · [Source code]({SOURCE_URL})"
	visitor_badge = _visitor_badge(VISITOR_COUNTER_ID)
	if visitor_badge:
	footer += f"\n\n{visitor_badge}"
	gr.Markdown(footer)

	# Hidden inputs preserve programmatic evaluation overrides. The interface
	# uses the editable defaults above without exposing expert controls.
	bigram_scoring_method = gr.Textbox(value=BIGRAM_SCORING_METHOD, visible=False)
	bigram_remasking_confidence = gr.Textbox(value=BIGRAM_REMASKING_CONFIDENCE, visible=False)
	pmi_min_count = gr.Number(value=PMI_MIN_COUNT, precision=0, visible=False)
	pmi_clip = gr.Number(value=PMI_CLIP, visible=False)
	if not (SHOW_BIGRAM_EVALUATION_CONTROLS and BIGRAM_ENABLED):
	bigram_eos_delay_penalty = gr.Number(value=BIGRAM_ENDING_DELAY_STRENGTH, visible=False)
	bigram_model_path = gr.Textbox(value=BIGRAM_PATH, visible=False)
	generation_seed = gr.Number(value=-1, precision=0, visible=False)
	measure_decoding = gr.Checkbox(value=False, visible=False)

	inputs = [
	question,
	system_prompt,
	max_new_tokens,
	num_steps,
	block_length,
	temperature,
	top_k,
	pause_per_step,
	trajectory_color_mode,
	remasking_strategy,
	revisable_tokens,
	ar_check_interval,
	delay_eos_eot,
	early_stopping,
	bigram_model_path,
	bigram_scoring_method,
	pmi_min_count,
	pmi_clip,
	generation_seed,
	measure_decoding,
	bigram_eos_delay_penalty,
	bigram_remasking_confidence,
	repetition_penalty_strength,
	adaptive_parallel_decoding,
	sequence_sampling,
	count_all_candidate_bigrams,
	bigram_weight,
	repetition_risk_penalty,
	decay_repetition_penalty,
	apd_aggressiveness,
	ngram_guidance,
	]
	def update_ar_controls(strategy, apd=False):
	"""Expose AR-only controls and clear revision when another strategy is selected."""
	enabled = strategy == "Autoregressive-confidence" and not apd
	return gr.update(interactive=enabled, value=False), gr.update(interactive=enabled)

	remasking_strategy.change(
	update_ar_controls,
	inputs=[remasking_strategy, adaptive_parallel_decoding],
	outputs=[revisable_tokens, ar_check_interval],
	queue=False,
	)
	adaptive_parallel_decoding.change(
	update_ar_controls, inputs=[remasking_strategy, adaptive_parallel_decoding],
	outputs=[revisable_tokens, ar_check_interval], queue=False,
	)
	def update_apd_aggressiveness(enabled):
	return gr.update(interactive=bool(enabled))

	adaptive_parallel_decoding.change(
	update_apd_aggressiveness, inputs=[adaptive_parallel_decoding],
	outputs=[apd_aggressiveness], queue=False,
	)
	def update_repetition_control(strategy):
	return gr.update(visible=strategy == "PMI Viterbi + Repetition Penalty")

	remasking_strategy.change(
	update_repetition_control, inputs=[remasking_strategy],
	outputs=[repetition_penalty_strength], queue=False,
	)
	run_event = run.click(show_loading, outputs=[status, trajectory], queue=False)
	run_event.then(generate, inputs=inputs, outputs=[status, trajectory])
	submit_event = question.submit(show_loading, outputs=[status, trajectory], queue=False)
	submit_event.then(generate, inputs=inputs, outputs=[status, trajectory])

	demo.queue(default_concurrency_limit=1).launch(ssr_mode=False)