Initial commit

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

app.py

@@ -0,0 +1,1496 @@
+import html
+import json
+import random
+import re
+import time
+
+import gradio as gr
+
+MAX_DEMO_LEN = 20
+
+ROSA_CODE = """
+def rosa_qkv(qqq, kkk, vvv):
+    n=len(qqq); out=[-1]*n
+    for i in range(n):
+        for w in range(i+1,0,-1):
+            t=qqq[i+1-w:i+1]
+            for j in range(i-w,-1,-1):
+                if kkk[j:j+w]==t:
+                    out[i]=vvv[j+w]
+                    break
+            if out[i]!=-1:
+                break
+    return out
+""".strip("\n")
+
+
+KEYWORDS = {
+    "def",
+    "for",
+    "in",
+    "if",
+    "break",
+    "return",
+    "assert",
+}
+BUILTINS = {"len", "range"}
+KEYWORD_RE = re.compile(r"\b(" + "|".join(sorted(KEYWORDS)) + r")\b")
+BUILTIN_RE = re.compile(r"\b(" + "|".join(sorted(BUILTINS)) + r")\b")
+NUMBER_RE = re.compile(r"(?<![\w.])(-?\d+)(?![\w.])")
+
+
+def split_comment(line: str) -> tuple[str, str]:
+    in_single = False
+    in_double = False
+    escaped = False
+    for index, ch in enumerate(line):
+        if escaped:
+            escaped = False
+            continue
+        if ch == "\\":
+            escaped = True
+            continue
+        if ch == "'" and not in_double:
+            in_single = not in_single
+            continue
+        if ch == '"' and not in_single:
+            in_double = not in_double
+            continue
+        if ch == "#" and not in_single and not in_double:
+            return line[:index], line[index:]
+    return line, ""
+
+
+def tokenize_strings(code: str) -> list[tuple[str, str]]:
+    segments: list[tuple[str, str]] = []
+    index = 0
+    while index < len(code):
+        ch = code[index]
+        if ch in ("'", '"'):
+            quote = ch
+            start = index
+            index += 1
+            escaped = False
+            while index < len(code):
+                if escaped:
+                    escaped = False
+                    index += 1
+                    continue
+                if code[index] == "\\":
+                    escaped = True
+                    index += 1
+                    continue
+                if code[index] == quote:
+                    index += 1
+                    break
+                index += 1
+            segments.append(("string", code[start:index]))
+            continue
+        start = index
+        while index < len(code) and code[index] not in ("'", '"'):
+            index += 1
+        segments.append(("text", code[start:index]))
+    return segments
+
+
+def highlight_text_segment(text: str) -> str:
+    escaped = html.escape(text)
+    escaped = KEYWORD_RE.sub(r'<span class="tok-keyword">\1</span>', escaped)
+    escaped = BUILTIN_RE.sub(r'<span class="tok-builtin">\1</span>', escaped)
+    escaped = NUMBER_RE.sub(r'<span class="tok-number">\1</span>', escaped)
+    return escaped
+
+
+def highlight_python_line(line: str) -> str:
+    code, comment = split_comment(line)
+    segments = tokenize_strings(code)
+    rendered: list[str] = []
+    for kind, text in segments:
+        if kind == "string":
+            rendered.append('<span class="tok-string">{}</span>'.format(html.escape(text)))
+        else:
+            rendered.append(highlight_text_segment(text))
+    if comment:
+        rendered.append('<span class="tok-comment">{}</span>'.format(html.escape(comment)))
+    return "".join(rendered)
+
+
+def build_code_html(code: str) -> str:
+    lines = code.splitlines()
+    rendered = ['<div id="rosa-code" class="rosa-code">']
+    marker_t = "__TOK_T__"
+    marker_k = "__TOK_K__"
+    marker_v = "__TOK_V__"
+    for index, line in enumerate(lines, start=1):
+        line_with_markers = line
+        if index == LINE_TRY:
+            line_with_markers = line_with_markers.replace("kkk[j:j+w]", marker_k, 1)
+            line_with_markers = line_with_markers.replace("t", marker_t, 1)
+        if index == LINE_ASSIGN:
+            line_with_markers = line_with_markers.replace("vvv[j+w]", marker_v, 1)
+        highlighted = highlight_python_line(line_with_markers)
+        highlighted = highlighted.replace(
+            marker_t, '<span class="code-token" data-token="t">t</span>'
+        )
+        highlighted = highlighted.replace(
+            marker_k, '<span class="code-token" data-token="k">kkk[j:j+w]</span>'
+        )
+        highlighted = highlighted.replace(
+            marker_v, '<span class="code-token" data-token="v">vvv[j+w]</span>'
+        )
+        rendered.append(
+            '<div class="code-line" data-line="{line}">'
+            '<span class="line-no">{line}</span>'
+            '<span class="line-text">{text}</span>'
+            "</div>".format(line=index, text=highlighted)
+        )
+    rendered.append("</div>")
+    return "\n".join(rendered)
+
+
+def find_line_number(lines: list[str], needle: str, fallback: int = 1) -> int:
+    for index, line in enumerate(lines, start=1):
+        if needle in line:
+            return index
+    return fallback
+
+
+ROSA_LINES = ROSA_CODE.splitlines()
+def find_line_after(lines: list[str], start_line: int, needle: str, fallback: int) -> int:
+    start_index = min(max(start_line, 0), len(lines))
+    for index in range(start_index, len(lines)):
+        if needle in lines[index]:
+            return index + 1
+    return fallback
+
+
+LINE_FOR_I = find_line_number(ROSA_LINES, "for i in range(n):")
+LINE_FOR_W = find_line_number(ROSA_LINES, "for w in range(i+1,0,-1):")
+LINE_T_ASSIGN = find_line_number(ROSA_LINES, "t=qqq[i+1-w:i+1]")
+LINE_FOR_J = find_line_number(ROSA_LINES, "for j in range(i-w,-1,-1):")
+LINE_TRY = find_line_number(ROSA_LINES, "if kkk[j:j+w]==t:")
+LINE_ASSIGN = find_line_number(ROSA_LINES, "out[i]=vvv[j+w]")
+LINE_BREAK_INNER = find_line_after(ROSA_LINES, LINE_ASSIGN, "break", LINE_ASSIGN)
+LINE_IF_OUT = find_line_number(ROSA_LINES, "if out[i]!=-1:")
+LINE_BREAK_OUTER = find_line_after(ROSA_LINES, LINE_IF_OUT, "break", LINE_IF_OUT)
+LINE_RETURN = find_line_number(ROSA_LINES, "return out")
+CODE_HTML = build_code_html(ROSA_CODE)
+
+
+def parse_bits(text: str, name: str) -> list[int]:
+    cleaned = re.sub(r"[,\s]+", "", (text or "").strip())
+    if not cleaned:
+        raise gr.Error(f"{name} cannot be empty")
+    if re.search(r"[^01]", cleaned):
+        raise gr.Error(f"{name} must contain only 0/1 (spaces or commas allowed)")
+    return [int(c) for c in cleaned]
+
+
+def rosa_steps(q: list[int], k: list[int], v: list[int]) -> tuple[list[dict], list[int]]:
+    n = len(q)
+    out = [-1] * n
+    steps: list[dict] = []
+    for i in range(n):
+        steps.append({"phase": "loop_i", "i": i, "line": LINE_FOR_I})
+        matched = False
+        match_j = -1
+        match_w = -1
+        for w in range(i + 1, 0, -1):
+            steps.append({"phase": "loop_w", "i": i, "w": w, "line": LINE_FOR_W})
+            t = q[i + 1 - w : i + 1]
+            t_str = "".join(str(x) for x in t)
+            steps.append({"phase": "assign_t", "i": i, "w": w, "t": t_str, "line": LINE_T_ASSIGN})
+            for j in range(i - w, -1, -1):
+                steps.append({"phase": "loop_j", "i": i, "w": w, "j": j, "line": LINE_FOR_J})
+                is_match = k[j : j + w] == t
+                steps.append(
+                    {
+                        "phase": "try",
+                        "i": i,
+                        "w": w,
+                        "j": j,
+                        "t": t_str,
+                        "matched": is_match,
+                        "line": LINE_TRY,
+                    }
+                )
+                if is_match:
+                    value = v[j + w]
+                    out[i] = value
+                    match_j = j
+                    match_w = w
+                    steps.append(
+                        {
+                            "phase": "assign",
+                            "i": i,
+                            "w": w,
+                            "j": j,
+                            "t": t_str,
+                            "v_index": j + w,
+                            "value": value,
+                            "line": LINE_ASSIGN,
+                        }
+                    )
+                    steps.append(
+                        {
+                            "phase": "break_inner",
+                            "i": i,
+                            "w": w,
+                            "j": j,
+                            "line": LINE_BREAK_INNER,
+                        }
+                    )
+                    matched = True
+                    break
+            if matched:
+                steps.append(
+                    {
+                        "phase": "check",
+                        "i": i,
+                        "w": match_w,
+                        "j": match_j,
+                        "matched": True,
+                        "line": LINE_IF_OUT,
+                    }
+                )
+                if i == n - 1:
+                    steps.append(
+                        {
+                            "phase": "break_outer",
+                            "i": i,
+                            "w": match_w,
+                            "j": match_j,
+                            "line": LINE_BREAK_OUTER,
+                        }
+                    )
+                break
+        if not matched:
+            out[i] = -1
+        steps.append(
+            {
+                "phase": "output",
+                "i": i,
+                "value": out[i],
+                "matched": matched,
+            }
+        )
+        if i == n - 1:
+            steps.append({"phase": "return", "line": LINE_RETURN})
+    return steps, out
+
+
+def format_out(out: list[int]) -> str:
+    if any(value < 0 or value > 1 for value in out):
+        return " ".join(str(value) for value in out)
+    return "".join(str(value) for value in out)
+
+
+def build_payload(q: list[int], k: list[int], v: list[int]) -> tuple[str, str]:
+    steps, out = rosa_steps(q, k, v)
+    payload = {
+        "q": q,
+        "k": k,
+        "v": v,
+        "steps": steps,
+        "out": out,
+        "run_id": time.time_ns(),
+    }
+    return json.dumps(payload, ensure_ascii=False), format_out(out)
+
+
+def on_demo(q_text: str, k_text: str, v_text: str) -> tuple[str, str]:
+    q = parse_bits(q_text, "q")
+    k = parse_bits(k_text, "k")
+    v = parse_bits(v_text, "v")
+    if not (len(q) == len(k) == len(v)):
+        raise gr.Error("q, k, v must have the same length")
+    if len(q) > MAX_DEMO_LEN:
+        raise gr.Error(f"For smooth playback, length should be <= {MAX_DEMO_LEN}")
+    return build_payload(q, k, v)
+
+
+def on_random(length: int) -> tuple[str, str, str]:
+    length = max(1, int(length))
+    q = [random.randint(0, 1) for _ in range(length)]
+    k = [random.randint(0, 1) for _ in range(length)]
+    v = [random.randint(0, 1) for _ in range(length)]
+    return "".join(str(x) for x in q), "".join(str(x) for x in k), "".join(str(x) for x in v)
+
+
+CSS = """
+.page-header {
+  text-align: center;
+  margin-bottom: 18px;
+  color: #0f172a;
+}
+.page-title {
+  font-size: 30px;
+  font-weight: 700;
+  letter-spacing: 0.4px;
+  margin-bottom: 6px;
+}
+.page-subtitle {
+  font-size: 14px;
+  color: #475569;
+}
+.rosa-shell {
+  display: flex;
+  gap: 24px;
+  align-items: flex-start;
+  justify-content: center;
+  flex-wrap: wrap;
+  position: relative;
+}
+.rosa-pane {
+  flex: 3 1 520px;
+  min-width: 320px;
+}
+.rosa-code-pane {
+  flex: 2 1 320px;
+  min-width: 300px;
+}
+#rosa-vis {
+  --rosa-blue: #3b82f6;
+  --rosa-sky: #38bdf8;
+  --rosa-violet: #a855f7;
+  --rosa-amber: #f59e0b;
+  --rosa-cyan: #06b6d4;
+  --rosa-green: #22c55e;
+  --rosa-green-soft: rgba(34, 197, 94, 0.16);
+  --rosa-red: #ef4444;
+  --rosa-red-soft: rgba(239, 68, 68, 0.14);
+  --rosa-blue-soft: rgba(59, 130, 246, 0.08);
+  --rosa-violet-soft: rgba(168, 85, 247, 0.08);
+}
+#rosa-vis .rosa-card {
+  background: #ffffff;
+  border-radius: 18px;
+  padding: 24px;
+  color: #0f172a;
+  box-shadow: none;
+  border: 1px solid #e2e8f0;
+  text-align: center;
+  position: relative;
+}
+#rosa-vis .rosa-rows {
+  display: flex;
+  flex-direction: column;
+  gap: 20px;
+  align-items: center;
+}
+#rosa-vis .rosa-row {
+  display: flex;
+  align-items: center;
+  gap: 14px;
+  flex-wrap: wrap;
+  justify-content: center;
+  width: 100%;
+}
+#rosa-vis .rosa-row.k-row {
+  margin-bottom: 0;
+}
+#rosa-vis .row-label {
+  min-width: 36px;
+  font-size: 12px;
+  color: #475569;
+  text-transform: uppercase;
+  letter-spacing: 0.2px;
+  text-align: center;
+}
+#rosa-vis .row-cells {
+  display: flex;
+  flex-wrap: wrap;
+  gap: 6px;
+  justify-content: center;
+  max-width: 100%;
+  min-width: 0;
+}
+#rosa-vis .cell {
+  width: 30px;
+  height: 30px;
+  border-radius: 8px;
+  background: #f1f5f9;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  font-weight: 600;
+  transition: transform 0.18s ease, box-shadow 0.18s ease, background 0.18s ease;
+  color: #0f172a;
+  box-shadow: none;
+}
+#rosa-vis .cell.active {
+  background: var(--rosa-blue);
+  color: #f8fafc;
+  box-shadow: none;
+}
+#rosa-vis .cell.suffix {
+  background: var(--rosa-sky);
+  color: #f8fafc;
+}
+#rosa-vis .cell.k-window {
+  background: var(--rosa-violet);
+  color: #f8fafc;
+}
+#rosa-vis .cell.v-pick {
+  background: var(--rosa-amber);
+  color: #1f2937;
+}
+#rosa-vis .cell.out,
+#rosa-vis .cell.out-fixed {
+  background: var(--rosa-amber);
+  color: #1f2937;
+}
+#rosa-vis .cell.filled {
+  box-shadow: none;
+}
+#rosa-vis .rosa-overlay {
+  position: absolute;
+  inset: 0;
+  pointer-events: none;
+  z-index: 5;
+}
+#rosa-vis .overlay-svg {
+  position: absolute;
+  inset: 0;
+  width: 100%;
+  height: 100%;
+  pointer-events: none;
+}
+#rosa-vis .overlay-box-layer {
+  position: absolute;
+  inset: 0;
+  pointer-events: none;
+}
+#rosa-vis .overlay-box {
+  position: absolute;
+  border: 2px solid rgba(59, 130, 246, 0.7);
+  border-radius: 10px;
+  box-shadow: 0 6px 16px rgba(15, 23, 42, 0.08);
+  box-sizing: border-box;
+  background: transparent;
+  transition: border-color 0.18s ease, background 0.18s ease;
+}
+#rosa-vis .overlay-hover-layer {
+  position: absolute;
+  inset: 0;
+  pointer-events: auto;
+  z-index: 6;
+}
+#rosa-vis .overlay-hover-box {
+  position: absolute;
+  background: transparent;
+  pointer-events: auto;
+  cursor: pointer;
+}
+#rosa-vis .overlay-box[data-label="t"] {
+  --overlay-label: "t";
+}
+#rosa-vis .overlay-box[data-label="kkk[j:j+w]"] {
+  --overlay-label: "kkk[j:j+w]";
+}
+#rosa-vis .overlay-box.t-box {
+  border-color: var(--rosa-blue);
+  background: var(--rosa-blue-soft);
+}
+#rosa-vis .overlay-box.k-box {
+  border-color: var(--rosa-violet);
+  background: var(--rosa-violet-soft);
+}
+#rosa-vis .overlay-box.try-match {
+  border-color: var(--rosa-green);
+  background: var(--rosa-green-soft);
+}
+#rosa-vis .overlay-box.try-miss {
+  border-color: var(--rosa-red);
+  background: var(--rosa-red-soft);
+}
+#rosa-vis .overlay-ray {
+  position: absolute;
+  inset: 0;
+  width: 100%;
+  height: 100%;
+  pointer-events: none;
+}
+#rosa-vis .overlay-ray-line {
+  stroke: var(--rosa-cyan);
+  stroke-width: 2;
+  fill: none;
+  stroke-linecap: round;
+}
+#rosa-vis .overlay-label {
+  position: absolute;
+  top: -12px;
+  right: 0;
+  font-size: 11px;
+  padding: 0;
+  background: transparent;
+  color: #0f172a;
+  letter-spacing: 0.2px;
+  white-space: nowrap;
+  display: none;
+}
+#rosa-vis .overlay-label.t-label {
+  color: #1d4ed8;
+}
+#rosa-vis .overlay-label.k-label {
+  color: #6d28d9;
+}
+.v-float {
+  position: fixed;
+  z-index: 9999;
+  border-radius: 8px;
+  background: #f59e0b;
+  color: #1f2937;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  font-weight: 600;
+  box-shadow: 0 10px 24px rgba(15, 23, 42, 0.2);
+  pointer-events: none;
+  transition: transform 0.35s ease, opacity 0.35s ease;
+  opacity: 1;
+}
+#rosa-vis .rosa-legend {
+  display: flex;
+  flex-wrap: wrap;
+  gap: 12px;
+  margin-bottom: 12px;
+  font-size: 12px;
+  color: #475569;
+  justify-content: center;
+}
+#rosa-vis .legend-item {
+  display: inline-flex;
+  align-items: center;
+  gap: 6px;
+}
+#rosa-vis .legend-dot {
+  width: 12px;
+  height: 12px;
+  border-radius: 4px;
+}
+#rosa-vis .legend-suffix {
+  background: var(--rosa-sky);
+}
+#rosa-vis .legend-window {
+  background: var(--rosa-violet);
+}
+#rosa-vis .legend-match {
+  background: var(--rosa-green);
+}
+#rosa-vis .legend-v {
+  background: var(--rosa-amber);
+}
+#rosa-vis .legend-out {
+  background: var(--rosa-amber);
+}
+#rosa-vis .index-row {
+  padding-bottom: 8px;
+  border-bottom: 1px dashed #e2e8f0;
+}
+#rosa-vis .index-cells {
+  gap: 6px;
+}
+#rosa-vis .index-cell {
+  width: 30px;
+  text-align: center;
+  font-size: 11px;
+  color: #64748b;
+  font-variant-numeric: tabular-nums;
+}
+#rosa-code {
+  background: #ffffff;
+  border: 1px solid #e2e8f0;
+  border-radius: 12px;
+  padding: 12px;
+  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
+  font-size: 12px;
+  color: #0f172a;
+  max-height: 520px;
+  overflow: auto;
+}
+#rosa-code .code-line {
+  display: flex;
+  gap: 10px;
+  padding: 2px 6px;
+  border-radius: 6px;
+}
+#rosa-code .line-no {
+  width: 28px;
+  text-align: right;
+  color: #94a3b8;
+  user-select: none;
+}
+#rosa-code .line-text {
+  white-space: pre;
+  flex: 1;
+}
+#rosa-code .code-line.active {
+  background: #dbeafe;
+  color: #1d4ed8;
+}
+#rosa-code .code-line.active .line-no {
+  color: #1d4ed8;
+}
+#rosa-code .tok-keyword {
+  color: #7c3aed;
+  font-weight: 600;
+}
+#rosa-code .tok-builtin {
+  color: #0ea5e9;
+}
+#rosa-code .tok-number {
+  color: #f97316;
+}
+#rosa-code .tok-string {
+  color: #10b981;
+}
+#rosa-code .tok-comment {
+  color: #64748b;
+  font-style: italic;
+}
+#rosa-code .code-token {
+  border-bottom: 1px dashed #94a3b8;
+  padding: 0 1px;
+}
+#rosa-code .code-token.active {
+  background: #fef3c7;
+  border-radius: 4px;
+  border-bottom-color: #f59e0b;
+}
+"""
+
+JS = """
+() => {
+  if (window.__rosaDemoInit) return;
+  window.__rosaDemoInit = true;
+
+  const rootId = "rosa-vis";
+  const stepsId = "steps_json";
+  const speedId = "speed_slider";
+  const baseDelay = 650;
+  const transferPauseMs = 200;
+  const transferMs = 750;
+  const debugAnim = true;
+  let runToken = 0;
+  let linkLayer = null;
+  let activeTokenEl = null;
+
+  function getAppRoot() {
+    const app = document.querySelector("gradio-app");
+    if (app && app.shadowRoot) return app.shadowRoot;
+    return document;
+  }
+
+  function findInputById(id, selector) {
+    const root = getAppRoot();
+    const direct = root.querySelector(`#${id}`) || document.getElementById(id);
+    if (!direct) return null;
+    if (direct.matches && direct.matches(selector)) return direct;
+    const nested = direct.querySelector(selector);
+    if (nested) return nested;
+    if (direct.shadowRoot) {
+      const shadowEl = direct.shadowRoot.querySelector(selector);
+      if (shadowEl) return shadowEl;
+    }
+    return null;
+  }
+
+  function getStepsBox() {
+    return findInputById(stepsId, "textarea, input");
+  }
+
+  function getSpeed() {
+    const slider = findInputById(speedId, 'input[type="range"], input');
+    const value = slider ? parseFloat(slider.value) : 1;
+    if (!Number.isFinite(value) || value <= 0) return 1;
+    return value;
+  }
+
+  function sleep(ms) {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+  }
+
+  function getOverlayHost() {
+    const root = getAppRoot();
+    if (root === document) return document.body;
+    return root;
+  }
+
+  function animateTransfer(fromEl, toEl, value, onFinish) {
+    if (!fromEl || !toEl) {
+      if (debugAnim) {
+        console.warn("[ROSA] animateTransfer missing element", {
+          hasFrom: !!fromEl,
+          hasTo: !!toEl,
+        });
+      }
+      return;
+    }
+    const fromRect = fromEl.getBoundingClientRect();
+    const toRect = toEl.getBoundingClientRect();
+    if (!fromRect || !toRect) {
+      if (debugAnim) {
+        console.warn("[ROSA] animateTransfer missing rect", {
+          fromRect,
+          toRect,
+        });
+      }
+      return;
+    }
+    const bubble = document.createElement("div");
+    bubble.className = "v-float";
+    bubble.textContent = value != null ? String(value) : fromEl.textContent;
+    bubble.style.position = "fixed";
+    bubble.style.zIndex = "9999";
+    bubble.style.borderRadius = "8px";
+    bubble.style.background = "#f59e0b";
+    bubble.style.color = "#1f2937";
+    bubble.style.display = "flex";
+    bubble.style.alignItems = "center";
+    bubble.style.justifyContent = "center";
+    bubble.style.fontWeight = "600";
+    bubble.style.boxShadow = "0 10px 24px rgba(15, 23, 42, 0.2)";
+    bubble.style.pointerEvents = "none";
+    bubble.style.transform = "translate(0px, 0px) scale(1)";
+    bubble.style.left = `${fromRect.left}px`;
+    bubble.style.top = `${fromRect.top}px`;
+    bubble.style.width = `${fromRect.width}px`;
+    bubble.style.height = `${fromRect.height}px`;
+    const host = getOverlayHost();
+    if (debugAnim) {
+      console.log("[ROSA] animateTransfer", {
+        from: {
+          left: fromRect.left,
+          top: fromRect.top,
+          width: fromRect.width,
+          height: fromRect.height,
+        },
+        to: {
+          left: toRect.left,
+          top: toRect.top,
+          width: toRect.width,
+          height: toRect.height,
+        },
+        host: host === document.body ? "document.body" : "shadowRoot",
+        value,
+      });
+      if (fromRect.width === 0 || fromRect.height === 0 || toRect.width === 0 || toRect.height === 0) {
+        console.warn("[ROSA] animateTransfer zero rect", {
+          fromRect,
+          toRect,
+        });
+      }
+    }
+    host.appendChild(bubble);
+    const dx = toRect.left - fromRect.left;
+    const dy = toRect.top - fromRect.top;
+    const toTransform = `translate(${dx}px, ${dy}px) scale(0.95)`;
+    const duration = transferMs;
+    const startAnimation = () => {
+      if (bubble.animate) {
+        const anim = bubble.animate(
+          [
+            { transform: "translate(0px, 0px) scale(1)" },
+            { transform: toTransform },
+          ],
+          { duration, easing: "ease-out", fill: "forwards" }
+        );
+        anim.addEventListener("finish", () => {
+          if (onFinish) onFinish();
+          bubble.remove();
+        });
+        return;
+      }
+      bubble.style.transition = `transform ${duration}ms ease`;
+      bubble.style.willChange = "transform";
+      requestAnimationFrame(() => {
+        bubble.style.transform = toTransform;
+      });
+      setTimeout(() => {
+        if (onFinish) onFinish();
+        bubble.remove();
+      }, duration + 80);
+    };
+    bubble.getBoundingClientRect();
+    requestAnimationFrame(() => {
+      requestAnimationFrame(startAnimation);
+    });
+  }
+
+  function getCodeLines() {
+    const root = getAppRoot();
+    const container = root.querySelector("#rosa-code") || document.getElementById("rosa-code");
+    if (!container) return {};
+    const lines = {};
+    container.querySelectorAll(".code-line").forEach((line) => {
+      const index = parseInt(line.dataset.line, 10);
+      if (Number.isFinite(index)) {
+        lines[index] = line;
+      }
+    });
+    return lines;
+  }
+
+  function resetCodeHighlight(codeLines) {
+    Object.values(codeLines).forEach((line) => {
+      line.classList.remove("active");
+    });
+  }
+
+  function setActiveCodeLine(state, line) {
+    if (!state.codeLines) return;
+    if (Number.isInteger(state.activeLine) && state.codeLines[state.activeLine]) {
+      state.codeLines[state.activeLine].classList.remove("active");
+    }
+    if (Number.isInteger(line) && state.codeLines[line]) {
+      state.codeLines[line].classList.add("active");
+      state.activeLine = line;
+    }
+  }
+
+  function buildRow(label, values, className) {
+    const row = document.createElement("div");
+    row.className = `rosa-row ${className || ""}`;
+    const labelEl = document.createElement("div");
+    labelEl.className = "row-label";
+    labelEl.textContent = label;
+    const cellsEl = document.createElement("div");
+    cellsEl.className = "row-cells";
+    const cells = values.map((val, idx) => {
+      const cell = document.createElement("div");
+      cell.className = "cell";
+      cell.textContent = String(val);
+      cell.dataset.idx = String(idx);
+      cellsEl.appendChild(cell);
+      return cell;
+    });
+    row.appendChild(labelEl);
+    row.appendChild(cellsEl);
+    return { row, cells };
+  }
+
+  function getCodeToken(rootEl, tokenKey) {
+    const root = rootEl || getAppRoot();
+    return root.querySelector(`[data-token="${tokenKey}"]`);
+  }
+
+  function ensureLinkLayer(shellEl) {
+    const host = shellEl || getAppRoot().querySelector("#rosa-shell") || document.body;
+    if (linkLayer && linkLayer.host === host) return linkLayer;
+    if (linkLayer && linkLayer.svg) {
+      linkLayer.svg.remove();
+    }
+    const svg = document.createElementNS("http://www.w3.org/2000/svg", "svg");
+    svg.setAttribute("width", "100%");
+    svg.setAttribute("height", "100%");
+    svg.style.position = "absolute";
+    svg.style.left = "0";
+    svg.style.top = "0";
+    svg.style.pointerEvents = "none";
+    svg.style.zIndex = "10000";
+    const defs = document.createElementNS("http://www.w3.org/2000/svg", "defs");
+    const marker = document.createElementNS("http://www.w3.org/2000/svg", "marker");
+    marker.setAttribute("id", "rosa-link-arrow");
+    marker.setAttribute("markerWidth", "8");
+    marker.setAttribute("markerHeight", "8");
+    marker.setAttribute("refX", "6");
+    marker.setAttribute("refY", "3");
+    marker.setAttribute("orient", "auto");
+    const markerPath = document.createElementNS("http://www.w3.org/2000/svg", "path");
+    markerPath.setAttribute("d", "M0,0 L6,3 L0,6 Z");
+    markerPath.setAttribute("fill", "#94a3b8");
+    marker.appendChild(markerPath);
+    defs.appendChild(marker);
+    svg.appendChild(defs);
+    const path = document.createElementNS("http://www.w3.org/2000/svg", "path");
+    path.setAttribute("stroke", "#94a3b8");
+    path.setAttribute("stroke-width", "1.5");
+    path.setAttribute("fill", "none");
+    path.setAttribute("marker-end", "url(#rosa-link-arrow)");
+    path.style.display = "none";
+    svg.appendChild(path);
+    host.appendChild(svg);
+    linkLayer = { svg, path, host };
+    return linkLayer;
+  }
+
+  function showLink(state, fromEl, tokenKey) {
+    if (!state || !state.shellEl) return;
+    const tokenEl = getCodeToken(state.shellEl, tokenKey);
+    if (!fromEl || !tokenEl) return;
+    const shellRect = state.shellEl.getBoundingClientRect();
+    const fromRect = fromEl.getBoundingClientRect();
+    const toRect = tokenEl.getBoundingClientRect();
+    const link = ensureLinkLayer(state.shellEl);
+    let startX = fromRect.right - shellRect.left;
+    let endX = toRect.left - shellRect.left;
+    if (toRect.left < fromRect.left) {
+      startX = fromRect.left - shellRect.left;
+      endX = toRect.right - shellRect.left;
+    }
+    const startY = fromRect.top + fromRect.height / 2 - shellRect.top;
+    const endY = toRect.top + toRect.height / 2 - shellRect.top;
+    link.path.setAttribute("d", `M ${startX} ${startY} L ${endX} ${endY}`);
+    link.path.style.display = "block";
+    if (activeTokenEl && activeTokenEl !== tokenEl) {
+      activeTokenEl.classList.remove("active");
+    }
+    tokenEl.classList.add("active");
+    activeTokenEl = tokenEl;
+  }
+
+  function hideLink() {
+    if (linkLayer) {
+      linkLayer.path.style.display = "none";
+    }
+    if (activeTokenEl) {
+      activeTokenEl.classList.remove("active");
+      activeTokenEl = null;
+    }
+  }
+
+  function getRangeRect(cells, start, end) {
+    const rects = [];
+    for (let idx = start; idx <= end; idx += 1) {
+      const cell = cells[idx];
+      if (!cell) continue;
+      rects.push(cell.getBoundingClientRect());
+    }
+    if (!rects.length) return null;
+    let left = rects[0].left;
+    let right = rects[0].right;
+    let top = rects[0].top;
+    let bottom = rects[0].bottom;
+    rects.forEach((rect) => {
+      left = Math.min(left, rect.left);
+      right = Math.max(right, rect.right);
+      top = Math.min(top, rect.top);
+      bottom = Math.max(bottom, rect.bottom);
+    });
+    return {
+      left,
+      top,
+      right,
+      bottom,
+      width: right - left,
+      height: bottom - top,
+    };
+  }
+
+  function toLocalRect(rect, containerRect) {
+    return {
+      left: rect.left - containerRect.left,
+      top: rect.top - containerRect.top,
+      width: rect.width,
+      height: rect.height,
+      right: rect.right - containerRect.left,
+      bottom: rect.bottom - containerRect.top,
+    };
+  }
+
+  function padRect(rect, pad, bounds) {
+    const left = Math.max(0, rect.left - pad);
+    const top = Math.max(0, rect.top - pad);
+    const right = Math.min(bounds.width, rect.left + rect.width + pad);
+    const bottom = Math.min(bounds.height, rect.top + rect.height + pad);
+    return {
+      left,
+      top,
+      width: Math.max(0, right - left),
+      height: Math.max(0, bottom - top),
+      right,
+      bottom,
+    };
+  }
+
+  function createOverlayBox(layer, rect, label, className, labelClass) {
+    const box = document.createElement("div");
+    box.className = `overlay-box ${className || ""}`;
+    box.style.left = `${rect.left}px`;
+    box.style.top = `${rect.top}px`;
+    box.style.width = `${rect.width}px`;
+    box.style.height = `${rect.height}px`;
+    const labelEl = document.createElement("div");
+    labelEl.className = `overlay-label ${labelClass || ""}`;
+    labelEl.textContent = label;
+    box.appendChild(labelEl);
+    box.dataset.label = label;
+    layer.appendChild(box);
+  }
+
+  function clearHoverBoxes(state) {
+    if (state.overlayHover) {
+      state.overlayHover.innerHTML = "";
+    }
+    state.hoverBoxes = [];
+    hideLink();
+  }
+
+  function clearOverlay(state) {
+    if (state.overlayBoxes) {
+      state.overlayBoxes.innerHTML = "";
+    }
+    if (state.rayLine) {
+      state.rayLine.style.display = "none";
+    }
+    clearHoverBoxes(state);
+  }
+
+  function updateOverlay(state, step) {
+    const hasI = Number.isInteger(step.i);
+    const hasW = Number.isInteger(step.w);
+    const hasJ = Number.isInteger(step.j);
+    if (step.phase === "loop_i") {
+      state.lastOverlay = null;
+    }
+    const showOverlay = [
+      "loop_w",
+      "assign_t",
+      "loop_j",
+      "try",
+      "assign",
+      "break_inner",
+      "check",
+      "break_outer",
+    ].includes(step.phase);
+    let overlay = null;
+    if (hasI && hasW) {
+      const sameWindow =
+        state.lastOverlay &&
+        state.lastOverlay.i === step.i &&
+        state.lastOverlay.w === step.w;
+      const jValue = hasJ ? step.j : (sameWindow ? state.lastOverlay.j : null);
+      state.lastOverlay = { i: step.i, w: step.w, j: jValue };
+      overlay = state.lastOverlay;
+    } else if (state.lastOverlay) {
+      overlay = state.lastOverlay;
+    }
+    clearOverlay(state);
+    if (!showOverlay) {
+      return;
+    }
+    if (!overlay) return;
+    const tStart = overlay.i - overlay.w + 1;
+    const tEnd = overlay.i;
+    const kStart = overlay.j;
+    const kEnd = Number.isInteger(overlay.j) ? overlay.j + overlay.w - 1 : null;
+    const tRect = getRangeRect(state.qCells, tStart, tEnd);
+    const kRect =
+      Number.isInteger(kStart) && Number.isInteger(kEnd)
+        ? getRangeRect(state.kCells, kStart, kEnd)
+        : null;
+    const vIndex = Number.isInteger(kStart) ? kStart + overlay.w : null;
+    const vCell = Number.isInteger(vIndex) ? state.vCells[vIndex] : null;
+    const vRect = vCell ? vCell.getBoundingClientRect() : null;
+    const cardRect = state.cardEl.getBoundingClientRect();
+    const cardStyle = window.getComputedStyle(state.cardEl);
+    const borderLeft = parseFloat(cardStyle.borderLeftWidth) || 0;
+    const borderTop = parseFloat(cardStyle.borderTopWidth) || 0;
+    const borderRight = parseFloat(cardStyle.borderRightWidth) || 0;
+    const borderBottom = parseFloat(cardStyle.borderBottomWidth) || 0;
+    const originLeft = cardRect.left + borderLeft;
+    const originTop = cardRect.top + borderTop;
+    const boxPad = 4;
+    const bounds = {
+      width: cardRect.width - borderLeft - borderRight,
+      height: cardRect.height - borderTop - borderBottom,
+    };
+    const toOverlayRect = (rect) => ({
+      left: rect.left - originLeft,
+      top: rect.top - originTop,
+      width: rect.width,
+      height: rect.height,
+      right: rect.right - originLeft,
+      bottom: rect.bottom - originTop,
+    });
+    const isTry = step.phase === "try";
+    const tryClass = isTry ? (step.matched ? "try-match" : "try-miss") : "";
+
+    if (tRect) {
+      const local = padRect(toOverlayRect(tRect), boxPad, bounds);
+      createOverlayBox(state.overlayBoxes, local, "t", `t-box ${tryClass}`, "t-label");
+      const hoverBox = document.createElement("div");
+      hoverBox.className = "overlay-hover-box";
+      hoverBox.style.left = `${local.left}px`;
+      hoverBox.style.top = `${local.top}px`;
+      hoverBox.style.width = `${local.width}px`;
+      hoverBox.style.height = `${local.height}px`;
+      hoverBox.addEventListener("mouseenter", () => showLink(state, hoverBox, "t"));
+      hoverBox.addEventListener("mouseleave", hideLink);
+      state.overlayHover.appendChild(hoverBox);
+      state.hoverBoxes.push(hoverBox);
+    }
+
+    if (kRect) {
+      const local = padRect(toOverlayRect(kRect), boxPad, bounds);
+      createOverlayBox(state.overlayBoxes, local, "kkk[j:j+w]", `k-box ${tryClass}`, "k-label");
+      const hoverBox = document.createElement("div");
+      hoverBox.className = "overlay-hover-box";
+      hoverBox.style.left = `${local.left}px`;
+      hoverBox.style.top = `${local.top}px`;
+      hoverBox.style.width = `${local.width}px`;
+      hoverBox.style.height = `${local.height}px`;
+      hoverBox.addEventListener("mouseenter", () => showLink(state, hoverBox, "k"));
+      hoverBox.addEventListener("mouseleave", hideLink);
+      state.overlayHover.appendChild(hoverBox);
+      state.hoverBoxes.push(hoverBox);
+    }
+
+    if (step.phase === "assign" && kRect && vRect && state.rayLine) {
+      const kLocal = padRect(toOverlayRect(kRect), boxPad, bounds);
+      const vLocal = toOverlayRect(vRect);
+      const startX = kLocal.right;
+      const startY = kLocal.bottom;
+      const endX = vLocal.left + vLocal.width / 2;
+      const endY = vLocal.top + vLocal.height / 2;
+      state.rayLine.setAttribute("x1", startX);
+      state.rayLine.setAttribute("y1", startY);
+      state.rayLine.setAttribute("x2", endX);
+      state.rayLine.setAttribute("y2", endY);
+      state.rayLine.style.display = "block";
+    }
+  }
+
+  function render(data) {
+    const root = getAppRoot().querySelector(`#${rootId}`);
+    if (!root) return null;
+    root.innerHTML = "";
+
+    const card = document.createElement("div");
+    card.className = "rosa-card";
+
+    const legend = document.createElement("div");
+    legend.className = "rosa-legend";
+    legend.innerHTML = `
+      <span class="legend-item"><span class="legend-dot legend-suffix"></span>Current suffix</span>
+      <span class="legend-item"><span class="legend-dot legend-window"></span>k window</span>
+      <span class="legend-item"><span class="legend-dot legend-match"></span>Match</span>
+      <span class="legend-item"><span class="legend-dot legend-v"></span>Read v</span>
+      <span class="legend-item"><span class="legend-dot legend-out"></span>Output</span>
+    `;
+    card.appendChild(legend);
+
+    const rowsWrap = document.createElement("div");
+    rowsWrap.className = "rosa-rows";
+
+    const indexRow = document.createElement("div");
+    indexRow.className = "rosa-row index-row";
+    const indexLabel = document.createElement("div");
+    indexLabel.className = "row-label";
+    indexLabel.textContent = "#";
+    const indexCells = document.createElement("div");
+    indexCells.className = "row-cells index-cells";
+    data.q.forEach((_, idx) => {
+      const cell = document.createElement("div");
+      cell.className = "index-cell";
+      cell.textContent = String(idx);
+      indexCells.appendChild(cell);
+    });
+    indexRow.appendChild(indexLabel);
+    indexRow.appendChild(indexCells);
+    rowsWrap.appendChild(indexRow);
+
+    const qRow = buildRow("q", data.q, "q-row");
+    const kRow = buildRow("k", data.k, "k-row");
+    const vRow = buildRow("v", data.v, "v-row");
+    const outRow = buildRow("out", data.q.map(() => "."), "out-row");
+
+    rowsWrap.appendChild(qRow.row);
+    rowsWrap.appendChild(kRow.row);
+    rowsWrap.appendChild(vRow.row);
+    rowsWrap.appendChild(outRow.row);
+
+    card.appendChild(rowsWrap);
+    const overlay = document.createElement("div");
+    overlay.className = "rosa-overlay";
+    const overlayRay = document.createElementNS("http://www.w3.org/2000/svg", "svg");
+    overlayRay.classList.add("overlay-ray");
+    const rayDefs = document.createElementNS("http://www.w3.org/2000/svg", "defs");
+    const rayMarker = document.createElementNS("http://www.w3.org/2000/svg", "marker");
+    rayMarker.setAttribute("id", "rosa-ray-head");
+    rayMarker.setAttribute("markerWidth", "8");
+    rayMarker.setAttribute("markerHeight", "8");
+    rayMarker.setAttribute("refX", "6");
+    rayMarker.setAttribute("refY", "3");
+    rayMarker.setAttribute("orient", "auto");
+    const rayMarkerPath = document.createElementNS("http://www.w3.org/2000/svg", "path");
+    rayMarkerPath.setAttribute("d", "M0,0 L6,3 L0,6 Z");
+    rayMarkerPath.setAttribute("fill", "var(--rosa-cyan)");
+    rayMarker.appendChild(rayMarkerPath);
+    rayDefs.appendChild(rayMarker);
+    overlayRay.appendChild(rayDefs);
+    const rayLine = document.createElementNS("http://www.w3.org/2000/svg", "line");
+    rayLine.classList.add("overlay-ray-line");
+    rayLine.setAttribute("marker-end", "url(#rosa-ray-head)");
+    rayLine.style.display = "none";
+    overlayRay.appendChild(rayLine);
+    const overlayBoxes = document.createElement("div");
+    overlayBoxes.className = "overlay-box-layer";
+    overlay.appendChild(overlayBoxes);
+    overlay.appendChild(overlayRay);
+    card.appendChild(overlay);
+    const overlayHover = document.createElement("div");
+    overlayHover.className = "overlay-hover-layer";
+    card.appendChild(overlayHover);
+    root.appendChild(card);
+
+    const codeLines = getCodeLines();
+    resetCodeHighlight(codeLines);
+
+    return {
+      shellEl: root.closest("#rosa-shell") || getAppRoot().querySelector("#rosa-shell"),
+      cardEl: card,
+      qCells: qRow.cells,
+      kCells: kRow.cells,
+      vCells: vRow.cells,
+      outCells: outRow.cells,
+      outFixed: new Set(),
+      outPending: new Set(),
+      overlayBoxes,
+      rayLine,
+      overlayHover,
+      hoverBoxes: [],
+      lastOverlay: null,
+      codeLines,
+      activeLine: null,
+    };
+  }
+
+  function clearHighlights(state) {
+    const holdActive =
+      state.vHold &&
+      Number.isInteger(state.vHold.index) &&
+      performance.now() < state.vHold.until;
+    const holdIndex = holdActive ? state.vHold.index : null;
+    const all = [...state.qCells, ...state.kCells, ...state.vCells, ...state.outCells];
+    all.forEach((cell) => {
+      cell.classList.remove("active", "suffix", "k-window", "v-pick", "out");
+    });
+    if (holdActive && state.vCells[holdIndex]) {
+      state.vCells[holdIndex].classList.add("v-pick");
+    }
+  }
+
+  function applyStep(state, step) {
+    clearHighlights(state);
+    const { qCells, kCells, vCells, outCells } = state;
+    if (Number.isInteger(step.line)) {
+      setActiveCodeLine(state, step.line);
+    }
+
+    const showWindow = [
+      "loop_w",
+      "assign_t",
+      "loop_j",
+      "try",
+      "assign",
+      "break_inner",
+      "check",
+      "break_outer",
+    ].includes(step.phase);
+    if (showWindow && Number.isInteger(step.i) && qCells[step.i]) {
+      qCells[step.i].classList.add("active");
+    }
+    if (showWindow && Number.isInteger(step.w)) {
+      for (let idx = step.i - step.w + 1; idx <= step.i; idx += 1) {
+        if (qCells[idx]) qCells[idx].classList.add("suffix");
+      }
+    }
+    if (showWindow && Number.isInteger(step.j)) {
+      for (let idx = step.j; idx <= step.j + step.w - 1; idx += 1) {
+        if (kCells[idx]) kCells[idx].classList.add("k-window");
+      }
+    }
+    updateOverlay(state, step);
+
+    if (step.phase === "try") {
+      return;
+    }
+
+    if (step.phase === "assign") {
+      if (debugAnim) {
+        console.log("[ROSA] assign step", {
+          i: step.i,
+          j: step.j,
+          w: step.w,
+          v_index: step.v_index,
+          value: step.value,
+        });
+      }
+      if (Number.isInteger(step.v_index) && vCells[step.v_index]) {
+        vCells[step.v_index].classList.add("v-pick");
+      }
+      if (
+        Number.isInteger(step.v_index) &&
+        Number.isInteger(step.i) &&
+        vCells[step.v_index] &&
+        outCells[step.i]
+      ) {
+        const holdToken = state.runToken;
+        if (state.outPending) {
+          state.outPending.add(step.i);
+        }
+        const totalWait = transferPauseMs + transferMs;
+        if (state.outPending) {
+          state.outPending.add(step.i);
+        }
+        const startTransfer = () => {
+          if (state.runToken !== holdToken) return;
+          animateTransfer(vCells[step.v_index], outCells[step.i], step.value, () => {
+            if (state.runToken !== holdToken) return;
+            if (state.outPending) {
+              state.outPending.delete(step.i);
+            }
+            if (state.outFixed) {
+              state.outFixed.add(step.i);
+            }
+            const outCell = outCells[step.i];
+            if (outCell) {
+              outCell.textContent = String(step.value);
+              outCell.classList.add("out-fixed", "filled");
+            }
+          });
+        };
+        if (transferPauseMs > 0) {
+          setTimeout(startTransfer, transferPauseMs);
+        } else {
+          startTransfer();
+        }
+        state.vHold = {
+          index: step.v_index,
+          until: performance.now() + totalWait,
+        };
+        const holdIndex = step.v_index;
+        setTimeout(() => {
+          if (state.runToken !== holdToken) return;
+          if (!state.vHold || state.vHold.index !== holdIndex) return;
+          if (performance.now() < state.vHold.until) return;
+          const cell = state.vCells[holdIndex];
+          if (cell) cell.classList.remove("v-pick");
+        }, totalWait + 60);
+        return totalWait;
+      }
+      return 0;
+    }
+
+    if (step.phase === "break_inner") {
+      return;
+    }
+
+    if (step.phase === "check") {
+      return;
+    }
+
+    if (step.phase === "break_outer") {
+      return;
+    }
+
+    if (step.phase === "output") {
+      if (
+        (state.outPending && state.outPending.has(step.i)) ||
+        (state.outFixed && state.outFixed.has(step.i))
+      ) {
+        return;
+      }
+      if (outCells[step.i]) {
+        outCells[step.i].textContent = String(step.value);
+        outCells[step.i].classList.add("out", "out-fixed", "filled");
+        if (state.outFixed) {
+          state.outFixed.add(step.i);
+        }
+      }
+      return;
+    }
+
+    if (step.phase === "return") {
+      return;
+    }
+  }
+
+  async function play(state, steps, token) {
+    if (!steps || !steps.length) return;
+    for (let idx = 0; idx < steps.length; idx += 1) {
+      if (token !== runToken) return;
+      const extraWait = applyStep(state, steps[idx]);
+      const delay = baseDelay / getSpeed();
+      const waitMs = Math.max(delay, Number.isFinite(extraWait) ? extraWait : 0);
+      await sleep(waitMs);
+    }
+  }
+
+  function start(data) {
+    runToken += 1;
+    const token = runToken;
+    const state = render(data);
+    if (!state) return;
+    state.runToken = token;
+    state.vHold = null;
+    if (state.outFixed) state.outFixed.clear();
+    if (state.outPending) state.outPending.clear();
+    play(state, data.steps, token);
+  }
+
+  let lastValue = "";
+  setInterval(() => {
+    const box = getStepsBox();
+    if (!box) return;
+    const value = box.value || "";
+    if (value && value !== lastValue) {
+      lastValue = value;
+      try {
+        const data = JSON.parse(value);
+        start(data);
+      } catch (err) {
+        console.error("Invalid ROSA steps payload", err);
+      }
+    }
+  }, 300);
+}
+"""
+
+
+with gr.Blocks(css=CSS, js=JS) as demo:
+    gr.HTML(
+        '<div class="page-header">'
+        '<div class="page-title">ROSA QKV Demo</div>'
+        '<div class="page-subtitle">Enter or randomize q/k/v (0/1 only). Click Start Demo to see suffix matching and retrieval.</div>'
+        "</div>"
+    )
+
+    with gr.Row():
+        q_text = gr.Textbox(label="q sequence", value="01010101010101010101", lines=1)
+        k_text = gr.Textbox(label="k sequence", value="10100110100110100110", lines=1)
+        v_text = gr.Textbox(label="v sequence", value="11001100110011001100", lines=1)
+
+    with gr.Row():
+        length = gr.Slider(4, 20, value=20, step=1, label="Random length")
+        random_btn = gr.Button("Randomize")
+        demo_btn = gr.Button("Start Demo", variant="primary")
+        speed = gr.Slider(
+            0.1,
+            5.0,
+            value=2.0,
+            step=0.05,
+            label="Playback speed",
+            elem_id="speed_slider",
+            interactive=True,
+        )
+
+    out_text = gr.Textbox(label="Output", interactive=False)
+    steps_json = gr.Textbox(visible=False, elem_id="steps_json")
+
+    gr.HTML(
+        f'<div id="rosa-shell" class="rosa-shell">'
+        f'<div class="rosa-pane"><div id="rosa-vis"></div></div>'
+        f'<div class="rosa-code-pane">{CODE_HTML}</div>'
+        f"</div>"
+    )
+
+    random_btn.click(on_random, inputs=[length], outputs=[q_text, k_text, v_text])
+    demo_btn.click(on_demo, inputs=[q_text, k_text, v_text], outputs=[steps_json, out_text])
+    demo.load(on_demo, inputs=[q_text, k_text, v_text], outputs=[steps_json, out_text])
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1 @@
+gradio>=5.0.0,<6.0.0