Sync playback speed and add quick-code panel

app.py

@@ -10,7 +10,7 @@ MAX_DEMO_LEN = 20
 GRADIO_MAJOR = int((gr.__version__ or "0").split(".", maxsplit=1)[0])
 
 ROSA_CODE = """
-def rosa_qkv(qqq, kkk, vvv):
+def rosa_qkv_naive(qqq, kkk, vvv):
     n=len(qqq); out=[-1]*n
     for i in range(n):
         for w in range(i+1,0,-1):
@@ -22,7 +22,34 @@ def rosa_qkv(qqq, kkk, vvv):
             if out[i]!=-1:
                 break
     return out
-""".strip("\n")
+""".strip(
+    "\n"
+)
+
+ROSA_QUICK_CODE = """
+def rosa_qkv_ref_minus1(qqq, kkk, vvv): # note: input will never contain "-1"
+    n=len(qqq); y=[-1]*n; s=2*n+1; t=[None]*s; f=[-1]*s; m=[0]*s; r=[-1]*s; t[0]={}; g=0; u=1; w=h=0; assert n==len(kkk)==len(vvv)
+    for i,(q,k) in enumerate(zip(qqq,kkk)):
+        p,x=w,h
+        while p!=-1 and q not in t[p]: x=m[p] if x>m[p] else x; p=f[p]
+        p,x=(t[p][q],x+1) if p!=-1 else (0,0); v=p
+        while f[v]!=-1 and m[f[v]]>=x: v=f[v]
+        while v!=-1 and (m[v]<=0 or r[v]<0): v=f[v]
+        y[i]=vvv[r[v]+1] if v!=-1 else -1; w,h=p,x; j=u; u+=1; t[j]={}; m[j]=m[g]+1; p=g
+        while p!=-1 and k not in t[p]: t[p][k]=j; p=f[p]
+        if p==-1: f[j]=0
+        else:
+            d=t[p][k]
+            if m[p]+1==m[d]: f[j]=d
+            else:
+                b=u; u+=1; t[b]=t[d].copy(); m[b]=m[p]+1; f[b]=f[d]; r[b]=r[d]; f[d]=f[j]=b
+                while p!=-1 and t[p][k]==d: t[p][k]=b; p=f[p]
+        v=g=j
+        while v!=-1 and r[v]<i: r[v]=i; v=f[v]
+    return y
+""".strip(
+    "\n"
+)
 
 
 KEYWORDS = {
@@ -30,11 +57,16 @@ KEYWORDS = {
     "for",
     "in",
     "if",
+    "else",
+    "while",
     "break",
     "return",
     "assert",
+    "None",
+    "True",
+    "False",
 }
-BUILTINS = {"len", "range"}
+BUILTINS = {"len", "range", "zip", "enumerate"}
 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.])")
@@ -130,15 +162,9 @@ def build_code_html(code: str) -> str:
         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>'
-        )
+        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>'
@@ -149,6 +175,21 @@ def build_code_html(code: str) -> str:
     return "\n".join(rendered)
 
 
+def build_plain_code_html(code: str, block_id: str) -> str:
+    lines = code.splitlines()
+    rendered = [f'<div id="{block_id}" class="rosa-code">']
+    for index, line in enumerate(lines, start=1):
+        highlighted = highlight_python_line(line)
+        rendered.append(
+            '<div class="code-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:
@@ -157,6 +198,8 @@ def find_line_number(lines: list[str], needle: str, fallback: int = 1) -> int:
 
 
 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)):
@@ -176,6 +219,7 @@ 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)
+QUICK_CODE_HTML = build_plain_code_html(ROSA_QUICK_CODE, "rosa-code-quick")
 
 
 def parse_bits(text: str, name: str) -> list[int]:
@@ -350,6 +394,28 @@ CSS = """
   flex: 2 1 320px;
   min-width: 300px;
 }
+.quick-code-details {
+  margin-top: 8px;
+}
+.quick-code-details > summary {
+  cursor: pointer;
+  user-select: none;
+  padding: 8px 10px;
+  border: 1px dashed #cbd5e1;
+  border-radius: 12px;
+  background: #f8fafc;
+  color: #0f172a;
+  font-weight: 600;
+}
+.quick-code-details[open] > summary {
+  margin-bottom: 10px;
+}
+.quick-code-details > summary::-webkit-details-marker {
+  display: none;
+}
+.quick-code-details .rosa-code {
+  max-height: 420px;
+}
 #rosa-vis {
   --rosa-blue: #3b82f6;
   --rosa-sky: #38bdf8;
@@ -599,7 +665,7 @@ CSS = """
   color: #64748b;
   font-variant-numeric: tabular-nums;
 }
-#rosa-code {
+.rosa-code {
   background: #ffffff;
   border: 1px solid #e2e8f0;
   border-radius: 12px;
@@ -610,51 +676,53 @@ CSS = """
   max-height: 520px;
   overflow: auto;
 }
-#rosa-code .code-line {
+.rosa-code .code-line {
   display: flex;
   gap: 10px;
   padding: 2px 6px;
   border-radius: 6px;
 }
-#rosa-code .line-no {
-  width: 28px;
-  text-align: right;
+.rosa-code .line-no {
+  flex: 0 0 36px;
+  width: 36px;
+  text-align: left;
   color: #94a3b8;
   user-select: none;
+  font-variant-numeric: tabular-nums;
 }
-#rosa-code .line-text {
+.rosa-code .line-text {
   white-space: pre;
   flex: 1;
 }
-#rosa-code .code-line.active {
+.rosa-code .code-line.active {
   background: #dbeafe;
   color: #1d4ed8;
 }
-#rosa-code .code-line.active .line-no {
+.rosa-code .code-line.active .line-no {
   color: #1d4ed8;
 }
-#rosa-code .tok-keyword {
+.rosa-code .tok-keyword {
   color: #7c3aed;
   font-weight: 600;
 }
-#rosa-code .tok-builtin {
+.rosa-code .tok-builtin {
   color: #0ea5e9;
 }
-#rosa-code .tok-number {
+.rosa-code .tok-number {
   color: #f97316;
 }
-#rosa-code .tok-string {
+.rosa-code .tok-string {
   color: #10b981;
 }
-#rosa-code .tok-comment {
+.rosa-code .tok-comment {
   color: #64748b;
   font-style: italic;
 }
-#rosa-code .code-token {
+.rosa-code .code-token {
   border-bottom: 1px dashed #94a3b8;
   padding: 0 1px;
 }
-#rosa-code .code-token.active {
+.rosa-code .code-token.active {
   background: #fef3c7;
   border-radius: 4px;
   border-bottom-color: #f59e0b;
@@ -703,7 +771,8 @@ JS_FUNC = """
 
   function getSpeed() {
     const slider = findInputById(speedId, 'input[type="range"], input');
-    const value = slider ? parseFloat(slider.value) : 1;
+    // If the speed slider isn't mounted yet (e.g. on first load), default to x2.
+    const value = slider ? parseFloat(slider.value) : 2;
     if (!Number.isFinite(value) || value <= 0) return 1;
     return value;
   }
@@ -718,7 +787,7 @@ JS_FUNC = """
     return root;
   }
 
-  function animateTransfer(fromEl, toEl, value, onFinish) {
+  function animateTransfer(fromEl, toEl, value, durationMs, onFinish) {
     if (!fromEl || !toEl) {
       if (debugAnim) {
         console.warn("[ROSA] animateTransfer missing element", {
@@ -787,7 +856,8 @@ JS_FUNC = """
     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 duration =
+      Number.isFinite(durationMs) && durationMs > 0 ? durationMs : transferMs;
     const startAnimation = () => {
       if (bubble.animate) {
         const anim = bubble.animate(
@@ -1164,11 +1234,11 @@ JS_FUNC = """
     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>
+      <span class="legend-item"><span class="legend-dot legend-suffix"></span>Current suffix (t)</span>
+      <span class="legend-item"><span class="legend-dot legend-window"></span>k window (kkk[j:j+w])</span>
+      <span class="legend-item"><span class="legend-dot legend-match"></span>Match (kkk[j:j+w]==t)</span>
+      <span class="legend-item"><span class="legend-dot legend-v"></span>Read v (vvv[j+w])</span>
+      <span class="legend-item"><span class="legend-dot legend-out"></span>Output (out)</span>
     `;
     card.appendChild(legend);
 
@@ -1328,17 +1398,20 @@ JS_FUNC = """
         vCells[step.v_index] &&
         outCells[step.i]
       ) {
+        const speed = getSpeed();
         const holdToken = state.runToken;
         if (state.outPending) {
           state.outPending.add(step.i);
         }
-        const totalWait = transferPauseMs + transferMs;
+        const pauseMs = transferPauseMs / speed;
+        const durationMs = transferMs / speed;
+        const totalWait = pauseMs + durationMs;
         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, () => {
+          animateTransfer(vCells[step.v_index], outCells[step.i], step.value, durationMs, () => {
             if (state.runToken !== holdToken) return;
             if (state.outPending) {
               state.outPending.delete(step.i);
@@ -1353,8 +1426,8 @@ JS_FUNC = """
             }
           });
         };
-        if (transferPauseMs > 0) {
-          setTimeout(startTransfer, transferPauseMs);
+        if (pauseMs > 0) {
+          setTimeout(startTransfer, pauseMs);
         } else {
           startTransfer();
         }
@@ -1458,8 +1531,8 @@ demo_context = gr.Blocks(css=CSS, js=JS_BOOT)
 with demo_context 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 class="page-title">RWKV-8 ROSA-QKV-1bit Demo</div>'
+        '<div class="page-subtitle">This is using naive algorithm (not suffix automaton). Enter or randomize q/k/v (0/1 only), then click [Start Demo].</div>'
         "</div>"
     )
 
@@ -1474,7 +1547,7 @@ with demo_context as demo:
         demo_btn = gr.Button("Start Demo", variant="primary")
         speed = gr.Slider(
             0.1,
-            5.0,
+            10.0,
             value=2.0,
             step=0.05,
             label="Playback speed",
@@ -1488,7 +1561,13 @@ with demo_context as demo:
     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 class="rosa-code-pane">'
+        f"{CODE_HTML}"
+        f'<details class="quick-code-details">'
+        f"<summary>Fast version (click to expand)</summary>"
+        f"{QUICK_CODE_HTML}"
+        f"</details>"
+        f"</div>"
         f"</div>"
     )