Update app.py

app.py

@@ -2,7 +2,8 @@
 import json
 from pathlib import Path
 import threading, time
-
+import anywidget
+import traitlets as t
 import solara
 import pandas as pd
 import plotly.graph_objects as go
@@ -236,52 +237,92 @@ def AutoPredictWatcher():
     solara.use_effect(effect, [text, auto])
     return solara.Text("", style={"display": "none"})
 
-# ---------- Predictions list ----------
+class HoverList(anywidget.AnyWidget):
+    """
+    Renders the prediction rows in the browser and streams hover/click
+    events back to Python via synced traitlets.
+    """
+    # Browser code: builds the list and wires events
+    _esm = """
+    export function render({ model, el }) {
+      const make = () => {
+        const items = model.get('items') || [];
+        el.innerHTML = "";
+        const wrap = document.createElement('div');
+        wrap.style.display = 'flex';
+        wrap.style.flexDirection = 'column';
+        items.forEach(({tid, label}) => {
+          const btn = document.createElement('button');
+          btn.textContent = label;
+          btn.className = 'rowbtn';           // your existing CSS
+          btn.setAttribute('type', 'button');
+          btn.setAttribute('role', 'button');
+          btn.setAttribute('tabindex', '0');
+
+          // hover → preview
+          const preview = () => {
+            model.set('hovered_id', tid);
+            model.save_changes();
+          };
+          btn.addEventListener('mouseenter', preview);
+          btn.addEventListener('mouseover',  preview);
+          btn.addEventListener('mousemove',  preview);
+          btn.addEventListener('focus',      preview);
+
+          // click → append
+          btn.addEventListener('click', () => {
+            model.set('clicked_id', tid);
+            model.save_changes();
+          });
+
+          wrap.appendChild(btn);
+        });
+        el.appendChild(wrap);
+      };
+
+      // initial render
+      make();
+
+      // re-render when items change
+      model.on('change:items', make);
+    }
+    """
+    # Data flowing between JS and Python
+    items       = t.List(trait=t.Dict()).tag(sync=True)   # [{tid:int, label:str}, ...]
+    hovered_id  = t.Int(allow_none=True).tag(sync=True)
+    clicked_id  = t.Int(allow_none=True).tag(sync=True)
+
+
 @solara.component
 def PredictionsList():
-    df = preds_rx.value
-    with solara.Column(gap="6px", style={"maxWidth": "720px"}):
-        solara.Markdown("### Prediction")
-        solara.Text(
-            " #  probs    token  predicted next token",
-            style={
-                "color": "var(--muted)",
-                "fontFamily": 'ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace',
-            },
-        )
+    token_probs = token_probs_rx.value or []
+    selected_token_id = selected_token_id_rx.value
 
-        for i, row in df.iterrows():
-            tid = int(row["id"])
-            prob = row["probs"]
-            tok_disp = display_token_from_id(tid)
-            row_label = fmt_row(i, prob, tid, tok_disp)
+    with solara.Column():
+        for tid, prob in token_probs[:10]:
+            row_label = f"{id_to_token.get(tid, tid)} ({prob:.2f})"
+            is_selected = tid == selected_token_id
 
-            # Wrapper DIV catches pointer/hover; Button handles click (and also hooks hover)
             with solara.Div(
-                classes=["rowbtn"],  # styling on the wrapper
-                style={"justifyContent": "flex-start", "width": "100%"},
-                attributes={"tabindex": "0", "role": "button"},
-                # ---- HOVER = preview neighborhood (wrapper) ----
-                on_mouse_enter=lambda *args, tid=tid: preview_token(tid),
-                on_mouse_over=lambda *args, tid=tid: preview_token(tid),
-                on_mouse_move=lambda *args, tid=tid: preview_token(tid),
-                on_pointer_enter=lambda *args, tid=tid: preview_token(tid),
-                on_pointer_move=lambda *args, tid=tid: preview_token(tid),
-                on_focus=lambda *args, tid=tid: preview_token(tid),   # keyboard focus
+                style={
+                    "background": "#333" if is_selected else "transparent",
+                    "padding": "2px",
+                    "borderRadius": "4px",
+                },
+                on_mouse_enter=lambda tid=tid: preview_token(tid),
+                on_mouse_leave=lambda tid=tid: preview_token(None),
             ):
                 solara.Button(
                     row_label,
-                    classes=[],  # keep wrapper styled; leave button unstyled
-                    style={"justifyContent": "flex-start", "width": "100%"},
-                    # ---- CLICK = append token to text ----
                     on_click=lambda *args, tid=tid: append_token(tid),
-                    # Redundant hover hooks on the button too (helps on some builds)
-                    on_mouse_enter=lambda *args, tid=tid: preview_token(tid),
-                    on_mouse_over=lambda *args, tid=tid: preview_token(tid),
-                    on_mouse_move=lambda *args, tid=tid: preview_token(tid),
-                    on_pointer_enter=lambda *args, tid=tid: preview_token(tid),
-                    on_pointer_move=lambda *args, tid=tid: preview_token(tid),
-                    on_focus=lambda *args, tid=tid: preview_token(tid),
+                    style={
+                        "justifyContent": "flex-start",
+                        "width": "100%",
+                        "textAlign": "left",
+                        "color": "#38bdf8" if is_selected else "white",
+                    },
+                    on_mouse_enter=lambda tid=tid: preview_token(tid),
+                    on_mouse_leave=lambda tid=tid: preview_token(None),
                 )
 
 # ---------- Page ----------