src/diffusion_lm/playground.py

"""Local browser playground that streams the token-unmasking process."""

from __future__ import annotations

import argparse
import hashlib
import html
import secrets
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

import torch
from torch import Tensor
from tokenizers import Tokenizer

from diffusion_lm.diffusion import UnmaskStep, iterative_unmask_steps
from diffusion_lm.model import DiffusionTransformer, format_parameter_count
from diffusion_lm.sample import load_model
from diffusion_lm.tokenizer import load_tokenizer, special_token_id, special_token_ids
from diffusion_lm.train import resolve_device


@dataclass(frozen=True)
class GenerationSettings:
    prompt: str = ""
    generation_length: int = 64
    steps: int = 64
    temperature: float = 0.8
    strategy: Literal["ancestral", "confidence"] = "confidence"
    seed: int = 1337

    def __post_init__(self) -> None:
        if self.generation_length <= 0:
            raise ValueError("La longitud debe ser mayor que cero.")
        if not 1 <= self.steps <= 512:
            raise ValueError("Los pasos deben estar entre 1 y 512.")
        if not 0.0 <= self.temperature <= 5.0:
            raise ValueError("La temperatura debe estar entre 0 y 5.")
        if self.strategy not in {"ancestral", "confidence"}:
            raise ValueError("La estrategia debe ser ancestral o confidence.")
        if not 0 <= self.seed < 2**63:
            raise ValueError("La semilla debe estar entre 0 y 2^63-1.")


@dataclass(frozen=True)
class PlaygroundUpdate:
    state: UnmaskStep
    prompt_tokens: int
    partial_text: str
    final_text: str
    token_html: str
    elapsed_seconds: float
    step_seconds: float


def _seed_generation(device: torch.device, seed: int) -> None:
    torch.manual_seed(seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)
    elif device.type == "mps" and hasattr(torch.mps, "manual_seed"):
        torch.mps.manual_seed(seed)


def _synchronize(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    elif device.type == "mps":
        torch.mps.synchronize()


def _token_label(tokenizer: Tokenizer, token_id: int, mask_token_id: int) -> str:
    raw = tokenizer.id_to_token(token_id) or f"#{token_id}"
    if token_id == mask_token_id:
        return "MASK"
    return (
        raw.replace("Ġ", "▁")
        .replace("Ċ", "↵")
        .replace("ĉ", "⇥")
        .replace("\n", "↵")
    ) or "∅"


def render_token_grid(
    tokenizer: Tokenizer,
    token_ids: list[int],
    *,
    mask_token_id: int,
    prompt_tokens: int,
    previous_token_ids: list[int] | None,
) -> str:
    """Render escaped token chips for a single sample."""

    chips: list[str] = []
    for position, token_id in enumerate(token_ids):
        if position < prompt_tokens:
            state = "prompt"
        elif token_id == mask_token_id:
            state = "mask"
        elif previous_token_ids is not None and previous_token_ids[position] == mask_token_id:
            state = "new"
        else:
            state = "revealed"

        raw = tokenizer.id_to_token(token_id) or f"token {token_id}"
        label = html.escape(_token_label(tokenizer, token_id, mask_token_id))
        title = html.escape(f"posición {position} · id {token_id} · {raw}", quote=True)
        chips.append(
            f'<span class="token-chip token-{state}" title="{title}">{label}</span>'
        )

    return (
        '<section class="token-stage" aria-label="Estado actual de los tokens">'
        '<div class="token-grid">'
        + "".join(chips)
        + "</div>"
        '<div class="token-legend" aria-label="Leyenda">'
        '<span><i class="legend-dot legend-prompt"></i>prompt</span>'
        '<span><i class="legend-dot legend-new"></i>recién revelado</span>'
        '<span><i class="legend-dot legend-mask"></i>máscara</span>'
        "</div></section>"
    )


class PlaygroundEngine:
    """Own one loaded model and serialize interactive generations."""

    def __init__(self, model: DiffusionTransformer, tokenizer_path: str | Path) -> None:
        self.model = model.eval()
        self.tokenizer_path = Path(tokenizer_path)
        self.tokenizer = load_tokenizer(self.tokenizer_path)
        self._lock = threading.Lock()
        self._validate_tokenizer()

    @property
    def device(self) -> torch.device:
        return next(self.model.parameters()).device

    def _validate_tokenizer(self) -> None:
        tokenizer_hash = hashlib.sha256(self.tokenizer_path.read_bytes()).hexdigest()
        model_hash = getattr(self.model, "tokenizer_sha256", None)
        if model_hash is not None and tokenizer_hash != model_hash:
            raise ValueError("El tokenizer no coincide con el usado para entrenar el checkpoint.")
        if self.tokenizer.get_vocab_size(with_added_tokens=True) != self.model.config.vocab_size:
            raise ValueError("El vocabulario del tokenizer no coincide con el checkpoint.")
        if special_token_id(self.tokenizer, "mask") != self.model.config.mask_token_id:
            raise ValueError("El id de [MASK] no coincide con el checkpoint.")

    def info(self) -> dict[str, str | int | bool]:
        return {
            "parameters": self.model.num_parameters,
            "parameters_human": format_parameter_count(self.model.num_parameters),
            "device": str(self.device),
            "context": self.model.config.max_seq_len,
            "vocab_size": self.model.config.vocab_size,
            "tokenizer": self.tokenizer_path.name,
            "mps_fp64_fallback": self.device.type == "mps",
        }

    def _prepare(self, settings: GenerationSettings) -> tuple[Tensor, int, tuple[int, ...]]:
        prompt_ids = (
            self.tokenizer.encode(settings.prompt, add_special_tokens=False).ids
            if settings.prompt
            else []
        )
        role_ids = special_token_ids(self.tokenizer)
        reserved_ids = set(role_ids.values())
        encountered = reserved_ids.intersection(prompt_ids)
        if encountered:
            raise ValueError(
                "El prompt contiene tokens especiales reservados. Escribí texto normal sin "
                "los sentinels internos del modelo."
            )

        total_length = len(prompt_ids) + settings.generation_length
        if total_length > self.model.config.max_seq_len:
            available = self.model.config.max_seq_len - len(prompt_ids)
            raise ValueError(
                f"El prompt usa {len(prompt_ids)} tokens y deja {max(0, available)} para generar; "
                f"solicitaste {settings.generation_length}."
            )

        input_ids = torch.full(
            (1, total_length),
            self.model.config.mask_token_id,
            dtype=torch.long,
            device=self.device,
        )
        if prompt_ids:
            input_ids[0, : len(prompt_ids)] = torch.tensor(prompt_ids, device=self.device)

        blocked = tuple(role_ids[role] for role in ("pad", "unk", "bos", "mask"))
        return input_ids, len(prompt_ids), blocked

    def _decode_final(self, token_ids: list[int], prompt_tokens: int) -> str:
        eos_id = special_token_id(self.tokenizer, "eos")
        if eos_id in token_ids[prompt_tokens:]:
            token_ids = token_ids[: token_ids.index(eos_id, prompt_tokens)]
        return self.tokenizer.decode(token_ids, skip_special_tokens=True)

    def stream(self, settings: GenerationSettings):
        """Yield one UI update per reverse-diffusion pass."""

        input_ids, prompt_tokens, blocked = self._prepare(settings)
        with self._lock:
            _seed_generation(self.device, settings.seed)
            started = time.perf_counter()
            # Measures engine work per pass; consumer time between yields is excluded.
            pass_started = started
            previous_ids: list[int] | None = None
            for state in iterative_unmask_steps(
                self.model,
                input_ids,
                self.model.config.mask_token_id,
                steps=settings.steps,
                temperature=settings.temperature,
                strategy=settings.strategy,
                blocked_token_ids=blocked,
            ):
                token_ids = state.tokens[0].detach().cpu().tolist()
                partial_text = self.tokenizer.decode(token_ids, skip_special_tokens=False)
                final_text = (
                    self._decode_final(token_ids, prompt_tokens)
                    if state.masked_remaining == 0
                    else ""
                )
                token_html = render_token_grid(
                    self.tokenizer,
                    token_ids,
                    mask_token_id=self.model.config.mask_token_id,
                    prompt_tokens=prompt_tokens,
                    previous_token_ids=previous_ids,
                )
                if state.masked_remaining == 0:
                    _synchronize(self.device)
                now = time.perf_counter()
                yield PlaygroundUpdate(
                    state=state,
                    prompt_tokens=prompt_tokens,
                    partial_text=partial_text,
                    final_text=final_text,
                    token_html=token_html,
                    elapsed_seconds=now - started,
                    step_seconds=now - pass_started,
                )
                previous_ids = token_ids
                pass_started = time.perf_counter()


PLAYGROUND_CSS = """
:root {
  --playground-accent: #7c3aed;
  --playground-accent-soft: rgba(124, 58, 237, 0.14);
  --playground-teal: #0f766e;
  --playground-border: rgba(100, 116, 139, 0.22);
}
.gradio-container { max-width: 1440px !important; }
.playground-header {
  padding: 22px 24px; border: 1px solid var(--playground-border); border-radius: 18px;
  background: linear-gradient(135deg, rgba(124,58,237,.10), rgba(15,118,110,.06));
  box-shadow: 0 1px 2px rgba(30,41,59,.05), 0 14px 34px rgba(71,85,105,.08);
}
.playground-header h1 { margin: 0; font-size: clamp(1.6rem, 3vw, 2.4rem); text-wrap: balance; }
.playground-header p { margin: 8px 0 0; color: var(--body-text-color-subdued); text-wrap: pretty; }
.model-strip { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 16px; }
.model-pill {
  padding: 7px 10px; border-radius: 999px; border: 1px solid var(--playground-border);
  background: var(--block-background-fill); font-variant-numeric: tabular-nums; font-size: .82rem;
}
.control-panel, .output-panel {
  border: 1px solid var(--playground-border) !important; border-radius: 18px !important;
  padding: 16px !important; box-shadow: 0 1px 2px rgba(30,41,59,.04), 0 10px 28px rgba(71,85,105,.06);
}
.token-stage { min-height: 220px; display: flex; flex-direction: column; justify-content: space-between; }
.token-grid { display: flex; flex-wrap: wrap; align-content: flex-start; gap: 7px; padding: 8px 2px 18px; }
.token-chip {
  display: inline-flex; min-height: 32px; align-items: center; padding: 5px 8px; border-radius: 9px;
  border: 1px solid transparent; font-family: ui-monospace, SFMono-Regular, Menlo, monospace;
  font-size: .82rem; font-variant-numeric: tabular-nums; transition: transform 160ms ease-out, opacity 180ms ease-out;
}
.token-chip:hover { transform: translateY(-1px); }
.token-prompt { color: #075985; background: rgba(14,165,233,.12); border-color: rgba(14,165,233,.26); }
.token-revealed { background: rgba(15,118,110,.10); border-color: rgba(15,118,110,.18); }
.token-new { color: #5b21b6; background: var(--playground-accent-soft); border-color: rgba(124,58,237,.35); }
.token-mask { color: var(--body-text-color-subdued); background: rgba(100,116,139,.08); border: 1px dashed rgba(100,116,139,.30); opacity: .7; }
.token-legend { display: flex; flex-wrap: wrap; gap: 14px; color: var(--body-text-color-subdued); font-size: .78rem; }
.token-legend span { display: inline-flex; align-items: center; gap: 6px; }
.legend-dot { width: 9px; height: 9px; border-radius: 50%; display: inline-block; }
.legend-prompt { background: #0ea5e9; } .legend-new { background: #7c3aed; } .legend-mask { background: #94a3b8; }
#generate-button, #stop-button { min-height: 44px; transition: transform 150ms ease-out; }
#generate-button:active, #stop-button:active { transform: scale(.98); }
@media (prefers-reduced-motion: reduce) { .token-chip, #generate-button, #stop-button { transition: none; } }
"""


def _model_header(engine: PlaygroundEngine) -> str:
    info = engine.info()
    warning = (
        " · MPS usa CPU para Gumbel fp64 cuando temperatura > 0"
        if info["mps_fp64_fallback"]
        else ""
    )
    return (
        '<header class="playground-header">'
        "<h1>Mini Diffusion LM Playground</h1>"
        "<p>Observá cómo el modelo transforma máscaras en texto usando contexto bidireccional."
        f"{html.escape(warning)}</p>"
        '<div class="model-strip">'
        f'<span class="model-pill">{info["parameters_human"]} parámetros</span>'
        f'<span class="model-pill">{html.escape(str(info["device"]))}</span>'
        f'<span class="model-pill">contexto {info["context"]}</span>'
        f'<span class="model-pill">vocabulario {info["vocab_size"]}</span>'
        f'<span class="model-pill">{html.escape(str(info["tokenizer"]))}</span>'
        "</div></header>"
    )


def build_playground(engine: PlaygroundEngine):
    """Build a Gradio Blocks app without importing Gradio for base-package users."""

    try:
        import gradio as gr
    except ImportError as exc:  # pragma: no cover - exercised by CLI environments.
        raise RuntimeError(
            'Falta Gradio. Instalalo con: pip install -e ".[playground]"'
        ) from exc

    max_context = engine.model.config.max_seq_len
    default_length = min(64, max_context)
    theme = gr.themes.Soft(primary_hue="violet", secondary_hue="teal", neutral_hue="slate")
    with gr.Blocks(
        title="Mini Diffusion LM Playground",
        analytics_enabled=False,
        fill_width=True,
    ) as demo:
        gr.HTML(_model_header(engine))
        with gr.Row():
            with gr.Column(scale=4, elem_classes="control-panel"):
                gr.Markdown("## Configuración")
                prompt = gr.Textbox(
                    label="Prompt (opcional)",
                    placeholder="Ej.: Once upon a time…",
                    lines=6,
                    max_lines=10,
                )
                with gr.Row():
                    length = gr.Slider(
                        minimum=1,
                        maximum=max_context,
                        value=default_length,
                        step=1,
                        label="Tokens a generar",
                    )
                    steps = gr.Slider(
                        minimum=1,
                        maximum=256,
                        value=min(64, max_context),
                        step=1,
                        label="Pasos de difusión",
                    )
                with gr.Accordion("Opciones avanzadas", open=False):
                    strategy = gr.Radio(
                        choices=[
                            ("Confianza · revela los tokens más seguros", "confidence"),
                            ("Ancestral · transición probabilística", "ancestral"),
                        ],
                        value="confidence",
                        label="Estrategia",
                    )
                    temperature = gr.Slider(
                        minimum=0.0,
                        maximum=2.0,
                        value=0.8,
                        step=0.05,
                        label="Temperatura",
                    )
                    seed = gr.Number(
                        value=0,
                        precision=0,
                        minimum=0,
                        maximum=2**31 - 1,
                        label='Semilla (0 = aleatoria en cada generación)',
                    )
                with gr.Row():
                    generate_button = gr.Button(
                        "Generar",
                        variant="primary",
                        elem_id="generate-button",
                    )
                    stop_button = gr.Button(
                        "Detener",
                        variant="stop",
                        elem_id="stop-button",
                    )

            with gr.Column(scale=7, elem_classes="output-panel"):
                status = gr.Markdown(
                    "### Listo\nConfigurá una muestra y presioná **Generar**."
                )
                token_view = gr.HTML(
                    '<div class="token-stage"><p>Los tokens aparecerán acá.</p></div>'
                )
                output = gr.Textbox(
                    label="Texto actual",
                    lines=7,
                    interactive=False,
                )
                metrics = gr.Markdown("`Esperando una generación`", elem_classes="metrics")

        def stream_generation(
            prompt_value: str,
            length_value: float,
            steps_value: float,
            strategy_value: str,
            temperature_value: float,
            seed_value: float,
        ):
            clicked = time.perf_counter()
            try:
                resolved_seed = int(seed_value) or secrets.randbelow(2**31 - 1) + 1
                settings = GenerationSettings(
                    prompt=prompt_value or "",
                    generation_length=int(length_value),
                    steps=int(steps_value),
                    strategy=strategy_value,  # type: ignore[arg-type]
                    temperature=float(temperature_value),
                    seed=resolved_seed,
                )
                for update in engine.stream(settings):
                    generated = settings.generation_length - update.state.masked_remaining
                    percent = 100.0 * generated / settings.generation_length
                    total_seconds = time.perf_counter() - clicked
                    average_step = update.elapsed_seconds / max(1, update.state.step)
                    tokens_per_second = (
                        generated / update.elapsed_seconds if update.elapsed_seconds > 0 else 0.0
                    )
                    if update.state.masked_remaining == 0:
                        status_text = (
                            f'### Completado en {total_seconds:.2f} s\n'
                            f'{generated} tokens en {update.state.step} pasos · '
                            f'{tokens_per_second:.1f} tok/s · '
                            f'{average_step * 1000:.0f} ms/paso promedio'
                        )
                    else:
                        status_text = (
                            f"### Paso {update.state.step}/{update.state.total_steps}\n"
                            f"{update.state.masked_remaining} máscaras restantes · "
                            f"{percent:.0f}% revelado"
                        )
                    visible_text = (
                        update.final_text
                        if update.state.masked_remaining == 0
                        else update.partial_text
                    )
                    metrics_text = (
                        f'`{total_seconds:.2f} s desde el clic` · '
                        f'`modelo {update.elapsed_seconds:.2f} s` · '
                        f'`paso {update.step_seconds * 1000:.0f} ms` · '
                        f'`prom. {average_step * 1000:.0f} ms/paso` · '
                        f'`{tokens_per_second:.1f} tok/s` · '
                        f'`{update.prompt_tokens} tokens de prompt` · '
                        f'`seed {settings.seed}`'
                    )
                    yield update.token_html, status_text, visible_text, metrics_text
            except ValueError as exc:
                raise gr.Error(str(exc)) from exc

        generation_event = generate_button.click(
            fn=stream_generation,
            inputs=[prompt, length, steps, strategy, temperature, seed],
            outputs=[token_view, status, output, metrics],
            show_progress="minimal",
            scroll_to_output=False,
            concurrency_limit=1,
            concurrency_id="diffusion-model",
            trigger_mode="once",
            stream_every=0.1,
            api_visibility="private",
        )
        stop_button.click(
            fn=lambda: "### Generación detenida",
            outputs=status,
            cancels=[generation_event],
            queue=False,
            api_visibility="private",
        )

    demo = demo.queue(max_size=8, default_concurrency_limit=1)
    # Gradio 6 moved presentation arguments from Blocks() to launch().
    demo._mini_diffusion_theme = theme
    return demo


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, required=True, help="checkpoint local confiable")
    parser.add_argument("--tokenizer", type=Path, required=True, help="tokenizer usado al entrenar")
    parser.add_argument("--device", default="auto", help="auto, cpu, mps, cuda…")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--no-browser", action="store_true", help="no abrir el navegador")
    return parser


def main() -> None:
    args = _build_parser().parse_args()
    if not args.checkpoint.is_file():
        raise SystemExit(f"Checkpoint inexistente: {args.checkpoint}")
    if not args.tokenizer.is_file():
        raise SystemExit(f"Tokenizer inexistente: {args.tokenizer}")
    if not 1 <= args.port <= 65535:
        raise SystemExit("El puerto debe estar entre 1 y 65535")

    device = resolve_device(args.device)
    print(f"Cargando {args.checkpoint} en {device}…")
    model = load_model(args.checkpoint, device)
    engine = PlaygroundEngine(model, args.tokenizer)
    demo = build_playground(engine)
    print(f"Playground: http://{args.host}:{args.port}")
    print("Usá únicamente checkpoints locales confiables.")
    demo.launch(
        server_name=args.host,
        server_port=args.port,
        inbrowser=not args.no_browser,
        share=False,
        show_error=True,
        strict_cors=True,
        max_threads=4,
        footer_links=[],
        enable_monitoring=False,
        ssr_mode=False,
        pwa=False,
        theme=demo._mini_diffusion_theme,
        css=PLAYGROUND_CSS,
    )


if __name__ == "__main__":
    main()