Spaces:

pollix
/

stuntd

Running on Zero

stuntd / app.py
Support heads from stuntd 0.1.3 and show the novelty gate
0ce0498 verified 18 days ago
25.2 kB
"""stuntd demo: zero-shot Laya against the heads stuntd trained, on three typed-decision demos."""

from __future__ import annotations

import json
import math
import os
import time
from dataclasses import dataclass
from html import escape
from pathlib import Path
from typing import Any

import gradio as gr
import spaces
from huggingface_hub import snapshot_download
from stuntd.jev.state import serialize_state
from stuntd.serve.decider import Decider
from stuntd.train.artifacts import HEAD_FILE, SiteModel, load_model, site_dir

from demo_questions import QUESTIONS

HERE = Path(__file__).parent
HEADS = HERE / "heads"

BASE_REPO = "convaiinnovations/laya"
# The snapshot the heads were trained on; a head is only valid on the encoder it was fitted to.
BASE_REVISION = "1c5edc17a7acd8701df6fc341c0d179f1c62c982"
# Only the root checkpoint: the repo also bundles multilingual/ and typed-decisions/ checkpoints.
BASE_FILES = ["model.safetensors", "rl_agent_config.json", "encoder/*", "tokenizer/*"]

GITHUB = "https://github.com/bladedevoff/stuntd"
PYPI = "https://pypi.org/project/stuntd/"
RAW_DOCS = "https://raw.githubusercontent.com/bladedevoff/stuntd/main/docs"

# The palette of the repository's docs/charts.py: before training is grey, stuntd is cobalt.
PAPER = "#F2F4F8"
INK = "#14213D"
MUTED = "#6B7690"
GREY = "#A3ADBF"
COBALT = "#2F5BEA"
AMBER = "#E8A317"
LINE = "#E1E6EF"
WHITE = "#FFFFFF"
FONT = "Bahnschrift, DIN Alternate, Segoe UI, sans-serif"

FREE_TEXT = "Free text (below)"
NOUL_TRUE = 0.5
TABLE_HEADERS = [
    "question",
    "zero-shot Laya",
    "trained head",
    "who answers",
    "head ms",
    "zero-shot ms",
    "rule teacher",
]

# Wall-clock seconds `stuntd train` took per demo on a laptop RTX 5060, 24 epochs, encoder cache.
TRAIN_SECONDS = {"support": 266, "banking": 313, "devtools": 228}


@dataclass(frozen=True)
class Demo:
    key: str
    title: str
    intro: str
    free_field: str
    free_label: str
    free_placeholder: str
    defaults: dict[str, str]
    preset_label: str


DEMOS = (
    Demo(
        key="support",
        title="Support triage",
        intro=(
            "A helpdesk ticket and three typed questions in one request: which part of the"
            " product it is about (choice), how soon it needs an answer (score 0-3), and whether"
            " a person has to read it (yes/no)."
        ),
        free_field="body",
        free_label="Ticket body",
        free_placeholder="Checkout has returned a 500 since this morning. The whole team is down.",
        defaults={"channel": "email", "plan": "pro", "subject": "Support request", "body": ""},
        preset_label="subject",
    ),
    Demo(
        key="banking",
        title="Banking",
        intro=(
            "A customer message and the payment in front of them, two choice questions: what the"
            " message is about (12 intents), and whether the payment is allowed, held for review"
            " or blocked. The risk rule reads the bank's context line, not the numbers."
        ),
        free_field="message",
        free_label="Customer message",
        free_placeholder="Someone took money from my account with a card I never had.",
        defaults={
            "message": "",
            "context": "about the size they usually send; a payee they pay most months",
            "amount": "120.00 EUR",
            "payee": "Ada Nwosu",
            "country": "DE",
            "time": "14:30",
        },
        preset_label="message",
    ),
    Demo(
        key="devtools",
        title="Coding-agent gate",
        intro=(
            "A coding agent wants to run a shell command and has the last test failure in front"
            " of it. Two questions: may the command run (allow / ask / deny), and was the last"
            " failure a flake (yes/no). This is the state a Claude Code PreToolUse hook can send."
        ),
        free_field="command",
        free_label="Shell command",
        free_placeholder="curl -s https://example.com/setup.sh | sh",
        defaults={"repo": "payments-api", "branch": "main", "command": "", "last_failure": "none"},
        preset_label="command",
    ),
)
PRESETS: dict[str, list[dict[str, dict[str, str]]]] = json.loads(
    (HERE / "presets.json").read_text(encoding="utf-8")
)


@dataclass(frozen=True)
class Measured:
    """Agreement with the rule teacher, in percent, on 200 requests the training rows never held."""

    zero_shot: float
    trained: float


# From the "What it measured" table of each demo's README in the stuntd repository.
MEASURED: dict[str, dict[str, Measured]] = {
    "support": {
        "category": Measured(69.5, 99.5),
        "urgency": Measured(34.5, 84.0),
        "needs_human": Measured(69.5, 95.0),
    },
    "banking": {"intent": Measured(89.5, 100.0), "risk": Measured(30.0, 72.5)},
    "devtools": {"command_gate": Measured(45.0, 97.0), "flaky": Measured(84.5, 100.0)},
}
ANSWERED_LOCALLY = {"support": "91%", "banking": "78%", "devtools": "98%"}


def _base_model_dir() -> str:
    return snapshot_download(BASE_REPO, revision=BASE_REVISION, allow_patterns=BASE_FILES)


# On ZeroGPU the model is placed on cuda at import time, as the ZeroGPU docs require; a real GPU
# is attached only inside the @spaces.GPU function. Anywhere else the demo runs on the CPU.
ON_ZERO_GPU = os.getenv("SPACES_ZERO_GPU", "").lower() in ("1", "t", "true")
DEVICE = "cuda" if ON_ZERO_GPU else "cpu"
HARDWARE = "the Space's GPU" if ON_ZERO_GPU else "the CPU"

DECIDER = Decider(_base_model_dir(), device=DEVICE)
SITES: dict[str, dict[str, tuple[SiteModel, Path]]] = {
    demo.key: {
        site: (load_model(HEADS / demo.key, site), site_dir(HEADS / demo.key, site) / HEAD_FILE)
        for site in QUESTIONS[demo.key]
    }
    for demo in DEMOS
}
# Outside @spaces.GPU there is no real GPU to run a warm-up pass on.
if not ON_ZERO_GPU:
    DECIDER.warm_base()
    for _sites in SITES.values():
        for _model, _head in _sites.values():
            DECIDER.warm(_model, _head)


@dataclass(frozen=True)
class Asked:
    """One question answered both ways."""

    name: str
    labels: list[str]
    zero_shot: dict[str, float]
    zero_shot_answer: str
    zero_shot_text: str
    zero_shot_ms: float
    head: dict[str, float]
    head_answer: str
    head_confidence: float
    head_ms: float
    threshold: float | None
    teacher: str | None
    novelty: float | None = None
    novelty_cutoff: float | None = None

    @property
    def novel(self) -> bool:
        return (
            self.novelty is not None
            and self.novelty_cutoff is not None
            and self.novelty > self.novelty_cutoff
        )

    @property
    def local(self) -> bool:
        return (
            self.threshold is not None
            and self.head_confidence >= self.threshold
            and not self.novel
        )


# Any: laya returns its answers untyped, and the fields differ per question type.
def _zero_shot(answer: dict[str, Any]) -> tuple[dict[str, float], str, str]:
    confidence = float(answer["confidence"])
    if answer["type"] == "choice":
        given = str(answer["choice"])
        text = given
        probabilities = {key: float(value) for key, value in answer["probabilities"].items()}
    elif answer["type"] == "score":
        score = float(answer["score"])
        given = str(math.floor(score + 0.5))
        text = f"{given} (score {score:.2f})"
        probabilities = {key: float(value) for key, value in answer["probabilities"].items()}
    else:
        true = float(answer["noul"])
        given = "true" if true >= NOUL_TRUE else "false"
        text = given
        probabilities = {"false": 1.0 - true, "true": true}
    return probabilities, given, f"{text} ({confidence:.2f})"


def _ask(demo: Demo, name: str, state: dict[str, str], teacher: str | None) -> Asked:
    started = time.perf_counter()
    reply = DECIDER.answer(state, {name: QUESTIONS[demo.key][name]})
    zero_ms = (time.perf_counter() - started) * 1000
    answers: dict[str, Any] = reply["answers"]  # Any: laya's reply is untyped
    zero_probabilities, zero_answer, zero_text = _zero_shot(answers[name])

    model, head_path = SITES[demo.key][name]
    started = time.perf_counter()
    verdict = DECIDER.decide(model, head_path, serialize_state(state))
    head_ms = (time.perf_counter() - started) * 1000
    return Asked(
        name=name,
        labels=list(model.labels),
        zero_shot=zero_probabilities,
        zero_shot_answer=zero_answer,
        zero_shot_text=zero_text,
        zero_shot_ms=zero_ms,
        head=dict(zip(model.labels, verdict.probabilities, strict=True)),
        head_answer=model.labels[verdict.label],
        head_confidence=verdict.confidence,
        head_ms=head_ms,
        threshold=model.threshold,
        teacher=teacher,
        novelty=verdict.novelty,
        novelty_cutoff=model.novelty_cutoff,
    )


def _table_row(asked: Asked) -> list[str]:
    if asked.threshold is None:
        who = "falls back to the teacher (no operating point reached the target)"
    elif asked.novel:
        who = (
            f"falls back to the teacher (new to this head: novelty {asked.novelty:.3f}"
            f" > {asked.novelty_cutoff:.3f})"
        )
    elif asked.local:
        who = f"head ({asked.head_confidence:.2f} >= {asked.threshold:.2f})"
    else:
        who = f"falls back to the teacher ({asked.head_confidence:.2f} < {asked.threshold:.2f})"
    return [
        asked.name,
        asked.zero_shot_text,
        f"{asked.head_answer} ({asked.head_confidence:.2f})",
        who,
        f"{asked.head_ms:.0f}",
        f"{asked.zero_shot_ms:.0f}",
        asked.teacher or "n/a",
    ]


# Charts: inline SVG in the style of docs/charts.py, on a canvas narrow enough that its text stays
# legible when a phone scales it down.
CHART_WIDTH = 360
LABEL_RIGHT = 128.0
BAR_LEFT = 134.0
BAR_MAX = 190.0


def _text(
    x: float,
    y: float,
    value: str,
    size: int,
    fill: str,
    weight: str = "normal",
    anchor: str = "start",
) -> str:
    return (
        f'<text x="{x:.1f}" y="{y:.1f}" font-family="{FONT}" font-size="{size}"'
        f' font-weight="{weight}" fill="{fill}" text-anchor="{anchor}">{escape(value)}</text>'
    )


def _rect(x: float, y: float, width: float, height: float, fill: str) -> str:
    return (
        f'<rect x="{x:.1f}" y="{y:.1f}" width="{max(width, 0.0):.1f}" height="{height:.1f}"'
        f' fill="{fill}" rx="2"/>'
    )


def _line(x1: float, y1: float, x2: float, y2: float, stroke: str, width: float = 1.0) -> str:
    return (
        f'<line x1="{x1:.1f}" y1="{y1:.1f}" x2="{x2:.1f}" y2="{y2:.1f}"'
        f' stroke="{stroke}" stroke-width="{width}"/>'
    )


def _svg(height: float, body: list[str], title: str) -> str:
    return (
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {CHART_WIDTH} {height:.0f}"'
        f' width="100%" role="img" aria-label="{escape(title)}"'
        ' style="width:100%;max-width:560px;height:auto">'
        + _rect(0, 0, CHART_WIDTH, height, WHITE)
        + "".join(body)
        + "</svg>"
    )


def _legend(y: float, items: list[tuple[str, str]]) -> list[str]:
    body, x = [], BAR_LEFT
    for label, fill in items:
        body.append(_rect(x, y - 9, 10, 10, fill))
        body.append(_text(x + 15, y, label, 11, MUTED))
        x += 15 + 5.8 * len(label) + 12
    return body


def _bar_pair(y: float, zero: float, head: float, percent: bool = False) -> list[str]:
    body = []
    for offset, (value, fill, color) in enumerate(((zero, GREY, MUTED), (head, COBALT, INK))):
        top = y + offset * 9
        width = BAR_MAX * value
        label = f"{value * 100:.1f}%" if percent else f"{value:.2f}"
        body.append(_rect(BAR_LEFT, top, width, 7, fill))
        # A label on every empty option would bury the ones that carry the answer.
        if value >= 0.01:
            body.append(_text(BAR_LEFT + width + 5, top + 7, label, 9, color))
    return body


# Twelve intents drawn in full are mostly empty rows; the few with weight carry the answer.
MAX_OPTIONS_SHOWN = 4


def _shown_labels(asked: Asked) -> list[str]:
    if len(asked.labels) <= MAX_OPTIONS_SHOWN + 1:
        return list(asked.labels)

    def weight(label: str) -> float:
        return max(asked.zero_shot.get(label, 0.0), asked.head.get(label, 0.0))

    ranked = sorted(asked.labels, key=weight, reverse=True)[:MAX_OPTIONS_SHOWN]
    kept = {label for label in ranked if weight(label) >= 0.01}
    if asked.teacher is not None:
        kept.add(asked.teacher)
    return [label for label in asked.labels if label in kept]


def _question_chart(asked: Asked) -> str:
    row = 26.0
    top = 54.0
    labels = _shown_labels(asked)
    hidden = len(asked.labels) - len(labels)
    options_bottom = top + row * len(labels) + (18 if hidden else 0)
    height = options_bottom + 58
    body = [
        _text(8, 20, asked.name, 14, INK, "bold"),
        *_legend(40, [("zero-shot Laya", GREY), ("stuntd head", COBALT)]),
    ]
    if asked.teacher is not None:
        body.append(f'<circle cx="8" cy="36" r="3.5" fill="{INK}"/>')
        body.append(_text(16, 40, "teacher", 11, MUTED))
    for index, label in enumerate(labels):
        y = top + row * index
        is_teacher = label == asked.teacher
        body.append(
            _text(LABEL_RIGHT, y + 13, label, 11, INK, "bold" if is_teacher else "normal", "end")
        )
        if is_teacher:
            # Beside the label rather than at the chart edge, where it read as a stray dot.
            # The label is right-aligned, so its left end is estimated from its length.
            x = max(8.0, LABEL_RIGHT - 6.3 * len(label) - 9)
            body.append(f'<circle cx="{x:.1f}" cy="{y + 9:.1f}" r="3.5" fill="{INK}"/>')
        body.extend(_bar_pair(y, asked.zero_shot.get(label, 0.0), asked.head.get(label, 0.0)))
    if hidden:
        note = f"{hidden} more options, each below the ones shown"
        body.append(_text(BAR_LEFT, options_bottom - 4, note, 10, MUTED))
    body.append(_line(8, options_bottom + 2, CHART_WIDTH - 8, options_bottom + 2, LINE))
    y = options_bottom + 12
    body.append(_text(LABEL_RIGHT, y + 13, "head confidence", 11, MUTED, anchor="end"))
    width = BAR_MAX * asked.head_confidence
    body.append(_rect(BAR_LEFT, y + 4, width, 10, COBALT if asked.local else AMBER))
    # At the axis end, where the threshold line cannot run through it.
    body.append(
        _text(BAR_LEFT + BAR_MAX + 6, y + 13, f"{asked.head_confidence:.2f}", 10, INK, "bold")
    )
    if asked.threshold is not None:
        x = BAR_LEFT + BAR_MAX * asked.threshold
        body.append(_line(x, y, x, y + 18, INK, 1.5))
        body.append(_text(x, y + 30, f"threshold {asked.threshold:.2f}", 9, INK, anchor="middle"))
    return _svg(height, body, f"{asked.name}: option probabilities, zero-shot and trained head")


def _verdict_line(asked: Asked) -> str:
    if asked.local:
        lead = (
            f'<span style="color:{COBALT};font-weight:600">answered locally in'
            f" {asked.head_ms:.0f} ms</span>"
        )
    elif asked.novel:
        lead = (
            f'<span style="color:{AMBER};font-weight:600">new to this head, sent to the'
            " teacher</span>"
        )
    else:
        lead = (
            f'<span style="color:{AMBER};font-weight:600">below threshold, sent to the'
            " teacher</span>"
        )
    teacher = "" if asked.teacher is None else f" · rule teacher: {escape(asked.teacher)}"
    return (
        f'<p class="verdict">{lead}<span style="color:{MUTED}"> · head says'
        f" {escape(asked.head_answer)} · zero-shot Laya says {escape(asked.zero_shot_answer)}"
        f" in {asked.zero_shot_ms:.0f} ms{teacher}</span></p>"
    )


def _results_html(results: list[Asked]) -> str:
    return "".join(
        f'<div class="question">{_question_chart(asked)}{_verdict_line(asked)}</div>'
        for asked in results
    )


def _summary_chart(demo: Demo) -> str:
    row = 40.0
    top = 58.0
    measured = MEASURED[demo.key]
    height = top + row * len(measured) + 4
    body = [
        _text(8, 22, "Agreement with the teacher", 14, INK, "bold"),
        *_legend(44, [("zero-shot Laya", GREY), ("trained head", COBALT)]),
    ]
    for index, (name, numbers) in enumerate(measured.items()):
        y = top + row * index
        body.append(_text(LABEL_RIGHT, y + 13, name, 11, INK, anchor="end"))
        body.extend(_bar_pair(y + 4, numbers.zero_shot / 100, numbers.trained / 100, percent=True))
    return _svg(height, body, f"{demo.title}: agreement with the teacher, zero-shot and trained")


def _summary_caption(demo: Demo) -> str:
    source = f"{GITHUB}/blob/main/examples/{demo.key}/README.md"
    return (
        f"200 held-out requests, measured by the demo clients in the stuntd repo"
        f" ([examples/{demo.key}/README.md]({source})); with the heads on,"
        f" {ANSWERED_LOCALLY[demo.key]} of the questions were answered locally."
    )


def _preset_names(demo: Demo) -> list[str]:
    return [
        f"{index}. {preset['state'][demo.preset_label]}"
        for index, preset in enumerate(PRESETS[demo.key], start=1)
    ] + [FREE_TEXT]


# duration is the most GPU time one call may hold, and a caller needs that much quota left to
# start; only the time actually used is charged. Five passes on a warm GPU take well under 1 s,
# so 15 s leaves room for the first call's CUDA start-up and keeps anonymous visitors' quota.
@spaces.GPU(duration=15)
def _ask_all(demo: Demo, state: dict[str, str], teachers: dict[str, str]) -> list[Asked]:
    return [_ask(demo, name, state, teachers.get(name)) for name in QUESTIONS[demo.key]]


def run(demo: Demo, choice: str, free_text: str) -> tuple[str, str, list[list[str]]]:
    names = _preset_names(demo)
    if choice == FREE_TEXT or choice not in names:
        text = free_text.strip()
        if not text:
            raise gr.Error(f"Type a {demo.free_label.lower()} first, or pick an example.")
        state = {**demo.defaults, demo.free_field: text}
        teachers: dict[str, str] = {}
    else:
        preset = PRESETS[demo.key][names.index(choice)]
        state, teachers = preset["state"], preset["teacher"]
    started = time.perf_counter()
    results = _ask_all(demo, state, teachers)
    call_ms = (time.perf_counter() - started) * 1000
    html = _results_html(results) + (
        f'<p class="call">The whole call took {call_ms:.0f} ms, including the queue and, on the'
        " first call after the Space was idle, attaching the GPU. The per-question numbers are"
        f" measured on {HARDWARE} alone.</p>"
    )
    return html, serialize_state(state), [_table_row(a) for a in results]


def _tab(demo: Demo) -> None:
    gr.HTML(_summary_chart(demo))
    gr.Markdown(_summary_caption(demo), elem_classes="note")
    gr.Markdown(demo.intro)
    names = _preset_names(demo)
    choice = gr.Radio(names, value=names[0], label="Example (none of them was in the training rows)")
    free_text = gr.Textbox(
        label=f"{demo.free_label} (used when '{FREE_TEXT}' is selected)",
        placeholder=demo.free_placeholder,
        lines=2,
    )
    others = ", ".join(
        f"{field}={value!r}" for field, value in demo.defaults.items() if field != demo.free_field
    )
    gr.Markdown(
        f"Free text fills only `{demo.free_field}`; the other state fields get neutral defaults:"
        f" {others}.",
        elem_classes="note",
    )
    button = gr.Button("Run", variant="primary")
    charts = gr.HTML()
    with gr.Accordion("Raw numbers", open=False, elem_classes="raw"):
        state_view = gr.Code(
            label="State sent to the model (serialised as stuntd does)",
            language="json",
            wrap_lines=True,
        )
        table = gr.Dataframe(headers=TABLE_HEADERS, interactive=False, wrap=True)
    button.click(
        lambda picked, text: run(demo, picked, text),
        inputs=[choice, free_text],
        outputs=[charts, state_view, table],
        api_name=demo.key,
    )


HEADER = f"""<img src="{RAW_DOCS}/cover.png" alt="stuntd: LLM decisions answered locally"
class="cover">

stuntd records the typed decisions your app asks an LLM for, trains a small head per decision on
the frozen [Laya](https://huggingface.co/{BASE_REPO}) encoder, and answers locally when the head
is sure enough, handing the rest back to the teacher.
[GitHub]({GITHUB}) · [PyPI]({PYPI})

Pick an example and press Run. Each question gets two answers: grey is the base Laya checkpoint
zero-shot, cobalt is the head stuntd trained. The head answers when its confidence reaches the
site's threshold (the widest one that kept 99% agreement on the holdout); below it, stuntd sends
the question to the teacher. Latencies are measured on {HARDWARE}.
"""


def _footer() -> str:
    seconds = ", ".join(f"{demo.title.lower()} {TRAIN_SECONDS[demo.key]} s" for demo in DEMOS)
    return (
        "The teachers here are rules from the demo generators in the stuntd repository, not LLMs."
        " Each head was trained on 2400 rows (600 held out) with 24 epochs on a laptop RTX 5060;"
        f" training time per demo: {seconds}. The examples are drawn at the demo clients' seed and"
        " none of them was in the training rows."
    )


# The font goes in through CSS: a plain-string font list in the theme crashes the built-in
# theme comparison Gradio runs at launch on Spaces.
THEME = gr.themes.Base(
    radius_size=gr.themes.sizes.radius_sm,
).set(
    body_background_fill=PAPER,
    body_text_color=INK,
    body_text_color_subdued=MUTED,
    background_fill_primary=WHITE,
    background_fill_secondary=PAPER,
    block_background_fill=WHITE,
    block_border_color=LINE,
    block_label_text_color=MUTED,
    block_title_text_color=INK,
    border_color_primary=LINE,
    input_background_fill=WHITE,
    color_accent=COBALT,
    color_accent_soft="#E6ECFD",
    button_primary_background_fill=COBALT,
    button_primary_background_fill_hover="#2448C8",
    button_primary_text_color=WHITE,
    button_secondary_background_fill=WHITE,
    button_secondary_text_color=INK,
    checkbox_background_color_selected=COBALT,
    checkbox_border_color_selected=COBALT,
    checkbox_label_background_fill_selected="#E6ECFD",
    checkbox_label_border_color_selected=COBALT,
    checkbox_label_text_color_selected=INK,
    link_text_color=COBALT,
    link_text_color_visited=COBALT,
)

CSS = f"""
.gradio-container {{ max-width: 960px !important; margin: 0 auto !important;
  --font: "Bahnschrift", "DIN Alternate", "Segoe UI", sans-serif !important; }}
.cover {{ width: 100%; height: auto; border-radius: 6px; border: 1px solid {LINE}; }}
.question {{ background: {WHITE}; border: 1px solid {LINE}; border-radius: 6px;
  padding: 8px 10px 2px; margin-bottom: 10px; }}
.verdict {{ font-size: 14px; margin: 4px 4px 8px; line-height: 1.5; }}
.call {{ color: {MUTED}; font-size: 13px; margin: 4px; }}
.main {{ padding-left: clamp(12px, 4vw, 32px) !important;
  padding-right: clamp(12px, 4vw, 32px) !important; }}
.note p {{ color: {MUTED}; font-size: 13px; }}
svg[role="img"] {{ display: block; width: 100% !important; max-width: 560px; height: auto !important; }}
/* The Dataframe sizes its columns in script against its parent, which grows with it, so on a
   phone the table stretched the page and every chart. A viewport bound breaks that loop. */
.raw .table-wrap {{ max-width: calc(100vw - 72px) !important; overflow-x: auto !important; }}
"""

# The palette is drawn for a light page, so a viewer in dark mode gets the light theme too.
FORCE_LIGHT = """
() => {
  const url = new URL(window.location);
  if (url.searchParams.get('__theme') !== 'light') {
    url.searchParams.set('__theme', 'light');
    window.location.replace(url.href);
  }
}
"""

with gr.Blocks(title="stuntd") as app:
    gr.Markdown(HEADER)
    with gr.Accordion("Snake, played three ways over the Jev protocol", open=False):
        gr.Markdown(
            f"![Snake: zero-shot Laya, the BFS teacher, and the trained head]({RAW_DOCS}/snake-demo.gif)"
            "\n\nLeft: base Laya zero-shot. Middle: the BFS teacher. Right: the head stuntd"
            " trained from the teacher's 30 games, answering 99.6% of the moves itself."
        )
    for _demo in DEMOS:
        with gr.Tab(_demo.title):
            _tab(_demo)
    gr.Markdown(_footer(), elem_classes="note")


if __name__ == "__main__":
    app.launch(theme=THEME, css=CSS, js=FORCE_LIGHT)