Big refactoring

- Speed gains
- Improved dendogram figures
- Better information while chatting with models or loading datasets
- Faster overall ui
- Probin UI imrpovements
- Default values changed for better user experiennce
- Code structure refactoring

.env.example

@@ -25,3 +25,8 @@ ARTIFACTS_DIR=artifacts
 # PERSONA_UI_STORE_CACHE_ENTRIES=4
 # PERSONA_UI_VECTOR_CACHE_ENTRIES=4
 # PERSONA_UI_PREPARED_CACHE_ENTRIES=8
+# PERSONA_UI_FIGURE_STATE_ENTRIES=2
+# PERSONA_UI_PREPARED_STATE_ENTRIES=4
+# PERSONA_UI_PROBE_CACHE_ENTRIES=8
+# PERSONA_UI_PROBE_SWEEP_CACHE_ENTRIES=4
+# PERSONA_UI_PROBE_DERIVED_CACHE_ENTRIES=12

README.md

@@ -118,6 +118,8 @@ ARTIFACTS_DIR=...      # Optional: where persona vectors are read from (default:
 PERSONA_VECTORS_HUB_REPO=...  # Optional: default Analysis/Probing Hub dataset repo
 PERSONA_UI_VECTOR_CACHE_ENTRIES=4     # Optional: loaded analysis datasets kept warm
 PERSONA_UI_PREPARED_CACHE_ENTRIES=8   # Optional: prepared projections / k-means groups kept warm
+PERSONA_UI_FIGURE_STATE_ENTRIES=2     # Optional: recent rendered Analysis figures kept in-session
+PERSONA_UI_PREPARED_STATE_ENTRIES=4   # Optional: recent projection-ready markers kept in-session
 ```
 
 The app picks up this file automatically via `load_dotenv()` on startup.
@@ -153,4 +155,4 @@ The store classes are `PersonaVectorStore` (local) and `HFPersonaVectorStore`
 
 ## Analysis responsiveness
 
-The Analysis tab keeps a small bounded cache of loaded vector datasets and prepared projection data. Once a projection has been computed, recoloring it by persona, attribute, or k-means group reuses the same coordinates; nearby Hub interactions also keep metadata warm instead of re-scanning after every figure. Tune `PERSONA_UI_VECTOR_CACHE_ENTRIES` if RAM is tight or you regularly switch among many selections, and `PERSONA_UI_PREPARED_CACHE_ENTRIES` if you revisit several projection configurations in one session.
+The Analysis tab keeps small bounded caches of loaded vector datasets, prepared projection data, and a tiny MRU window of rendered figures. Once a projection has been computed, recoloring it by persona, attribute, or k-means group reuses the same coordinates; nearby method switches can reuse the last couple of figures instead of rebuilding immediately, while the caps keep RAM bounded. Tune `PERSONA_UI_VECTOR_CACHE_ENTRIES` if RAM is tight or you regularly switch among many selections, `PERSONA_UI_PREPARED_CACHE_ENTRIES` if you revisit several projection configurations in one session, and `PERSONA_UI_FIGURE_STATE_ENTRIES` if you want more or less method-switch warmth. Probe loading, probe sweeps, and per-trace probe outputs are bounded separately via `PERSONA_UI_PROBE_CACHE_ENTRIES`, `PERSONA_UI_PROBE_SWEEP_CACHE_ENTRIES`, and `PERSONA_UI_PROBE_DERIVED_CACHE_ENTRIES`; the derived-output cache defaults to a wider MRU window because those tensors are small compared with traced activations and are cheap wins to keep warm.

app.py

@@ -4,11 +4,7 @@ from dataclasses import dataclass
 import streamlit as st
 from dotenv import load_dotenv
 
-from utils.analysis_sources import (
-    DEFAULT_COMPARE_MODEL,
-    DEFAULT_HUB_REPO,
-    SOURCE_HUB,
-)
+from utils.analysis_sources import DEFAULT_COMPARE_MODEL, DEFAULT_HUB_REPO, SOURCE_HUB
 from utils.helpers import DATASET_SOURCES, session_key, widget_key
 from utils.preload import preload_once
 from utils.runtime import list_remote_models
@@ -60,21 +56,34 @@ def _hub_metadata_preload_calls() -> tuple[
     calls: list[tuple[str, tuple[str, str, str, str | None]]] = []
 
     def add(repo: str, model: str, mask_strategy: str, variant: str | None) -> None:
-        calls.append((
-            "utils.analysis_sources:prefetch_hub_metadata",
-            (repo, model, mask_strategy, variant),
-        ))
+        calls.append(
+            (
+                "utils.analysis_sources:prefetch_hub_metadata",
+                (repo, model, mask_strategy, variant),
+            )
+        )
+
+    shared_source = st.session_state.get("source:last_source", SOURCE_HUB)
+    shared_mask_strategy = st.session_state.get(
+        "source:last_mask_strategy", "answer_mean"
+    )
 
-    analysis_source = st.session_state.get("analysis:last_source", SOURCE_HUB)
+    analysis_source = st.session_state.get("analysis:last_source", shared_source)
     if analysis_source == SOURCE_HUB:
-        repo = st.session_state.get("analysis:hub_repo", DEFAULT_HUB_REPO)
+        repo = st.session_state.get(
+            "analysis:hub_repo",
+            st.session_state.get("source:hub_repo", DEFAULT_HUB_REPO),
+        )
         mask_strategy = st.session_state.get(
             "analysis:last_mask_strategy",
-            "answer_mean",
+            shared_mask_strategy,
         )
         model = st.session_state.get(
             widget_key("load", "hub_model", repo, mask_strategy),
-            st.session_state.get("analysis:hub_model_fallback", DEFAULT_COMPARE_MODEL),
+            st.session_state.get(
+                "analysis:hub_model_fallback",
+                st.session_state.get("source:hub_model", DEFAULT_COMPARE_MODEL),
+            ),
         )
         variant = st.session_state.get(
             "analysis:last_projection_variant",
@@ -82,16 +91,22 @@ def _hub_metadata_preload_calls() -> tuple[
         )
         add(repo, model, mask_strategy, variant)
 
-    probe_source = st.session_state.get(widget_key("probe", "source"), SOURCE_HUB)
+    probe_source = st.session_state.get(widget_key("probe", "source"), shared_source)
     if probe_source == SOURCE_HUB:
-        repo = st.session_state.get("probe:hub_repo", DEFAULT_HUB_REPO)
+        repo = st.session_state.get(
+            "probe:hub_repo",
+            st.session_state.get("source:hub_repo", DEFAULT_HUB_REPO),
+        )
         mask_strategy = st.session_state.get(
             "probe:last_mask_strategy",
-            "answer_mean",
+            shared_mask_strategy,
         )
         model = st.session_state.get(
             widget_key("probe", "hub_model", repo, mask_strategy),
-            st.session_state.get("probe:hub_model_fallback", DEFAULT_COMPARE_MODEL),
+            st.session_state.get(
+                "probe:hub_model_fallback",
+                st.session_state.get("source:hub_model", DEFAULT_COMPARE_MODEL),
+            ),
         )
         add(repo, model, mask_strategy, st.session_state.get("probe:variant"))
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "persona-ui"
-version = "0.4.0"
+version = "0.5.0"
 description = "Streamlit UI for persona-vectors"
 readme = "README.md"
 requires-python = ">=3.12"

state.py

@@ -21,9 +21,19 @@ class ChatState(TypedDict):
 
 
 def chat_session_key(model_name: str, dataset_source: str) -> str:
-    """Build the session-state key for a chat context."""
+    """Build the session-state key for a chat conversation.
 
-    return session_key("chat_state", model_name, dataset_source)
+    A model/backend switch changes *how* the next turn is generated, not which
+    conversation the user is looking at. Keeping the model out of the key means
+    toggling local/remote execution (or selecting another model) no longer makes
+    an existing thread appear to vanish behind a fresh empty state.
+
+    ``model_name`` stays in the signature for call-site compatibility and to
+    make the intent explicit where chat state is requested.
+    """
+
+    _ = model_name
+    return session_key("chat_state", dataset_source)
 
 
 def default_chat_state() -> ChatState:

tabs/analysis/_shared.py

@@ -261,6 +261,7 @@ def _render_persona_count_controls(
     *,
     default_count: int,
     include_assistant_default: bool,
+    max_count_limit: int | None = None,
 ) -> tuple[int, bool]:
     count_key = widget_key(
         "load",
@@ -280,11 +281,16 @@ def _render_persona_count_controls(
     )
 
     if options.regular_ids:
+        max_count = (
+            min(max_count_limit, len(options.regular_ids))
+            if max_count_limit is not None
+            else len(options.regular_ids)
+        )
         persona_count = st.slider(
             "Personas",
             min_value=0 if options.assistant_id is not None else 1,
-            max_value=len(options.regular_ids),
-            value=default_count,
+            max_value=max_count,
+            value=min(default_count, max_count),
             key=count_key,
             help="Use the first N available non-assistant personas.",
         )
@@ -310,6 +316,7 @@ def _select_artifact_personas(
     remember_key: str,
     default_all: bool = False,
     default_count_limit: int | None = None,
+    max_count_limit: int | None = None,
 ) -> list[str]:
     empty_message = _personas_empty_message(variants)
     options = _load_persona_options(
@@ -336,6 +343,7 @@ def _select_artifact_personas(
         options,
         default_count=default_count,
         include_assistant_default=include_assistant_default,
+        max_count_limit=max_count_limit,
     )
 
     persona_ids = options.regular_ids[:persona_count]
@@ -361,6 +369,48 @@ def _select_artifact_personas(
     return persona_ids
 
 
+def _render_persona_select_controls(
+    options: PersonaOptions,
+    widget_scope: str,
+    *,
+    max_selections: int | None = None,
+) -> list[str]:
+    select_key = widget_key("load", "persona_select", widget_scope)
+    assistant_key = widget_key("load", "persona_select_assistant", widget_scope)
+
+    label_map = {
+        persona_id: f"{options.persona_names.get(persona_id, persona_id)} ({persona_id})"
+        for persona_id in options.regular_ids
+    }
+    sorted_labels = sorted(label_map.values())
+    selected_labels = st.multiselect(
+        "Select personas",
+        options=sorted_labels,
+        key=select_key,
+        placeholder="Search and select personas...",
+        max_selections=max_selections,
+    )
+    label_to_id = {label: persona_id for persona_id, label in label_map.items()}
+    selected_ids = [label_to_id[label] for label in selected_labels]
+
+    if options.assistant_id is not None:
+        include_assistant = st.checkbox(
+            "Include Assistant persona",
+            key=assistant_key,
+        )
+        if include_assistant:
+            selected_ids.append(options.assistant_id)
+
+    st.session_state[_persona_names_state_key(widget_scope)] = dict(
+        options.persona_names
+    )
+
+    if not selected_ids:
+        st.info("Select at least one persona.")
+
+    return selected_ids
+
+
 def _render_save_buttons(
     figs: list[object],
     filenames: list[str],
@@ -398,6 +448,7 @@ def _render_mask_strategy_select(scope: str) -> MaskStrategy:
     return render_mask_strategy_select(
         key=widget_key("load", "mask_strategy", scope),
         last_key=_LAST_MASK_STRATEGY_KEY,
+        remember_key="source:last_mask_strategy",
         help_text="Which extracted activation set to load.",
     )
 
@@ -410,6 +461,8 @@ def _select_single_variant_samples(
     remember_key: str,
     variant_remember_key: str,
     default_count_limit: int,
+    max_count_limit: int | None = None,
+    allow_specific_personas: bool = False,
 ) -> tuple[str, list[str], str, list[int]] | None:
     variants = available_variants(store, mask_strategy)
     if not variants:
@@ -425,14 +478,41 @@ def _select_single_variant_samples(
         default=default_variant,
         format_func=prompt_variant_label,
     )
-    persona_ids = _select_artifact_personas(
-        store,
-        [variant],
-        mask_strategy,
-        widget_scope=f"{scope}:{store_id(store)}",
-        remember_key=remember_key,
-        default_count_limit=default_count_limit,
-    )
+    widget_scope = f"{scope}:{store_id(store)}"
+    select_specific = False
+    if allow_specific_personas:
+        select_specific = st.toggle(
+            "Select specific personas",
+            value=False,
+            key=widget_key("load", "select_specific_personas", scope, store_id(store)),
+            help="Search and select specific personas instead of using the first N.",
+        )
+
+    if select_specific:
+        options = _load_persona_options(
+            store,
+            [variant],
+            mask_strategy,
+            empty_message=_personas_empty_message([variant]),
+        )
+        if options is None:
+            st.session_state.pop(_persona_names_state_key(widget_scope), None)
+            return None
+        persona_ids = _render_persona_select_controls(
+            options,
+            widget_scope,
+            max_selections=max_count_limit,
+        )
+    else:
+        persona_ids = _select_artifact_personas(
+            store,
+            [variant],
+            mask_strategy,
+            widget_scope=widget_scope,
+            remember_key=remember_key,
+            default_count_limit=default_count_limit,
+            max_count_limit=max_count_limit,
+        )
     if not persona_ids:
         return None
 

tabs/analysis/_state.py

@@ -4,7 +4,7 @@ import streamlit as st
 from persona_data.synth_persona import BASELINE_PERSONA_ID
 from persona_vectors.attributes import DEFAULT_MAX_ATTRIBUTE_CATEGORIES
 
-from utils.helpers import slugify, widget_key
+from utils.helpers import env_int, slugify, widget_key
 
 
 def _filename(*parts: str) -> str:
@@ -30,11 +30,15 @@ _LAST_LAYER_FRAMES_KEY = "analysis:last_layer_frames"
 
 _DEFAULT_LAYER_FRAMES = 16
 _DEFAULT_PERSONA_LIMITS = {
-    "similarity": 120,
+    "similarity": 20,
     "pca": 500,
     "umap": 500,
     "isomap": 500,
-    "dendro": 160,
+    "dendro": 20,
+}
+_MAX_PERSONA_COUNTS = {
+    "similarity": 100,
+    "dendro": 100,
 }
 _MAX_SIMILARITY_CELLS = 4_000_000
 _MAX_PAIR_TRAJECTORY_TRACES = 500
@@ -136,28 +140,38 @@ def _sequence_to_list(value: object) -> list[object] | None:
 
 
 _TRACKED_STATE_KEYS_KEY = "analysis:_tracked_state_keys"
+_FIGURE_STATE_ENTRIES = env_int("PERSONA_UI_FIGURE_STATE_ENTRIES", 2)
+_PREPARED_STATE_ENTRIES = env_int("PERSONA_UI_PREPARED_STATE_ENTRIES", 4)
 
 
-def _clear_old_load_states(current_key: str, suffix: str) -> None:
-    # Only one heavy figure state should live at a time. We track
-    # the keys we create per suffix so eviction is O(1) instead of scanning
-    # all of session_state on every rerun. Every such key is passed through
-    # this function before it is set, so the registry stays authoritative.
-    tracked: dict[str, set[str]] = st.session_state.setdefault(
+def _touch_load_state(current_key: str, suffix: str, *, max_entries: int) -> None:
+    # Keep a tiny MRU window of heavy state instead of scanning all of
+    # session_state or retaining every figure forever. This makes nearby
+    # method-switching feel warm while still giving RAM a hard ceiling.
+    tracked: dict[str, list[str]] = st.session_state.setdefault(
         _TRACKED_STATE_KEYS_KEY, {}
     )
-    for key in tracked.get(suffix, ()):
-        if key != current_key:
-            st.session_state.pop(key, None)
-    tracked[suffix] = {current_key}
+    keys = [key for key in tracked.get(suffix, []) if key != current_key]
+    keys.append(current_key)
+    while len(keys) > max(1, max_entries):
+        st.session_state.pop(keys.pop(0), None)
+    tracked[suffix] = keys
 
 
 def _clear_old_figure_states(current_key: str) -> None:
-    _clear_old_load_states(current_key, "_fig_state")
+    _touch_load_state(
+        current_key,
+        "_fig_state",
+        max_entries=_FIGURE_STATE_ENTRIES,
+    )
 
 
 def _clear_old_prepared_states(current_key: str) -> None:
-    _clear_old_load_states(current_key, "_projection_ready")
+    _touch_load_state(
+        current_key,
+        "_projection_ready",
+        max_entries=_PREPARED_STATE_ENTRIES,
+    )
 
 
 def _store_figure_state(key: str, value: object) -> None:

tabs/analysis/cosine.py

@@ -4,9 +4,6 @@ import streamlit as st
 from persona_vectors.extraction import MaskStrategy
 from persona_vectors.plots import plot_layer_similarity
 
-from utils.analysis_sources import Store, available_variants, store_id
-from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
-
 from tabs.analysis._shared import (
     _load_variant_vectors,
     _plotly_chart,
@@ -21,6 +18,8 @@ from tabs.analysis._state import (
     _filename,
     _store_figure_state,
 )
+from utils.analysis_sources import Store, available_variants, store_id
+from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
 
 
 def _render_cosine_selection(

tabs/analysis/dendrogram.py

@@ -1,15 +1,10 @@
+from copy import deepcopy
+
+import plotly.graph_objects as go
 import streamlit as st
 from persona_vectors.extraction import MaskStrategy
 from persona_vectors.plots import plot_persona_dendrogram
-
-from utils.analysis_sources import (
-    Store,
-    available_variants,
-    store_cache_parts,
-    store_id,
-    store_layers_cached,
-)
-from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
+from plotly.subplots import make_subplots
 
 from tabs.analysis._shared import (
     _load_persona_options,
@@ -17,60 +12,113 @@ from tabs.analysis._shared import (
     _plotly_chart,
     _release_vector_memory,
     _render_layer_frame_controls,
+    _render_persona_select_controls,
     _render_save_buttons,
     _select_artifact_personas,
 )
 from tabs.analysis._state import (
     _DEFAULT_PERSONA_LIMITS,
-    PersonaOptions,
+    _MAX_PERSONA_COUNTS,
     _clear_old_figure_states,
     _filename,
     _persona_names_state_key,
     _personas_empty_message,
     _store_figure_state,
 )
+from utils.analysis_sources import (
+    Store,
+    available_variants,
+    store_cache_parts,
+    store_id,
+    store_layers_cached,
+)
+from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
 
 _LAST_DENDRO_PERSONAS_KEY = "analysis:last_personas:dendro"
 _DENDRO_LINKAGE_OPTIONS = ["ward", "complete", "average", "single"]
 
 
-def _render_persona_select_controls(
-    options: PersonaOptions,
-    widget_scope: str,
-) -> list[str]:
-    select_key = widget_key("load", "persona_select", widget_scope)
-    assistant_key = widget_key("load", "persona_select_assistant", widget_scope)
-
-    label_map = {
-        pid: f"{options.persona_names.get(pid, pid)} ({pid})"
-        for pid in options.regular_ids
-    }
-    sorted_labels = sorted(label_map.values())
-    selected_labels = st.multiselect(
-        "Select personas",
-        options=sorted_labels,
-        key=select_key,
-        placeholder="Search and select personas...",
+def _comparison_dendrogram_figure(
+    fig_a: go.Figure,
+    fig_b: go.Figure,
+    *,
+    title_a: str,
+    title_b: str,
+) -> go.Figure:
+    """Merge two layered dendrograms so one slider drives both panels."""
+    combined = make_subplots(
+        rows=1,
+        cols=2,
+        subplot_titles=(title_a, title_b),
+        shared_yaxes=True,
+        horizontal_spacing=0.05,
     )
-    label_to_id = {v: k for k, v in label_map.items()}
-    selected_ids = [label_to_id[lbl] for lbl in selected_labels]
+    for trace in fig_a.data:
+        combined.add_trace(deepcopy(trace), row=1, col=1)
+    for trace in fig_b.data:
+        combined.add_trace(deepcopy(trace), row=1, col=2)
 
-    if options.assistant_id is not None:
-        include_assistant = st.checkbox(
-            "Include Assistant persona",
-            key=assistant_key,
+    frames: list[go.Frame] = []
+    for frame_a, frame_b in zip(fig_a.frames, fig_b.frames, strict=True):
+        right_data = []
+        for trace in frame_b.data:
+            copied = deepcopy(trace)
+            copied.update(xaxis="x2", yaxis="y2")
+            right_data.append(copied)
+        frame_xaxis = frame_a.layout.xaxis.to_plotly_json()
+        frame_xaxis2 = frame_b.layout.xaxis.to_plotly_json()
+        frame_xaxis2["matches"] = None
+        frame_xaxis2["anchor"] = "y2"
+        frame_yaxis = frame_a.layout.yaxis.to_plotly_json()
+        frame_yaxis2 = frame_b.layout.yaxis.to_plotly_json()
+        frame_yaxis2["matches"] = "y"
+        frame_yaxis2["anchor"] = "x2"
+        frames.append(
+            go.Frame(
+                name=frame_a.name,
+                data=[*deepcopy(frame_a.data), *right_data],
+                layout={
+                    "title": {"text": f"Dendrogram comparison - Layer {frame_a.name}"},
+                    "xaxis": frame_xaxis,
+                    "xaxis2": frame_xaxis2,
+                    "yaxis": frame_yaxis,
+                    "yaxis2": frame_yaxis2,
+                },
+            )
         )
-        if include_assistant:
-            selected_ids.append(options.assistant_id)
 
-    st.session_state[_persona_names_state_key(widget_scope)] = dict(
-        options.persona_names
+    y_ranges = [
+        fig_a.layout.yaxis.range,
+        fig_b.layout.yaxis.range,
+    ]
+    max_y = max(float(axis_range[1]) for axis_range in y_ranges if axis_range)
+    first_layer = fig_a.frames[0].name if fig_a.frames else ""
+    combined.frames = frames
+    combined.update_layout(
+        title={
+            "text": f"Dendrogram comparison - Layer {first_layer}",
+            "font": {"size": 24},
+            "y": 0.98,
+            "yanchor": "top",
+        },
+        template="plotly_white",
+        height=750,
+        margin=dict(t=140, b=260),
+        updatemenus=fig_a.layout.updatemenus,
+        sliders=fig_a.layout.sliders,
     )
-
-    if not selected_ids:
-        st.info("Select at least one persona.")
-
-    return selected_ids
+    left_xaxis = fig_a.layout.xaxis.to_plotly_json()
+    right_xaxis = fig_b.layout.xaxis.to_plotly_json()
+    right_xaxis["matches"] = None
+    right_xaxis["anchor"] = "y2"
+    combined.update_layout(xaxis=left_xaxis, xaxis2=right_xaxis)
+    combined.update_xaxes(tickangle=-45, automargin=True)
+    combined.update_yaxes(
+        title_text=fig_a.layout.yaxis.title.text,
+        range=[0.0, max_y],
+        automargin=True,
+    )
+    return combined
 
 
 def _render_dendrogram_analysis(
@@ -132,6 +180,7 @@ def _render_dendrogram_analysis(
         persona_ids = _render_persona_select_controls(
             options,
             widget_scope=f"dendro:{store_id(store)}",
+            max_selections=_MAX_PERSONA_COUNTS["dendro"],
         )
         if not persona_ids:
             return
@@ -143,6 +192,7 @@ def _render_dendrogram_analysis(
             widget_scope=f"dendro:{store_id(store)}",
             remember_key=_LAST_DENDRO_PERSONAS_KEY,
             default_count_limit=_DEFAULT_PERSONA_LIMITS["dendro"],
+            max_count_limit=_MAX_PERSONA_COUNTS["dendro"],
         )
         if not persona_ids:
             return
@@ -221,7 +271,6 @@ def _render_dendrogram_analysis(
                 title=f"Dendrogram — {prompt_variant_label(variant_a)}",
             )
             fig_a.update_layout(height=750)
-            del samples_a
             fig_b = None
             if variant_a != variant_b:
                 progress.progress(60, text="Building second dendrogram…")
@@ -236,10 +285,26 @@ def _render_dendrogram_analysis(
                 )
                 fig_b.update_layout(height=750)
                 del samples_b
+            del samples_a
+            comparison_fig = None
+            if fig_b is not None and layered_mode:
+                comparison_fig = _comparison_dendrogram_figure(
+                    fig_a,
+                    fig_b,
+                    title_a=prompt_variant_label(variant_a),
+                    title_b=prompt_variant_label(variant_b),
+                )
             progress.progress(90, text="Storing figure state…")
             _store_figure_state(
                 fig_key,
-                (fig_a, fig_b, len(persona_ids), variant_a, variant_b),
+                (
+                    None if comparison_fig is not None else fig_a,
+                    None if comparison_fig is not None else fig_b,
+                    comparison_fig,
+                    len(persona_ids),
+                    variant_a,
+                    variant_b,
+                ),
             )
             progress.progress(100, text="Done.")
         except Exception as exc:
@@ -250,8 +315,16 @@ def _render_dendrogram_analysis(
             progress.empty()
 
     if fig_key in st.session_state:
-        fig_a, fig_b, n_personas, va, vb = st.session_state[fig_key]
-        if fig_b is not None:
+        saved = st.session_state[fig_key]
+        if len(saved) == 5:
+            # Drop pre-refactor state so hot-reloaded sessions do not unpack the
+            # old two-figure payload shape.
+            st.session_state.pop(fig_key, None)
+            return
+        fig_a, fig_b, comparison_fig, n_personas, va, vb = saved
+        if comparison_fig is not None:
+            _plotly_chart(comparison_fig)
+        elif fig_b is not None:
             col_a, col_b = st.columns(2)
             with col_a:
                 st.subheader(prompt_variant_label(va))
@@ -262,14 +335,22 @@ def _render_dendrogram_analysis(
         else:
             _plotly_chart(fig_a)
 
-        figs = [fig_a] + ([fig_b] if fig_b else [])
-        filenames = [
-            _filename("dendro", store.model_name, mask_strategy.value, va),
-            *(
-                [_filename("dendro", store.model_name, mask_strategy.value, vb)]
-                if fig_b
-                else []
-            ),
-        ]
+        figs = (
+            [comparison_fig]
+            if comparison_fig is not None
+            else [fig_a] + ([fig_b] if fig_b else [])
+        )
+        filenames = (
+            [_filename("dendro_compare", store.model_name, mask_strategy.value, va, vb)]
+            if comparison_fig is not None
+            else [
+                _filename("dendro", store.model_name, mask_strategy.value, va),
+                *(
+                    [_filename("dendro", store.model_name, mask_strategy.value, vb)]
+                    if fig_b
+                    else []
+                ),
+            ]
+        )
         _render_save_buttons(figs, filenames, "dendro")
         st.success(f"Generated dendrogram(s) for {n_personas} persona(s).")

tabs/analysis/layered.py

@@ -2,10 +2,7 @@ from collections.abc import Callable
 
 import plotly.graph_objects as go
 import streamlit as st
-from persona_vectors.attributes import (
-    attribute_color_kwargs,
-    attribute_display_label,
-)
+from persona_vectors.attributes import attribute_color_kwargs, attribute_display_label
 from persona_vectors.extraction import MaskStrategy
 from persona_vectors.plots import (
     build_layered_figure,
@@ -13,19 +10,6 @@ from persona_vectors.plots import (
     build_similarity_figures,
 )
 
-from utils.analysis_metadata import (
-    synth_persona_attribute_names,
-    synth_persona_dataset_cached,
-)
-from utils.analysis_sources import (
-    Store,
-    kmeans_groups_cached,
-    projection_data_cached,
-    store_cache_parts,
-    store_id,
-)
-from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
-
 from tabs.analysis._shared import (
     _gray_out_unselected_personas,
     _load_persona_vectors,
@@ -61,6 +45,18 @@ from tabs.analysis._state import (
     _remembered_selectbox,
     _store_figure_state,
 )
+from utils.analysis_metadata import (
+    synth_persona_attribute_names,
+    synth_persona_dataset_cached,
+)
+from utils.analysis_sources import (
+    Store,
+    kmeans_groups_cached,
+    projection_data_cached,
+    store_cache_parts,
+    store_id,
+)
+from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
 
 
 def _render_pair_trajectory_control(
@@ -446,6 +442,8 @@ def _render_layered_figure_analysis(
     n_components: int = 2,
     remember_key: str = _LAST_PROJECTION_PERSONAS_KEY,
     default_count_limit: int = 500,
+    max_count_limit: int | None = None,
+    allow_specific_personas: bool = False,
 ) -> None:
     """Render a single-variant layered analysis: select → button → figure(s).
 
@@ -463,6 +461,8 @@ def _render_layered_figure_analysis(
             else _LAST_SIMILARITY_VARIANT_KEY
         ),
         default_count_limit=default_count_limit,
+        max_count_limit=max_count_limit,
+        allow_specific_personas=allow_specific_personas,
     )
     if selected is None:
         return

tabs/analysis_core.py

@@ -1,27 +1,4 @@
-from pathlib import Path
-
 import streamlit as st
-from persona_data.environment import get_artifacts_dir
-from persona_vectors.extraction import MaskStrategy
-
-from utils.analysis_sources import (
-    DEFAULT_COMPARE_MODEL,
-    DEFAULT_HUB_REPO,
-    SOURCE_HUB,
-    SOURCE_LOCAL,
-    SOURCES,
-    Store,
-    activation_store_cached,
-    hub_models_by_mask_strategy,
-    local_model_matches,
-    local_model_options_cached,
-)
-from utils.helpers import (
-    ANALYSIS_HELP_TEXT,
-    ANALYSIS_MODES,
-    prompt_variant_label,
-    widget_key,
-)
 
 from tabs.analysis._shared import _render_mask_strategy_select
 from tabs.analysis._state import (
@@ -29,153 +6,18 @@ from tabs.analysis._state import (
     _LAST_PROJECTION_DIMS_KEY,
     _LAST_SIMILARITY_PERSONAS_KEY,
     _LAST_SOURCE_KEY,
+    _MAX_PERSONA_COUNTS,
 )
 from tabs.analysis.cosine import _render_cosine_similarity
 from tabs.analysis.dendrogram import _render_dendrogram_analysis
 from tabs.analysis.layered import _render_layered_figure_analysis
-
-
-def _render_source_select() -> str:
-    last_source = st.session_state.get(_LAST_SOURCE_KEY, SOURCE_HUB)
-    source = st.segmented_control(
-        "Source",
-        options=SOURCES,
-        default=last_source if last_source in SOURCES else SOURCE_HUB,
-        key=widget_key("load", "source"),
-        label_visibility="collapsed",
-    )
-    if source is None:
-        source = SOURCE_HUB
-    st.session_state[_LAST_SOURCE_KEY] = source
-    return source
-
-
-def _render_hub_model_select(
-    repo_id: str,
-    mask_strategy: MaskStrategy,
-) -> str:
-    fallback_model = st.session_state.get(
-        "analysis:hub_model_fallback",
-        DEFAULT_COMPARE_MODEL,
-    )
-    try:
-        models_by_strategy = hub_models_by_mask_strategy(repo_id)
-    except Exception as exc:
-        st.warning(f"Could not load Hub configs for `{repo_id}`: {exc}")
-        return st.text_input(
-            "Hub model",
-            value=fallback_model,
-            key="analysis:hub_model_fallback",
-            help="Analysis-only model id to use if Hub config discovery is unavailable.",
-        )
-
-    model_options = models_by_strategy.get(mask_strategy, [])
-    if not model_options:
-        st.warning(
-            f"No Hub vector configs found for `{mask_strategy.value}` in `{repo_id}`."
-        )
-        return st.text_input(
-            "Hub model",
-            value=fallback_model,
-            key="analysis:hub_model_fallback",
-            help="Analysis-only model id to use for this Hub repo.",
-        )
-
-    previous_model = st.session_state.get(
-        widget_key("load", "hub_model", repo_id, mask_strategy.value),
-        fallback_model,
-    )
-    default_model = (
-        previous_model if previous_model in model_options else model_options[0]
-    )
-
-    return st.selectbox(
-        "Hub model",
-        options=model_options,
-        index=model_options.index(default_model),
-        key=widget_key("load", "hub_model", repo_id, mask_strategy.value),
-        help="Models with vectors in the selected Hub repo and mask strategy.",
-    )
-
-
-def _render_local_model_select(
-    artifacts_root: str,
-    mask_strategy: MaskStrategy,
-) -> str:
-    fallback_model = st.session_state.get("analysis:local_model", DEFAULT_COMPARE_MODEL)
-    model_options = local_model_options_cached(artifacts_root, mask_strategy.value)
-    if not model_options:
-        return st.text_input(
-            "Local model",
-            value=fallback_model,
-            key="analysis:local_model",
-            help="Analysis-only local model id or path.",
-        )
-
-    custom = st.toggle(
-        "Custom local model",
-        value=False,
-        key="analysis:local_model_custom_enabled",
-        help="Enter a model id/path manually instead of choosing from activation directories.",
-    )
-    if custom:
-        return st.text_input(
-            "Local model",
-            value=fallback_model,
-            key="analysis:local_model",
-            help="Analysis-only local model id or path.",
-        )
-
-    previous_model = st.session_state.get("analysis:local_model_select", fallback_model)
-    if not any(local_model_matches(previous_model, option) for option in model_options):
-        previous_model = fallback_model
-    default_model = next(
-        (
-            option
-            for option in model_options
-            if local_model_matches(option, previous_model)
-        ),
-        model_options[0],
-    )
-    selected = st.selectbox(
-        "Local model",
-        options=model_options,
-        index=model_options.index(default_model),
-        key="analysis:local_model_select",
-        help="Models discovered under the selected artifacts root.",
-    )
-    st.session_state["analysis:local_model"] = selected
-    return selected
-
-
-def _build_store(source: str, mask_strategy: MaskStrategy) -> Store:
-    if source == SOURCE_HUB:
-        repo = st.text_input(
-            "Hub repo",
-            value=st.session_state.get("analysis:hub_repo", DEFAULT_HUB_REPO),
-            key="analysis:hub_repo",
-            help="Hugging Face dataset published by `scripts/push_to_hf.py`.",
-        )
-        hub_model_name = _render_hub_model_select(repo, mask_strategy)
-        return activation_store_cached(
-            SOURCE_HUB,
-            repo,
-            hub_model_name,
-            mask_strategy.value,
-        )
-    artifacts_root = st.text_input(
-        "Artifacts root",
-        value=str(get_artifacts_dir() / "activations"),
-        key="analysis:artifacts_root",
-    )
-    artifacts_root = str(Path(artifacts_root).expanduser())
-    local_model_name = _render_local_model_select(artifacts_root, mask_strategy)
-    return activation_store_cached(
-        SOURCE_LOCAL,
-        artifacts_root,
-        local_model_name,
-        mask_strategy.value,
-    )
+from utils.helpers import (
+    ANALYSIS_HELP_TEXT,
+    ANALYSIS_MODES,
+    prompt_variant_label,
+    widget_key,
+)
+from utils.source_controls import render_source_select, render_store_select
 
 
 def render_analysis_tab() -> None:
@@ -186,7 +28,7 @@ def render_analysis_tab() -> None:
         "Analyse persona vectors by cosine similarity, PCA, UMAP, Isomap, or hierarchical clustering."
     )
 
-    source = _render_source_select()
+    source = render_source_select(widget_scope="load", last_source_key=_LAST_SOURCE_KEY)
 
     analysis_mode = st.segmented_control(
         "Analysis mode",
@@ -201,7 +43,18 @@ def render_analysis_tab() -> None:
 
     with st.expander("Source settings", expanded=True):
         mask_strategy = _render_mask_strategy_select(analysis_mode)
-        store = _build_store(source, mask_strategy)
+        store = render_store_select(
+            source,
+            mask_strategy,
+            state_prefix="analysis",
+            widget_scope="load",
+            artifacts_root_key="analysis:artifacts_root",
+            model_label="Hub model",
+            local_model_label="Local model",
+            allow_custom_local_model=True,
+            repo_help="Hugging Face dataset published by `scripts/push_to_hf.py`.",
+            fallback_help="Analysis-only model id to use if Hub config discovery is unavailable.",
+        )
 
     if analysis_mode == "Cosine similarity":
         _render_cosine_similarity(store, mask_strategy)
@@ -219,6 +72,8 @@ def render_analysis_tab() -> None:
             include_pair_trajectories=True,
             remember_key=_LAST_SIMILARITY_PERSONAS_KEY,
             default_count_limit=_DEFAULT_PERSONA_LIMITS["similarity"],
+            max_count_limit=_MAX_PERSONA_COUNTS["similarity"],
+            allow_specific_personas=True,
         )
         return
 

tabs/chat.py

@@ -15,6 +15,8 @@ from tabs.chat_shared import (
     generate_chat_reply_result,
     hydrate_chat_state,
     load_chat_personas,
+    mark_model_loaded,
+    model_load_status,
     render_chat_selection,
 )
 from tabs.chat_ui import (
@@ -25,7 +27,7 @@ from tabs.chat_ui import (
 )
 from utils.chat import build_chat_messages, resolve_system_prompt
 from utils.chat_export import save_chat_export
-from utils.helpers import session_key, widget_key
+from utils.helpers import format_ndif_status, session_key, widget_key
 from utils.runtime import cached_model
 
 if TYPE_CHECKING:
@@ -94,9 +96,26 @@ def _handle_single_chat_generation(
     chat_log,
 ) -> None:
     messages = build_chat_messages(active_system_prompt, chat_state["messages"])
+    status_box = st.empty()
+
+    def _show_phase(text: str) -> None:
+        status_box.caption(text)
+
+    def _show_ndif_status(job_id: str, status_name: str, description: str) -> None:
+        status_box.caption(
+            format_ndif_status(
+                job_id,
+                status_name,
+                description,
+                completed_detail="Downloading result...",
+            )
+        )
 
     with st.spinner("Generating reply..."):
+        _show_phase(model_load_status(model_name))
         model = cached_model(model_name=model_name)
+        mark_model_loaded(model_name)
+        _show_phase("Submitting to NDIF..." if remote else "Generating locally...")
 
         def _show_error(exc: Exception) -> None:
             with chat_log:
@@ -108,15 +127,19 @@ def _handle_single_chat_generation(
             messages=messages,
             remote=remote,
             generation=generation,
+            on_status=_show_ndif_status if remote else None,
             on_error=_show_error,
         )
         if error is not None:
+            status_box.empty()
             if pending_action == "new_user_prompt" and chat_state["messages"]:
                 chat_state["messages"].pop()
             return
         if reply is None:
+            status_box.empty()
             return
 
+    status_box.empty()
     chat_state["messages"].append({"role": "assistant", "content": reply.text})
     st.rerun()
 

tabs/chat_shared.py

@@ -23,6 +23,9 @@ class ChatSelection:
     changed: bool
 
 
+_LOADED_MODEL_NAMES_KEY = session_key("chat", "loaded_model_names")
+
+
 def load_chat_personas(dataset_source: str) -> list[PersonaData] | None:
     personas_file_key = session_key("extract", "personas_file")
     qa_file_key = session_key("extract", "qa_file")
@@ -84,12 +87,27 @@ def render_chat_selection(
     return ChatSelection(selected_persona, prompt_mode, changed)
 
 
+def model_load_status(model_name: str) -> str:
+    """Return an honest coarse-grained loading label for the current session."""
+
+    loaded_names = st.session_state.setdefault(_LOADED_MODEL_NAMES_KEY, set())
+    return "Using cached model..." if model_name in loaded_names else "Loading model..."
+
+
+def mark_model_loaded(model_name: str) -> None:
+    """Remember that this session has already requested a model once."""
+
+    loaded_names = st.session_state.setdefault(_LOADED_MODEL_NAMES_KEY, set())
+    loaded_names.add(model_name)
+
+
 def generate_chat_reply_result(
     *,
     model: object,
     messages: list[dict[str, str]],
     remote: bool,
     generation: GenerationConfig,
+    on_status: Callable[[str, str, str], None] | None = None,
     on_error: Callable[[Exception], None] | None = None,
 ) -> tuple[ChatReply | None, Exception | None]:
     try:
@@ -98,6 +116,7 @@ def generate_chat_reply_result(
                 model=model,
                 messages=messages,
                 remote=remote,
+                on_status=on_status,
                 **generation.to_generate_kwargs(),
             ),
             None,

tabs/chat_ui.py

@@ -16,6 +16,7 @@ from utils.helpers import (
 
 if TYPE_CHECKING:
     from persona_data.synth_persona import PersonaData
+
     from utils.contrast import TokenContrast
 
 GENERATION_DEFAULTS = {

tabs/compare_chat.py

@@ -14,7 +14,7 @@ from tabs.chat_shared import (
 from utils.chat import ChatReply, build_chat_messages, resolve_system_prompt
 from utils.chat_export import save_chat_export
 from utils.contrast import compute_contrast, compute_contrast_pair
-from utils.helpers import persona_label, session_key, widget_key
+from utils.helpers import format_ndif_status, persona_label, session_key, widget_key
 from utils.runtime import cached_model
 
 from .chat_ui import (
@@ -142,15 +142,40 @@ def _generate_panels(
     spinner_label: str,
 ) -> list[ChatReply | Exception]:
     results: list[ChatReply | Exception] = []
+    status_box = st.empty()
     with st.spinner(spinner_label):
         for panel in panels:
+            panel_label = panel.side.title()
+            status_box.caption(
+                f"{panel_label}: {'Submitting to NDIF...' if remote else 'Generating locally...'}"
+            )
+
+            def _show_ndif_status(
+                job_id: str,
+                status_name: str,
+                description: str,
+                *,
+                label: str = panel_label,
+            ) -> None:
+                status_box.caption(
+                    format_ndif_status(
+                        job_id,
+                        status_name,
+                        description,
+                        prefix=label,
+                        completed_detail="Downloading result...",
+                    )
+                )
+
             reply, error = generate_chat_reply_result(
                 model=model,
                 messages=build_chat_messages(panel.prompt, panel.state["messages"]),
                 remote=remote,
                 generation=generation,
+                on_status=_show_ndif_status if remote else None,
             )
             results.append(reply if error is None else error)
+    status_box.empty()
     return results
 
 

tabs/extract.py

@@ -20,7 +20,7 @@ from utils.datasets import (
     warm_qa_in_background,
 )
 from utils.helpers import (
-    NDIF_STATUS_ICONS,
+    format_ndif_status,
     persona_label,
     prompt_variant_label,
     session_key,
@@ -353,8 +353,7 @@ def _run_extraction_plan(
     ndif_status_box = st.empty()
 
     def _on_ndif_status(job_id: str, status_name: str, description: str) -> None:
-        icon = NDIF_STATUS_ICONS.get(status_name, "•")
-        ndif_status_box.caption(f"{icon} `{job_id}` **{status_name}** — {description}")
+        ndif_status_box.caption(format_ndif_status(job_id, status_name, description))
 
     with st.spinner("Loading model..."):
         model = cached_model(model_name=model_name)

tabs/probe.py

@@ -11,43 +11,28 @@ is a thin Streamlit wrapper around them.
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from pathlib import Path
-
 import streamlit as st
-from persona_data.environment import get_artifacts_dir
 from persona_vectors.analysis import LayeredSamples
 from persona_vectors.attributes import attribute_display_label
 from persona_vectors.extraction import MaskStrategy
 from persona_vectors.plots import plot_metric_comparison, plot_metric_over_layers
 from persona_vectors.probes import (
     AttributeLabels,
-    attribute_probe_labels,
     default_probe_kinds,
-    filter_attribute_samples_min_count,
     infer_probe_task,
     layer_matrix,
     save_probe_artifact,
     shuffle_label_baseline,
-    sweep_attribute,
 )
 
+from tabs.probe_sweep import SweepInputs, cached_sweep
 from utils.analysis_metadata import (
     synth_persona_attribute_names,
     synth_persona_dataset_cached,
 )
 from utils.analysis_sources import (
-    DEFAULT_COMPARE_MODEL,
-    DEFAULT_HUB_REPO,
-    SOURCE_HUB,
-    SOURCE_LOCAL,
-    SOURCES,
     Store,
-    activation_store_cached,
     available_variants,
-    hub_models_by_mask_strategy,
-    load_persona_vectors_cached,
-    local_model_options_cached,
     persona_names_cached,
     personas_cached,
     store_cache_parts,
@@ -55,6 +40,7 @@ from utils.analysis_sources import (
 )
 from utils.controls import render_mask_strategy_select
 from utils.helpers import widget_key
+from utils.source_controls import render_source_select, render_store_select
 
 # ---------------------------------------------------------------------------
 # Constants and config
@@ -78,94 +64,6 @@ _SECONDARY_METRIC = {
 }
 
 
-@dataclass(frozen=True)
-class _SweepInputs:
-    source: str
-    location: str
-    model_name: str
-    mask_value: str
-    variant: str
-    persona_ids: tuple[str, ...]
-    attributes: tuple[str, ...]
-    task: str
-    probe_kinds: tuple[str, ...]
-    n_pca_components: int | None
-    layers: tuple[int, ...]
-    min_class_count: int
-    seed: int
-
-
-# ---------------------------------------------------------------------------
-# Source / store selection (slim mirror of the analysis tab pattern)
-# ---------------------------------------------------------------------------
-
-
-def _select_source() -> str:
-    key = widget_key("probe", "source")
-    source = st.segmented_control(
-        "Source",
-        options=SOURCES,
-        default=st.session_state.get(key, SOURCE_HUB),
-        key=key,
-        label_visibility="collapsed",
-    )
-    return source or SOURCE_HUB
-
-
-def _select_store(source: str, mask_strategy: MaskStrategy) -> Store:
-    if source == SOURCE_HUB:
-        repo = st.text_input(
-            "Hub repo",
-            value=st.session_state.get("probe:hub_repo", DEFAULT_HUB_REPO),
-            key="probe:hub_repo",
-        )
-        models = hub_models_by_mask_strategy(repo).get(mask_strategy, [])
-        if not models:
-            st.warning(
-                f"No Hub vector configs for `{mask_strategy.value}` in `{repo}`."
-            )
-            model_name = st.text_input(
-                "Model",
-                value=st.session_state.get("probe:hub_model_fallback", DEFAULT_COMPARE_MODEL),
-                key="probe:hub_model_fallback",
-            )
-        else:
-            previous = st.session_state.get(
-                widget_key("probe", "hub_model", repo, mask_strategy.value),
-                models[0],
-            )
-            model_name = st.selectbox(
-                "Model",
-                options=models,
-                index=models.index(previous) if previous in models else 0,
-                key=widget_key("probe", "hub_model", repo, mask_strategy.value),
-            )
-        return activation_store_cached(SOURCE_HUB, repo, model_name, mask_strategy.value)
-
-    root = st.text_input(
-        "Artifacts root",
-        value=str(get_artifacts_dir() / "activations"),
-        key="probe:local_root",
-    )
-    root = str(Path(root).expanduser())
-    models = local_model_options_cached(root, mask_strategy.value)
-    if models:
-        previous = st.session_state.get("probe:local_model", models[0])
-        model_name = st.selectbox(
-            "Model",
-            options=models,
-            index=models.index(previous) if previous in models else 0,
-            key="probe:local_model",
-        )
-    else:
-        model_name = st.text_input(
-            "Model",
-            value=st.session_state.get("probe:local_model_fallback", DEFAULT_COMPARE_MODEL),
-            key="probe:local_model_fallback",
-        )
-    return activation_store_cached(SOURCE_LOCAL, root, model_name, mask_strategy.value)
-
-
 def _select_variant(store: Store, mask_strategy: MaskStrategy) -> str | None:
     variants = available_variants(store, mask_strategy)
     if not variants:
@@ -184,7 +82,9 @@ def _select_personas(
     store: Store, variant: str, mask_strategy: MaskStrategy
 ) -> list[str]:
     source, location, model_name = store_cache_parts(store)
-    all_ids = personas_cached(source, location, model_name, mask_strategy.value, (variant,))
+    all_ids = personas_cached(
+        source, location, model_name, mask_strategy.value, (variant,)
+    )
     if not all_ids:
         st.info("No personas found for this variant.")
         return []
@@ -225,7 +125,12 @@ def _select_personas(
     st.session_state["probe:persona_count"] = count
     persona_ids = regular[:count]
     persona_names_cached(
-        source, location, model_name, mask_strategy.value, (variant,), tuple(persona_ids)
+        source,
+        location,
+        model_name,
+        mask_strategy.value,
+        (variant,),
+        tuple(persona_ids),
     )
     st.caption(f"Probing {len(persona_ids)} of {len(regular)} non-assistant personas.")
     return persona_ids
@@ -323,13 +228,15 @@ def _select_layers(num_layers: int) -> list[int]:
     )
     if not fast:
         return list(range(num_layers))
-    return sorted({
-        0,
-        num_layers // 4,
-        num_layers // 2,
-        (3 * num_layers) // 4,
-        num_layers - 1,
-    })
+    return sorted(
+        {
+            0,
+            num_layers // 4,
+            num_layers // 2,
+            (3 * num_layers) // 4,
+            num_layers - 1,
+        }
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -337,66 +244,12 @@ def _select_layers(num_layers: int) -> list[int]:
 # ---------------------------------------------------------------------------
 
 
-@st.cache_resource(show_spinner=False)
-def _cached_sweep(
-    inputs: _SweepInputs,
-) -> tuple[
-    dict[str, list[dict[str, object]]],
-    dict[str, tuple[AttributeLabels, LayeredSamples]],
-]:
-    samples = load_persona_vectors_cached(
-        inputs.source, inputs.location, inputs.model_name,
-        inputs.mask_value, inputs.variant, inputs.persona_ids,
-    )
-    dataset = synth_persona_dataset_cached()
-    # The min-count filter drops personas per attribute, so each attribute keeps
-    # its own (labels, samples) pair for the downstream selectivity/save tools.
-    per_attr: dict[str, tuple[AttributeLabels, LayeredSamples]] = {}
-
-    def _labels_and_samples(attribute: str) -> tuple[AttributeLabels, LayeredSamples]:
-        if attribute not in per_attr:
-            labels = attribute_probe_labels(
-                dataset, attribute, list(inputs.persona_ids), task=inputs.task,  # type: ignore[arg-type]
-            )
-            probe_samples, labels = filter_attribute_samples_min_count(
-                samples, labels, min_count=inputs.min_class_count
-            )
-            per_attr[attribute] = (labels, probe_samples)
-        return per_attr[attribute]
-
-    def _sweep(attribute: str, n_pca: int | None) -> list[dict[str, object]]:
-        labels, probe_samples = _labels_and_samples(attribute)
-        return sweep_attribute(
-            probe_samples, labels,
-            layers=list(inputs.layers),
-            probe_kinds=list(inputs.probe_kinds),  # type: ignore[arg-type]
-            n_pca_components=n_pca,
-            seed=inputs.seed,
-        )
-
-    def _sweep_all(n_pca: int | None) -> list[dict[str, object]]:
-        rows: list[dict[str, object]] = []
-        for attribute in inputs.attributes:
-            rows.extend(_sweep(attribute, n_pca))
-        return rows
-
-    if inputs.n_pca_components is not None:
-        # Always overlay the compressed sweep against full activations.
-        rows_by_label = {
-            "full": _sweep_all(None),
-            f"pca{inputs.n_pca_components}": _sweep_all(inputs.n_pca_components),
-        }
-    else:
-        rows_by_label = {"full": _sweep_all(None)}
-    return rows_by_label, per_attr
-
-
 def _show_sweep(
     rows_by_label: dict[str, list[dict[str, object]]],
     per_attr: dict[str, tuple[AttributeLabels, LayeredSamples]],
     attributes: tuple[str, ...],
     task: str,
-    inputs: _SweepInputs,
+    inputs: SweepInputs,
 ) -> None:
     primary = _PRIMARY_METRIC[task]
     secondary = _SECONDARY_METRIC.get(task)
@@ -442,8 +295,7 @@ def _show_sweep(
         for label, label_rows in rows_by_label.items():
             for attribute in attributes:
                 attr_rows = [
-                    row for row in label_rows
-                    if row.get("attribute") == attribute
+                    row for row in label_rows if row.get("attribute") == attribute
                 ]
                 label_best = _best_row(attr_rows)
                 if label_best is None:
@@ -451,22 +303,23 @@ def _show_sweep(
                 summary_row: dict[str, object] = {}
                 if multi_attr:
                     summary_row["attribute"] = attribute
-                summary_row.update({
-                    "features": label,
-                    "best_layer": label_best["layer"],
-                    "probe": label_best["probe_kind"],
-                    primary: round(float(label_best[primary]), 3),
-                    f"baseline_{primary}": round(
-                        float(label_best.get(f"baseline_{primary}", float("nan"))), 3
-                    ),
-                })
+                summary_row.update(
+                    {
+                        "features": label,
+                        "best_layer": label_best["layer"],
+                        "probe": label_best["probe_kind"],
+                        primary: round(float(label_best[primary]), 3),
+                        f"baseline_{primary}": round(
+                            float(label_best.get(f"baseline_{primary}", float("nan"))),
+                            3,
+                        ),
+                    }
+                )
                 summary_rows.append(summary_row)
         if summary_rows:
             st.dataframe(summary_rows, width="stretch", hide_index=True)
 
-    feature_desc = (
-        f" · pca{inputs.n_pca_components}" if inputs.n_pca_components else ""
-    )
+    feature_desc = f" · pca{inputs.n_pca_components}" if inputs.n_pca_components else ""
 
     best_attr = str(best["attribute"])
     labels, samples = per_attr[best_attr]
@@ -495,7 +348,7 @@ def _render_selectivity_control(
     labels: AttributeLabels,
     samples: LayeredSamples,
     task: str,
-    inputs: _SweepInputs,
+    inputs: SweepInputs,
 ) -> None:
     if task == "numeric":
         return  # selectivity control is classification-only
@@ -507,14 +360,18 @@ def _render_selectivity_control(
             "dataset artifacts, not the property."
         )
         n_repeats = st.slider(
-            "Shuffle repeats", min_value=3, max_value=15, value=5,
+            "Shuffle repeats",
+            min_value=3,
+            max_value=15,
+            value=5,
             key="probe:shuffle_repeats",
         )
         if st.button("Run selectivity control", key="probe:run_shuffle"):
             with st.spinner("Running shuffled-label control..."):
                 X = layer_matrix(samples, int(best["layer"]))
                 shuffled = shuffle_label_baseline(
-                    X, labels.y,
+                    X,
+                    labels.y,
                     task=task,  # type: ignore[arg-type]
                     layer=int(best["layer"]),
                     probe_kind=best["probe_kind"],  # type: ignore[arg-type]
@@ -539,7 +396,7 @@ def _render_save_artifact(
     labels: AttributeLabels,
     samples: LayeredSamples,
     task: str,
-    inputs: _SweepInputs,
+    inputs: SweepInputs,
 ) -> None:
     def synced_default(key: str, default: str) -> str:
         default_key = f"{key}:default"
@@ -575,7 +432,9 @@ def _render_save_artifact(
         if st.button("Save", key="probe:save_artifact"):
             X = layer_matrix(samples, int(best["layer"]))
             directory = save_probe_artifact(
-                X=X, y=labels.y, labels=labels,
+                X=X,
+                y=labels.y,
+                labels=labels,
                 task=task,  # type: ignore[arg-type]
                 probe_kind=best["probe_kind"],  # type: ignore[arg-type]
                 n_pca_components=inputs.n_pca_components,
@@ -601,14 +460,21 @@ def _render_save_artifact(
 def render_probing_tab() -> None:
     st.title("Probing")
 
-    source = _select_source()
+    source = render_source_select(widget_scope="probe")
     with st.expander("Source", expanded=True):
         mask_strategy = render_mask_strategy_select(
             key=widget_key("probe", "mask_strategy"),
             last_key="probe:last_mask_strategy",
+            remember_key="source:last_mask_strategy",
             help_text="Which extracted activation set to load.",
         )
-        store = _select_store(source, mask_strategy)
+        store = render_store_select(
+            source,
+            mask_strategy,
+            state_prefix="probe",
+            widget_scope="probe",
+            artifacts_root_key="probe:local_root",
+        )
         variant = _select_variant(store, mask_strategy)
         if variant is None:
             return
@@ -644,13 +510,19 @@ def render_probing_tab() -> None:
         min_class_count = _MIN_CLASS_COUNT
         seed = 0
 
-    inputs = _SweepInputs(
-        source=source, location=location, model_name=model_name,
-        mask_value=mask_strategy.value, variant=variant,
-        persona_ids=tuple(persona_ids), attributes=tuple(attributes), task=task,
+    inputs = SweepInputs(
+        source=source,
+        location=location,
+        model_name=model_name,
+        mask_value=mask_strategy.value,
+        variant=variant,
+        persona_ids=tuple(persona_ids),
+        attributes=tuple(attributes),
+        task=task,
         probe_kinds=tuple(probe_kinds),
         n_pca_components=n_pca_components,
-        layers=tuple(layers), min_class_count=min_class_count,
+        layers=tuple(layers),
+        min_class_count=min_class_count,
         seed=int(seed),
     )
 
@@ -659,7 +531,7 @@ def render_probing_tab() -> None:
     if run:
         with st.spinner("Evaluating probes across layers..."):
             try:
-                sweep, per_attr = _cached_sweep(inputs)
+                sweep, per_attr = cached_sweep(inputs)
             except Exception as exc:
                 st.error(f"Sweep failed: {exc}")
                 st.session_state.pop(state_key, None)
@@ -674,6 +546,9 @@ def render_probing_tab() -> None:
         else:
             sweep, per_attr, result_inputs = saved_result
             _show_sweep(
-                sweep, per_attr, result_inputs.attributes,
-                result_inputs.task, result_inputs,
+                sweep,
+                per_attr,
+                result_inputs.attributes,
+                result_inputs.task,
+                result_inputs,
             )

tabs/probe_sweep.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import streamlit as st
+from persona_vectors.analysis import LayeredSamples
+from persona_vectors.probes import (
+    AttributeLabels,
+    attribute_probe_labels,
+    filter_attribute_samples_min_count,
+    sweep_attribute,
+)
+
+from utils.analysis_metadata import synth_persona_dataset_cached
+from utils.analysis_sources import load_persona_vectors_cached
+from utils.helpers import env_int
+
+_SWEEP_CACHE_ENTRIES = env_int("PERSONA_UI_PROBE_SWEEP_CACHE_ENTRIES", 4)
+
+
+@dataclass(frozen=True)
+class SweepInputs:
+    source: str
+    location: str
+    model_name: str
+    mask_value: str
+    variant: str
+    persona_ids: tuple[str, ...]
+    attributes: tuple[str, ...]
+    task: str
+    probe_kinds: tuple[str, ...]
+    n_pca_components: int | None
+    layers: tuple[int, ...]
+    min_class_count: int
+    seed: int
+
+
+@st.cache_resource(show_spinner=False, max_entries=_SWEEP_CACHE_ENTRIES)
+def cached_sweep(
+    inputs: SweepInputs,
+) -> tuple[
+    dict[str, list[dict[str, object]]],
+    dict[str, tuple[AttributeLabels, LayeredSamples]],
+]:
+    samples = load_persona_vectors_cached(
+        inputs.source,
+        inputs.location,
+        inputs.model_name,
+        inputs.mask_value,
+        inputs.variant,
+        inputs.persona_ids,
+    )
+    dataset = synth_persona_dataset_cached()
+    per_attr: dict[str, tuple[AttributeLabels, LayeredSamples]] = {}
+
+    def labels_and_samples(attribute: str) -> tuple[AttributeLabels, LayeredSamples]:
+        if attribute not in per_attr:
+            labels = attribute_probe_labels(
+                dataset,
+                attribute,
+                list(inputs.persona_ids),
+                task=inputs.task,  # type: ignore[arg-type]
+            )
+            probe_samples, labels = filter_attribute_samples_min_count(
+                samples,
+                labels,
+                min_count=inputs.min_class_count,
+            )
+            per_attr[attribute] = (labels, probe_samples)
+        return per_attr[attribute]
+
+    def sweep_one(attribute: str, n_pca: int | None) -> list[dict[str, object]]:
+        labels, probe_samples = labels_and_samples(attribute)
+        return sweep_attribute(
+            probe_samples,
+            labels,
+            layers=list(inputs.layers),
+            probe_kinds=list(inputs.probe_kinds),  # type: ignore[arg-type]
+            n_pca_components=n_pca,
+            seed=inputs.seed,
+        )
+
+    def sweep_all(n_pca: int | None) -> list[dict[str, object]]:
+        rows: list[dict[str, object]] = []
+        for attribute in inputs.attributes:
+            rows.extend(sweep_one(attribute, n_pca))
+        return rows
+
+    rows_by_label = {"full": sweep_all(None)}
+    if inputs.n_pca_components is not None:
+        rows_by_label[f"pca{inputs.n_pca_components}"] = sweep_all(
+            inputs.n_pca_components
+        )
+    return rows_by_label, per_attr

tabs/probe_ui.py

@@ -6,7 +6,15 @@ import streamlit as st
 import torch
 
 from utils.chat import build_chat_messages
-from utils.helpers import session_key, widget_key
+from utils.helpers import env_int, session_key, widget_key
+from utils.probe_files import (
+    DEFAULT_LOCAL_PROBE_DIR,
+    DEFAULT_PROBE_REPO,
+    list_local_probe_files,
+    list_probe_files,
+    model_probe_dir_name,
+    parse_probe_filename,
+)
 from utils.probe_overlay import (
     attach_overlays,
     build_classification_overlays,
@@ -15,24 +23,25 @@ from utils.probe_overlay import (
 )
 from utils.probe_trace import ConversationTrace, trace_conversation
 from utils.probes import (
-    DEFAULT_LOCAL_PROBE_DIR,
-    DEFAULT_PROBE_REPO,
     LoadedProbe,
-    list_local_probe_files,
-    list_probe_files,
     load_local_probe,
     load_probe,
     load_probe_from_bytes,
-    model_probe_dir_name,
-    parse_probe_filename,
 )
 from utils.runtime import cached_model
+from utils.selection_controls import remembered_segmented_control
 
 _LAST_SOURCE_KEY = session_key("probe", "last_source")
 _LAST_LOCAL_FILE_KEY = session_key("probe", "last_local_file")
 _LAST_HUB_FILE_KEY = session_key("probe", "last_hub_file")
 
 _PROBE_SOURCES = ("Local artifact", "Hugging Face repo", "Upload .pt")
+_DERIVED_CACHE_TRACKER_KEY = session_key("probe", "derived_cache_keys")
+# Keep enough room for the three retained traces plus a few recently explored
+# probes per trace. Derived outputs are much smaller than the trace activations
+# themselves, so this avoids needless recomputation without reopening
+# unbounded growth.
+_DERIVED_CACHE_ENTRIES = env_int("PERSONA_UI_PROBE_DERIVED_CACHE_ENTRIES", 12)
 
 
 # ---------------------------------------------------------------------------
@@ -62,23 +71,16 @@ def _default_file(files: list[str], remembered: str | None) -> str:
     return files[0]
 
 
-def _render_probe_selector(
-    *, context_key: str, model_name: str
-) -> LoadedProbe | None:
+def _render_probe_selector(*, context_key: str, model_name: str) -> LoadedProbe | None:
     """Inline source + file selector. Returns the loaded probe or None."""
-    source_key = widget_key(context_key, "probe_source")
-    if source_key not in st.session_state:
-        st.session_state[source_key] = st.session_state.get(
-            _LAST_SOURCE_KEY, _PROBE_SOURCES[0]
-        )
-    source = st.segmented_control(
+    source = remembered_segmented_control(
         "Probe source",
         options=_PROBE_SOURCES,
-        key=source_key,
+        key=widget_key(context_key, "probe_source"),
+        remember_key=_LAST_SOURCE_KEY,
+        default=_PROBE_SOURCES[0],
         label_visibility="collapsed",
     )
-    source = source or _PROBE_SOURCES[0]
-    st.session_state[_LAST_SOURCE_KEY] = source
 
     if source == "Local artifact":
         return _render_local_probe(context_key=context_key, model_name=model_name)
@@ -87,9 +89,7 @@ def _render_probe_selector(
     return _render_upload_probe(context_key=context_key)
 
 
-def _render_local_probe(
-    *, context_key: str, model_name: str
-) -> LoadedProbe | None:
+def _render_local_probe(*, context_key: str, model_name: str) -> LoadedProbe | None:
     root_dir = st.text_input(
         "Probe directory",
         value=st.session_state.get(
@@ -118,9 +118,7 @@ def _render_local_probe(
         return None
 
 
-def _render_hub_probe(
-    *, context_key: str, model_name: str
-) -> LoadedProbe | None:
+def _render_hub_probe(*, context_key: str, model_name: str) -> LoadedProbe | None:
     repo_id = st.text_input(
         "Probe repo",
         value=st.session_state.get(
@@ -249,15 +247,43 @@ def _validate(
 # ---------------------------------------------------------------------------
 
 
+def _store_derived_cache(key: str, value: object) -> None:
+    """Store one derived probe result while keeping a small MRU window."""
+
+    tracked = st.session_state.setdefault(_DERIVED_CACHE_TRACKER_KEY, [])
+    if not isinstance(tracked, list):
+        tracked = []
+    tracked = [existing for existing in tracked if existing != key]
+    tracked.append(key)
+    while len(tracked) > _DERIVED_CACHE_ENTRIES:
+        st.session_state.pop(tracked.pop(0), None)
+    st.session_state[_DERIVED_CACHE_TRACKER_KEY] = tracked
+    st.session_state[key] = value
+
+
+def _get_derived_cache(key: str) -> object | None:
+    """Return a derived probe result and refresh its MRU position."""
+
+    cached = st.session_state.get(key)
+    if cached is None:
+        return None
+    tracked = st.session_state.get(_DERIVED_CACHE_TRACKER_KEY)
+    if isinstance(tracked, list) and key in tracked:
+        tracked = [existing for existing in tracked if existing != key]
+        tracked.append(key)
+        st.session_state[_DERIVED_CACHE_TRACKER_KEY] = tracked
+    return cached
+
+
 def _classification_predictions(
     probe: LoadedProbe, activations: torch.Tensor, cache_key: str
 ) -> tuple[torch.Tensor, torch.Tensor]:
     full_key = widget_key("probe_predictions", cache_key, str(id(probe)))
-    cached = st.session_state.get(full_key)
+    cached = _get_derived_cache(full_key)
     if cached is not None:
         return cached
     _, probs, predicted = probe.run_batch(activations)
-    st.session_state[full_key] = (probs, predicted)
+    _store_derived_cache(full_key, (probs, predicted))
     return probs, predicted
 
 
@@ -265,11 +291,11 @@ def _regression_values(
     probe: LoadedProbe, activations: torch.Tensor, cache_key: str
 ) -> torch.Tensor:
     full_key = widget_key("probe_values", cache_key, str(id(probe)))
-    cached = st.session_state.get(full_key)
+    cached = _get_derived_cache(full_key)
     if cached is not None:
         return cached
     values = probe.predict_batch(activations)
-    st.session_state[full_key] = values
+    _store_derived_cache(full_key, values)
     return values
 
 
@@ -297,9 +323,7 @@ def _apply_overlays(
         probs, predicted = _classification_predictions(
             probe, trace.activations, trace.cache_key
         )
-        binary = probs.shape[1] == 1 or (
-            probs.shape[1] == 2 and len(probe.labels) == 2
-        )
+        binary = probs.shape[1] == 1 or (probs.shape[1] == 2 and len(probe.labels) == 2)
         overlays = build_classification_overlays(
             trace=trace,
             probs=probs,
@@ -332,9 +356,7 @@ def render_probe_inspector(
     def _conversation_sig() -> int:
         return hash(
             tuple(
-                (m.get("role"), m.get("content"))
-                for m in messages
-                if m.get("content")
+                (m.get("role"), m.get("content")) for m in messages if m.get("content")
             )
         )
 
@@ -349,9 +371,7 @@ def render_probe_inspector(
             st.caption("Probe overlay shows up after the first assistant reply.")
             return
 
-        probe = _render_probe_selector(
-            context_key=context_key, model_name=model_name
-        )
+        probe = _render_probe_selector(context_key=context_key, model_name=model_name)
         if probe is None:
             _reset()
             return

tests/test_datasets.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+from utils import datasets
+
+
+class _Progress:
+    def __init__(self) -> None:
+        self.updates: list[tuple[float, str | None]] = []
+
+    def progress(self, value: float, *, text: str | None = None) -> None:
+        self.updates.append((value, text))
+
+
+def test_download_missing_startup_files_only_fetches_uncached_files(monkeypatch):
+    warnings: list[str] = []
+    progress = _Progress()
+    downloads: list[tuple[str, str, str]] = []
+
+    monkeypatch.setattr(
+        datasets,
+        "_is_cached",
+        lambda _repo, filename: filename == "already.jsonl",
+    )
+    monkeypatch.setattr(datasets.st, "warning", warnings.append)
+    monkeypatch.setattr(
+        datasets.st,
+        "progress",
+        lambda value, *, text=None: progress,
+    )
+    monkeypatch.setattr(
+        datasets,
+        "hf_hub_download",
+        lambda repo, filename, *, repo_type: downloads.append(
+            (repo, filename, repo_type)
+        ),
+    )
+
+    datasets._download_missing_startup_files_if_needed(
+        "org/repo",
+        ("already.jsonl", "missing.jsonl"),
+        "Example",
+    )
+
+    assert warnings and "First-time setup for Example" in warnings[0]
+    assert downloads == [("org/repo", "missing.jsonl", "dataset")]
+    assert progress.updates[-1] == (1.0, "Downloaded missing.jsonl (1/1)")
+
+
+def test_download_missing_startup_files_stays_quiet_when_cached(monkeypatch):
+    monkeypatch.setattr(datasets, "_is_cached", lambda *_args: True)
+
+    def unexpected(*_args, **_kwargs):
+        raise AssertionError("cold-download UI should not render for warm cache")
+
+    monkeypatch.setattr(datasets.st, "warning", unexpected)
+    monkeypatch.setattr(datasets.st, "progress", unexpected)
+    monkeypatch.setattr(datasets, "hf_hub_download", unexpected)
+
+    datasets._download_missing_startup_files_if_needed(
+        "org/repo",
+        ("cached.jsonl",),
+        "Example",
+    )
+
+
+def test_prepare_nemotron_prefetches_first_parquet_shard(monkeypatch):
+    calls: list[tuple[str, tuple[str, ...], str]] = []
+    monkeypatch.setattr(
+        datasets,
+        "list_repo_files",
+        lambda *_args, **_kwargs: (
+            "README.md",
+            "data/train-00001-of-00002.parquet",
+            "data/train-00000-of-00002.parquet",
+        ),
+    )
+    monkeypatch.setattr(
+        datasets,
+        "_download_missing_startup_files_if_needed",
+        lambda repo, filenames, label: calls.append((repo, filenames, label)),
+    )
+
+    datasets._prepare_nemotron_startup_download(
+        datasets.DatasetSource.NEMOTRON_USA.value,
+        "Nemotron USA",
+    )
+
+    assert calls == [
+        (
+            "nvidia/Nemotron-Personas-USA",
+            ("data/train-00000-of-00002.parquet",),
+            "Nemotron USA",
+        )
+    ]
+
+
+def test_warm_qa_makes_synth_qa_download_visible_before_thread(monkeypatch):
+    calls: list[tuple[str, tuple[str, ...], str]] = []
+    started: list[bool] = []
+
+    class DummySynth:
+        def prefetch_qa(self) -> None:
+            pass
+
+    class DummyThread:
+        def __init__(self, *args, **kwargs) -> None:
+            pass
+
+        def start(self) -> None:
+            started.append(True)
+
+    monkeypatch.setattr(datasets, "SynthPersonaDataset", DummySynth)
+    monkeypatch.setattr(
+        datasets,
+        "_download_missing_startup_files_if_needed",
+        lambda repo, filenames, label: calls.append((repo, filenames, label)),
+    )
+    monkeypatch.setattr(datasets.threading, "Thread", DummyThread)
+
+    datasets.warm_qa_in_background(DummySynth())
+
+    assert calls == [
+        (
+            "implicit-personalization/synth-persona",
+            ("dataset_qa.jsonl",),
+            "SynthPersona QA",
+        )
+    ]
+    assert started == [True]

tests/test_probe_cache_bounds.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import torch
+
+from tabs import probe_ui
+from utils import probe_trace
+
+
+def test_store_derived_cache_evicts_oldest(monkeypatch):
+    session_state: dict[str, object] = {}
+    monkeypatch.setattr(probe_ui.st, "session_state", session_state)
+    monkeypatch.setattr(probe_ui, "_DERIVED_CACHE_ENTRIES", 2)
+
+    probe_ui._store_derived_cache("k1", 1)
+    probe_ui._store_derived_cache("k2", 2)
+    probe_ui._store_derived_cache("k3", 3)
+
+    assert "k1" not in session_state
+    assert session_state["k2"] == 2
+    assert session_state["k3"] == 3
+    assert session_state[probe_ui._DERIVED_CACHE_TRACKER_KEY] == ["k2", "k3"]
+
+
+def test_get_derived_cache_refreshes_recently_used_entry(monkeypatch):
+    session_state: dict[str, object] = {}
+    monkeypatch.setattr(probe_ui.st, "session_state", session_state)
+    monkeypatch.setattr(probe_ui, "_DERIVED_CACHE_ENTRIES", 2)
+
+    probe_ui._store_derived_cache("k1", 1)
+    probe_ui._store_derived_cache("k2", 2)
+
+    assert probe_ui._get_derived_cache("k1") == 1
+    probe_ui._store_derived_cache("k3", 3)
+
+    assert "k1" in session_state
+    assert "k2" not in session_state
+    assert session_state[probe_ui._DERIVED_CACHE_TRACKER_KEY] == ["k1", "k3"]
+
+
+def test_trace_eviction_drops_derived_results(monkeypatch):
+    session_state: dict[str, object] = {}
+    monkeypatch.setattr(probe_trace.st, "session_state", session_state)
+    monkeypatch.setattr(probe_trace, "_MAX_CACHED_TRACES", 1)
+
+    trace = probe_trace.ConversationTrace(
+        cache_key="old",
+        model_name="m",
+        remote=False,
+        prompt_text="p",
+        prompt_hash="h",
+        layer=0,
+        location="post_reasoning",
+        input_ids=torch.tensor([1]),
+        activations=torch.zeros((1, 1)),
+        tokens=["x"],
+        assistant_spans=[],
+        is_special=torch.tensor([False]),
+    )
+    old_prediction_key = "probe_predictions::old::probe"
+    kept_prediction_key = "probe_predictions::new::probe"
+    session_state[probe_trace._DERIVED_CACHE_TRACKER_KEY] = [
+        old_prediction_key,
+        kept_prediction_key,
+    ]
+    session_state[old_prediction_key] = object()
+    session_state[kept_prediction_key] = object()
+
+    probe_trace._store_cached_trace("old", trace)
+    probe_trace._store_cached_trace(
+        "new",
+        probe_trace.ConversationTrace(
+            **{**trace.__dict__, "cache_key": "new"},
+        ),
+    )
+
+    assert old_prediction_key not in session_state
+    assert kept_prediction_key in session_state
+    assert session_state[probe_trace._DERIVED_CACHE_TRACKER_KEY] == [
+        kept_prediction_key
+    ]

tests/test_probe_sweep.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import torch
+from persona_vectors.analysis import LayeredSamples
+from persona_vectors.probes import AttributeLabels
+
+from tabs import probe_sweep
+
+
+def test_cached_sweep_keeps_per_attribute_samples_and_full_plus_pca(monkeypatch):
+    samples = LayeredSamples(
+        vectors=torch.zeros((3, 2, 4)),
+        labels=["p0", "p1", "p2"],
+        hover_text=["p0", "p1", "p2"],
+    )
+    sweep_calls: list[tuple[str, int | None]] = []
+
+    monkeypatch.setattr(
+        probe_sweep,
+        "load_persona_vectors_cached",
+        lambda *args: samples,
+    )
+    monkeypatch.setattr(
+        probe_sweep,
+        "synth_persona_dataset_cached",
+        lambda: SimpleNamespace(),
+    )
+
+    def labels_for(_dataset, attribute, _persona_ids, *, task):
+        return AttributeLabels(
+            attribute_name=attribute,
+            task=task,
+            y=torch.tensor([0, 1, 0]).numpy(),
+            labels=["a", "b", "a"],
+            class_names=["a", "b"],
+        )
+
+    monkeypatch.setattr(probe_sweep, "attribute_probe_labels", labels_for)
+
+    def filtered(input_samples, labels, *, min_count):
+        assert min_count == 2
+        return input_samples, labels
+
+    monkeypatch.setattr(
+        probe_sweep,
+        "filter_attribute_samples_min_count",
+        filtered,
+    )
+
+    def sweep(input_samples, labels, *, layers, probe_kinds, n_pca_components, seed):
+        assert input_samples is samples
+        assert layers == [0, 1]
+        assert probe_kinds == ["logistic_regression"]
+        assert seed == 0
+        sweep_calls.append((labels.attribute_name, n_pca_components))
+        return [
+            {
+                "attribute": labels.attribute_name,
+                "layer": 0,
+                "probe_kind": probe_kinds[0],
+                "balanced_accuracy": 0.5,
+            }
+        ]
+
+    monkeypatch.setattr(probe_sweep, "sweep_attribute", sweep)
+
+    inputs = probe_sweep.SweepInputs(
+        source="src",
+        location="loc",
+        model_name="model",
+        mask_value="answer_mean",
+        variant="templated",
+        persona_ids=("p0", "p1", "p2"),
+        attributes=("sex", "gender"),
+        task="binary",
+        probe_kinds=("logistic_regression",),
+        n_pca_components=2,
+        layers=(0, 1),
+        min_class_count=2,
+        seed=0,
+    )
+
+    rows_by_label, per_attr = probe_sweep.cached_sweep.__wrapped__(inputs)
+
+    assert list(rows_by_label) == ["full", "pca2"]
+    assert [row["attribute"] for row in rows_by_label["full"]] == ["sex", "gender"]
+    assert set(per_attr) == {"sex", "gender"}
+    assert sweep_calls == [
+        ("sex", None),
+        ("gender", None),
+        ("sex", 2),
+        ("gender", 2),
+    ]

tests/test_probes.py

@@ -11,17 +11,16 @@ two correctness fixes:
 
 import pytest
 import torch
-
 from persona_vectors.probes import ProbeArtifact
+
+from utils.probe_files import parse_probe_filename
 from utils.probes import (
     LoadedProbe,
     _LinearProbe,
     _loaded_probe_from_artifact,
     _normalize_labels,
-    parse_probe_filename,
 )
 
-
 # --------------------------------------------------------------------------- #
 # parse_probe_filename
 # --------------------------------------------------------------------------- #
@@ -123,9 +122,7 @@ def test_normalize_batch_pca_only_applies_pca():
     probe = _probe(
         2,
         pca_mean=torch.ones(3),
-        pca_components=torch.tensor(
-            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
-        ),
+        pca_components=torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]),
     )
     batch = torch.tensor([[2.0, 4.0, 9.0]])
     out = probe._normalize_batch(batch)

tests/test_state.py

@@ -0,0 +1,16 @@
+from state import chat_session_key
+
+
+def test_chat_session_key_is_stable_across_model_switches() -> None:
+    dataset = "HuggingFace: synth-persona"
+
+    assert chat_session_key("google/gemma-2-2b-it", dataset) == chat_session_key(
+        "google/gemma-2-9b-it",
+        dataset,
+    )
+
+
+def test_chat_session_key_still_separates_datasets() -> None:
+    model = "google/gemma-2-2b-it"
+
+    assert chat_session_key(model, "dataset-a") != chat_session_key(model, "dataset-b")

utils/analysis_sources.py

@@ -7,8 +7,8 @@ from persona_vectors.analysis import (
     load_analysis_dataset,
 )
 from persona_vectors.artifacts import (
-    PersonaVectorStore,
     HFPersonaVectorStore,
+    PersonaVectorStore,
     discover_activation_models,
     model_dir_name,
 )

utils/chat.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Literal
@@ -185,6 +186,7 @@ def generate_chat_reply(
     top_k: int = 50,
     repetition_penalty: float = 1.0,
     seed: int | None = None,
+    on_status: Callable[[str, str, str], None] | None = None,
 ) -> ChatReply:
     """Generate one assistant reply from a full chat history.
 
@@ -228,9 +230,16 @@ def generate_chat_reply(
         generation_kwargs["repetition_penalty"] = repetition_penalty
     # `remote` is captured by nnsight's RemoteableMixin.trace() and is NOT
     # forwarded to the underlying model's generate
+    backend = _build_remote_backend(model, on_status) if remote else None
+
     with (
         _seeded_rng(seed if do_sample and not remote else None),
-        model.generate(prompt, remote=remote, **generation_kwargs) as tracer,
+        model.generate(
+            prompt,
+            remote=remote,
+            backend=backend,
+            **generation_kwargs,
+        ) as tracer,
     ):
         generated = tracer.result.save()
 
@@ -247,3 +256,34 @@ def generate_chat_reply(
         text=text,
         generated_ids=generated_ids.detach().cpu(),
     )
+
+
+def _build_remote_backend(
+    model: StandardizedTransformer,
+    on_status: Callable[[str, str, str], None] | None,
+):
+    """Build an NDIF backend that can surface lifecycle updates to callers."""
+
+    if on_status is None:
+        return None
+
+    from nnsight.intervention.backends.remote import JobStatusDisplay, RemoteBackend
+
+    class _CallbackJobStatusDisplay(JobStatusDisplay):
+        def update(
+            self,
+            job_id: str = "",
+            status_name: str = "",
+            description: str = "",
+        ):
+            super().update(job_id, status_name, description)
+            if status_name:
+                on_status(job_id, status_name, description)
+
+    backend = RemoteBackend(model.to_model_key())
+    backend.CONNECT_TIMEOUT = 300.0
+    backend.status_display = _CallbackJobStatusDisplay(
+        enabled=True,
+        verbose=backend.verbose,
+    )
+    return backend

utils/contrast.py

@@ -247,9 +247,7 @@ def render_contrast_html(result: TokenContrast) -> str:
     # those render as blank lines before the first word. Drop leading
     # whitespace-only tokens (and left-trim the first visible one) so the
     # contrast starts at real content. Display-only — weights stay aligned.
-    items = list(
-        zip(result.tokens, result.weights, result.raw_diffs, strict=True)
-    )
+    items = list(zip(result.tokens, result.weights, result.raw_diffs, strict=True))
     start = 0
     while start < len(items) and not items[start][0].strip():
         start += 1

utils/controls.py

@@ -7,8 +7,12 @@ def render_mask_strategy_select(
     key: str,
     last_key: str,
     help_text: str,
+    remember_key: str | None = None,
 ) -> MaskStrategy:
-    last_strategy = st.session_state.get(last_key, MaskStrategy.ANSWER_MEAN.value)
+    last_strategy = st.session_state.get(
+        remember_key,
+        st.session_state.get(last_key, MaskStrategy.ANSWER_MEAN.value),
+    )
     strategies = list(MaskStrategy)
     selected = st.selectbox(
         "Mask strategy",
@@ -26,4 +30,6 @@ def render_mask_strategy_select(
         help=help_text,
     )
     st.session_state[last_key] = selected.value
+    if remember_key is not None:
+        st.session_state[remember_key] = selected.value
     return selected

utils/datasets.py

@@ -7,6 +7,7 @@ from tempfile import mkdtemp
 from typing import Any
 
 import streamlit as st
+from huggingface_hub import hf_hub_download, list_repo_files, try_to_load_from_cache
 from persona_data.nemotron_personas import (
     NemotronPersonasFranceDataset,
     NemotronPersonasUSADataset,
@@ -16,6 +17,17 @@ from persona_data.synth_persona import SynthPersonaDataset
 
 from .helpers import DatasetSource
 
+_SYNTH_PERSONA_REPO = "implicit-personalization/synth-persona"
+_SYNTH_PERSONA_STARTUP_FILES = (
+    "implicit_shared_mc_bank.json",
+    "dataset_personas.jsonl",
+)
+_SYNTH_PERSONA_QA_FILE = "dataset_qa.jsonl"
+_NEMOTRON_REPOS = {
+    DatasetSource.NEMOTRON_FRANCE.value: "nvidia/Nemotron-Personas-France",
+    DatasetSource.NEMOTRON_USA.value: "nvidia/Nemotron-Personas-USA",
+}
+
 
 @st.cache_resource(show_spinner=False)
 def _cached_dataset(cls: type) -> Any:
@@ -39,13 +51,19 @@ def warm_qa_in_background(dataset: Any) -> None:
     warm = getattr(dataset, "prefetch_qa", None)
     if warm is None:
         return  # persona-only dataset (e.g. Nemotron) has no QA
+    if isinstance(dataset, SynthPersonaDataset):
+        # Extract will need QA soon. Make the one-time large transfer explicit,
+        # then leave the CPU-heavy parse on the existing background thread.
+        _download_missing_startup_files_if_needed(
+            _SYNTH_PERSONA_REPO,
+            (_SYNTH_PERSONA_QA_FILE,),
+            "SynthPersona QA",
+        )
     with _qa_warm_lock:
         if getattr(dataset, "_qa_warm_started", False):
             return
         dataset._qa_warm_started = True
-    threading.Thread(
-        target=warm, name="persona-ui-warm-qa", daemon=True
-    ).start()
+    threading.Thread(target=warm, name="persona-ui-warm-qa", daemon=True).start()
 
 
 @st.cache_resource(show_spinner=False)
@@ -118,12 +136,19 @@ def load_dataset(
     """Load the selected dataset source for the UI."""
 
     if dataset_source == DatasetSource.SYNTH_PERSONA.value:
+        _download_missing_startup_files_if_needed(
+            _SYNTH_PERSONA_REPO,
+            _SYNTH_PERSONA_STARTUP_FILES,
+            "SynthPersona",
+        )
         return _cached_dataset(SynthPersonaDataset), "SynthPersona"
 
     if dataset_source == DatasetSource.NEMOTRON_FRANCE.value:
+        _prepare_nemotron_startup_download(dataset_source, "Nemotron France")
         return _cached_dataset(NemotronPersonasFranceDataset), "Nemotron France"
 
     if dataset_source == DatasetSource.NEMOTRON_USA.value:
+        _prepare_nemotron_startup_download(dataset_source, "Nemotron USA")
         return _cached_dataset(NemotronPersonasUSADataset), "Nemotron USA"
 
     if personas_file is None or qa_file is None:
@@ -132,3 +157,60 @@ def load_dataset(
     personas_path = _uploaded_file_to_temp_path(personas_file, stem="personas")
     qa_path = _uploaded_file_to_temp_path(qa_file, stem="qa")
     return _cached_local_dataset(str(personas_path), str(qa_path)), "Local upload"
+
+
+def _is_cached(repo_id: str, filename: str) -> bool:
+    """Return whether a Hub dataset file already exists in the local HF cache."""
+
+    cached = try_to_load_from_cache(repo_id, filename, repo_type="dataset")
+    return isinstance(cached, str)
+
+
+def _download_missing_startup_files_if_needed(
+    repo_id: str,
+    filenames: tuple[str, ...],
+    label: str,
+) -> None:
+    """Make first-time Hub downloads visible before dataset construction blocks.
+
+    Hugging Face handles byte-level transfer internally. We expose file-level
+    progress here, which is the useful unit this UI can know in advance.
+    """
+
+    missing = tuple(
+        filename for filename in filenames if not _is_cached(repo_id, filename)
+    )
+    if not missing:
+        return
+
+    st.warning(
+        f"First-time setup for {label}: downloading dataset files from Hugging Face. "
+        "Later loads should use the local cache."
+    )
+    progress = st.progress(0.0, text=f"Preparing {label} download…")
+    total = len(missing)
+    for index, filename in enumerate(missing, start=1):
+        progress.progress(
+            (index - 1) / total,
+            text=f"Downloading {filename} ({index}/{total})",
+        )
+        hf_hub_download(repo_id, filename, repo_type="dataset")
+        progress.progress(
+            index / total,
+            text=f"Downloaded {filename} ({index}/{total})",
+        )
+
+
+def _prepare_nemotron_startup_download(dataset_source: str, label: str) -> None:
+    """Prefetch the first parquet shard used by the default Nemotron sample."""
+
+    repo_id = _NEMOTRON_REPOS[dataset_source]
+    parquet_files = tuple(
+        sorted(
+            filename
+            for filename in list_repo_files(repo_id, repo_type="dataset")
+            if filename.startswith("data/train-") and filename.endswith(".parquet")
+        )
+    )
+    if parquet_files:
+        _download_missing_startup_files_if_needed(repo_id, (parquet_files[0],), label)

utils/helpers.py

@@ -64,6 +64,26 @@ NDIF_STATUS_ICONS = {
 }
 
 
+def format_ndif_status(
+    job_id: str,
+    status_name: str,
+    description: str,
+    *,
+    prefix: str | None = None,
+    completed_detail: str | None = None,
+) -> str:
+    """Build the shared one-line NDIF status label used across the UI."""
+
+    icon = NDIF_STATUS_ICONS.get(status_name, "•")
+    detail = (
+        completed_detail
+        if completed_detail is not None and status_name == "COMPLETED"
+        else description
+    )
+    label = f"{icon} `{job_id}` **{status_name}** — {detail}"
+    return f"{prefix}: {label}" if prefix else label
+
+
 def slugify(value: str) -> str:
     """Convert a string to a filesystem-safe slug."""
 

utils/probe_files.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+import streamlit as st
+
+PROBE_FILENAME_RE = re.compile(
+    r"^cognitive_map_probe_layer(?P<layer>\d+)_(?P<model_type>[a-z0-9]+)_"
+    r"(?P<location>pre_reasoning|post_reasoning)_all_(?P<scope>general|size\d+)\.pt$"
+)
+
+PERSONA_PROBE_DIR_RE = re.compile(
+    r"^(?P<probe_kind>[a-z_]+?)(?:_pca(?P<pca>\d+))?_layer(?P<layer>\d+)$"
+)
+
+DEFAULT_PROBE_REPO = "project-telos/cognitive_map_probes"
+DEFAULT_LOCAL_PROBE_DIR = os.environ.get("PERSONA_PROBES_DIR", "artifacts/probes")
+
+
+@dataclass(frozen=True)
+class ProbeFileMetadata:
+    filename: str
+    layer: int | None
+    model_type: str
+    location: str | None
+    scope: str | None
+    label: str
+    model_name: str | None = None
+    attribute_name: str | None = None
+
+
+def model_probe_dir_name(model_name: str) -> str:
+    return model_name.replace("/", "__")
+
+
+def parse_probe_filename(filename: str) -> ProbeFileMetadata:
+    path = Path(filename)
+    match = PROBE_FILENAME_RE.match(path.name)
+    if match:
+        layer = int(match.group("layer"))
+        model_type = match.group("model_type")
+        location = match.group("location")
+        scope = match.group("scope")
+        scope_label = scope.replace("size", "size ")
+        return ProbeFileMetadata(
+            filename=filename,
+            layer=layer,
+            model_type=model_type,
+            location=location,
+            scope=scope,
+            label=(
+                f"Layer {layer} - {model_type.upper()} - "
+                f"{location.replace('_', ' ')} - {scope_label}"
+            ),
+        )
+
+    parent_match = PERSONA_PROBE_DIR_RE.match(path.parent.name)
+    if parent_match and path.name in {"probe.json", "weights.safetensors"}:
+        layer = int(parent_match.group("layer"))
+        probe_kind = parent_match.group("probe_kind")
+        pca = parent_match.group("pca")
+        scope = f"pca{pca}" if pca else None
+        attribute = path.parent.parent.name or "attribute"
+        model_name = path.parts[0].replace("__", "/") if len(path.parts) >= 5 else None
+        label = f"{attribute} - layer {layer} - {probe_kind}"
+        if pca:
+            label += f" (pca{pca})"
+        return ProbeFileMetadata(
+            filename=filename,
+            layer=layer,
+            model_type=probe_kind,
+            location=None,
+            scope=scope,
+            label=label,
+            model_name=model_name,
+            attribute_name=attribute,
+        )
+
+    return ProbeFileMetadata(
+        filename=filename,
+        layer=None,
+        model_type="unknown",
+        location=None,
+        scope=None,
+        label=path.stem.replace("_", " "),
+    )
+
+
+@st.cache_data(show_spinner=False, ttl=300)
+def list_probe_files(repo_id: str) -> list[str]:
+    from huggingface_hub import list_repo_files
+
+    return _dedupe_probe_entries(list_repo_files(repo_id, repo_type="model"))
+
+
+@st.cache_data(show_spinner=False, ttl=30)
+def list_local_probe_files(root_dir: str) -> list[str]:
+    root = Path(root_dir).expanduser()
+    if not root.is_dir():
+        return []
+    files = _dedupe_probe_entries(
+        [
+            str(path.relative_to(root))
+            for path in root.rglob("*")
+            if path.is_file()
+            and path.name in {"probe.pt", "probe.json", "weights.safetensors"}
+        ]
+    )
+    return sorted(files, key=_probe_sort_key)
+
+
+@st.cache_data(show_spinner=False, ttl=300)
+def download_probe_file(repo_id: str, filename: str) -> str:
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id, filename, repo_type="model")
+
+
+@st.cache_data(show_spinner=False, ttl=300)
+def download_probe_json_and_weights(repo_id: str, filename: str) -> tuple[str, str]:
+    from huggingface_hub import hf_hub_download
+
+    metadata_path = hf_hub_download(repo_id, filename, repo_type="model")
+    weights_name = str(Path(filename).with_name("weights.safetensors"))
+    weights_path = hf_hub_download(repo_id, weights_name, repo_type="model")
+    return metadata_path, weights_path
+
+
+def _probe_sort_key(filename: str) -> tuple[str, str, int, str]:
+    metadata = parse_probe_filename(filename)
+    return (
+        metadata.model_name or "",
+        metadata.attribute_name or "",
+        metadata.layer if metadata.layer is not None else 10**9,
+        filename,
+    )
+
+
+def _dedupe_probe_entries(files: list[str]) -> list[str]:
+    by_dir: dict[str, set[str]] = {}
+    standalone: list[str] = []
+    for filename in files:
+        path = Path(filename)
+        if path.name in {"probe.pt", "probe.json", "weights.safetensors"}:
+            by_dir.setdefault(str(path.parent), set()).add(path.name)
+        elif filename.endswith(".pt"):
+            standalone.append(filename)
+
+    entries = list(standalone)
+    for directory, names in by_dir.items():
+        selected = (
+            "probe.json"
+            if "probe.json" in names
+            else "probe.pt"
+            if "probe.pt" in names
+            else "weights.safetensors"
+        )
+        entries.append(str(Path(directory) / selected))
+    return sorted(entries, key=_probe_sort_key)

utils/probe_overlay.py

@@ -124,18 +124,14 @@ def build_regression_overlays(
     return overlays
 
 
-def attach_overlays(
-    messages: list[dict], overlays: list[ProbeOverlay]
-) -> None:
+def attach_overlays(messages: list[dict], overlays: list[ProbeOverlay]) -> None:
     """Attach one overlay to each assistant message, in order.
 
     Requires a 1:1 match. If the counts don't line up (e.g. the chat template
     doesn't mark assistant tokens), clear overlays so the caller can show a
     clear status instead of painting the wrong message.
     """
-    assistant_idxs = [
-        i for i, m in enumerate(messages) if m.get("role") == "assistant"
-    ]
+    assistant_idxs = [i for i, m in enumerate(messages) if m.get("role") == "assistant"]
     clear_overlays(messages)
     if not assistant_idxs or len(overlays) != len(assistant_idxs):
         return
@@ -189,8 +185,7 @@ def _tooltip(probs_row: list[float], labels: list[str | None]) -> str:
         # Single-output sigmoid: synthesize the complementary class so the
         # hover shows both label probabilities, not just one.
         return escape(
-            f"{positive_label} {positive:.2f} · "
-            f"not {positive_label} {1 - positive:.2f}"
+            f"{positive_label} {positive:.2f} · not {positive_label} {1 - positive:.2f}"
         )
     ranked = sorted(enumerate(probs_row), key=lambda item: item[1], reverse=True)
     parts = [f"{_label_for(labels, idx)} {prob:.2f}" for idx, prob in ranked]

utils/probe_trace.py

@@ -11,6 +11,7 @@ from persona_data.prompts import normalize_messages, supports_system_role
 from utils.chat import decode_token, format_generation_prompt, resolve_saved_tensor
 
 _TRACE_CACHE_KEY = "probe:trace_cache"
+_DERIVED_CACHE_TRACKER_KEY = "probe:derived_cache_keys"
 _MAX_CACHED_TRACES = 3
 
 
@@ -92,9 +93,7 @@ def trace_conversation(
 
     n_tokens = int(input_ids.shape[0])
     assistant_spans = _clip_spans(
-        _assistant_spans_from_offsets(
-            model.tokenizer, prompt_text, messages, n_tokens
-        ),
+        _assistant_spans_from_offsets(model.tokenizer, prompt_text, messages, n_tokens),
         n_tokens,
     )
     if not assistant_spans and assistant_mask_seq is not None:
@@ -182,6 +181,30 @@ def _store_cached_trace(cache_key: str, trace: ConversationTrace) -> None:
     while len(cache) > _MAX_CACHED_TRACES:
         oldest_key = next(iter(cache))
         cache.pop(oldest_key, None)
+        _drop_derived_results_for_trace(oldest_key)
+
+
+def _drop_derived_results_for_trace(cache_key: str) -> None:
+    """Remove probe predictions tied to a trace that just aged out."""
+
+    prefixes = (
+        f"probe_predictions::{cache_key}::",
+        f"probe_values::{cache_key}::",
+    )
+    tracked = st.session_state.get(_DERIVED_CACHE_TRACKER_KEY)
+    if isinstance(tracked, list):
+        kept: list[str] = []
+        for key in tracked:
+            if isinstance(key, str) and key.startswith(prefixes):
+                st.session_state.pop(key, None)
+            else:
+                kept.append(key)
+        st.session_state[_DERIVED_CACHE_TRACKER_KEY] = kept
+        return
+
+    for key in list(st.session_state):
+        if isinstance(key, str) and key.startswith(prefixes):
+            st.session_state.pop(key, None)
 
 
 def _compute_assistant_mask(
@@ -302,9 +325,7 @@ def _assistant_spans_from_prefixes(
         for i, message in enumerate(messages):
             if message.get("role") != "assistant":
                 continue
-            prefix_ids = apply(
-                messages[:i], tokenize=True, add_generation_prompt=True
-            )
+            prefix_ids = apply(messages[:i], tokenize=True, add_generation_prompt=True)
             through_ids = apply(
                 messages[: i + 1], tokenize=True, add_generation_prompt=False
             )
@@ -332,9 +353,7 @@ def _flatten_ids(value: object) -> list[int] | None:
         return None
 
 
-def _clip_spans(
-    spans: list[tuple[int, int]], n_tokens: int
-) -> list[tuple[int, int]]:
+def _clip_spans(spans: list[tuple[int, int]], n_tokens: int) -> list[tuple[int, int]]:
     clipped: list[tuple[int, int]] = []
     for start, end in spans:
         s = max(0, min(start, n_tokens))

utils/probes.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 
 import io
-import os
-import re
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -13,33 +11,14 @@ import torch.nn as nn
 import torch.nn.functional as F
 from persona_vectors.probes import ProbeArtifact, load_probe_artifact
 
-PROBE_FILENAME_RE = re.compile(
-    r"^cognitive_map_probe_layer(?P<layer>\d+)_(?P<model_type>[a-z0-9]+)_"
-    r"(?P<location>pre_reasoning|post_reasoning)_all_(?P<scope>general|size\d+)\.pt$"
+from utils.helpers import env_int
+from utils.probe_files import (
+    download_probe_file,
+    download_probe_json_and_weights,
+    parse_probe_filename,
 )
 
-# persona-vectors layout: .../{model}/{mask}/{variant}/{attribute}/{probe_kind}[_pca{K}]_layer{N}/weights.safetensors
-PERSONA_PROBE_DIR_RE = re.compile(
-    r"^(?P<probe_kind>[a-z_]+?)(?:_pca(?P<pca>\d+))?_layer(?P<layer>\d+)$"
-)
-
-DEFAULT_PROBE_REPO = "project-telos/cognitive_map_probes"
-DEFAULT_LOCAL_PROBE_DIR = os.environ.get(
-    "PERSONA_PROBES_DIR",
-    "artifacts/probes",
-)
-
-
-@dataclass(frozen=True)
-class ProbeFileMetadata:
-    filename: str
-    layer: int | None
-    model_type: str
-    location: str | None
-    scope: str | None
-    label: str
-    model_name: str | None = None
-    attribute_name: str | None = None
+_PROBE_CACHE_ENTRIES = env_int("PERSONA_UI_PROBE_CACHE_ENTRIES", 8)
 
 
 @dataclass(frozen=True)
@@ -195,9 +174,7 @@ class LoadedProbe:
             predicted = torch.argmax(probs, dim=-1)
         return logits, probs, predicted
 
-    def _forward_batch(
-        self, batch: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    def _forward_batch(self, batch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         normalized = self._normalize_batch(batch)
         with torch.no_grad():
             logits = self.model(normalized).detach().cpu()
@@ -233,104 +210,7 @@ class LoadedProbe:
         return batch
 
 
-def model_probe_dir_name(model_name: str) -> str:
-    return model_name.replace("/", "__")
-
-
-def parse_probe_filename(filename: str) -> ProbeFileMetadata:
-    path = Path(filename)
-    match = PROBE_FILENAME_RE.match(path.name)
-    if match:
-        layer = int(match.group("layer"))
-        model_type = match.group("model_type")
-        location = match.group("location")
-        scope = match.group("scope")
-        scope_label = scope.replace("size", "size ")
-        return ProbeFileMetadata(
-            filename=filename,
-            layer=layer,
-            model_type=model_type,
-            location=location,
-            scope=scope,
-            label=(
-                f"Layer {layer} - {model_type.upper()} - "
-                f"{location.replace('_', ' ')} - {scope_label}"
-            ),
-        )
-
-    # persona-vectors layout: parent dir holds probe_kind[_pca{K}]_layer{N},
-    # and the dir above that is the attribute name.
-    parent_match = PERSONA_PROBE_DIR_RE.match(path.parent.name)
-    if parent_match and path.name in {"probe.json", "weights.safetensors"}:
-        layer = int(parent_match.group("layer"))
-        probe_kind = parent_match.group("probe_kind")
-        pca = parent_match.group("pca")
-        scope = f"pca{pca}" if pca else None
-        attribute = path.parent.parent.name or "attribute"
-        model_name = path.parts[0].replace("__", "/") if len(path.parts) >= 5 else None
-        label = f"{attribute} - layer {layer} - {probe_kind}"
-        if pca:
-            label += f" (pca{pca})"
-        return ProbeFileMetadata(
-            filename=filename,
-            layer=layer,
-            model_type=probe_kind,
-            location=None,
-            scope=scope,
-            label=label,
-            model_name=model_name,
-            attribute_name=attribute,
-        )
-
-    return ProbeFileMetadata(
-        filename=filename,
-        layer=None,
-        model_type="unknown",
-        location=None,
-        scope=None,
-        label=path.stem.replace("_", " "),
-    )
-
-
-@st.cache_data(show_spinner=False, ttl=300)
-def list_probe_files(repo_id: str) -> list[str]:
-    from huggingface_hub import list_repo_files
-
-    files = list_repo_files(repo_id, repo_type="model")
-    return _dedupe_probe_entries(files)
-
-
-@st.cache_data(show_spinner=False, ttl=30)
-def list_local_probe_files(root_dir: str) -> list[str]:
-    root = Path(root_dir).expanduser()
-    if not root.is_dir():
-        return []
-    files = _dedupe_probe_entries([
-        str(path.relative_to(root))
-        for path in root.rglob("*")
-        if path.is_file() and path.name in {"probe.pt", "probe.json", "weights.safetensors"}
-    ])
-    return sorted(files, key=_probe_sort_key)
-
-
-@st.cache_data(show_spinner=False, ttl=300)
-def download_probe_file(repo_id: str, filename: str) -> str:
-    from huggingface_hub import hf_hub_download
-
-    return hf_hub_download(repo_id, filename, repo_type="model")
-
-
-@st.cache_data(show_spinner=False, ttl=300)
-def download_probe_json_and_weights(repo_id: str, filename: str) -> tuple[str, str]:
-    from huggingface_hub import hf_hub_download
-
-    metadata_path = hf_hub_download(repo_id, filename, repo_type="model")
-    weights_name = str(Path(filename).with_name("weights.safetensors"))
-    weights_path = hf_hub_download(repo_id, weights_name, repo_type="model")
-    return metadata_path, weights_path
-
-
-@st.cache_resource(show_spinner=False)
+@st.cache_resource(show_spinner=False, max_entries=_PROBE_CACHE_ENTRIES)
 def load_probe(repo_id: str, filename: str) -> LoadedProbe:
     if filename.endswith("probe.json"):
         metadata_path, weights_path = download_probe_json_and_weights(repo_id, filename)
@@ -346,7 +226,7 @@ def load_probe(repo_id: str, filename: str) -> LoadedProbe:
     )
 
 
-@st.cache_resource(show_spinner=False)
+@st.cache_resource(show_spinner=False, max_entries=_PROBE_CACHE_ENTRIES)
 def load_local_probe(root_dir: str, filename: str) -> LoadedProbe:
     root = Path(root_dir).expanduser()
     path = (root / filename).resolve()
@@ -370,7 +250,7 @@ def load_local_probe(root_dir: str, filename: str) -> LoadedProbe:
     )
 
 
-@st.cache_resource(show_spinner=False)
+@st.cache_resource(show_spinner=False, max_entries=_PROBE_CACHE_ENTRIES)
 def load_probe_from_bytes(filename: str, data: bytes) -> LoadedProbe:
     return _load_probe_payload(
         filename=filename,
@@ -432,16 +312,20 @@ def _load_probe_payload(
             _optional_str(payload.get("attribute_name")) or metadata.attribute_name
         ),
         feature_space=(
-            (f"pca{payload['n_pca_components']}"
-             if payload.get("n_pca_components")
-             else None)
+            (
+                f"pca{payload['n_pca_components']}"
+                if payload.get("n_pca_components")
+                else None
+            )
             or _optional_str(payload.get("feature_space"))
             or metadata.scope
         ),
         task=_optional_str(payload.get("task")),
         probe_kind=_optional_str(payload.get("probe_kind")),
         scaler_mean=_as_cpu_tensor(payload.get("scaler_mean")),
-        scaler_std=_as_cpu_tensor(_first_present(payload, "scaler_std", "scaler_scale")),
+        scaler_std=_as_cpu_tensor(
+            _first_present(payload, "scaler_std", "scaler_scale")
+        ),
         pca_mean=_as_cpu_tensor(payload.get("pca_mean")),
         pca_components=_as_cpu_tensor(payload.get("pca_components")),
     )
@@ -617,39 +501,6 @@ def _first_present(payload: dict[str, Any], *keys: str) -> Any:
     return None
 
 
-def _probe_sort_key(filename: str) -> tuple[str, str, int, str]:
-    metadata = parse_probe_filename(filename)
-    return (
-        metadata.model_name or "",
-        metadata.attribute_name or "",
-        metadata.layer if metadata.layer is not None else 10**9,
-        filename,
-    )
-
-
-def _dedupe_probe_entries(files: list[str]) -> list[str]:
-    by_dir: dict[str, set[str]] = {}
-    standalone: list[str] = []
-    for filename in files:
-        path = Path(filename)
-        if path.name in {"probe.pt", "probe.json", "weights.safetensors"}:
-            by_dir.setdefault(str(path.parent), set()).add(path.name)
-        elif filename.endswith(".pt"):
-            standalone.append(filename)
-
-    entries = list(standalone)
-    for directory, names in by_dir.items():
-        selected = (
-            "probe.json"
-            if "probe.json" in names
-            else "probe.pt"
-            if "probe.pt" in names
-            else "weights.safetensors"
-        )
-        entries.append(str(Path(directory) / selected))
-    return sorted(entries, key=_probe_sort_key)
-
-
 def _normalize_labels(raw_labels: Any, num_classes: int) -> list[str | None]:
     if isinstance(raw_labels, (list, tuple)):
         labels = [str(label) for label in raw_labels[:num_classes]]

utils/selection_controls.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import streamlit as st
+
+
+def remembered_segmented_control(
+    label: str,
+    *,
+    options: Sequence[str],
+    key: str,
+    remember_key: str | None = None,
+    default: str | None = None,
+    label_visibility: str = "visible",
+) -> str:
+    """Render a segmented control with one small, reusable memory pattern."""
+    fallback = default or options[0]
+    remembered = st.session_state.get(
+        remember_key,
+        st.session_state.get(key, fallback),
+    )
+    selected = (
+        st.segmented_control(
+            label,
+            options=options,
+            default=remembered if remembered in options else fallback,
+            key=key,
+            label_visibility=label_visibility,
+        )
+        or fallback
+    )
+    if remember_key is not None:
+        st.session_state[remember_key] = selected
+    return selected

utils/source_controls.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import streamlit as st
+from persona_data.environment import get_artifacts_dir
+from persona_vectors.extraction import MaskStrategy
+
+from utils.analysis_sources import (
+    DEFAULT_COMPARE_MODEL,
+    DEFAULT_HUB_REPO,
+    SOURCE_HUB,
+    SOURCE_LOCAL,
+    SOURCES,
+    Store,
+    activation_store_cached,
+    hub_models_by_mask_strategy,
+    local_model_matches,
+    local_model_options_cached,
+)
+from utils.helpers import widget_key
+from utils.selection_controls import remembered_segmented_control
+
+_SHARED_SOURCE_KEY = "source:last_source"
+_SHARED_HUB_REPO_KEY = "source:hub_repo"
+_SHARED_HUB_MODEL_KEY = "source:hub_model"
+_SHARED_LOCAL_ROOT_KEY = "source:local_root"
+_SHARED_LOCAL_MODEL_KEY = "source:local_model"
+
+
+def render_source_select(
+    *,
+    widget_scope: str,
+    last_source_key: str | None = None,
+) -> str:
+    key = widget_key(widget_scope, "source")
+    if last_source_key is not None and last_source_key not in st.session_state:
+        shared_source = st.session_state.get(_SHARED_SOURCE_KEY)
+        if shared_source is not None:
+            st.session_state[last_source_key] = shared_source
+    selected = remembered_segmented_control(
+        "Source",
+        options=SOURCES,
+        key=key,
+        remember_key=last_source_key or _SHARED_SOURCE_KEY,
+        default=SOURCE_HUB,
+        label_visibility="collapsed",
+    )
+    st.session_state[_SHARED_SOURCE_KEY] = selected
+    if last_source_key is not None:
+        st.session_state[last_source_key] = selected
+    return selected
+
+
+def _render_hub_model_select(
+    *,
+    state_prefix: str,
+    widget_scope: str,
+    repo_id: str,
+    mask_strategy: MaskStrategy,
+    model_label: str,
+    fallback_help: str,
+    selection_help: str,
+) -> str:
+    fallback_key = f"{state_prefix}:hub_model_fallback"
+    fallback_model = st.session_state.get(
+        fallback_key,
+        st.session_state.get(_SHARED_HUB_MODEL_KEY, DEFAULT_COMPARE_MODEL),
+    )
+    try:
+        models_by_strategy = hub_models_by_mask_strategy(repo_id)
+    except Exception as exc:
+        st.warning(f"Could not load Hub configs for `{repo_id}`: {exc}")
+        model = st.text_input(
+            model_label,
+            value=fallback_model,
+            key=fallback_key,
+            help=fallback_help,
+        )
+        st.session_state[_SHARED_HUB_MODEL_KEY] = model
+        return model
+
+    model_options = models_by_strategy.get(mask_strategy, [])
+    if not model_options:
+        st.warning(
+            f"No Hub vector configs found for `{mask_strategy.value}` in `{repo_id}`."
+        )
+        model = st.text_input(
+            model_label,
+            value=fallback_model,
+            key=fallback_key,
+            help=fallback_help,
+        )
+        st.session_state[_SHARED_HUB_MODEL_KEY] = model
+        return model
+
+    select_key = widget_key(widget_scope, "hub_model", repo_id, mask_strategy.value)
+    previous_model = st.session_state.get(
+        select_key,
+        st.session_state.get(_SHARED_HUB_MODEL_KEY, fallback_model),
+    )
+    default_model = (
+        previous_model if previous_model in model_options else model_options[0]
+    )
+    selected = st.selectbox(
+        model_label,
+        options=model_options,
+        index=model_options.index(default_model),
+        key=select_key,
+        help=selection_help,
+    )
+    st.session_state[fallback_key] = selected
+    st.session_state[_SHARED_HUB_MODEL_KEY] = selected
+    return selected
+
+
+def _render_local_model_select(
+    *,
+    state_prefix: str,
+    artifacts_root: str,
+    mask_strategy: MaskStrategy,
+    allow_custom_toggle: bool,
+    model_label: str,
+) -> str:
+    fallback_key = f"{state_prefix}:local_model"
+    fallback_model = st.session_state.get(
+        fallback_key,
+        st.session_state.get(_SHARED_LOCAL_MODEL_KEY, DEFAULT_COMPARE_MODEL),
+    )
+    model_options = local_model_options_cached(artifacts_root, mask_strategy.value)
+    if not model_options:
+        model = st.text_input(model_label, value=fallback_model, key=fallback_key)
+        st.session_state[_SHARED_LOCAL_MODEL_KEY] = model
+        return model
+
+    if allow_custom_toggle:
+        custom = st.toggle(
+            "Custom local model",
+            value=False,
+            key=f"{state_prefix}:local_model_custom_enabled",
+            help="Enter a model id/path manually instead of choosing from activation directories.",
+        )
+        if custom:
+            model = st.text_input("Local model", value=fallback_model, key=fallback_key)
+            st.session_state[_SHARED_LOCAL_MODEL_KEY] = model
+            return model
+
+    select_key = f"{state_prefix}:local_model_select"
+    previous_model = st.session_state.get(
+        select_key,
+        st.session_state.get(_SHARED_LOCAL_MODEL_KEY, fallback_model),
+    )
+    if not any(local_model_matches(previous_model, option) for option in model_options):
+        previous_model = fallback_model
+    default_model = next(
+        (
+            option
+            for option in model_options
+            if local_model_matches(option, previous_model)
+        ),
+        model_options[0],
+    )
+    selected = st.selectbox(
+        model_label,
+        options=model_options,
+        index=model_options.index(default_model),
+        key=select_key,
+        help="Models discovered under the selected artifacts root.",
+    )
+    st.session_state[fallback_key] = selected
+    st.session_state[_SHARED_LOCAL_MODEL_KEY] = selected
+    return selected
+
+
+def render_store_select(
+    source: str,
+    mask_strategy: MaskStrategy,
+    *,
+    state_prefix: str,
+    widget_scope: str,
+    artifacts_root_key: str,
+    model_label: str = "Model",
+    local_model_label: str = "Model",
+    allow_custom_local_model: bool = False,
+    repo_help: str | None = None,
+    fallback_help: str = "Model id to use if Hub config discovery is unavailable.",
+) -> Store:
+    if source == SOURCE_HUB:
+        repo_key = f"{state_prefix}:hub_repo"
+        repo = st.text_input(
+            "Hub repo",
+            value=st.session_state.get(
+                repo_key,
+                st.session_state.get(_SHARED_HUB_REPO_KEY, DEFAULT_HUB_REPO),
+            ),
+            key=repo_key,
+            help=repo_help,
+        )
+        st.session_state[_SHARED_HUB_REPO_KEY] = repo
+        model_name = _render_hub_model_select(
+            state_prefix=state_prefix,
+            widget_scope=widget_scope,
+            repo_id=repo,
+            mask_strategy=mask_strategy,
+            model_label=model_label,
+            fallback_help=fallback_help,
+            selection_help="Models with vectors in the selected Hub repo and mask strategy.",
+        )
+        return activation_store_cached(
+            SOURCE_HUB, repo, model_name, mask_strategy.value
+        )
+
+    root = st.text_input(
+        "Artifacts root",
+        value=st.session_state.get(
+            _SHARED_LOCAL_ROOT_KEY,
+            str(get_artifacts_dir() / "activations"),
+        ),
+        key=artifacts_root_key,
+    )
+    root = str(Path(root).expanduser())
+    st.session_state[_SHARED_LOCAL_ROOT_KEY] = root
+    model_name = _render_local_model_select(
+        state_prefix=state_prefix,
+        artifacts_root=root,
+        mask_strategy=mask_strategy,
+        allow_custom_toggle=allow_custom_local_model,
+        model_label=local_model_label,
+    )
+    return activation_store_cached(SOURCE_LOCAL, root, model_name, mask_strategy.value)

uv.lock

@@ -464,11 +464,11 @@ wheels = [
 
 [[package]]
 name = "decorator"
-version = "5.3.0"
+version = "5.3.1"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/5c/50/a39dd7ab407e93978dfa07d109b7d633e37958c89f30cbcec061b77b3ebc/decorator-5.3.0.tar.gz", hash = "sha256:95fda3122972c847cf0ff7e0ce2829bf25136f2526b627b3da85b60ca5f485c0", size = 58431, upload-time = "2026-05-17T06:59:57.258Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/60/8b/32f9823da46cde7df2087faa08cd98d01b908f8dcab982cdba9c84e85355/decorator-5.3.1.tar.gz", hash = "sha256:4cbcdd55a6efadb9dbea26b858f4fb3264567b52d69ca0d25b721b553f60ea82", size = 58084, upload-time = "2026-05-18T06:03:28.057Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/d5/6f/f8d0bba4dc2a69817d74f640d504650241ebf2f9f7263426f1b953b344d4/decorator-5.3.0-py3-none-any.whl", hash = "sha256:f8c2d71ede92f073144ddd7f3e9fbbc3bd0f2f29522c9d75ee648d66553834f4", size = 11104, upload-time = "2026-05-17T06:59:54.676Z" },
+    { url = "https://files.pythonhosted.org/packages/05/7f/798705f5296a58ca505d600456748d1be48078eac8a7050d8a98bc9edb89/decorator-5.3.1-py3-none-any.whl", hash = "sha256:f47fe6fdbd2edd623ecfe36875d37aba411624e2670dd395dddae1358689bb3c", size = 10365, upload-time = "2026-05-18T06:03:26.517Z" },
 ]
 
 [[package]]
@@ -1585,7 +1585,7 @@ wheels = [
 
 [[package]]
 name = "persona-ui"
-version = "0.4.0"
+version = "0.5.0"
 source = { virtual = "." }
 dependencies = [
     { name = "catppuccin" },
@@ -2145,11 +2145,11 @@ wheels = [
 
 [[package]]
 name = "python-multipart"
-version = "0.0.28"
+version = "0.0.29"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/82/54/a85eb421fbdd5007bc5af39d0f4ed9fa609e0fedbfdc2adcf0b34526870e/python_multipart-0.0.28.tar.gz", hash = "sha256:8550da197eac0f7ab748961fc9509b999fa2662ea25cef857f05249f6893c0f8", size = 45314, upload-time = "2026-05-10T11:05:16.596Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/4e/fe/70bd71a6738b09a0bdf6480ca6436b167469ca4578b2a0efbe390b4b0e70/python_multipart-0.0.29.tar.gz", hash = "sha256:643e93849196645e2dbdd81a0f8829a23123ad7f797a84a364c6fb3563f18904", size = 45678, upload-time = "2026-05-17T17:29:47.654Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/f3/a2/43bbc5860b5034e2af4ef99a0e04d726ff329c43e192ef3abaa8d7ecfce5/python_multipart-0.0.28-py3-none-any.whl", hash = "sha256:10faac07eb966c3f48dc415f9dee46c04cb10d58d30a35677db8027c825ed9b6", size = 29438, upload-time = "2026-05-10T11:05:15.052Z" },
+    { url = "https://files.pythonhosted.org/packages/8f/cb/769cfc37177252872a45a71f3fbdde9d51b471a3f3c14bfe95dde3407386/python_multipart-0.0.29-py3-none-any.whl", hash = "sha256:2ddcc971cef266225f54f552d8fa10bcfbb1f14446caec199060daac59ff2d69", size = 29640, upload-time = "2026-05-17T17:29:45.69Z" },
 ]
 
 [[package]]