Update app.py

app.py

@@ -43,8 +43,8 @@ hr{ border-color:var(--border); }
 # ------------------ 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([])     # [(tok_display, sim), ...]
+selected_token_id_rx = solara.reactive(None)  # the token we’re currently highlighting
+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)
@@ -59,11 +59,11 @@ neighbors = {}
 ids_set = set()
 
 if COORDS_PATH.exists() and NEIGH_PATH.exists():
-    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], ...] } }
+    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.")
+    notice_rx.set("Embedding files not found — add assets/embeddings/*.json`.")
 
 # ------------------ Helpers ------------------
 def display_token_from_id(tid: int) -> str:
@@ -109,10 +109,17 @@ def predict_top10(prompt: str) -> pd.DataFrame:
     return df
 
 def on_predict():
+    """Update predictions. Keep current highlight unless there is none yet."""
     df = predict_top10(text_rx.value)
     preds_rx.set(df)
-    if len(df) > 0:
-        preview_token(int(df.iloc[0]["id"]))  # highlight top-1 by default
+    if len(df) == 0:
+        return
+    # Only auto-select the top-1 if we don't have a selection yet.
+    if selected_token_id_rx.value is None:
+        preview_token(int(df.iloc[0]["id"]))
+    else:
+        # Keep highlighting whatever the user last hovered/clicked.
+        fig_rx.set(highlight(int(selected_token_id_rx.value)))
 
 # ------------------ Plotly figure ------------------
 def base_scatter():
@@ -183,8 +190,9 @@ 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 appended token
-    on_predict()                  # refresh predictions
+    # lock highlight to the clicked token and keep it after re-predict
+    preview_token(int(token_id))
+    on_predict()  # refresh next tokens but do NOT override highlight anymore
 
 # ------------------ Auto-predict on typing (debounced) ------------------
 @solara.component
@@ -232,18 +240,21 @@ def PredictionsList():
             tok_disp = display_token_from_id(tid)
             label = fmt_row(i, prob, tid, tok_disp)
 
-            # Use Div so hover works across Solara versions; note (event, ...) signatures
+            # Wrap a Button (click) inside a Div (hover). Div events pass (event, ...),
+            # Button click is reliable across Solara versions.
             with solara.Div(
-                classes=["rowbtn"],
-                style={"justifyContent": "flex-start"},
                 attributes={"tabindex": "0", "role": "button"},
-                on_click=lambda e=None, tid=tid: append_token(tid),
                 on_mouse_enter=lambda e=None, tid=tid: preview_token(tid),
                 on_mouse_over=lambda e=None, tid=tid: preview_token(tid),
                 on_mouse_move=lambda e=None, tid=tid: preview_token(tid),
                 on_focus=lambda e=None, tid=tid: preview_token(tid),
             ):
-                solara.Text(label)
+                solara.Button(
+                    label,
+                    on_click=lambda tid=tid: append_token(tid),
+                    classes=["rowbtn"],
+                    style={"justifyContent": "flex-start", "width": "100%"},
+                )
 
 # ------------------ Page ------------------
 @solara.component