Massive speedup update + download plot fixing

app.py

@@ -1,4 +1,5 @@
 import os
+import threading
 
 import streamlit as st
 from dotenv import load_dotenv
@@ -16,6 +17,22 @@ _TABS = ["Chat", "Compare", "Extract"]
 _TAB_ICONS = [":material/chat:", ":material/search:", ":material/tune:"]
 
 
+def _preload_default_model() -> None:
+    """Background-warm the default local model so the first chat is instant."""
+    try:
+        import torch
+
+        torch.set_grad_enabled(False)
+        from utils.runtime import cached_model
+
+        cached_model(DEFAULT_MODEL)
+    except Exception:
+        pass
+
+
+threading.Thread(target=_preload_default_model, daemon=True).start()
+
+
 def _remote_model_input(remote_models: list[str]) -> str:
     """Return the active remote model id, picking from running NDIF deployments or a custom value."""
 
@@ -78,8 +95,7 @@ def _sidebar_controls() -> tuple[bool, str, str, str]:
     from utils.runtime import list_remote_models
 
     with st.sidebar:
-        st.markdown("# Persona UI")
-        st.caption("Chat, extract, and compare persona runs.")
+        st.markdown("## Persona UI")
 
         if "sidebar__active_tab" not in st.session_state:
             st.session_state["sidebar__active_tab"] = "Chat"

tabs/chat.py

@@ -18,18 +18,20 @@ from utils.chat import (
     resolve_system_prompt,
 )
 from utils.chat_export import save_chat_export
-from utils.datasets import load_dataset
+from utils.datasets import load_persona_list
 from utils.helpers import widget_key
 from utils.runtime import cached_model
 
 _LAST_PERSONA_ID_KEY = "chat:last_persona_id"
 _LAST_PROMPT_MODE_KEY = "chat:last_prompt_mode"
 _LAST_COMPARE_MODE_KEY = "chat:last_compare_mode"
+_LAST_PROBE_ENABLED_KEY = "chat:last_probe_enabled"
+_LAST_TOKEN_CONTRAST_KEY = "chat:last_token_contrast"
 
 
 def _load_personas(dataset_source: str) -> list[PersonaData] | None:
     try:
-        dataset, dataset_status = load_dataset(
+        personas, dataset_status = load_persona_list(
             dataset_source,
             personas_file=st.session_state.get("extract__personas_file"),
             qa_file=st.session_state.get("extract__qa_file"),
@@ -40,7 +42,6 @@ def _load_personas(dataset_source: str) -> list[PersonaData] | None:
         st.info("Check the selected dataset source or upload both JSONL files.")
         return None
 
-    personas = list(dataset)
     if not personas:
         st.warning("No personas found in the selected dataset.")
         st.info("Try a different dataset source or upload a non-empty personas file.")
@@ -152,6 +153,8 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
         context_key,
         remote,
         last_compare_mode_key=_LAST_COMPARE_MODE_KEY,
+        last_probe_enabled_key=_LAST_PROBE_ENABLED_KEY,
+        last_token_contrast_key=_LAST_TOKEN_CONTRAST_KEY,
     )
     if tools.compare_mode:
         render_compare_mode(

tabs/chat_ui.py

@@ -21,6 +21,21 @@ GENERATION_DEFAULTS = {
     "repetition_penalty": 1.0,
 }
 
+_LAST_GEN_PREFIX = "chat:last_gen:"
+
+
+def _persisted_key(context_key: str, name: str, default) -> str:
+    """Per-context widget key, seeded from the last cross-context value."""
+    last_key = f"{_LAST_GEN_PREFIX}{name}"
+    key = widget_key(context_key, name)
+    if key not in st.session_state:
+        st.session_state[key] = st.session_state.get(last_key, default)
+    return key
+
+
+def _remember(name: str, value) -> None:
+    st.session_state[f"{_LAST_GEN_PREFIX}{name}"] = value
+
 
 @dataclass(frozen=True)
 class GenerationConfig:
@@ -119,6 +134,8 @@ def render_advanced_settings(
     remote: bool,
     *,
     last_compare_mode_key: str,
+    last_probe_enabled_key: str = "",
+    last_token_contrast_key: str = "",
 ) -> tuple[GenerationConfig, ChatTools]:
     """Render the Advanced expander: tool toggles + generation settings."""
     with st.expander("Advanced", expanded=False):
@@ -130,12 +147,23 @@ def render_advanced_settings(
                 last_compare_mode_key, False
             )
 
+        probe_key = widget_key(context_key, "probe_enabled")
+        if probe_key not in st.session_state:
+            st.session_state[probe_key] = st.session_state.get(
+                last_probe_enabled_key, False
+            )
+
+        token_contrast_key = widget_key(context_key, "token_contrast")
+        if token_contrast_key not in st.session_state:
+            st.session_state[token_contrast_key] = st.session_state.get(
+                last_token_contrast_key, False
+            )
+
         tools_col1, tools_col2, tools_col3 = st.columns(3)
         with tools_col1:
             probe_enabled = st.toggle(
                 "Probe tools",
-                value=False,
-                key=widget_key(context_key, "probe_enabled"),
+                key=probe_key,
                 help="Trace chat activations and run compatible `.pt` probes on tapped tokens.",
             )
         with tools_col2:
@@ -147,8 +175,7 @@ def render_advanced_settings(
         with tools_col3:
             token_contrast = st.toggle(
                 "Token contrast",
-                value=False,
-                key=widget_key(context_key, "token_contrast"),
+                key=token_contrast_key,
                 disabled=not compare_mode,
                 help=(
                     "Color each generated token by how characteristic it is of each persona. "
@@ -158,95 +185,122 @@ def render_advanced_settings(
                 ),
             )
         st.session_state[last_compare_mode_key] = compare_mode
+        if last_probe_enabled_key:
+            st.session_state[last_probe_enabled_key] = probe_enabled
+        if last_token_contrast_key:
+            st.session_state[last_token_contrast_key] = token_contrast
 
         st.divider()
         st.caption("Generation")
+        generation = _render_generation_fragment(context_key, remote)
 
-        config_col1, config_col2 = st.columns([2, 1])
-        with config_col1:
-            max_new_tokens = st.slider(
-                "Max new tokens",
-                min_value=16,
-                max_value=512,
-                value=GENERATION_DEFAULTS["max_new_tokens"],
-                step=16,
-                key=widget_key(context_key, "max_new_tokens"),
-            )
-        with config_col2:
-            repetition_penalty = st.slider(
-                "Repetition penalty",
-                min_value=0.5,
-                max_value=2.0,
-                value=GENERATION_DEFAULTS["repetition_penalty"],
-                step=0.05,
-                key=widget_key(context_key, "repetition_penalty"),
-            )
+    tools = ChatTools(
+        probe_enabled=probe_enabled,
+        compare_mode=compare_mode,
+        token_contrast=token_contrast and compare_mode,
+    )
+    return generation, tools
 
-        use_sampling = st.checkbox(
-            "Random sampling",
-            value=False,
-            key=widget_key(context_key, "use_sampling"),
+
+@st.fragment
+def _render_generation_fragment(context_key: str, remote: bool) -> GenerationConfig:
+    """Render generation sliders inside a fragment so tweaks don't full-rerun."""
+    config_col1, config_col2 = st.columns([2, 1])
+    with config_col1:
+        max_new_tokens = st.slider(
+            "Max new tokens",
+            min_value=16,
+            max_value=512,
+            step=16,
+            key=_persisted_key(
+                context_key, "max_new_tokens", GENERATION_DEFAULTS["max_new_tokens"]
+            ),
+        )
+    with config_col2:
+        repetition_penalty = st.slider(
+            "Repetition penalty",
+            min_value=0.5,
+            max_value=2.0,
+            step=0.05,
+            key=_persisted_key(
+                context_key,
+                "repetition_penalty",
+                GENERATION_DEFAULTS["repetition_penalty"],
+            ),
         )
 
-        sampling_disabled = not use_sampling
-        sampling_col1, sampling_col2, sampling_col3 = st.columns(3)
-        with sampling_col1:
-            temperature = st.slider(
-                "Temperature",
-                min_value=0.01,
-                max_value=2.0,
-                value=GENERATION_DEFAULTS["temperature"],
-                step=0.01,
-                disabled=sampling_disabled,
-                key=widget_key(context_key, "temperature"),
-            )
-        with sampling_col2:
-            top_p = st.slider(
-                "Top-p",
-                min_value=0.01,
-                max_value=1.0,
-                value=GENERATION_DEFAULTS["top_p"],
-                step=0.01,
-                disabled=sampling_disabled,
-                key=widget_key(context_key, "top_p"),
-            )
-        with sampling_col3:
-            top_k = st.slider(
-                "Top-k (0 = off)",
+    use_sampling = st.checkbox(
+        "Random sampling",
+        key=_persisted_key(context_key, "use_sampling", False),
+    )
+
+    sampling_disabled = not use_sampling
+    sampling_col1, sampling_col2, sampling_col3 = st.columns(3)
+    with sampling_col1:
+        temperature = st.slider(
+            "Temperature",
+            min_value=0.01,
+            max_value=2.0,
+            step=0.01,
+            disabled=sampling_disabled,
+            key=_persisted_key(
+                context_key, "temperature", GENERATION_DEFAULTS["temperature"]
+            ),
+        )
+    with sampling_col2:
+        top_p = st.slider(
+            "Top-p",
+            min_value=0.01,
+            max_value=1.0,
+            step=0.01,
+            disabled=sampling_disabled,
+            key=_persisted_key(context_key, "top_p", GENERATION_DEFAULTS["top_p"]),
+        )
+    with sampling_col3:
+        top_k = st.slider(
+            "Top-k (0 = off)",
+            min_value=0,
+            max_value=100,
+            step=1,
+            disabled=sampling_disabled,
+            key=_persisted_key(context_key, "top_k", GENERATION_DEFAULTS["top_k"]),
+        )
+
+    seed_disabled = sampling_disabled or remote
+    seed_enabled = st.checkbox(
+        "Fix seed",
+        disabled=seed_disabled,
+        key=_persisted_key(context_key, "seed_enabled", False),
+    )
+    seed = None
+    if seed_enabled:
+        seed = int(
+            st.number_input(
+                "Seed",
                 min_value=0,
-                max_value=100,
-                value=GENERATION_DEFAULTS["top_k"],
+                max_value=2_147_483_647,
                 step=1,
-                disabled=sampling_disabled,
-                key=widget_key(context_key, "top_k"),
+                disabled=seed_disabled,
+                key=_persisted_key(context_key, "seed", 0),
             )
-
-        seed_disabled = sampling_disabled or remote
-        seed_enabled = st.checkbox(
-            "Fix seed",
-            value=False,
-            disabled=seed_disabled,
-            key=widget_key(context_key, "seed_enabled"),
         )
-        seed = None
-        if seed_enabled:
-            seed = int(
-                st.number_input(
-                    "Seed",
-                    min_value=0,
-                    max_value=2_147_483_647,
-                    value=0,
-                    step=1,
-                    disabled=seed_disabled,
-                    key=widget_key(context_key, "seed"),
-                )
-            )
 
-        if remote:
-            st.caption("Seed is local-only and disabled for remote runs.")
+    if remote:
+        st.caption("Seed is local-only and disabled for remote runs.")
+
+    for name, value in (
+        ("max_new_tokens", max_new_tokens),
+        ("repetition_penalty", repetition_penalty),
+        ("use_sampling", use_sampling),
+        ("temperature", temperature),
+        ("top_p", top_p),
+        ("top_k", top_k),
+        ("seed_enabled", seed_enabled),
+    ):
+        _remember(name, value)
 
     do_sample = bool(use_sampling)
-    generation = GenerationConfig(
+    return GenerationConfig(
         max_new_tokens=int(max_new_tokens),
         do_sample=do_sample,
         temperature=float(temperature),
@@ -255,12 +309,6 @@ def render_advanced_settings(
         repetition_penalty=float(repetition_penalty),
         seed=seed if do_sample and seed is not None and not remote else None,
     )
-    tools = ChatTools(
-        probe_enabled=probe_enabled,
-        compare_mode=compare_mode,
-        token_contrast=token_contrast and compare_mode,
-    )
-    return generation, tools
 
 
 def render_chat_message(

tabs/compare.py

@@ -14,7 +14,6 @@ from persona_vectors.plots import (
     build_pair_similarity_figure,
     plot_layer_similarity,
     save_plot_html,
-    save_plot_png,
 )
 
 from utils.helpers import (
@@ -147,19 +146,13 @@ def _render_save_buttons(
     filenames: list[str],
     key_suffix: str,
 ) -> None:
-    """Render Save HTML / Save PNG column buttons for one or more figures."""
-    col1, col2 = st.columns(2)
-    with col1:
-        if st.button("Save HTML", key=widget_key("load", "save_html", key_suffix)):
+    """Render the Save HTML button for one or more figures."""
+    if st.button("Save HTML", key=widget_key("load", "save_html", key_suffix)):
+        try:
             paths = [save_plot_html(fig, fn) for fig, fn in zip(figs, filenames)]
             st.success(f"Saved {len(paths)} HTML file(s) to `artifacts/plots`.")
-    with col2:
-        if st.button("Save PNG", key=widget_key("load", "save_png", key_suffix)):
-            try:
-                paths = [save_plot_png(fig, fn) for fig, fn in zip(figs, filenames)]
-                st.success(f"Saved {len(paths)} PNG file(s) to `artifacts/plots`.")
-            except Exception as exc:
-                st.error(f"Could not save PNG: {exc}")
+        except Exception as exc:
+            st.error(f"Could not save HTML: {exc}")
 
 
 def _render_mask_strategy_select(scope: str) -> MaskStrategy:
@@ -463,15 +456,7 @@ def _render_layered_figure_analysis(
         "persona_vector",
         persona_key,
     )
-    filename = _filename(
-        "compare",
-        scope,
-        store.model_name,
-        mask_strategy.value,
-        variant,
-        "persona_vector",
-        persona_key,
-    )
+    filename = scope
 
     if st.button(button_label, type="primary"):
         try:
@@ -487,6 +472,8 @@ def _render_layered_figure_analysis(
                 layers=selected_layers,
                 title=title_fn(variant),
             )
+            if figure_kind in {"umap", "pca"}:
+                main_fig.update_layout(height=700)
             extra_fig = (
                 build_pair_similarity_figure(
                     samples,

utils/datasets.py

@@ -46,6 +46,28 @@ def _uploaded_file_to_temp_path(uploaded_file: Any, stem: str) -> Path:
     return temp_path
 
 
+def load_persona_list(
+    dataset_source: str,
+    personas_file: Any = None,
+    qa_file: Any = None,
+) -> tuple[list, str]:
+    """Like ``load_dataset`` but returns ``(personas_list, status)``.
+
+    The list is memoized on the cached dataset instance so repeated reruns
+    don't pay for re-iteration.
+    """
+
+    dataset, status = load_dataset(dataset_source, personas_file, qa_file)
+    cached = getattr(dataset, "_persona_list_cache", None)
+    if cached is None:
+        cached = list(dataset)
+        try:
+            dataset._persona_list_cache = cached
+        except (AttributeError, TypeError):
+            pass
+    return cached, status
+
+
 def load_dataset(
     dataset_source: str,
     personas_file: Any = None,