Update app.py

Auto-predict
A tiny AutoPredictWatcher() component debounces text_rx changes (~250 ms) and calls on_predict() automatically. No more Predict button.

Click anywhere on a row
The top-10 is now a clean, custom list of full-width “row buttons” (styled to look like rows). Clicking anywhere in the row calls append_token(tid).

Clearer neighborhood
The Plotly figure now has three layers:

Pale cloud (all tokens): rgba(56,189,248,0.15)

Darker neighbors (10–20): rgba(56,189,248,0.75)

Bright selected token: rgba(34,211,238,1.0)

Hover preview + neighbor chips
Each row has on_mouse_enter that calls preview_token(tid), which updates the plot and populates a right-side chip list with the 10–20 closest neighbors and their cosine %.

app.py

@@ -1,65 +1,78 @@
 # app.py
 import json
-import random
 from pathlib import Path
+import threading, time
 
 import solara
-import pandas as pd
 import torch
 import torch.nn.functional as F
+import pandas as pd
+import plotly.graph_objects as go
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ---- Model (same as original Space) -----------------------------------------
+# ------------------ Model ------------------
 MODEL_ID = "Qwen/Qwen3-0.6B"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
 
-# ---- App state ---------------------------------------------------------------
-text_rx = solara.reactive("twinkle, twinkle, little ")
-top10_rx = solara.reactive(pd.DataFrame(columns=["probs", "next token ID", "predicted next token"]))
-selected_token_id_rx = solara.reactive(None)   # for neighborhood focus
-notice_rx = solara.reactive("Enter text to see predictions.")
-theme_css = solara.reactive("""
-<style>
-:root {
-  --primary: #38bdf8;   /* light blue */
-  --bg: #ffffff;        /* white */
-  --text: #000000;      /* black */
-  --muted: #6b7280;     /* gray-500 */
-  --border: #e5e7eb;    /* gray-200 */
+# ------------------ Theme ------------------
+theme_css = """
+:root{
+  --primary:#38bdf8;     /* light blue */
+  --bg:#ffffff;          /* white */
+  --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); }
-table td, table th { border-color: var(--border) !important; }
-.solara-dataframe .MuiTableCell-root { font-size: 14px; }
-.btn-primary { background: var(--primary); color: #000; border: 1px solid var(--primary); padding: 6px 10px; border-radius: 8px; }
-.badge { display:inline-block; padding:2px 8px; border:1px solid var(--border); border-radius:999px; color:var(--text); }
-</style>
-""")
-
-# ---- Load embedding assets (your files) --------------------------------------
+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;
+}
+.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)
+
+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
+
+# ------------------ Embedding assets ------------------
 ASSETS = Path("assets/embeddings")
 COORDS_PATH = ASSETS / "pca_top5k_coords.json"
-NEIGH_PATH = ASSETS / "neighbors_top5k_k40.json"
+NEIGH_PATH  = ASSETS / "neighbors_top5k_k40.json"
 
 coords = {}
 neighbors = {}
 ids_set = set()
 
 if COORDS_PATH.exists() and NEIGH_PATH.exists():
-    with COORDS_PATH.open("r", encoding="utf-8") as f:
-        coords = json.load(f)            # {token_id: [x, y], ...}
-    with NEIGH_PATH.open("r", encoding="utf-8") as f:
-        neighbors = json.load(f)         # {"neighbors": {token_id: [[nid, sim], ...]}}
+    coords = json.loads(COORDS_PATH.read_text("utf-8"))
+    neighbors = json.loads(NEIGH_PATH.read_text("utf-8"))
     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 to enable the map.")
 
-# ---- Helpers -----------------------------------------------------------------
+# ------------------ Predict ------------------
 def predict_top10(prompt: str) -> pd.DataFrame:
     if not prompt:
-        return pd.DataFrame(columns=["probs", "next token ID", "predicted next token"])
-
+        return pd.DataFrame(columns=["probs", "id", "tok"])
     tokens = tokenizer.encode(prompt, return_tensors="pt")
     out = model.generate(
         tokens,
@@ -67,148 +80,180 @@ def predict_top10(prompt: str) -> pd.DataFrame:
         output_scores=True,
         return_dict_in_generate=True,
         pad_token_id=tokenizer.eos_token_id,
-        do_sample=False,  # greedy (deterministic)
+        do_sample=False,
         temperature=0.0,
         top_k=1,
         top_p=1.0,
     )
-    scores = F.softmax(out.scores[0], dim=-1)      # [1, vocab]
-    top_10 = torch.topk(scores, 10)
-
-    df = pd.DataFrame()
-    df["probs"] = top_10.values[0].detach().cpu().numpy()
-    df["probs"] = [f"{p:.2%}" for p in df["probs"]]
-    ids = [int(top_10.indices[0][i].detach().cpu().item()) for i in range(10)]
-    df["next token ID"] = ids
-    df["predicted next token"] = [tokenizer.decode([i]) for i in ids]
+    scores = torch.softmax(out.scores[0], dim=-1)          # [1, vocab]
+    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]
+    df = pd.DataFrame({"probs": probs, "id": ids, "tok": toks})
+    df["probs"] = df["probs"].map(lambda p: f"{p:.2%}")
     return df
 
-def get_neighbor_list(token_id: int, k: int = 18):
-    if not ids_set or token_id not in ids_set:
-        return []
-    raw = neighbors.get("neighbors", {}).get(str(token_id), [])
-    # raw item is [nid, sim]; keep top k
-    return raw[:k]
-
-# ---- Plot (Plotly scatter) ---------------------------------------------------
-# We generate a static "all points" scatter once, then reuse it with highlights.
-import plotly.graph_objects as go
+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
 
+# ------------------ Plotly figure ------------------
 def base_scatter():
-    if not coords:
-        return go.Figure().update_layout(
-            height=440, margin=dict(l=10, r=10, t=10, b=10),
-            paper_bgcolor="white", plot_bgcolor="white",
-        )
-    # unpack coordinates
-    xs, ys, tids = [], [], []
-    for tid_str, pt in coords.items():
-        xs.append(pt[0]); ys.append(pt[1]); tids.append(int(tid_str))
     fig = go.Figure()
-    fig.add_trace(go.Scattergl(
-        x=xs, y=ys, mode="markers",
-        marker=dict(size=3, opacity=0.85),
-        text=[f"id {t}" for t in tids],
-        hoverinfo="skip",  # keep hover minimal; we’ll show neighbors explicitly
-    ))
+    if coords:
+        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)"),
+            hoverinfo="skip",
+        ))
     fig.update_layout(
-        height=440, margin=dict(l=10, r=10, t=10, b=10),
+        height=460, margin=dict(l=10, r=10, t=10, b=10),
         paper_bgcolor="white", plot_bgcolor="white",
         xaxis=dict(visible=False), yaxis=dict(visible=False),
+        showlegend=False,
     )
     return fig
 
-base_fig = base_scatter()
-fig_rx = solara.reactive(base_fig)
+fig_rx = solara.reactive(base_scatter())
 
-def highlight(token_id: int):
-    """Return a new figure with neighbors + target highlighted."""
-    fig = base_fig.to_dict()  # detach copy
-    fig = go.Figure(fig)
+def get_neighbor_list(token_id: int, k: int = 18):
+    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], ... ]
 
+def highlight(token_id: int):
+    """Return figure with neighbors + target highlighted and update neighbor chip list."""
+    fig = base_scatter()
     if not coords or token_id not in ids_set:
+        neighbor_list_rx.set([])
         return fig
 
-    # Target
-    tx, ty = coords[str(token_id)]
-    fig.add_trace(go.Scattergl(
-        x=[tx], y=[ty], mode="markers",
-        marker=dict(size=8, line=dict(width=1), symbol="circle"),
-        name="target",
-    ))
-
-    # Neighbors
-    nbrs = get_neighbor_list(token_id)
+    # neighbors
+    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]
         fig.add_trace(go.Scattergl(
             x=nx, y=ny, mode="markers",
-            marker=dict(size=6, symbol="circle-open"),
-            name="neighbors",
+            marker=dict(size=6, color="rgba(56,189,248,0.75)", symbol="circle"),
+            hoverinfo="skip",
         ))
-    fig.update_layout(showlegend=False)
+        # update chip list
+        chips = []
+        for nid, sim in nbrs:
+            chips.append((tokenizer.decode([int(nid)]), float(sim)))
+        neighbor_list_rx.set(chips)
+    else:
+        neighbor_list_rx.set([])
+
+    # target
+    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)),
+        hoverinfo="skip",
+    ))
     return fig
 
-# ---- UI actions --------------------------------------------------------------
-def on_append_cell(column, row_index):
-    # append chosen next token to the text input
-    df = top10_rx.value
-    if row_index < len(df):
-        token_id = int(df.iloc[row_index]["next token ID"])
-        decoded = tokenizer.decode([token_id])
-        text_rx.set(text_rx.value + decoded)
-        selected_token_id_rx.set(token_id)
-        # Update plot
-        fig_rx.set(highlight(token_id))
+def preview_token(token_id: int):
+    selected_token_id_rx.set(int(token_id))
+    fig_rx.set(highlight(int(token_id)))
 
-cell_actions = [solara.CellAction(icon="mdi-plus", name="Append & highlight", on_click=on_append_cell)]
+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
 
-def on_predict():
-    df = predict_top10(text_rx.value)
-    top10_rx.set(df)
-    notice_rx.set("Click a candidate to append it and highlight its neighborhood.")
-    # also set selected to the top-1 for convenience
-    if len(df) > 0:
-        tid = int(df.iloc[0]["next token ID"])
-        selected_token_id_rx.set(tid)
-        fig_rx.set(highlight(tid))
-
-def on_show_neighborhood():
-    # take last token in the prompt (if any), otherwise do nothing
-    ids = tokenizer.encode(text_rx.value)
-    if ids:
-        token_id = int(ids[-1])
-        selected_token_id_rx.set(token_id)
-        fig_rx.set(highlight(token_id))
-
-# ---- Page --------------------------------------------------------------------
+# ------------------ Auto-predict on typing (debounced) ------------------
+@solara.component
+def AutoPredictWatcher():
+    text = text_rx.value
+    auto = auto_running_rx.value
+
+    def effect():
+        if not auto:
+            return
+        cancelled = False
+        snap = text
+
+        def worker():
+            time.sleep(0.25)  # debounce ~250ms
+            if not cancelled and snap == text_rx.value:
+                on_predict()
+
+        threading.Thread(target=worker, daemon=True).start()
+
+        def cleanup():
+            nonlocal cancelled
+            cancelled = True
+        return cleanup
+
+    solara.use_effect(effect, [text, auto])
+    return solara.Text("", style={"display":"none"})  # nothing visible
+
+# ------------------ UI: custom clickable/hoverable list ------------------
+@solara.component
+def PredictionsList():
+    df = preds_rx.value
+    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>')
+        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"],
+                          )
+
+# ------------------ Page ------------------
 @solara.component
 def Page():
-    solara.Style(theme_css.value)
+    solara.Style(theme_css)
 
     with solara.Column(margin=12, gap="16px"):
         solara.Markdown("# Next-Token Predictor + Semantic Neighborhood")
         solara.Markdown(
-            "Type text, then **Predict** to see the next-token distribution. "
-            "Click a candidate to append it and highlight its **semantic neighborhood**."
+            "Type text to see predictions update automatically. "
+            "Click a candidate to append it and highlight its **semantic neighborhood**. "
+            "Hover a candidate to preview its neighborhood."
         )
-        with solara.Row(gap="8px"):
-            solara.InputText("Enter text", value=text_rx, continuous_update=True, style={"minWidth": "520px"})
-            solara.Button("Predict", on_click=on_predict, classes=["btn-primary"])
-            solara.Button("Show neighborhood of last token", on_click=on_show_neighborhood)
+
+        # Input row (no Predict button needed anymore)
+        solara.InputText("Enter text", value=text_rx, continuous_update=True, style={"minWidth":"520px"})
 
         solara.Markdown(f"*{notice_rx.value}*")
 
-        # Top-10 table
-        solara.Markdown("### Prediction")
-        solara.DataFrame(top10_rx.value, items_per_page=10, cell_actions=cell_actions)
+        with solara.Row(align_items="start", gap="24px"):
+            with solara.Column():
+                PredictionsList()
+
+            with solara.Column():
+                solara.Markdown("### Semantic Neighborhood")
+                if not coords:
+                    solara.Markdown("> Embedding map unavailable – add `assets/embeddings/*.json`.")
+                else:
+                    solara.FigurePlotly(fig_rx.value)
+
+                # Neighbor chips
+                if neighbor_list_rx.value:
+                    solara.Markdown("**Nearest neighbors:**")
+                    with solara.Row(wrap=True):
+                        for tok, sim in neighbor_list_rx.value:
+                            solara.HTML(tag="span", unsafe_innerHTML=f'<span class="badge">{tok} &nbsp; {(sim*100):.1f}%</span>')
 
-        # Neighborhood panel
-        solara.Markdown("### Semantic Neighborhood")
-        if not coords:
-            solara.Markdown("> Embedding map unavailable – add `assets/embeddings/*.json`.")
-        else:
-            solara.FigurePlotly(fig_rx.value)
+        AutoPredictWatcher()  # invisible component that triggers auto-predict
 
+# Seed initial predictions
+on_predict()
 Page()