Update app.py

app.py

@@ -1,4 +1,5 @@
 # app.py
+# app.py
 import json
 from pathlib import Path
 import threading, time
@@ -30,10 +31,12 @@ hr{ border-color:var(--border); }
 .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;
+  width:100%; padding:10px 12px; border-radius:12px;
   border:1px solid var(--border); background:#fff; color:var(--text);
+  display:flex; justify-content:flex-start; align-items:center;
+  text-align:left;
   font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
-  letter-spacing: .2px;
+  letter-spacing:.2px;
 }
 .rowbtn:hover{ background:#f7fbff; border-color:#c3e8fb; cursor:pointer; }
 """
@@ -42,7 +45,7 @@ hr{ border-color:var(--border); }
 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([])     # [(tok, sim), ...]
+neighbor_list_rx = solara.reactive([])     # [(tok_display, sim), ...]
 notice_rx = solara.reactive("Click a candidate (or hover to preview).")
 auto_running_rx = solara.reactive(True)
 last_hovered_id_rx = solara.reactive(None)
@@ -64,15 +67,21 @@ 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}"
+def display_token_from_id(tid: int) -> str:
+    """Readable label for a single token id (no leading tokenizer markers)."""
+    toks = tokenizer.convert_ids_to_tokens([int(tid)], skip_special_tokens=True)
+    t = toks[0] if toks else ""
+    for lead in ("▁", "Ġ"):  # common leading markers
+        if t.startswith(lead):
+            t = t[len(lead):]
+    t = t.replace("\n", "↵")
+    if t.strip() == "":
+        return "␠"  # visible space marker for pure whitespace
+    return t
+
+def fmt_row(idx: int, prob: str, tid: int, tok_disp: str) -> str:
+    # left-justified simple columns
+    return f"{idx:<2}  {prob:<7}  {tid:<6}  {tok_disp}"
 
 # ------------------ Predict ------------------
 def predict_top10(prompt: str) -> pd.DataFrame:
@@ -95,7 +104,7 @@ def predict_top10(prompt: str) -> pd.DataFrame:
     topk = torch.topk(scores, 10)
     ids = [int(topk.indices[0, i]) for i in range(10)]
     probs = [float(topk.values[0, i]) for i in range(10)]
-    toks = [tokenizer.decode([i]) for i in ids]
+    toks = [tokenizer.decode([i]) for i in ids]  # used for append; display uses display_token_from_id
     df = pd.DataFrame({"probs": probs, "id": ids, "tok": toks})
     df["probs"] = df["probs"].map(lambda p: f"{p:.2%}")
     return df
@@ -113,7 +122,7 @@ def base_scatter():
         xs, ys = zip(*[coords[k] for k in coords.keys()])
         fig.add_trace(go.Scattergl(
             x=xs, y=ys, mode="markers",
-            marker=dict(size=3, opacity=1.0, color="rgba(56,189,248,0.15)"),
+            marker=dict(size=3, opacity=1.0, color="rgba(56,189,248,0.15)"),  # pale cloud
             hoverinfo="skip",
         ))
     fig.update_layout(
@@ -145,12 +154,12 @@ def highlight(token_id: int):
         ny = [coords[str(nid)][1] for nid, _ in nbrs]
         fig.add_trace(go.Scattergl(
             x=nx, y=ny, mode="markers",
-            marker=dict(size=6, color="rgba(56,189,248,0.75)", symbol="circle"),
+            marker=dict(size=6, color="rgba(56,189,248,0.75)", symbol="circle"),  # darker neighbors
             hoverinfo="skip",
         ))
         chips = []
         for nid, sim in nbrs:
-            chips.append((tokenizer.decode([int(nid)]), float(sim)))
+            chips.append((display_token_from_id(int(nid)), float(sim)))
         neighbor_list_rx.set(chips)
     else:
         neighbor_list_rx.set([])
@@ -158,7 +167,7 @@ def highlight(token_id: int):
     tx, ty = coords[str(token_id)]
     fig.add_trace(go.Scattergl(
         x=[tx], y=[ty], mode="markers",
-        marker=dict(size=10, color="rgba(34,211,238,1.0)", line=dict(width=1)),
+        marker=dict(size=10, color="rgba(34,211,238,1.0)", line=dict(width=1)),  # bright target
         hoverinfo="skip",
     ))
     return fig
@@ -172,10 +181,11 @@ def preview_token(token_id: int):
     fig_rx.set(highlight(token_id))
 
 def append_token(token_id: int):
+    # keep decode() here so spacing stays correct in the prompt
     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
+    preview_token(int(token_id))  # highlight appended token
+    on_predict()                  # refresh predictions
 
 # ------------------ Auto-predict on typing (debounced) ------------------
 @solara.component
@@ -190,7 +200,7 @@ def AutoPredictWatcher():
         snap = text
 
         def worker():
-            time.sleep(0.25)  # debounce ~250ms
+            time.sleep(0.25)  # ~250ms debounce
             if not cancelled and snap == text_rx.value:
                 on_predict()
 
@@ -211,20 +221,25 @@ def PredictionsList():
     with solara.Column(gap="6px", style={"maxWidth": "720px"}):
         solara.Markdown("### Prediction")
         solara.Text(
-            " #   probs    token  predicted next token",
+            " #  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 = fmt_row(i, prob, tid, tok)
+            tid = int(row["id"])
+            prob = row["probs"]
+            tok_disp = display_token_from_id(tid)  # clean label
+            label = fmt_row(i, prob, tid, tok_disp)
             solara.Button(
                 label,
                 on_click=lambda tid=tid: append_token(tid),
-                on_mouse_enter=lambda tid=tid: preview_token(tid),
+                on_mouse_enter=lambda tid=tid: preview_token(tid),  # newer solara
+                on_mouse_over=lambda tid=tid: preview_token(tid),   # older solara
+                on_focus=lambda tid=tid: preview_token(tid),        # keyboard
                 classes=["rowbtn"],
+                style={"justifyContent": "flex-start"},
             )
 
 # ------------------ Page ------------------
@@ -240,7 +255,7 @@ def Page():
             "Hover a candidate to preview its neighborhood."
         )
 
-        # Input (auto-predict is handled by watcher)
+        # Input (auto-predict handled by watcher)
         solara.InputText("Enter text", value=text_rx, continuous_update=True, style={"minWidth": "520px"})
         solara.Markdown(f"*{notice_rx.value}*")
 
@@ -262,11 +277,11 @@ def Page():
                         for tok, sim in neighbor_list_rx.value:
                             solara.HTML(
                                 tag="span",
-                                unsafe_innerHTML=f'<span class="badge">{show_token(tok)} &nbsp; {(sim*100):.1f}%</span>',
+                                unsafe_innerHTML=f'<span class="badge">{tok} &nbsp; {(sim*100):.1f}%</span>',
                             )
 
         AutoPredictWatcher()
 
 # Seed initial predictions and mount
 on_predict()
-Page()
+Page()