Update app.py

app.py

@@ -13,6 +13,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 # for robust hover/click from the browser
 import anywidget
 import traitlets as t
+import html  # for escaping token text in the HTML label
 
 
 # ---------- Model ----------
@@ -83,7 +84,7 @@ selected_token_id_rx = solara.reactive(None)
 neighbor_list_rx = solara.reactive([])
 last_hovered_id_rx = solara.reactive(None)
 auto_running_rx = solara.reactive(True)
-
+neigh_msg_rx = solara.reactive("")  # message shown when no neighborhood is available
 
 # ---------- Embedding assets ----------
 ASSETS = Path("assets/embeddings")
@@ -178,11 +179,20 @@ def get_neighbor_list(token_id: int, k: int = 20):
 
 def highlight(token_id: int):
     fig = base_scatter()
+
+    # Not in map (or missing map) → clear chips and show message
     if not coords or token_id not in ids_set:
         neighbor_list_rx.set([])
+        if not coords:
+            neigh_msg_rx.set("Embedding map unavailable – add `assets/embeddings/*.json`.")
+        else:
+            neigh_msg_rx.set("Neighborhood unavailable for this token (not in the top-5k set).")
         return fig
 
+    # In map → clear message and draw neighbors/target
+    neigh_msg_rx.set("")
     nbrs = get_neighbor_list(token_id, k=20)
+
     if nbrs:
         nx = [coords[str(nid)][0] for nid,_ in nbrs]
         ny = [coords[str(nid)][1] for nid,_ in nbrs]
@@ -328,7 +338,17 @@ def PredictionsList():
             tid = int(row["id"])
             prob = row["probs"]                 # already a formatted string like "4.12%"
             tok_disp = display_token_from_id(tid)
-            items.append({"tid": tid, "label": fmt_row(i, prob, tid, tok_disp)})
+            tok_safe = html.escape(tok_disp)   # protect the HTML label
+
+                label_html = (
+                    f'<div class="rowbtn-grid">'
+                    f'  <span class="c0">{i}</span>'
+                    f'  <span class="c1">{prob}</span>'
+                    f'  <span class="c2">{tid}</span>'
+                    f'  <span class="c3">{tok_safe}</span>'
+                    f'</div>'
+                )
+                items.append({"tid": tid, "label_html": label_html})   # <-- note label_html
 
         w = HoverList()
         w.items = items
@@ -382,6 +402,8 @@ def Page():
                         for tok, sim in neighbor_list_rx.value:
                             solara.HTML(tag="span",
                                 unsafe_innerHTML=f'<span class="badge">{tok} &nbsp; {(sim*100):.1f}%</span>')
+                elif neigh_msg_rx.value:
+                solara.Text(neigh_msg_rx.value, style={"color":"var(--muted)"})
 
         AutoPredictWatcher()