Update app.py

app.py

@@ -4,18 +4,21 @@ from pathlib import Path
 import threading, time
 
 import solara
-import torch
-import torch.nn.functional as F
 import pandas as pd
 import plotly.graph_objects as go
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ------------------ Model ------------------
+# ---------- versions (shows up in Space logs) ----------
+import plotly
+print("VERSIONS:", "solara", solara.__version__, "plotly", plotly.__version__, "torch", torch.__version__)
+
+# ---------- Model ----------
 MODEL_ID = "Qwen/Qwen3-0.6B"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
 
-# ------------------ Theme ------------------
+# ---------- Theme & layout fixes ----------
 theme_css = """
 :root{
   --primary:#38bdf8;     /* light blue */
@@ -24,18 +27,15 @@ theme_css = """
   --muted:#6b7280;       /* gray-500 */
   --border:#e5e7eb;      /* gray-200 */
 }
-body{ background:var(--bg); color:var(--text);}
-h1,h2,h3{ color:var(--text); }
-hr{ border-color:var(--border); }
+body{ background:var(--bg); color:var(--text); }
 .badge{ display:inline-block; padding:2px 8px; border:1px solid var(--border); border-radius:999px; margin:2px; }
 
-/* Keep the predictions list above the plot so it can receive pointer events */
+/* Make sure the prediction list can receive pointer events even if Plotly expands */
 .predictions-panel { position: relative; z-index: 5; }
-.plot-panel { position: relative; z-index: 1; }
-
-/* Safety: prevent any wide Plotly overlay from swallowing events on the left */
+.plot-panel        { position: relative; z-index: 1; }
 .plot-panel .js-plotly-plot { position: relative; z-index: 1; }
 
+/* Row style */
 .rowbtn{
   width:100%; padding:10px 12px; border-radius:12px;
   border:1px solid var(--border); background:#fff; color:var(--text);
@@ -47,16 +47,16 @@ hr{ border-color:var(--border); }
 .rowbtn:hover{ background:#f7fbff; border-color:#c3e8fb; }
 """
 
-# ------------------ App state ------------------
+# ---------- 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)   # currently highlighted token id
-neighbor_list_rx = solara.reactive([])         # [(tok_display, sim), ...]
+preds_rx = solara.reactive(pd.DataFrame(columns=["probs","id","tok"]))
+selected_token_id_rx = solara.reactive(None)
+neighbor_list_rx = solara.reactive([])
+last_hovered_id_rx = solara.reactive(None)
 notice_rx = solara.reactive("Click a candidate (or hover to preview).")
 auto_running_rx = solara.reactive(True)
-last_hovered_id_rx = solara.reactive(None)
 
-# ------------------ Embedding assets ------------------
+# ---------- Embedding assets ----------
 ASSETS = Path("assets/embeddings")
 COORDS_PATH = ASSETS / "pca_top5k_coords.json"
 NEIGH_PATH  = ASSETS / "neighbors_top5k_k40.json"
@@ -66,34 +66,31 @@ 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"))
+    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`.")
+    notice_rx.set("Embedding files not found — add assets/embeddings/*.json to enable the map.")
 
-# ------------------ Helpers ------------------
+# ---------- Helpers ----------
 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 ("▁", "Ġ"):
         if t.startswith(lead):
             t = t[len(lead):]
-    t = t.replace("\n", "↵")
+    t = t.replace("\n","↵")
     if t.strip() == "":
-        return "␠"  # visible space marker for pure whitespace
+        return "␠"
     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 ------------------
+# ---------- Predict ----------
 def predict_top10(prompt: str) -> pd.DataFrame:
     if not prompt:
-        return pd.DataFrame(columns=["probs", "id", "tok"])
-
+        return pd.DataFrame(columns=["probs","id","tok"])
     tokens = tokenizer.encode(prompt, return_tensors="pt")
     out = model.generate(
         tokens,
@@ -101,43 +98,41 @@ 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,
-        temperature=0.0,
-        top_k=1,
-        top_p=1.0,
+        do_sample=False, temperature=0.0, top_k=1, top_p=1.0,
     )
-    scores = torch.softmax(out.scores[0], dim=-1)          # [1, vocab]
+    scores = torch.softmax(out.scores[0], dim=-1)
     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]  # for append; display uses display_token_from_id
+    toks = [tokenizer.decode([i]) for i in ids]  # used for append only
     df = pd.DataFrame({"probs": probs, "id": ids, "tok": toks})
     df["probs"] = df["probs"].map(lambda p: f"{p:.2%}")
     return df
 
 def on_predict():
-    """Update predictions. Keep current highlight unless there is none yet."""
+    """Update predictions; keep current highlight unless none yet."""
     df = predict_top10(text_rx.value)
     preds_rx.set(df)
     if len(df) == 0:
         return
     if selected_token_id_rx.value is None:
-        preview_token(int(df.iloc[0]["id"]))
+        preview_token(int(df.iloc[0]["id"]))  # first time only
     else:
+        # keep the user's last selection/hover
         fig_rx.set(highlight(int(selected_token_id_rx.value)))
 
-# ------------------ Plotly figure ------------------
+# ---------- Plot ----------
 def base_scatter():
     fig = go.Figure()
     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)"),  # pale cloud
+            marker=dict(size=3, opacity=1.0, color="rgba(56,189,248,0.15)"),
             hoverinfo="skip",
         ))
     fig.update_layout(
-        height=460, 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,
@@ -153,7 +148,6 @@ def get_neighbor_list(token_id: int, k: int = 20):
     return raw[:k]
 
 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([])
@@ -161,16 +155,14 @@ def highlight(token_id: int):
 
     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]
+        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, color="rgba(56,189,248,0.75)", symbol="circle"),  # darker neighbors
+            marker=dict(size=6, color="rgba(56,189,248,0.75)"),
             hoverinfo="skip",
         ))
-        chips = []
-        for nid, sim in nbrs:
-            chips.append((display_token_from_id(int(nid)), float(sim)))
+        chips = [(display_token_from_id(int(nid)), float(sim)) for nid,sim in nbrs]
         neighbor_list_rx.set(chips)
     else:
         neighbor_list_rx.set([])
@@ -178,16 +170,15 @@ 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)),  # bright target
+        marker=dict(size=10, color="rgba(34,211,238,1.0)", line=dict(width=1)),
         hoverinfo="skip",
     ))
     return fig
 
 def preview_token(token_id: int):
-    print("preview ->", token_id)  # TEMP: check Logs
-    token_id = int(token_id)
-    # TEMP DEBUG: verify hover fires in Space logs
+    # DEBUG: confirm events reach Python
     print("preview ->", token_id)
+    token_id = int(token_id)
     if last_hovered_id_rx.value == token_id:
         return
     last_hovered_id_rx.set(token_id)
@@ -195,14 +186,14 @@ 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
-    print("append ->", token_id)   # TEMP: check Logs
+    # DEBUG
+    print("append ->", token_id)
     decoded = tokenizer.decode([int(token_id)])
     text_rx.set(text_rx.value + decoded)
-    preview_token(int(token_id))  # keep highlight on clicked token
-    on_predict()                  # refresh predictions, preserve selection
+    preview_token(int(token_id))
+    on_predict()
 
-# ------------------ Auto-predict on typing (debounced) ------------------
+# ---------- Auto-predict (debounced) ----------
 @solara.component
 def AutoPredictWatcher():
     text = text_rx.value
@@ -215,7 +206,7 @@ def AutoPredictWatcher():
         snap = text
 
         def worker():
-            time.sleep(0.25)  # ~250ms debounce
+            time.sleep(0.25)
             if not cancelled and snap == text_rx.value:
                 on_predict()
 
@@ -229,41 +220,39 @@ def AutoPredictWatcher():
     solara.use_effect(effect, [text, auto])
     return solara.Text("", style={"display": "none"})
 
-# ------------------ UI: rows as Div (hover + click here) ------------------
+# ---------- Predictions 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.Text(
             " #  probs    token  predicted next token",
             style={
-                "color": "var(--muted)",
-                "fontFamily": 'ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace',
+                "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"])
-            prob = row["probs"]
+            tid = int(row["id"]); prob = row["probs"]
             tok_disp = display_token_from_id(tid)
             label = fmt_row(i, prob, tid, tok_disp)
 
-            # One Div per row: both hover and click handlers live here.
+            # Use Div so pointer events are reliable; accept *args to handle any signature
             with solara.Div(
                 classes=["rowbtn"],
-                style={"justifyContent": "flex-start", "width": "100%"},
-                attributes={"tabindex": "0", "role": "button"},
-                on_click=lambda e=None, tid=tid: append_token(tid),            # click to append
-                on_mouse_enter=lambda e=None, tid=tid: preview_token(tid),     # hover preview
-                on_mouse_over=lambda e=None, tid=tid: preview_token(tid),
-                on_mouse_move=lambda e=None, tid=tid: preview_token(tid),
-                on_pointer_enter=lambda e=None, tid=tid: preview_token(tid),
-                on_focus=lambda e=None, tid=tid: preview_token(tid),
+                style={"justifyContent":"flex-start","width":"100%"},
+                attributes={"tabindex":"0","role":"button"},
+                on_click=lambda *args, tid=tid: append_token(tid),
+                on_mouse_enter=lambda *args, tid=tid: preview_token(tid),
+                on_mouse_over=lambda *args, tid=tid: preview_token(tid),
+                on_mouse_move=lambda *args, tid=tid: preview_token(tid),
+                on_pointer_enter=lambda *args, tid=tid: preview_token(tid),
+                on_focus=lambda *args, tid=tid: preview_token(tid),
             ):
                 solara.Text(label)
 
-
-# ------------------ Page ------------------
+# ---------- Page ----------
 @solara.component
 def Page():
     solara.Style(theme_css)
@@ -276,12 +265,10 @@ def Page():
             "Hover a candidate to preview its neighborhood."
         )
 
-        # Input (auto-predict handled by watcher)
-        solara.InputText("Enter text", value=text_rx, continuous_update=True, style={"minWidth": "520px"})
+        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.Row(gap="24px", style={"align-items":"flex-start"}):
             with solara.Column(classes=["predictions-panel"]):
                 PredictionsList()
 
@@ -294,15 +281,13 @@ def Page():
 
                 if neighbor_list_rx.value:
                     solara.Markdown("**Nearest neighbors:**")
-                    with solara.Row(style={"flex-wrap": "wrap"}):
+                    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">{tok} &nbsp; {(sim*100):.1f}%</span>')
 
         AutoPredictWatcher()
 
-# Seed initial predictions and mount
+# ---------- Kickoff ----------
 on_predict()
-Page()
+Page()