Spaces:
Sleeping
Sleeping
File size: 4,381 Bytes
74679d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """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>"
|