Spaces:
Sleeping
Sleeping
| """Model loading and the generation loop of the demo (independent of Gradio).""" | |
| from __future__ import annotations | |
| import os | |
| import time | |
| from dataclasses import dataclass | |
| from typing import Iterator | |
| from .render import Detokenizer, replay_html, stats_html, tokens_html | |
| # display name -> Hub repo with the heads (base model and tokenizer are pinned in its config) | |
| MODELS = { | |
| "Qwen3-8B-Base": "FlorianVal/DSSD-Qwen3-8B-Base", | |
| "Llama-3.1-8B": "FlorianVal/Llama-3.1-8B-DSSD", | |
| } | |
| MAX_PROMPT_TOKENS = 1536 | |
| MAX_NEW_TOKENS = 256 | |
| def _graphs() -> bool | str: | |
| # CUDA graphs are captured once per process; under ZeroGPU each call may run in a fresh GPU | |
| # worker, so the eager engine is the safe default there. DSSD_GRAPHS=1 forces them on. | |
| v = os.environ.get("DSSD_GRAPHS") | |
| if v is not None: | |
| return v == "1" | |
| return False if os.environ.get("SPACES_ZERO_GPU") else "auto" | |
| def load_models(backend: str | None = None) -> dict: | |
| """{display name: DSSD}. backend "hub" (default): the published heads and their base models | |
| on CUDA; "tiny": a random tiny model on CPU for local UI work.""" | |
| backend = backend or os.environ.get("DSSD_DEMO_BACKEND", "hub") | |
| if backend == "tiny": | |
| from .tiny import tiny_engine | |
| return {"Tiny random model (local test)": tiny_engine()} | |
| import torch | |
| from dssd import DSSD | |
| names = [n.strip() for n in os.environ.get("DSSD_DEMO_MODELS", ",".join(MODELS)).split(",") if n.strip()] | |
| device = "cuda" if (torch.cuda.is_available() or os.environ.get("SPACES_ZERO_GPU")) else "cpu" | |
| token = os.environ.get("HF_TOKEN") | |
| return {n: DSSD.from_pretrained(MODELS[n], device=device, graphs=_graphs(), token=token) for n in names} | |
| class Update: | |
| dssd: str # HTML of the DSSD output (coloured by source) | |
| ar: str # HTML of the autoregressive output ("" when not compared) | |
| stats: str # HTML stat cards (final update) or a progress line | |
| replay: str # hidden replay payload (final update) | |
| done: bool = False | |
| def generate_updates(engine, prompt: str, domain: str, max_new: int, compare: bool, | |
| min_interval: float = 0.05) -> Iterator[Update]: | |
| """Stream DSSD (then AR when compare) as HTML updates, at most one every min_interval seconds, | |
| then a final update with stats and the replay timelines.""" | |
| if not prompt or not prompt.strip(): | |
| raise ValueError("Please enter a prompt.") | |
| n = len(engine.tokenizer.encode(prompt, add_special_tokens=False)) | |
| if n > MAX_PROMPT_TOKENS: | |
| raise ValueError(f"The prompt is too long ({n} tokens; at most {MAX_PROMPT_TOKENS}).") | |
| max_new = max(1, min(int(max_new), MAX_NEW_TOKENS)) | |
| detok = {"dssd": Detokenizer(engine.tokenizer), "ar": Detokenizer(engine.tokenizer)} | |
| pieces = {"dssd": [], "ar": []} | |
| timeline = {"dssd": [], "ar": []} | |
| last = 0.0 | |
| def view(done=False, stats="", replay=""): | |
| return Update(tokens_html(pieces["dssd"], cursor=not done and not pieces["ar"]), | |
| tokens_html(pieces["ar"], colored=False, cursor=not done and bool(pieces["ar"])) | |
| if compare else "", | |
| stats, replay, done) | |
| for ev in engine.stream(prompt, max_new, domain=domain, ar=compare): | |
| if ev.done: | |
| for ph in ("dssd", "ar"): | |
| new = detok[ph].flush() | |
| pieces[ph] += new | |
| if timeline[ph]: | |
| timeline[ph] += [[timeline[ph][-1][0], t, s] for t, s in new] | |
| yield view(True, stats_html(ev.result), replay_html(timeline)) | |
| return | |
| new = detok[ev.phase].feed(ev.token_ids, ev.sources) | |
| pieces[ev.phase] += new | |
| timeline[ev.phase] += [[round(ev.t, 4), t, s] for t, s in new] | |
| now = time.perf_counter() | |
| if now - last >= min_interval: | |
| last = now | |
| phase = "Generating with DSSD…" if ev.phase == "dssd" else "Generating autoregressively for comparison…" | |
| yield view(stats=f'<div class="progress">{phase}</div>') | |