Update app.py

app.py

@@ -23,35 +23,29 @@ theme_css = """
   --text:#0b0f14;        /* near-black */
   --muted:#6b7280;       /* gray-500 */
   --border:#e5e7eb;      /* gray-200 */
-  --pale: rgba(56,189,248,0.15);   /* base dots */
-  --mid:  rgba(56,189,248,0.65);   /* neighbors */
-  --bright: rgba(34,211,238,1.0);  /* selected */
 }
 body{ background:var(--bg); color:var(--text);}
 h1,h2,h3{ color:var(--text); }
 hr{ border-color:var(--border); }
-.btn{ background:var(--primary); color:#000; border:1px solid var(--primary); padding:6px 10px; border-radius:10px; }
 .badge{ display:inline-block; padding:2px 8px; border:1px solid var(--border); border-radius:999px; margin:2px; }
+
 .rowbtn{
   width:100%; text-align:left; padding:10px 12px; border-radius:12px;
-  border:1px solid var(--border); background:#fff;
+  border:1px solid var(--border); background:#fff; color:var(--text);
+  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
+  letter-spacing: .2px;
 }
 .rowbtn:hover{ background:#f7fbff; border-color:#c3e8fb; cursor:pointer; }
-.rowgrid{ display:grid; grid-template-columns: 70px 120px 1fr; gap:10px; align-items:center; }
-.rowprob{ color:#111; font-weight:600; }
-.rowid{ color:var(--muted); font-variant-numeric: tabular-nums; }
-.rowtok{ color:#111; overflow:hidden; text-overflow:ellipsis; white-space:nowrap; }
-.listheader{ color:var(--muted); font-size:12px; margin-bottom:6px;}
 """
 
 # ------------------ App state ------------------
 text_rx = solara.reactive("twinkle, twinkle, little ")
 preds_rx = solara.reactive(pd.DataFrame(columns=["probs", "id", "tok"]))
 selected_token_id_rx = solara.reactive(None)
-neighbor_list_rx = solara.reactive([])  # list of (tok, sim)
-
+neighbor_list_rx = solara.reactive([])     # [(tok, sim), ...]
 notice_rx = solara.reactive("Click a candidate (or hover to preview).")
-auto_running_rx = solara.reactive(True)   # toggle if you ever want to disable auto-predict
+auto_running_rx = solara.reactive(True)
+last_hovered_id_rx = solara.reactive(None)
 
 # ------------------ Embedding assets ------------------
 ASSETS = Path("assets/embeddings")
@@ -63,16 +57,28 @@ neighbors = {}
 ids_set = set()
 
 if COORDS_PATH.exists() and NEIGH_PATH.exists():
-    coords = json.loads(COORDS_PATH.read_text("utf-8"))
-    neighbors = json.loads(NEIGH_PATH.read_text("utf-8"))
+    coords = json.loads(COORDS_PATH.read_text("utf-8"))               # { "tid": [x,y], ... }
+    neighbors = json.loads(NEIGH_PATH.read_text("utf-8"))             # { "neighbors": { "tid": [[nid,sim], ...] } }
     ids_set = set(map(int, coords.keys()))
 else:
     notice_rx.set("Embedding files not found — add assets/embeddings/*.json to enable the map.")
 
+# ------------------ Helpers ------------------
+def show_token(tok: str) -> str:
+    """Make invisible whitespace visible but readable in UI."""
+    if tok == "":
+        return "⟂"
+    return tok.replace(" ", "␠")
+
+def fmt_row(idx: int, prob: str, tid: int, tok: str) -> str:
+    tok_disp = show_token(tok)
+    return f"{idx:>2}  {prob:>6}  {tid:<6}  {tok_disp}"
+
 # ------------------ Predict ------------------
 def predict_top10(prompt: str) -> pd.DataFrame:
     if not prompt:
         return pd.DataFrame(columns=["probs", "id", "tok"])
+
     tokens = tokenizer.encode(prompt, return_tensors="pt")
     out = model.generate(
         tokens,
@@ -98,8 +104,7 @@ def on_predict():
     df = predict_top10(text_rx.value)
     preds_rx.set(df)
     if len(df) > 0:
-        tid = int(df.iloc[0]["id"])
-        preview_token(tid)  # highlight top-1 by default
+        preview_token(int(df.iloc[0]["id"]))  # highlight top-1 by default
 
 # ------------------ Plotly figure ------------------
 def base_scatter():
@@ -121,11 +126,11 @@ def base_scatter():
 
 fig_rx = solara.reactive(base_scatter())
 
-def get_neighbor_list(token_id: int, k: int = 18):
+def get_neighbor_list(token_id: int, k: int = 20):
     if not ids_set or token_id not in ids_set:
         return []
     raw = neighbors.get("neighbors", {}).get(str(token_id), [])
-    return raw[:k]  # [ [nid, sim], ... ]
+    return raw[:k]
 
 def highlight(token_id: int):
     """Return figure with neighbors + target highlighted and update neighbor chip list."""
@@ -134,7 +139,6 @@ def highlight(token_id: int):
         neighbor_list_rx.set([])
         return fig
 
-    # neighbors
     nbrs = get_neighbor_list(token_id, k=20)
     if nbrs:
         nx = [coords[str(nid)][0] for nid, _ in nbrs]
@@ -144,7 +148,6 @@ def highlight(token_id: int):
             marker=dict(size=6, color="rgba(56,189,248,0.75)", symbol="circle"),
             hoverinfo="skip",
         ))
-        # update chip list
         chips = []
         for nid, sim in nbrs:
             chips.append((tokenizer.decode([int(nid)]), float(sim)))
@@ -152,7 +155,6 @@ def highlight(token_id: int):
     else:
         neighbor_list_rx.set([])
 
-    # target
     tx, ty = coords[str(token_id)]
     fig.add_trace(go.Scattergl(
         x=[tx], y=[ty], mode="markers",
@@ -162,14 +164,18 @@ def highlight(token_id: int):
     return fig
 
 def preview_token(token_id: int):
-    selected_token_id_rx.set(int(token_id))
-    fig_rx.set(highlight(int(token_id)))
+    token_id = int(token_id)
+    if last_hovered_id_rx.value == token_id:
+        return
+    last_hovered_id_rx.set(token_id)
+    selected_token_id_rx.set(token_id)
+    fig_rx.set(highlight(token_id))
 
 def append_token(token_id: int):
     decoded = tokenizer.decode([int(token_id)])
     text_rx.set(text_rx.value + decoded)
     preview_token(int(token_id))  # highlight the one we appended
-    on_predict()  # refresh next top-10 for new text
+    on_predict()                  # refresh next top-10 for new text
 
 # ------------------ Auto-predict on typing (debounced) ------------------
 @solara.component
@@ -196,25 +202,30 @@ def AutoPredictWatcher():
         return cleanup
 
     solara.use_effect(effect, [text, auto])
-    return solara.Text("", style={"display":"none"})  # nothing visible
+    return solara.Text("", style={"display": "none"})
 
 # ------------------ UI: custom clickable/hoverable list ------------------
 @solara.component
 def PredictionsList():
     df = preds_rx.value
-    with solara.Column(gap="6px", style={"maxWidth":"720px"}):
+    with solara.Column(gap="6px", style={"maxWidth": "720px"}):
         solara.Markdown("### Prediction")
-        solara.HTML(tag="div", unsafe_innerHTML='<div class="listheader"># &nbsp;&nbsp; probs &nbsp;&nbsp; token id &nbsp;&nbsp; predicted next token</div>')
+        solara.Text(
+            " #   probs    token  predicted next token",
+            style={
+                "color": "var(--muted)",
+                "fontFamily": 'ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace',
+            },
+        )
         for i, row in df.iterrows():
             tid = int(row["id"]); tok = row["tok"]; prob = row["probs"]
-            label = f'<div class="rowgrid"><div class="rowprob">{i}&nbsp;&nbsp;{prob}</div><div class="rowid">{tid}</div><div class="rowtok">{tok}</div></div>'
-            # full-width button acts like a row; hover previews, click appends
-            solara.Button(None,
-                          on_click=lambda tid=tid: append_token(tid),
-                          on_mouse_enter=lambda tid=tid: preview_token(tid),
-                          html=label,
-                          classes=["rowbtn"],
-                          )
+            label = fmt_row(i, prob, tid, tok)
+            solara.Button(
+                label,
+                on_click=lambda tid=tid: append_token(tid),
+                on_mouse_enter=lambda tid=tid: preview_token(tid),
+                classes=["rowbtn"],
+            )
 
 # ------------------ Page ------------------
 @solara.component
@@ -229,11 +240,11 @@ def Page():
             "Hover a candidate to preview its neighborhood."
         )
 
-        # Input row (no Predict button needed anymore)
-        solara.InputText("Enter text", value=text_rx, continuous_update=True, style={"minWidth":"520px"})
-
+        # Input (auto-predict is handled by watcher)
+        solara.InputText("Enter text", value=text_rx, continuous_update=True, style={"minWidth": "520px"})
         solara.Markdown(f"*{notice_rx.value}*")
 
+        # Two columns
         with solara.Row(gap="24px", style={"align-items": "flex-start"}):
             with solara.Column():
                 PredictionsList()
@@ -245,15 +256,17 @@ def Page():
                 else:
                     solara.FigurePlotly(fig_rx.value)
 
-                # Neighbor chips
                 if neighbor_list_rx.value:
                     solara.Markdown("**Nearest neighbors:**")
                     with solara.Row(style={"flex-wrap": "wrap"}):
                         for tok, sim in neighbor_list_rx.value:
-                            solara.HTML(tag="span", unsafe_innerHTML=f'<span class="badge">{tok} &nbsp; {(sim*100):.1f}%</span>')
+                            solara.HTML(
+                                tag="span",
+                                unsafe_innerHTML=f'<span class="badge">{show_token(tok)} &nbsp; {(sim*100):.1f}%</span>',
+                            )
 
-        AutoPredictWatcher()  # invisible component that triggers auto-predict
+        AutoPredictWatcher()
 
-# Seed initial predictions
+# Seed initial predictions and mount
 on_predict()
 Page()