Update app.py

app.py

@@ -29,7 +29,8 @@ theme_css = """
 }
 
 /* Base */
-body{ background:var(--bg); color:var(--text); }
+body{ background:var(--bg); color:var(--text); margin:0;}
+h1{ margin:6px 0 8px; }   
 
 /* Two-column layout */
 .app-row { display:flex; align-items:flex-start; gap:16px; }        /* was 24px */
@@ -53,6 +54,18 @@ body{ background:var(--bg); color:var(--text); }
 }
 .rowbtn:hover{ background:#f7fbff; border-color:#c3e8fb; }
 
+/* New: 4-column grid inside each row button */
+.rowbtn-grid{
+  display:grid;
+  grid-template-columns: 28px 72px 72px 1fr; /* # | probs | tokenID | token */
+  column-gap:8px;
+  align-items:center;
+  width:100%;
+  font-family: var(--mono);
+  font-size:13px;
+  line-height:1.15;
+}
+
 /* Neighbor chips (smaller) */
 .badge{
   display:inline-block; padding:2px 6px;       /* was 2px 8px */
@@ -241,7 +254,7 @@ def AutoPredictWatcher():
 class HoverList(anywidget.AnyWidget):
     """
     Renders the prediction rows in the browser and streams hover/click events
-    back to Python via synced traitlets.
+    back to Python via synced traitlets. Supports HTML row labels via `label_html`.
     """
     _esm = """
     export function render({ model, el }) {
@@ -252,14 +265,20 @@ class HoverList(anywidget.AnyWidget):
         wrap.style.display = 'flex';
         wrap.style.flexDirection = 'column';
 
-        items.forEach(({tid, label}) => {
+        items.forEach((item) => {
+          const { tid, label, label_html } = item;
+
           const btn = document.createElement('button');
-          btn.textContent = label;
           btn.className = 'rowbtn';
           btn.setAttribute('type', 'button');
           btn.setAttribute('role', 'button');
           btn.setAttribute('tabindex', '0');
 
+          // Prefer HTML layout if provided; fall back to plain text
+          if (label_html) { btn.innerHTML = label_html; }
+          else            { btn.textContent = label || ""; }
+
+          // Hover → preview (bind several events for reliability)
           const preview = () => {
             model.set('hovered_id', tid);
             model.save_changes();
@@ -269,6 +288,7 @@ class HoverList(anywidget.AnyWidget):
           btn.addEventListener('mousemove',  preview);
           btn.addEventListener('focus',      preview);
 
+          // Click → append
           btn.addEventListener('click', () => {
             model.set('clicked_id', tid);
             model.save_changes();
@@ -284,11 +304,10 @@ class HoverList(anywidget.AnyWidget):
       model.on('change:items', renderList);
     }
     """
-    items      = t.List(trait=t.Dict()).tag(sync=True)   # [{tid:int, label:str}, ...]
+    items      = t.List(trait=t.Dict()).tag(sync=True)   # [{tid:int, label?:str, label_html?:str}, ...]
     hovered_id = t.Int(allow_none=True).tag(sync=True)
     clicked_id = t.Int(allow_none=True).tag(sync=True)
 
-
 # ---------- Predictions list (uses HoverList) ----------
 @solara.component
 def PredictionsList():
@@ -336,7 +355,7 @@ def PredictionsList():
 def Page():
     solara.Style(theme_css)
 
-    with solara.Column(margin=10, gap="12px"):
+    with solara.Column(margin=8, gap="10px"):
         solara.Markdown("# Next-Token Predictor + Semantic Neighborhood")
         solara.Markdown(
             "Type text to see AI's top predictions for the next token. "