src/diffusion_lm/inference.py

"""Interactive iterative denoising inference for saved LoRA adapters."""
from __future__ import annotations

import gc
from html import escape
import json
import math
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

import torch
import torch.nn.functional as F
from peft import PeftModel
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer

from .bigram import BigramViterbiSelector
from .apd import commitment_mask, confidence_threshold, resolve_mask_predictions, selected_token_probabilities
from .data import validate_mask_token
from .legacy_compat import install_legacy_pickle_modules, patch_legacy_lora_modules, restore_legacy_pickle_modules
from .modeling import TRAINABLE_BASE_STATE_FILENAME, forward_bidirectional, load_trainable_base_state


def find_adapters(outputs_dir: str | Path = "outputs") -> list[str]:
    """Return adapter directories relative to outputs_dir, newest first."""
    root = Path(outputs_dir).resolve()
    if not root.exists():
        return []
    paths = [path for path in root.rglob("adapter_config.json") if path.parent.is_dir()]
    return [str(path.parent.relative_to(root)) for path in sorted(paths, key=lambda p: p.stat().st_mtime, reverse=True)]


def _safe_adapter_path(outputs_dir: str | Path, selection: str) -> Path:
    """Resolve a selected adapter while preventing paths outside outputs_dir."""
    root = Path(outputs_dir).resolve()
    path = (root / selection).resolve()
    if root not in path.parents or not (path / "adapter_config.json").is_file():
        raise ValueError("Select a valid adapter directory below outputs/.")
    return path


def _precision_dtype(precision: str, device: torch.device) -> torch.dtype:
    """Map configured precision to a safe dtype for the selected device."""
    requested = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}.get(precision, torch.float32)
    # CPU inference with low-precision weights is not generally supported; MPS
    # has better float32 compatibility for interactive single-request inference.
    if device.type == "cpu":
        return torch.float32
    # T4-class CUDA GPUs have no native BF16 Tensor Core support.  BF16
    # quantized compute there is substantially slower than FP16, so retain the
    # saved run's preference only where the hardware can execute it natively.
    if device.type == "cuda" and requested == torch.bfloat16 and not torch.cuda.is_bf16_supported():
        return torch.float16
    return requested


def select_device(requested: str = "auto") -> torch.device:
    """Choose an available CUDA, MPS, or CPU device from a UI selection."""
    if requested == "auto":
        requested = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    if requested == "cuda" and not torch.cuda.is_available():
        raise ValueError("CUDA was requested but is unavailable.")
    if requested == "mps" and not torch.backends.mps.is_available():
        raise ValueError("MPS was requested but is unavailable.")
    return torch.device(requested)


@dataclass
class InferenceSession:
    model: torch.nn.Module
    tokenizer: Any
    device: torch.device
    adapter_path: Path
    config: dict[str, Any]
    mask_token_id: int
    quantization: str = "none"
    compute_dtype: str = "unknown"
    legacy_wrapper: bool = False
    prompt_format: str = "chat_template"
    llada: bool = False


def _load_adapter_path(
    adapter_path: str | Path,
    device_name: str = "auto",
    quantization: str | None = None,
    preflight: bool = True,
) -> InferenceSession:
    """Load an adapter directory that has already been resolved and validated."""
    adapter_path = Path(adapter_path).expanduser().resolve()
    if not (adapter_path / "adapter_config.json").is_file():
        raise ValueError(f"Adapter directory has no adapter_config.json: {adapter_path}")
    config_candidates = (
        adapter_path / "resolved_config.json",
        adapter_path / "lad_run_config.json",
        adapter_path.parent / "resolved_config.json",
    )
    run_config_path = next((path for path in config_candidates if path.is_file()), None)
    run_config = json.loads(run_config_path.read_text()) if run_config_path else {}
    adapter_config = json.loads((adapter_path / "adapter_config.json").read_text())
    base_model = adapter_config["base_model_name_or_path"]
    device = select_device(device_name)
    dtype = _precision_dtype(run_config.get("precision", "fp32"), device)
    requested_quantization = str(quantization or "auto").lower()
    resolved_quantization = str(run_config.get("quantization", "none") if requested_quantization == "auto" else requested_quantization).lower()
    if resolved_quantization in {"4-bit", "qlora"}:
        resolved_quantization = "4bit"
    if resolved_quantization not in {"none", "off", "false", "4bit"}:
        raise ValueError("Inference quantization must be 'auto', 'none', or '4bit'.")
    use_4bit = resolved_quantization == "4bit"
    if use_4bit and (adapter_path / TRAINABLE_BASE_STATE_FILENAME).is_file():
        raise ValueError(
            "4-bit adapter loading is unsupported for checkpoints with fully trained base layers; "
            "load this experiment with quantization='none' or merge it first"
        )
    compute_dtype = dtype
    cache_dir = run_config.get("base_model_cache_dir", "base_models")
    token = os.getenv("HF_TOKEN")
    tokenizer_name = run_config.get("tokenizer_name_or_path", base_model)
    tokenizer_kwargs: dict[str, Any] = {}
    if "mistral" in str(tokenizer_name).lower():
        tokenizer_kwargs["fix_mistral_regex"] = True
    tokenizer = AutoTokenizer.from_pretrained(
        tokenizer_name,
        use_fast=True,
        token=token,
        cache_dir=cache_dir,
        clean_up_tokenization_spaces=False,
        **tokenizer_kwargs,
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    load_kwargs: dict[str, Any] = dict(torch_dtype=dtype, token=token, cache_dir=cache_dir, trust_remote_code=False)
    if use_4bit:
        if device.type != "cuda":
            raise RuntimeError("4-bit inference requires an NVIDIA CUDA device.")
        try:
            from transformers import BitsAndBytesConfig
            import bitsandbytes  # noqa: F401
        except ImportError as exc:
            raise ImportError("4-bit inference requires bitsandbytes; install with `pip install -e '.[cuda]'`.") from exc
        compute_dtype = _precision_dtype(run_config.get("compute_dtype", run_config.get("precision", "bf16")), device)
        load_kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type=str(run_config.get("quantization_type", "nf4")),
            bnb_4bit_compute_dtype=compute_dtype,
            bnb_4bit_use_double_quant=bool(run_config.get("double_quant", True)),
        )
        # Quantized modules cannot subsequently be moved with model.to().
        load_kwargs["device_map"] = {"": device.index if device.index is not None else 0}
    base = AutoModelForCausalLM.from_pretrained(base_model, **load_kwargs)
    base.config.use_cache = False
    base.config.is_causal = False
    if hasattr(base.config, "use_bidirectional_attention"):
        base.config.use_bidirectional_attention = True
    adapter_load_kwargs: dict[str, Any] = {}
    if not use_4bit:
        # Stage adapter tensors on CPU before moving the assembled root module.
        # This is required by ZeroGPU, where CUDA is represented by a proxy at
        # module-import time but no physical GPU has been allocated yet.
        adapter_load_kwargs["device_map"] = {"": "cpu"}
        adapter_load_kwargs["torch_device"] = "cpu"
    model = PeftModel.from_pretrained(
        base,
        adapter_path,
        is_trainable=False,
        **adapter_load_kwargs,
    )
    load_trainable_base_state(model, adapter_path)
    norm_path = adapter_path / "normalization_state.pt"
    if norm_path.is_file():
        model.load_state_dict(torch.load(norm_path, map_location="cpu", weights_only=True), strict=False)
    if not use_4bit:
        model.to(device)
    model.eval()
    mask_info = validate_mask_token(tokenizer, str(run_config.get("mask_token", "MASK")))
    session = InferenceSession(model, tokenizer, device, adapter_path, run_config, mask_info["mask_token_id"], "4bit" if use_4bit else "none", str(compute_dtype).removeprefix("torch."))
    if preflight:
        preflight_session(session)
    return session


def load_session(adapter_selection: str, outputs_dir: str | Path = "outputs", device_name: str = "auto", quantization: str | None = None) -> InferenceSession:
    """Load a base model, saved LoRA adapter, tokenizer, and norm state."""
    adapter_path = _safe_adapter_path(outputs_dir, adapter_selection)
    return _load_adapter_path(adapter_path, device_name, quantization)


def load_hub_adapter_session(
    repo_id: str,
    device_name: str = "auto",
    quantization: str | None = None,
    revision: str | None = None,
    cache_dir: str | Path | None = None,
    preflight: bool = True,
) -> InferenceSession:
    """Download and load a BYOD adapter from the Hugging Face Hub."""
    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:
        raise ImportError(
            "Hub inference requires huggingface_hub; install it with `pip install huggingface-hub`."
        ) from exc
    adapter_path = snapshot_download(
        repo_id=repo_id,
        repo_type="model",
        revision=revision,
        cache_dir=str(cache_dir) if cache_dir is not None else None,
        token=os.getenv("HF_TOKEN"),
    )
    return _load_adapter_path(adapter_path, device_name, quantization, preflight=preflight)


def load_merged_session(
    model_path: str | Path,
    device_name: str = "auto",
    quantization: str | None = None,
    source_config: dict[str, Any] | None = None,
) -> InferenceSession:
    """Load a standalone model produced by ``merge_adapter.py``."""
    model_path = Path(model_path).expanduser().resolve()
    if not (model_path / "config.json").is_file():
        raise ValueError(f"Merged model directory has no config.json: {model_path}")

    run_config = dict(source_config or {})
    saved_run_config = model_path / "lad_run_config.json"
    if not run_config and saved_run_config.is_file():
        run_config = json.loads(saved_run_config.read_text())

    device = select_device(device_name)
    dtype = _precision_dtype(run_config.get("precision", "bf16"), device)
    requested_quantization = str(quantization or "none").lower()
    if requested_quantization == "auto":
        requested_quantization = "none"
    if requested_quantization in {"4-bit", "qlora"}:
        requested_quantization = "4bit"
    if requested_quantization not in {"none", "off", "false", "4bit"}:
        raise ValueError("Merged-model quantization must be 'none' or '4bit'.")

    token = os.getenv("HF_TOKEN")
    tokenizer = AutoTokenizer.from_pretrained(
        model_path,
        use_fast=True,
        token=token,
        clean_up_tokenization_spaces=False,
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    use_4bit = requested_quantization == "4bit"
    compute_dtype = dtype
    load_kwargs: dict[str, Any] = {
        "torch_dtype": dtype,
        "token": token,
        "trust_remote_code": False,
    }
    if use_4bit:
        if device.type != "cuda":
            raise RuntimeError("4-bit merged-model inference requires an NVIDIA CUDA device.")
        try:
            from transformers import BitsAndBytesConfig
            import bitsandbytes  # noqa: F401
        except ImportError as exc:
            raise ImportError("4-bit inference requires bitsandbytes; install with `pip install -e '.[cuda]'`.") from exc
        compute_dtype = _precision_dtype(run_config.get("compute_dtype", run_config.get("precision", "bf16")), device)
        load_kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type=str(run_config.get("quantization_type", "nf4")),
            bnb_4bit_compute_dtype=compute_dtype,
            bnb_4bit_use_double_quant=bool(run_config.get("double_quant", True)),
        )
        load_kwargs["device_map"] = {"": device.index if device.index is not None else 0}

    model = AutoModelForCausalLM.from_pretrained(model_path, **load_kwargs)
    model.config.use_cache = False
    model.config.is_causal = False
    if hasattr(model.config, "use_bidirectional_attention"):
        model.config.use_bidirectional_attention = True
    if not use_4bit:
        model.to(device)
    model.eval()

    mask_info = validate_mask_token(tokenizer, str(run_config.get("mask_token", "MASK")))
    session = InferenceSession(
        model=model,
        tokenizer=tokenizer,
        device=device,
        adapter_path=model_path,
        config=run_config,
        mask_token_id=mask_info["mask_token_id"],
        quantization="4bit" if use_4bit else "none",
        compute_dtype=str(compute_dtype).removeprefix("torch."),
    )
    preflight_session(session)
    return session


def _load_legacy_checkpoint_session(checkpoint: str | Path, tokenizer_name_or_path: str, device_name: str = "auto", source_config: dict[str, Any] | None = None) -> InferenceSession:
    """Load one trusted legacy full-object checkpoint from a local path."""
    checkpoint = Path(checkpoint).expanduser().resolve()
    if not checkpoint.is_file():
        raise ValueError(f"Legacy checkpoint does not exist: {checkpoint}")
    if not tokenizer_name_or_path.strip():
        raise ValueError("Legacy loading requires a tokenizer name or local tokenizer path.")
    # A full-object checkpoint can execute pickle code.  This loader is for
    # checkpoints the user trusts, including their locally archived model.
    previous_modules = install_legacy_pickle_modules()
    try:
        model = torch.load(checkpoint, map_location="cpu", weights_only=False)
    finally:
        restore_legacy_pickle_modules(previous_modules)
    if not isinstance(model, torch.nn.Module):
        raise ValueError(f"{checkpoint} is not a full torch.nn.Module checkpoint.")
    patch_legacy_lora_modules(model)

    token = os.getenv("HF_TOKEN")
    device = select_device(device_name)
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path.strip(), use_fast=True, token=token, clean_up_tokenization_spaces=False)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    model.to(device).eval()
    mask_info = validate_mask_token(tokenizer)
    session = InferenceSession(
        model=model,
        tokenizer=tokenizer,
        device=device,
        adapter_path=checkpoint,
        config=source_config or {"model_source": "local_legacy", "checkpoint": str(checkpoint), "tokenizer_name_or_path": tokenizer_name_or_path.strip()},
        mask_token_id=mask_info["mask_token_id"],
        quantization="none",
        compute_dtype=str(next((parameter.dtype for parameter in model.parameters() if parameter.is_floating_point()), torch.float32)).removeprefix("torch."),
        legacy_wrapper=True,
        prompt_format="legacy_llama",
    )
    preflight_session(session)
    return session


def load_local_legacy_session(checkpoint_path: str | Path, tokenizer_name_or_path: str, device_name: str = "auto") -> InferenceSession:
    """Load and preflight a trusted local legacy full-model checkpoint."""
    return _load_legacy_checkpoint_session(checkpoint_path, tokenizer_name_or_path, device_name)


def load_hosted_legacy_session(repo_id: str, filename: str, tokenizer_name_or_path: str, device_name: str = "auto") -> InferenceSession:
    """Load the trusted legacy full-model checkpoint hosted on Hugging Face.

    This exists for controlled comparisons: the checkpoint uses its original
    wrapper to construct full bidirectional attention, while decoding uses the
    current project's prompt and denoising loop.  Pickled checkpoints are only
    safe to load from a repository you trust.
    """
    if not repo_id.strip() or not filename.strip() or not tokenizer_name_or_path.strip():
        raise ValueError("Hosted legacy loading requires a repository ID, checkpoint filename, and tokenizer name.")
    try:
        from huggingface_hub import hf_hub_download
    except ImportError as exc:
        raise ImportError("Hosted model loading requires huggingface_hub, installed with transformers.") from exc

    token = os.getenv("HF_TOKEN")
    checkpoint = hf_hub_download(repo_id=repo_id.strip(), filename=filename.strip(), token=token)
    return _load_legacy_checkpoint_session(
        checkpoint,
        tokenizer_name_or_path,
        device_name,
        {"model_source": "huggingface_legacy", "repo_id": repo_id.strip(), "filename": filename.strip(), "tokenizer_name_or_path": tokenizer_name_or_path.strip()},
    )


def load_llada_session(repo_id: str = "GSAI-ML/LLaDA-8B-Instruct", device_name: str = "auto", *, revision: str | None = None, preflight: bool = True) -> InferenceSession:
    """Load LLaDA Instruct as a mask predictor for this app's denoising loop."""
    if not repo_id.strip():
        raise ValueError("LLaDA loading requires a Hugging Face repository ID.")
    device = select_device(device_name)
    if device.type not in {"cuda", "mps"}:
        raise ValueError("LLaDA-8B-Instruct requires CUDA or MPS inference; select a GPU-capable runtime.")
    # Apple MPS does not reliably support BF16 inference for this remote model.
    # FP16 is the practical MPS format; CUDA retains BF16 where available.
    dtype = torch.float16 if device.type == "mps" else _precision_dtype("bf16", device)
    token = os.getenv("HF_TOKEN")
    cache_dir = "base_models"
    try:
        tokenizer = AutoTokenizer.from_pretrained(
            repo_id.strip(), trust_remote_code=True, token=token, cache_dir=cache_dir, revision=revision,
        )
        model = AutoModel.from_pretrained(
            repo_id.strip(), trust_remote_code=True, torch_dtype=dtype, token=token, cache_dir=cache_dir, revision=revision,
        )
    except Exception as exc:
        raise RuntimeError(
            "Could not load LLaDA. Check the pinned model revision, remote-code files, "
            "and Transformers compatibility."
        ) from exc
    if tokenizer.pad_token_id == 126336:
        raise ValueError("LLaDA's pad token must differ from its fixed mask token (126336).")
    tokenizer.padding_side = "left"
    model.to(device).eval()
    session = InferenceSession(
        model=model,
        tokenizer=tokenizer,
        device=device,
        adapter_path=Path(repo_id.strip()),
        config={"model_source": "huggingface_llada", "repo_id": repo_id.strip(), "revision": revision},
        mask_token_id=126336,
        quantization="none",
        compute_dtype=str(dtype).removeprefix("torch."),
        prompt_format="llada",
        llada=True,
    )
    if preflight:
        preflight_session(session)
    return session


def _sample(
    logits: torch.Tensor,
    temperature: float,
    top_k: int,
    generator: torch.Generator | None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Sample tokens; greedy decoding retains full-vocabulary confidence."""
    if float(temperature) == 0.0:
        picked = logits.argmax(dim=-1)
        probabilities = F.softmax(logits.float(), dim=-1)
        return picked, probabilities.gather(-1, picked.unsqueeze(-1)).squeeze(-1)
    logits = logits / max(temperature, 1e-5)
    vocab_size = logits.shape[-1]
    k = min(max(int(top_k), 1), vocab_size)
    values, indices = torch.topk(logits, k, dim=-1)
    probabilities = F.softmax(values, dim=-1)
    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


def _sampling_token_confidence(
    logits: torch.Tensor,
    token_ids: torch.Tensor,
    temperature: float,
    top_k: int,
) -> torch.Tensor:
    """Return each existing token's probability under the top-k sampler."""
    if float(temperature) == 0.0:
        probabilities = F.softmax(logits.float(), dim=-1)
        return probabilities.gather(-1, token_ids.unsqueeze(-1)).squeeze(-1)
    scaled = logits / max(float(temperature), 1e-5)
    k = min(max(int(top_k), 1), scaled.shape[-1])
    values, indices = torch.topk(scaled, k, dim=-1)
    probabilities = F.softmax(values, dim=-1)
    matches = indices == token_ids.unsqueeze(-1)
    return (probabilities * matches).sum(dim=-1)


def _llada_gumbel_noise(logits: torch.Tensor, temperature: float) -> torch.Tensor:
    """Apply the float64 Gumbel-max transform used by official LLaDA decoding."""
    if float(temperature) == 0.0:
        return logits
    logits = logits.to(torch.float64)
    noise = torch.rand_like(logits, dtype=torch.float64)
    return logits.exp() / (-torch.log(noise)).pow(float(temperature))


def _apply_repetition_penalty(
    logits: torch.Tensor,
    answer_ids: torch.Tensor,
    penalty: float,
    mask_token_id: int,
    *,
    exclude_self: bool = False,
    excluded_token_ids: set[int] | None = None,
) -> torch.Tensor:
    """Reduce repeated-token probability weights by ``penalty ** count``.

    Only tokens already present in the generated answer are penalized; prompt
    tokens, MASK, and configured special tokens never contribute. Subtracting
    ``count * log(penalty)`` from a logit divides its unnormalized softmax
    probability by ``penalty ** count``. For the revisable denoise stream, a
    position's current token is excluded from its count, avoiding needless
    churn of unique predictions.
    """
    penalty = float(penalty)
    if penalty < 1.0:
        raise ValueError("repetition_penalty must be at least 1.0")
    if penalty == 1.0:
        return logits
    if logits.ndim != 3 or answer_ids.ndim != 2 or logits.shape[:2] != answer_ids.shape:
        raise ValueError("repetition penalty expects logits [batch, length, vocab] matching answer IDs")

    adjusted = logits.clone()
    vocabulary_size = adjusted.shape[-1]
    for batch_index in range(answer_ids.shape[0]):
        valid = answer_ids[batch_index]
        valid_mask = (
            (valid != int(mask_token_id))
            & (valid >= 0)
            & (valid < vocabulary_size)
        )
        excluded = set(excluded_token_ids or ())
        excluded.add(int(mask_token_id))
        if excluded:
            excluded_tensor = torch.tensor(
                sorted(excluded), device=valid.device, dtype=valid.dtype
            )
            valid_mask &= ~torch.isin(valid, excluded_tensor)
        valid = valid[valid_mask]
        if not len(valid):
            continue
        token_ids, counts = torch.unique(valid, return_counts=True)
        current = answer_ids[batch_index]
        scores = adjusted[batch_index, :, token_ids]
        exponents = counts[None, :].expand(logits.shape[1], -1)
        if exclude_self:
            exponents = exponents - (current[:, None] == token_ids[None, :]).to(
                dtype=exponents.dtype
            )
        log_penalty = torch.log(
            torch.tensor(penalty, device=logits.device, dtype=torch.float32)
        )
        penalized = scores.float() - exponents.float() * log_penalty
        # Penalty arithmetic stays in FP32 for numerical stability, then returns
        # to the model's native BF16/FP16 dtype for indexed assignment.
        adjusted[batch_index, :, token_ids] = penalized.to(dtype=adjusted.dtype)
    return adjusted


def _apply_eos_eot_prediction_penalty(
    logits: torch.Tensor,
    penalty: float,
    eos_token_id: int,
    eot_token_id: int | None = None,
) -> torch.Tensor:
    """Reduce EOS/EoT sampling weights without changing retention confidence.

    Subtracting ``log(penalty)`` from the selected logits divides their
    unnormalized probability weight by ``penalty``.  This is deliberately
    independent of the LLaDA-style delayed-retention option, which changes
    which sampled positions are retained or re-masked rather than what token
    is sampled in the first place.
    """
    penalty = float(penalty)
    if penalty < 1.0:
        raise ValueError("EOS/EOT prediction penalty must be at least 1.0")
    if penalty == 1.0:
        return logits

    token_ids = {int(eos_token_id)}
    if eot_token_id is not None:
        token_ids.add(int(eot_token_id))
    token_ids = {token_id for token_id in token_ids if 0 <= token_id < logits.shape[-1]}
    if not token_ids:
        return logits

    adjusted = logits.clone()
    log_penalty = torch.log(
        torch.tensor(penalty, device=logits.device, dtype=torch.float32)
    )
    indices = torch.tensor(sorted(token_ids), device=logits.device, dtype=torch.long)
    adjusted[..., indices] = (
        adjusted[..., indices].float() - log_penalty
    ).to(dtype=adjusted.dtype)
    return adjusted


def _llada_transfer_schedule(mask_count: int, steps: int) -> list[int]:
    """Distribute a linear-noise transfer budget uniformly across steps."""
    if mask_count < 0 or steps < 1:
        raise ValueError("mask_count must be non-negative and steps must be positive")
    base, remainder = divmod(mask_count, steps)
    return [base + int(index < remainder) for index in range(steps)]


def _block_step_plan(
    generation_length: int,
    steps: int,
    block_length: int | None,
) -> list[tuple[int, int, int, int, int]]:
    """Allocate a fixed total step budget across sequential answer blocks."""
    generation_length, steps = int(generation_length), int(steps)
    requested_block_length = generation_length if block_length is None else int(block_length)
    if generation_length < 1 or steps < 1 or requested_block_length < 1:
        raise ValueError("generation length, steps, and block length must be positive")
    effective_block_length = min(requested_block_length, generation_length)
    num_blocks = (generation_length + effective_block_length - 1) // effective_block_length
    if steps < num_blocks:
        raise ValueError(
            f"Denoising steps ({steps}) must be at least the number of blocks ({num_blocks})"
        )
    base, remainder = divmod(steps, num_blocks)
    plan = []
    for block_index in range(num_blocks):
        block_start = block_index * effective_block_length
        block_end = min(generation_length, block_start + effective_block_length)
        block_steps = base + int(block_index < remainder)
        for block_step in range(block_steps):
            plan.append((block_index, block_start, block_end, block_step, block_steps))
    return plan


def _remask_offsets(confidence: torch.Tensor, mask_probability: float, confidence_guided: bool) -> torch.Tensor:
    """Choose answer offsets to re-mask, preferring uncertain tokens when guided."""
    probability = max(0.0, min(1.0, float(mask_probability)))
    if confidence_guided:
        count = round(probability * len(confidence))
        return torch.argsort(confidence)[:count]
    return torch.where(torch.rand(len(confidence), device=confidence.device) < probability)[0]


def _autoregressive_token_confidence(
    session: InferenceSession,
    input_ids: torch.Tensor,
    answer_start: int,
) -> torch.Tensor:
    """Score answer tokens causally with the original model under the adapter.

    The returned value at answer offset ``i`` is
    ``p(x_i | prompt, x_{<i})``. The adapter is disabled only for this forward
    pass; no model weights are copied or reloaded.
    """
    disable_adapter = getattr(session.model, "disable_adapter", None)
    if disable_adapter is None:
        raise ValueError(
            "Autoregressive-confidence remasking requires an unmerged PEFT adapter "
            "so the original causal model can be restored without loading a second model."
        )
    if answer_start < 1 or answer_start >= input_ids.shape[1]:
        raise ValueError("answer_start must identify a non-empty answer after the prompt")
    attention_mask = torch.ones_like(input_ids, dtype=torch.long)
    with disable_adapter():
        outputs = session.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            use_cache=False,
        )
    # The logit immediately before each answer token predicts that token in a
    # causal LM. Work in FP32 for stable ranking of low-probability candidates.
    causal_logits = outputs.logits[:, answer_start - 1 : -1].float()
    targets = input_ids[:, answer_start:]
    target_logits = causal_logits.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
    return (target_logits - torch.logsumexp(causal_logits, dim=-1)).exp()[0]


def _native_eot_token_id(tokenizer: Any) -> int | None:
    """Return a tokenizer's native end-of-turn ID when it has one."""
    convert = getattr(tokenizer, "convert_tokens_to_ids", None)
    if convert is None:
        return None
    unknown = getattr(tokenizer, "unk_token_id", None)
    for token in ("<|eot_id|>", "<end_of_turn>", "<|end_of_turn|>"):
        token_id = convert(token)
        if token_id is not None and token_id != unknown and int(token_id) >= 0:
            return int(token_id)
    return None


def _ending_token_ids(tokenizer: Any) -> set[int]:
    """Recognize EOS and the backbone's native end-of-turn in every strategy."""
    ids = {int(tokenizer.eos_token_id)}
    eot = _native_eot_token_id(tokenizer)
    if eot is not None:
        ids.add(eot)
    return ids


def _first_ending_offset(tokens: list[int], ending_ids: set[int]) -> int | None:
    return next((offset for offset, token in enumerate(tokens) if token in ending_ids), None)


def forward_denoising(session: InferenceSession, input_ids: torch.Tensor, padding_mask: torch.Tensor) -> torch.Tensor:
    """Return denoising logits for either the current or legacy model wrapper."""
    if session.llada:
        # LLaDA caches rotary embeddings during preflight. Its remote model code
        # requires all later uses of those cached inference tensors to remain in
        # inference mode as well.
        with torch.inference_mode():
            outputs = session.model(input_ids, attention_mask=(~padding_mask).to(dtype=torch.long))
        return outputs.logits
    if session.legacy_wrapper:
        # The archived CustomTransformerModel builds its own full-attention
        # 4-D mask and passes use_cache=False to its inner Peft model. Passing
        # either argument here would duplicate the wrapper's keyword.
        outputs = session.model(input_ids=input_ids)
        return outputs["logits"] if isinstance(outputs, dict) else outputs.logits
    return forward_bidirectional(session.model, input_ids, padding_mask)


@torch.inference_mode()
def preflight_session(session: InferenceSession) -> tuple[int, int]:
    """Run one real forward pass and fail early if a loaded model is unusable."""
    prefix = _prompt_ids(session.tokenizer, "Reply with OK.", "You are a helpful assistant.", session.prompt_format)
    input_ids = torch.tensor([prefix + [session.mask_token_id]], device=session.device, dtype=torch.long)
    padding = torch.zeros_like(input_ids, dtype=torch.bool)
    try:
        logits = forward_denoising(session, input_ids, padding)
    except Exception as exc:
        source = "LLaDA" if session.llada else "legacy hosted checkpoint" if session.legacy_wrapper else "saved adapter"
        raise RuntimeError(f"Inference preflight failed for {source}; the model was not loaded for generation: {exc}") from exc
    if logits.ndim != 3 or logits.shape[:2] != input_ids.shape:
        raise RuntimeError(f"Inference preflight returned invalid logits shape {tuple(logits.shape)} for input shape {tuple(input_ids.shape)}")
    if not torch.isfinite(logits[:, -1]).all():
        raise RuntimeError("Inference preflight produced non-finite final-token logits.")
    return int(input_ids.shape[1]), int(logits.shape[-1])


def _prompt_ids(tokenizer: Any, question: str, system_prompt: str, prompt_format: str = "chat_template") -> list[int]:
    """Render system/user messages through a tokenizer’s native chat template."""
    from .data import apply_neutral_chat_template
    if not question.strip():
        raise ValueError("Enter a question or prompt.")
    if prompt_format == "legacy_llama":
        # The hosted historical checkpoint used a base Llama tokenizer with no
        # chat_template. Match the prompt layout from its original app while
        # still running the current project's denoising/sampling loop.
        prompt = (
            "<|begin_of_text|>\n"
            "<|start_header_id|>system<|end_header_id|>\n"
            f"{system_prompt}\n"
            "<|start_header_id|>user<|end_header_id|>\n"
            f"{question.strip()}\n"
            "<|start_header_id|>assistant<|end_header_id|>\n"
        )
        return list(tokenizer.encode(prompt, add_special_tokens=False))
    if prompt_format == "llada":
        content = f"{system_prompt}\n\n{question.strip()}" if system_prompt.strip() else question.strip()
        rendered = apply_neutral_chat_template(
            tokenizer,
            [{"role": "user", "content": content}],
            tokenize=True,
            add_generation_prompt=True,
        )
        if isinstance(rendered, str):
            rendered = tokenizer.encode(rendered, add_special_tokens=False)
        elif hasattr(rendered, "input_ids"):
            rendered = rendered.input_ids
        if rendered and isinstance(rendered[0], list):
            rendered = rendered[0]
        return list(rendered)
    if prompt_format != "chat_template":
        raise ValueError(f"Unknown prompt format: {prompt_format}")
    if not getattr(tokenizer, "chat_template", None):
        raise ValueError(f"{tokenizer.name_or_path} has no chat template; inference needs one to identify the answer boundary.")
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": question},
    ]
    try:
        rendered = apply_neutral_chat_template(tokenizer, messages, tokenize=True, add_generation_prompt=True)
    except Exception as exc:
        # Gemma's official template rejects a separate system role; preserve
        # the prompt by folding it into the user message, as training does.
        if exc.__class__.__name__ != "TemplateError" or "System role not supported" not in str(exc):
            raise
        rendered = apply_neutral_chat_template(tokenizer, [
            {"role": "user", "content": f"{system_prompt}\n\n{question}"},
        ], tokenize=True, add_generation_prompt=True)
    if isinstance(rendered, str):
        rendered = tokenizer.encode(rendered, add_special_tokens=False)
    elif hasattr(rendered, "input_ids"):
        rendered = rendered.input_ids
    if rendered and isinstance(rendered[0], list):
        rendered = rendered[0]
    if not all(isinstance(token, int) for token in rendered):
        raise ValueError(f"Tokenizer returned a non-integer chat-template encoding: {type(rendered).__name__}")
    return list(rendered)


@torch.inference_mode()
def llada_generate(
    session: InferenceSession,
    question: str,
    *,
    gen_length: int,
    steps: int,
    block_length: int | None = None,
    temperature: float = 0.0,
    cfg_scale: float = 0.0,
    remasking: str = "low_confidence",
    logits_eos_inf: bool = False,
    confidence_eos_eot_inf: bool = False,
    eot_token_id: int | None = None,
    system_prompt: str = "",
    seed: int = 1234,
    repetition_penalty: float = 1.0,
    eos_eot_prediction_penalty: float = 1.0,
) -> str:
    """Generate with the official LLaDA fixed-budget transfer algorithm.

    This intentionally bypasses ``denoise_stream``: official LLaDA predicts
    only still-masked positions, permanently transfers a fixed number per
    reverse step, and uses neither proportional unmasking nor a mask-ratio
    heuristic. The sampler is model-agnostic, so mask-only LAD adapters can use
    it with their native tokenizer and prompt format as well.
    """
    gen_length, steps = int(gen_length), int(steps)
    block_length = int(block_length or gen_length)
    if gen_length < 1 or steps < 1 or block_length < 1:
        raise ValueError("gen_length, steps, and block_length must be positive")
    if gen_length % block_length:
        raise ValueError("LLaDA gen_length must be divisible by block_length")
    num_blocks = gen_length // block_length
    if steps % num_blocks:
        raise ValueError("LLaDA steps must be divisible by the number of blocks")
    if remasking not in {"low_confidence", "random"}:
        raise ValueError("LLaDA remasking must be 'low_confidence' or 'random'")

    torch.manual_seed(int(seed))
    if session.device.type == "cuda":
        torch.cuda.manual_seed_all(int(seed))
    prefix = _prompt_ids(session.tokenizer, question, system_prompt, session.prompt_format)
    prompt_length = len(prefix)
    x = torch.full((1, prompt_length + gen_length), session.mask_token_id, dtype=torch.long, device=session.device)
    x[0, :prompt_length] = torch.tensor(prefix, dtype=torch.long, device=session.device)
    padding = torch.zeros_like(x, dtype=torch.bool)
    prompt_index = x != session.mask_token_id
    steps_per_block = steps // num_blocks
    eos_token_id = int(session.tokenizer.eos_token_id)

    for block in range(num_blocks):
        block_start = prompt_length + block * block_length
        block_end = block_start + block_length
        transfer_schedule = _llada_transfer_schedule(int((x[:, block_start:block_end] == session.mask_token_id).sum()), steps_per_block)
        for transfer_count in transfer_schedule:
            mask_index = x == session.mask_token_id
            if cfg_scale > 0.0:
                unconditional = x.clone()
                unconditional[prompt_index] = session.mask_token_id
                model_input = torch.cat([x, unconditional], dim=0)
                model_padding = torch.cat([padding, padding], dim=0)
                conditional_logits, unconditional_logits = forward_denoising(session, model_input, model_padding).chunk(2, dim=0)
                logits = unconditional_logits + (float(cfg_scale) + 1.0) * (conditional_logits - unconditional_logits)
            else:
                logits = forward_denoising(session, x, padding)
            logits[:, prompt_length:] = _apply_repetition_penalty(
                logits[:, prompt_length:],
                x[:, prompt_length:],
                repetition_penalty,
                session.mask_token_id,
                excluded_token_ids=set(getattr(session.tokenizer, "all_special_ids", [])),
            )
            logits[:, prompt_length:] = _apply_eos_eot_prediction_penalty(
                logits[:, prompt_length:],
                eos_eot_prediction_penalty,
                eos_token_id,
                eot_token_id,
            )
            if logits_eos_inf:
                logits = logits.clone()
                logits[..., eos_token_id] = -torch.inf
            predictions = torch.argmax(_llada_gumbel_noise(logits, temperature), dim=-1)
            if remasking == "low_confidence":
                probabilities = F.softmax(logits, dim=-1)
                confidence = probabilities.gather(-1, predictions.unsqueeze(-1)).squeeze(-1)
                if confidence_eos_eot_inf:
                    # Appendix B.4 delays EOS/EoT predictions by assigning
                    # them the lowest transfer confidence; they remain valid
                    # predictions and can still transfer in later steps.
                    special_prediction = predictions == eos_token_id
                    if eot_token_id is not None and 0 <= int(eot_token_id) < logits.shape[-1]:
                        special_prediction |= predictions == int(eot_token_id)
                    confidence = confidence.masked_fill(special_prediction, torch.finfo(confidence.dtype).min)
            else:
                confidence = torch.rand(predictions.shape, device=session.device)
            candidate = mask_index.clone()
            candidate[:, :block_start] = False
            candidate[:, block_end:] = False
            confidence = confidence.masked_fill(~candidate, -torch.inf)
            if transfer_count:
                transfer = torch.topk(confidence[0], k=int(transfer_count)).indices
                x[0, transfer] = predictions[0, transfer]

    answer = x[0, prompt_length:].tolist()
    return session.tokenizer.decode(answer, skip_special_tokens=True).strip()


def render_denoising_step(
    tokens: list[int],
    confidences: list[float],
    answer_start: int,
    tokenizer: Any,
    mask_token_id: int,
    step: int,
    total_steps: int,
    retained: set[int] | None = None,
    frozen: dict[int, int] | None = None,
    frozen_confidences: dict[int, float] | None = None,
    frozen_steps: dict[int, int] | None = None,
    color_mode: str = "Prediction probability",
) -> str:
    """Render one denoising state with optional token coloring."""
    ending_ids = _ending_token_ids(tokenizer)
    pieces = []
    answer = tokens[answer_start:]
    output_token_count = 0
    for offset, token in enumerate(answer):
        if token in ending_ids:
            break
        output_token_count += 1
        token_text = escape(tokenizer.decode([token], skip_special_tokens=False))
        if token == mask_token_id:
            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"token position {offset} · masked at iteration {step}"
        else:
            confidence = (
                frozen_confidences[offset]
                if frozen and offset in frozen and frozen_confidences and offset in frozen_confidences
                else float(confidences[offset]) if offset < len(confidences) else 0.0
            )
            confidence = max(0.0, min(1.0, confidence))
            predicted_step = (
                frozen_steps[offset]
                if frozen and offset in frozen and frozen_steps and offset in frozen_steps
                else step
            )
            if color_mode == "Prediction iteration":
                iteration_fraction = max(0.0, min(1.0, predicted_step / max(total_steps, 1)))
                lightness = 72 - round(42 * iteration_fraction)
                style = f"color:hsl(210,90%,{lightness}%);font-weight:500"
            elif color_mode == "Prediction probability":
                hue = int(confidence * 120)
                style = f"color:hsl({hue},90%,30%);font-weight:{'600' if confidence > .8 else '400'}"
            else:
                style = "color:inherit;font-weight:400"
            title = (
                f"token position {offset} · predicted at iteration {predicted_step} · "
                f"sampling probability {confidence:.1%}"
            )
        pieces.append(f"<span style='{style}' title='{title}'>{token_text}</span>")
    end_offset = _first_ending_offset(answer, ending_ids)
    visible_answer = answer[:end_offset] if end_offset is not None else answer
    # A plain textarea copies its value, avoiding rich-text serialization of
    # per-token color/bold spans. Preserve the complete decoded code verbatim.
    plain_text = (
        decode_denoising_state(visible_answer, tokenizer, mask_token_id)
        if mask_token_id in visible_answer
        else tokenizer.decode(visible_answer, skip_special_tokens=True, clean_up_tokenization_spaces=False)
    )
    copyable = (
        "<details style='margin-top:10px'><summary>Plain text / copy code</summary>"
        "<textarea readonly aria-label='Plain text output' spellcheck='false' "
        "style='box-sizing:border-box;width:100%;min-height:9em;margin-top:6px;"
        "font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;"
        "white-space:pre;tab-size:4'>" + escape(plain_text) + "</textarea></details>"
    )
    pct = int(100 * step / max(total_steps, 1))
    if color_mode == "Prediction iteration":
        legend = "Light-to-dark blue indicates earlier-to-later prediction iterations"
    elif color_mode == "Prediction probability":
        legend = "Red-to-green indicates lower-to-higher sampling probability"
    else:
        legend = "Token coloring is disabled"
    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'>{legend}; gray chips are masks. Hover over a token for its position, prediction iteration, and probability.</div>{copyable}</div>")


def decode_denoising_state(
    tokens: list[int],
    tokenizer: Any,
    mask_token_id: int,
    *,
    show_eos_tokens: bool = False,
) -> str:
    """Decode one answer state with unresolved positions and optional EOS shown."""
    eos_id = tokenizer.eos_token_id
    visible_end_ids = _ending_token_ids(tokenizer)
    end_offset = _first_ending_offset(tokens, visible_end_ids)
    if not show_eos_tokens and end_offset is not None:
        tokens = tokens[:end_offset]

    # Decode contiguous resolved spans so subword spacing remains natural, but
    # make adjacent mask tokens unambiguous and independent of the configured
    # mask vocabulary item (some model configs use markers such as `<?>`).
    pieces: list[str] = []
    resolved: list[int] = []

    def flush_resolved() -> None:
        if resolved:
            text = tokenizer.decode(
                resolved,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
            ).strip()
            if text:
                pieces.append(text)
            resolved.clear()

    for token in tokens:
        if token == mask_token_id:
            flush_resolved()
            pieces.append("MASK")
        elif show_eos_tokens and token in visible_end_ids:
            flush_resolved()
            marker = tokenizer.decode(
                [token],
                skip_special_tokens=False,
                clean_up_tokenization_spaces=False,
            ).strip()
            fallback = "<EOS>" if token == eos_id else "<EOT>"
            pieces.append(marker or fallback)
        else:
            resolved.append(token)
    flush_resolved()
    return " ".join(pieces)


def denoise_stream(session: InferenceSession, question: str, system_prompt: str, max_new_tokens: int, num_steps: int, noise_level: float, temperature: float, top_k: int, seed: int, permanent_unmask: bool = False, confidence_guided: bool = False, proportional_unmask: bool = True, early_stopping: bool = False, confidence_eos_eot_inf: bool = False, freeze_retained_tokens: bool = True, repetition_penalty: float = 1.0, eos_eot_prediction_penalty: float = 1.0, include_pre_remask_prediction: bool = False, block_length: int | None = None, trajectory_color_mode: str = "Prediction probability", autoregressive_guided: bool = False, revisable_tokens: bool = False, autoregressive_check_interval: int = 1, bigram_selector: BigramViterbiSelector | None = None, bigram_top_k: int | None = None, bigram_eos_delay_penalty: float = 10.0, adaptive_parallel_decoding: bool = False, apd_aggressiveness: float = 0.0):
    """Yield denoising states with retained positions and optional locked values.

    Delayed endings also receive a fading unary penalty in bigram mode only.
    With early stopping, a finalized EOS/EOT and resolved prefix finish the
    whole answer; a converged block without an ending advances to the next.
    """
    prefix = _prompt_ids(session.tokenizer, question, system_prompt, session.prompt_format)
    max_new_tokens, num_steps = int(max_new_tokens), int(num_steps)
    if max_new_tokens < 1 or num_steps < 1:
        raise ValueError("max_new_tokens and num_steps must both be at least 1.")
    if adaptive_parallel_decoding:
        confidence_threshold(0, 1, apd_aggressiveness)  # Validate before the first forward.
        # APD replaces the retention policy, not the selected-token scorer.
        permanent_unmask, freeze_retained_tokens, revisable_tokens = True, True, False
    if bigram_selector is not None:
        # A separate candidate count is retained for legacy offline callers;
        # the hosted interface uses the shared top_k control.
        bigram_top_k = int(top_k) if bigram_top_k is None else int(bigram_top_k)
        if not math.isfinite(bigram_eos_delay_penalty) or bigram_eos_delay_penalty < 1:
            raise ValueError("bigram_eos_delay_penalty must be finite and at least 1")
        if int(bigram_top_k) < 1:
            raise ValueError("bigram_top_k must be positive")
        if autoregressive_guided or revisable_tokens:
            raise ValueError("Bigram selection uses diffusion-confidence remasking without revisable tokens")
        confidence_guided = True
    if confidence_guided and autoregressive_guided:
        raise ValueError("Choose either diffusion-confidence or autoregressive-confidence remasking, not both.")
    if autoregressive_guided and not hasattr(session.model, "disable_adapter"):
        raise ValueError("Autoregressive-confidence remasking requires an unmerged PEFT adapter.")
    if revisable_tokens and not autoregressive_guided:
        raise ValueError("Revisable tokens are supported only with autoregressive-confidence remasking.")
    autoregressive_check_interval = int(autoregressive_check_interval)
    if autoregressive_check_interval < 1:
        raise ValueError("autoregressive_check_interval must be at least 1")
    step_plan = _block_step_plan(max_new_tokens, num_steps, block_length)
    num_blocks = step_plan[-1][0] + 1
    ids = prefix + [session.mask_token_id] * max_new_tokens
    answer_start = len(prefix)
    # Use the device's default RNG so this works consistently on CUDA, MPS, and
    # CPU; seed it once per request for reproducible interactive runs.
    torch.manual_seed(int(seed))
    if session.device.type == "cuda":
        torch.cuda.manual_seed_all(int(seed))
    padding = torch.zeros((1, len(ids)), device=session.device, dtype=torch.bool)
    last_confidence = 0.0
    retained: set[int] = set()
    frozen: dict[int, int] = {}
    frozen_confidences: dict[int, float] = {}
    frozen_steps: dict[int, int] = {}
    last_predictions: list[tuple[int, ...]] = []
    eot_token_id = _native_eot_token_id(session.tokenizer)
    ending_ids = _ending_token_ids(session.tokenizer)
    # EOS/EOT delaying and confidence-guided retention are independent. In
    # random mode, endings are still retained last, while ordinary tokens are
    # selected randomly.
    guided_retention = confidence_guided or autoregressive_guided
    skip_block_index: int | None = None
    forward_passes = 0
    for step, (block_index, block_start, block_end, block_step, block_steps) in enumerate(step_plan):
        if block_index == skip_block_index:
            continue
        if block_step == 0:
            last_predictions.clear()
        has_followup_step = block_step + 1 < block_steps
        ar_check_due = autoregressive_guided and not adaptive_parallel_decoding and has_followup_step and (
            (block_step + 1) % autoregressive_check_interval == 0
            or block_step + 2 == block_steps
        )
        tokens = torch.tensor([ids], device=session.device, dtype=torch.long)
        bigram_retention_scores = None
        with torch.inference_mode():
            answer_ids = tokens[:, answer_start:]
            logits = forward_denoising(session, tokens, padding)[:, answer_start:]
            forward_passes += 1
            logits = _apply_repetition_penalty(
                logits,
                answer_ids,
                repetition_penalty,
                session.mask_token_id,
                exclude_self=True,
                excluded_token_ids=set(getattr(session.tokenizer, "all_special_ids", [])),
            )[0]
            logits = _apply_eos_eot_prediction_penalty(
                logits,
                eos_eot_prediction_penalty,
                session.tokenizer.eos_token_id,
                eot_token_id,
            )
            if bigram_selector is None:
                sampled, confidence = _sample(logits, float(temperature), int(top_k), None)
            else:
                # Keep every actual sequence position, including intervening
                # frozen rows; compressing editable positions creates false bigrams.
                sampled = answer_ids[0].clone()
                confidence = torch.zeros(logits.shape[0], device=logits.device)
                current = answer_ids[0, block_start:block_end].clone()
                editable = torch.ones_like(current, dtype=torch.bool)
                for offset, token in frozen.items():
                    if block_start <= offset < block_end:
                        current[offset - block_start] = token
                        editable[offset - block_start] = False
                left_id = ids[answer_start + block_start - 1] if answer_start + block_start else None
                right_index = answer_start + block_end
                right_id = ids[right_index] if right_index < len(ids) and ids[right_index] != session.mask_token_id else None
                selection_logits = logits[block_start:block_end]
                if confidence_eos_eot_inf:
                    # Retention alone cannot alter a deterministic Viterbi
                    # winner. Fade this bias away by the final block step.
                    delay_penalty = float(bigram_eos_delay_penalty) ** (1.0 - (block_step + 1) / block_steps)
                    selection_logits = _apply_eos_eot_prediction_penalty(
                        selection_logits, delay_penalty, session.tokenizer.eos_token_id, eot_token_id,
                    )
                set_context = getattr(bigram_selector, "set_step_context", None)
                if set_context is not None:
                    set_context(block_step, block_steps)
                selected, selected_confidence = bigram_selector.select_block(
                    selection_logits, current, editable,
                    top_k=bigram_top_k,
                    temperature=float(temperature), left_id=left_id, right_id=right_id,
                )
                sampled[block_start:block_end] = selected
                confidence[block_start:block_end] = selected_confidence
                margins = getattr(bigram_selector, "last_path_margins", None)
                if margins is not None:
                    bigram_retention_scores = torch.zeros_like(confidence)
                    bigram_retention_scores[block_start:block_end] = margins
            if adaptive_parallel_decoding:
                confidence = confidence.float()
                if bigram_selector is None:
                    selection_logits = logits[block_start:block_end]
                    confidence[block_start:block_end] = selected_token_probabilities(
                        selection_logits, sampled[block_start:block_end],
                    )
                selected, selected_confidence = resolve_mask_predictions(
                    selection_logits, sampled[block_start:block_end], confidence[block_start:block_end],
                    answer_ids[0, block_start:block_end] == session.mask_token_id, session.mask_token_id,
                )
                sampled[block_start:block_end], confidence[block_start:block_end] = selected, selected_confidence
        if revisable_tokens:
            # Fill current masks while preserving visible tokens. The complete
            # candidate is scored below, after which any position in the block
            # may be selected for re-masking.
            for offset in range(block_start, block_end):
                if ids[answer_start + offset] == session.mask_token_id:
                    ids[answer_start + offset] = int(sampled[offset])
        else:
            ids[answer_start + block_start : answer_start + block_end] = sampled[block_start:block_end].tolist()
        if freeze_retained_tokens and not revisable_tokens:
            for offset, token in frozen.items():
                ids[answer_start + offset] = token
        retention_confidence = confidence if bigram_retention_scores is None else bigram_retention_scores
        completed_tokens = torch.tensor([ids], device=session.device, dtype=torch.long)
        if revisable_tokens and (confidence_guided or not ar_check_due):
            retention_confidence = _sampling_token_confidence(
                logits,
                completed_tokens[0, answer_start:],
                float(temperature),
                int(top_k),
            )
        if ar_check_due:
            retention_confidence = _autoregressive_token_confidence(
                session, completed_tokens, answer_start
            )
        special_prediction = completed_tokens[0, answer_start:] == session.tokenizer.eos_token_id
        if eot_token_id is not None:
            special_prediction |= completed_tokens[0, answer_start:] == eot_token_id
        if confidence_eos_eot_inf:
            retention_confidence = retention_confidence.masked_fill(
                special_prediction, torch.finfo(retention_confidence.dtype).min
            )
        predicted_text = decode_denoising_state(
            ids[answer_start:],
            session.tokenizer,
            session.mask_token_id,
            show_eos_tokens=include_pre_remask_prediction,
        )
        last_confidence = float(confidence[block_start:block_end].mean().cpu())
        # Compare the visible sampled answer before the next iteration's
        # re-masking. Tokens after the first EOS/EOT are not part of the answer and
        # must not prevent convergence. Excluding the ending itself still preserves
        # its position through the tuple length: moving it changes the prefix.
        prediction = ids[answer_start:]
        predicted_ending = _first_ending_offset(prediction, ending_ids)
        if predicted_ending is not None:
            prediction = prediction[:predicted_ending]
        last_predictions.append(tuple(prediction))
        if len(last_predictions) > 2:
            last_predictions.pop(0)
        stopped_early = early_stopping and not adaptive_parallel_decoding and len(last_predictions) == 2 and len(set(last_predictions)) == 1
        if session.mask_token_id in prediction[block_start:block_end]:
            stopped_early = False
        if (confidence_eos_eot_inf and predicted_ending is not None and has_followup_step
                and predicted_ending not in frozen):
            # Convergence must not bypass delayed retention by accepting an
            # ending which the normal remasking schedule would still discard.
            stopped_early = False
        # Progressively reduce corruption. Re-mask independently, retaining the
        # legacy schedule's initial noise_level and ending with a clean sample.
        apd_block_complete = False
        apd_newly_committed = 0
        if adaptive_parallel_decoding:
            pending = answer_ids[0, block_start:block_end] == session.mask_token_id
            chosen = commitment_mask(
                confidence[block_start:block_end], pending, block_step, block_steps,
                endings=special_prediction[block_start:block_end], delay_endings=confidence_eos_eot_inf,
                aggressiveness=apd_aggressiveness,
                minimum_tokens=(block_end - block_start + block_steps - 1) // block_steps,
            )
            offsets = (chosen.nonzero(as_tuple=True)[0] + block_start).tolist()
            apd_newly_committed = len(offsets)
            for offset in offsets:
                retained.add(offset)
                frozen[offset] = ids[answer_start + offset]
                frozen_confidences[offset] = float(confidence[offset].cpu())
                frozen_steps[offset] = forward_passes
            for offset in range(block_start, block_end):
                if offset not in retained:
                    ids[answer_start + offset] = session.mask_token_id
            apd_block_complete = all(offset in retained for offset in range(block_start, block_end))
        elif block_step + 1 < block_steps and not stopped_early:
            block_size = block_end - block_start
            mask_probability = max(0.0, min(1.0, float(noise_level) * (1.0 - (block_step + 1) / block_steps)))
            if revisable_tokens:
                remask_offsets = _remask_offsets(
                    retention_confidence[block_start:block_end],
                    mask_probability,
                    guided_retention,
                )
                if confidence_eos_eot_inf and not guided_retention:
                    ending_offsets = torch.where(special_prediction[block_start:block_end])[0]
                    remask_offsets = torch.unique(torch.cat((remask_offsets, ending_offsets)))
                for relative_offset in remask_offsets.tolist():
                    ids[answer_start + block_start + relative_offset] = session.mask_token_id
            elif permanent_unmask:
                keep_count = min(block_size, max(0, round((1.0 - mask_probability) * block_size)))
                retained_in_block = sum(block_start <= i < block_end for i in retained)
                needed = keep_count - retained_in_block
                candidates = [i for i in range(block_start, block_end) if i not in retained]
                if needed > 0 and candidates:
                    def random_retention_order(pool: list[int]) -> list[int]:
                        order = torch.randperm(len(pool), device=session.device).tolist()
                        randomized = [pool[i] for i in order]
                        if not confidence_eos_eot_inf:
                            return randomized
                        ordinary = [i for i in randomized if not bool(special_prediction[i])]
                        endings = [i for i in randomized if bool(special_prediction[i])]
                        return ordinary + endings

                    if proportional_unmask:
                        eos_positions = [i for i in range(block_start, block_end) if ids[answer_start + i] == session.tokenizer.eos_token_id]
                        boundary = min(eos_positions) if eos_positions else block_end
                        pools = [[i for i in candidates if i < boundary], [i for i in candidates if i >= boundary]]
                        target_normal = round(keep_count * (boundary - block_start) / block_size)
                        target_counts = [
                            max(0, target_normal - sum(block_start <= i < boundary for i in retained)),
                            max(0, keep_count - target_normal - sum(boundary <= i < block_end for i in retained)),
                        ]
                        chosen = []
                        for pool, target in zip(pools, target_counts):
                            if not pool or target <= 0:
                                continue
                            if guided_retention:
                                order = torch.argsort(retention_confidence, descending=True).tolist()
                                chosen.extend([i for i in order if i in pool][:target])
                            else:
                                chosen.extend(random_retention_order(pool)[:target])
                        if len(chosen) < needed:
                            remainder = [i for i in candidates if i not in chosen]
                            chosen.extend(remainder[: needed - len(chosen)])
                    elif guided_retention:
                        confidence_order = torch.argsort(retention_confidence, descending=True).tolist()
                        chosen = [i for i in confidence_order if i in candidates][:needed]
                    else:
                        chosen = random_retention_order(candidates)[:needed]
                    for offset in chosen:
                        retained.add(offset)
                        if freeze_retained_tokens:
                            frozen[offset] = ids[answer_start + offset]
                            frozen_confidences[offset] = float(confidence[offset].cpu())
                            frozen_steps[offset] = step + 1
                for offset in range(block_start, block_end):
                    if offset not in retained:
                        ids[answer_start + offset] = session.mask_token_id
            else:
                # Confidence-guided refinement keeps every token revisable, but
                # preferentially re-masks the least certain predictions. The
                # unguided mode retains the original random re-masking policy.
                remask_offsets = _remask_offsets(
                    retention_confidence[block_start:block_end],
                    mask_probability,
                    guided_retention,
                )
                if confidence_eos_eot_inf and not guided_retention:
                    ending_offsets = torch.where(special_prediction[block_start:block_end])[0]
                    remask_offsets = torch.unique(torch.cat((remask_offsets, ending_offsets)))
                for offset in remask_offsets.tolist():
                    ids[answer_start + block_start + offset] = session.mask_token_id
        current_answer = ids[answer_start:]
        ending_offset = _first_ending_offset(current_answer, ending_ids)
        visible_answer = current_answer[:ending_offset] if ending_offset is not None else current_answer
        answer_finished = (
            early_stopping and ending_offset is not None
            and session.mask_token_id not in visible_answer
            and (not has_followup_step or stopped_early or ending_offset in frozen)
        )
        remasked_text = decode_denoising_state(
            current_answer,
            session.tokenizer,
            session.mask_token_id,
            show_eos_tokens=include_pre_remask_prediction,
        )
        current_text = remasked_text
        if include_pre_remask_prediction:
            remask_label = (
                "State after re-mask"
                if has_followup_step and not stopped_early and not answer_finished
                else "State after re-mask (unchanged; final state)"
            )
            current_text = (
                f"Predicted (before re-mask):\n{predicted_text}\n"
                f"{remask_label}:\n{remasked_text}"
            )
        display_step = forward_passes if adaptive_parallel_decoding else step + 1
        status = f"Denoising step {display_step}/{num_steps} · {len(visible_answer)} output tokens · mean confidence {last_confidence:.3f}"
        if num_blocks > 1:
            status += f" · block {block_index + 1}/{num_blocks}"
        if permanent_unmask and not revisable_tokens:
            status += f" · retained {len(retained)} tokens"
        if revisable_tokens:
            status += " · revisable tokens"
        if autoregressive_guided and not adaptive_parallel_decoding:
            status += f" · AR-guided remasking every {autoregressive_check_interval} step(s)"
        if bigram_retention_scores is not None:
            if not adaptive_parallel_decoding:
                status += " · Viterbi-margin remasking"
        if adaptive_parallel_decoding:
            status += f" · APD threshold {confidence_threshold(block_step, block_steps, apd_aggressiveness):.3f} · committed {apd_newly_committed} · forwards {forward_passes}"
            if apd_block_complete:
                status += " · block committed"
        if stopped_early:
            status += " · block stopped early (same answer for 2 consecutive iterations)"
        if answer_finished:
            ending_name = "EOS" if current_answer[ending_offset] == session.tokenizer.eos_token_id else "EOT"
            status += f" · answer finished ({ending_name})"
            if block_index + 1 < num_blocks:
                status += " · remaining blocks skipped"
        html = render_denoising_step(
            ids,
            confidence.tolist(),
            answer_start,
            session.tokenizer,
            session.mask_token_id,
            display_step,
            num_steps,
            retained if permanent_unmask else None,
            frozen if permanent_unmask and freeze_retained_tokens else None,
            frozen_confidences if permanent_unmask and freeze_retained_tokens else None,
            frozen_steps if permanent_unmask and freeze_retained_tokens else None,
            trajectory_color_mode,
        )
        yield current_text, status, html
        if answer_finished:
            break
        if adaptive_parallel_decoding and apd_block_complete:
            if block_index + 1 == num_blocks:
                break
            skip_block_index = block_index
        if stopped_early:
            if num_blocks == 1:
                break
            skip_block_index = block_index
    answer = ids[answer_start:]
    final_ending = _first_ending_offset(answer, ending_ids)
    if final_ending is not None:
        answer = answer[:final_ending]
    text = session.tokenizer.decode(answer, skip_special_tokens=True).strip()
    return


def denoise(session: InferenceSession, question: str, system_prompt: str, max_new_tokens: int, num_steps: int, noise_level: float, temperature: float, top_k: int, seed: int, permanent_unmask: bool = False, confidence_guided: bool = False, proportional_unmask: bool = True, early_stopping: bool = False, progress: Callable[[float, str], None] | None = None, confidence_eos_eot_inf: bool = False, freeze_retained_tokens: bool = True, repetition_penalty: float = 1.0, eos_eot_prediction_penalty: float = 1.0, block_length: int | None = None, autoregressive_guided: bool = False, revisable_tokens: bool = False, autoregressive_check_interval: int = 1, adaptive_parallel_decoding: bool = False, apd_aggressiveness: float = 0.0) -> tuple[str, str]:
    """Run denoising to completion and return only the final text and status."""
    result = ("", "")
    for step, (text, status, _html) in enumerate(denoise_stream(session, question, system_prompt, max_new_tokens, num_steps, noise_level, temperature, top_k, seed, permanent_unmask, confidence_guided, proportional_unmask, early_stopping, confidence_eos_eot_inf, freeze_retained_tokens, repetition_penalty, eos_eot_prediction_penalty, False, block_length, "Prediction probability", autoregressive_guided, revisable_tokens, autoregressive_check_interval, adaptive_parallel_decoding=adaptive_parallel_decoding, apd_aggressiveness=apd_aggressiveness), start=1):
        result = (text, status)
        if progress:
            progress(step / int(num_steps), status)
    return result


def release_session(session: InferenceSession | None) -> None:
    """Free a loaded inference model and release backend allocator caches."""
    if session is None:
        return
    del session.model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    if torch.backends.mps.is_available():
        torch.mps.empty_cache()