Spaces:
Sleeping
Sleeping
| """HTML rendering for the demo: coloured tokens by source, legend, stats, replay payload.""" | |
| from __future__ import annotations | |
| import html | |
| import json | |
| # token source -> (legend label, tooltip) | |
| SOURCES = { | |
| "copy": ("Copied from context", "copied from the prompt or earlier output by n-gram lookup, then verified"), | |
| "head": ("Exit head (layer 4)", "drafted by the shallow exit head, then verified"), | |
| "star": ("Head top-4 candidate", "an alternative from the head's top-4, then verified"), | |
| "draft": ("Draft model", "drafted by a separate model, then verified"), | |
| "model": ("Full model", "produced by the full model (correction or bonus token of a verification)"), | |
| } | |
| class Detokenizer: | |
| """Incremental decoding into (text, source) pieces. Tokens that end inside a multi-byte | |
| character are held back until the character is complete; a piece takes the source of its | |
| first token.""" | |
| def __init__(self, tokenizer): | |
| self.tok = tokenizer | |
| self.ids: list[int] = [] | |
| self.text = "" | |
| self.pending: list[str] = [] | |
| def _emit(self, final: bool) -> list[tuple[str, str]]: | |
| cur = self.tok.decode(self.ids) | |
| if cur.endswith("�") and not final: | |
| return [] | |
| piece = cur[len(self.text):] | |
| src = self.pending[0] if self.pending else "model" | |
| self.text, self.pending = cur, [] | |
| return [(piece, src)] if piece else [] | |
| def feed(self, ids: list[int], sources: list[str]) -> list[tuple[str, str]]: | |
| out = [] | |
| for t, s in zip(ids, sources): | |
| self.ids.append(int(t)) | |
| self.pending.append(s) | |
| out += self._emit(final=False) | |
| return out | |
| def flush(self) -> list[tuple[str, str]]: | |
| return self._emit(final=True) if self.pending else [] | |
| def tokens_html(pieces, colored: bool = True, cursor: bool = False) -> str: | |
| out = [] | |
| for text, src in pieces: | |
| t = html.escape(text) | |
| if colored and src != "model" and src in SOURCES: | |
| out.append(f'<span class="tok tok-{src}" title="{html.escape(SOURCES[src][1])}">{t}</span>') | |
| else: | |
| out.append(t) | |
| if cursor: | |
| out.append('<span class="cursor"></span>') | |
| return f'<div class="dssd-text">{"".join(out)}</div>' | |
| def legend_html() -> str: | |
| chips = "".join( | |
| f'<span class="legend-item"><span class="tok tok-{s} swatch">Aa</span>' | |
| f"{html.escape(SOURCES[s][0])}</span>" | |
| for s in ("copy", "head", "star", "model")) | |
| return f'<div class="legend">{chips}</div>' | |
| def _card(value: str, label: str, cls: str = "") -> str: | |
| return f'<div class="stat {cls}"><div class="stat-value">{value}</div><div class="stat-label">{label}</div></div>' | |
| def stats_html(g) -> str: | |
| """Stat cards and the source-share bar for a DSSD Generation (with g.ar when compared).""" | |
| n = max(len(g.sources), 1) | |
| counts = {s: g.sources.count(s) for s in ("copy", "head", "star", "draft", "model")} | |
| cards = [] | |
| if g.ar is not None and g.speedup is not None: | |
| cards.append(_card(f"{g.speedup:.2f}×", "decoding speedup", "stat-main")) | |
| cards.append(_card(f"{g.decode_tokens_per_s:.1f}", "DSSD tokens/s")) | |
| if g.ar is not None: | |
| cards.append(_card(f"{g.ar.decode_tokens_per_s:.1f}", "autoregressive tokens/s")) | |
| cards.append(_card(f"{g.stats['tau']:.2f}", "tokens per verification")) | |
| if g.ar is not None: | |
| ok = bool(g.identical) | |
| cards.append(_card("✓ identical" if ok else "✗ outputs differ", "DSSD vs autoregressive", | |
| "stat-ok" if ok else "stat-bad")) | |
| bar = "".join( | |
| f'<div class="share tok-{s}" style="width:{100 * c / n:.2f}%" title="{SOURCES[s][0]}: {c} tokens"></div>' | |
| for s, c in counts.items() if c) | |
| labels = " · ".join(f"{SOURCES[s][0]} {round(100 * c / n)}%" for s, c in counts.items() if c) | |
| return (f'<div class="stats">{"".join(cards)}</div>' | |
| f'<div class="share-bar">{bar}</div><div class="share-labels">{html.escape(labels)}</div>') | |
| def replay_html(timelines: dict) -> str: | |
| """Hidden element holding {"dssd": [[t, text, source], ...], "ar": [...]} for the client-side | |
| side-by-side replay at the measured speed.""" | |
| data = html.escape(json.dumps(timelines, ensure_ascii=False), quote=True) | |
| return f"<div class=\"replay-data\" data-replay='{data}'></div>" | |