Updated to latest probing options

- Cleaned up repo
- Improved performance drastically updating to the latest versions of
the librarires + less reloading smarter caching and prefetchign

.env.example

@@ -23,3 +23,5 @@ ARTIFACTS_DIR=artifacts
 # Keep model cache at 1 unless you have enough RAM for multiple loaded models.
 # PERSONA_UI_MODEL_CACHE_ENTRIES=1
 # PERSONA_UI_STORE_CACHE_ENTRIES=4
+# PERSONA_UI_VECTOR_CACHE_ENTRIES=4
+# PERSONA_UI_PREPARED_CACHE_ENTRIES=8

README.md

@@ -116,6 +116,8 @@ NDIF_API_KEY=...       # Required for remote (NDIF) model execution
 HF_HOME=...            # Optional: HuggingFace cache directory
 ARTIFACTS_DIR=...      # Optional: where persona vectors are read from (default: ./artifacts)
 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
 ```
 
 The app picks up this file automatically via `load_dotenv()` on startup.
@@ -148,3 +150,7 @@ the Analysis/Probing tab's Local source path) at the tree you want to load.
 
 The store classes are `PersonaVectorStore` (local) and `HFPersonaVectorStore`
 (Hub) — same API, both imported by `utils/analysis_sources.py`.
+
+## 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.

pyproject.toml

@@ -5,7 +5,7 @@ description = "Streamlit UI for persona-vectors"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
-    "persona-vectors>=0.8.2",
+    "persona-vectors>=0.8.3",
     "datasets>=4.8.5",
     "huggingface-hub>=1.14.0",
     "streamlit>=1.44.0",

tabs/analysis/_shared.py

@@ -6,6 +6,19 @@ from persona_data.synth_persona import BASELINE_PERSONA_ID
 from persona_vectors.extraction import MaskStrategy
 from persona_vectors.plots import save_plot_html
 
+from tabs.analysis._state import (
+    _DEFAULT_LAYER_FRAMES,
+    _HIGHLIGHT_OTHER_COLOR,
+    _HIGHLIGHT_OTHER_LABEL,
+    _LAST_LAYER_FRAMES_KEY,
+    _LAST_MASK_STRATEGY_KEY,
+    PersonaOptions,
+    _is_assistant_persona,
+    _persona_names_state_key,
+    _personas_empty_message,
+    _remembered_selectbox,
+    _sequence_to_list,
+)
 from utils.analysis_sources import (
     Store,
     available_variants,
@@ -13,7 +26,6 @@ from utils.analysis_sources import (
     load_variant_vectors_cached,
     persona_names_cached,
     personas_cached,
-    release_hf_store_cache,
     store_cache_parts,
     store_id,
     store_layers_cached,
@@ -22,20 +34,6 @@ from utils.controls import render_mask_strategy_select
 from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
 from utils.theme import active_base, style_plotly_layer_controls
 
-from tabs.analysis._state import (
-    _DEFAULT_LAYER_FRAMES,
-    _HIGHLIGHT_OTHER_COLOR,
-    _HIGHLIGHT_OTHER_LABEL,
-    _LAST_LAYER_FRAMES_KEY,
-    _LAST_MASK_STRATEGY_KEY,
-    PersonaOptions,
-    _is_assistant_persona,
-    _persona_names_state_key,
-    _personas_empty_message,
-    _remembered_selectbox,
-    _sequence_to_list,
-)
-
 
 def _gray_out_unselected_personas(fig: go.Figure) -> None:
     def _gray_trace(trace: object) -> None:
@@ -118,8 +116,7 @@ def _load_variant_vectors(
     )
 
 
-def _release_vector_memory(store: Store, variants: list[str] | tuple[str, ...]) -> None:
-    release_hf_store_cache(store, variants)
+def _release_vector_memory() -> None:
     gc.collect()
 
 

tabs/analysis/_state.py

@@ -45,7 +45,7 @@ _CLUSTER_MODES = {
     "First selected layer": "first_layer",
     "Per layer": "per_layer",
 }
-_PROJECTION_COLOR_MODES = ["Persona", "K-means clusters", "Persona attribute"]
+_PROJECTION_COLOR_MODES = ["Persona attribute", "Persona", "K-means clusters"]
 _MAX_ATTRIBUTE_CATEGORIES = DEFAULT_MAX_ATTRIBUTE_CATEGORIES
 
 
@@ -87,7 +87,7 @@ class ProjectionColorConfig:
 @dataclass(frozen=True)
 class LayeredFigureStateKeys:
     figure: str
-    projection: str | None = None
+    prepared: str | None = None
 
 
 _HIGHLIGHT_OTHER_LABEL = "Other"
@@ -139,7 +139,7 @@ _TRACKED_STATE_KEYS_KEY = "analysis:_tracked_state_keys"
 
 
 def _clear_old_load_states(current_key: str, suffix: str) -> None:
-    # Only one heavy figure/projection state should live at a time. We track
+    # 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.
@@ -156,8 +156,8 @@ def _clear_old_figure_states(current_key: str) -> None:
     _clear_old_load_states(current_key, "_fig_state")
 
 
-def _clear_old_projection_states(current_key: str) -> None:
-    _clear_old_load_states(current_key, "_projection_state")
+def _clear_old_prepared_states(current_key: str) -> None:
+    _clear_old_load_states(current_key, "_projection_ready")
 
 
 def _store_figure_state(key: str, value: object) -> None:

tabs/analysis/cosine.py

@@ -78,22 +78,15 @@ def _build_cosine_figures(
     mask_strategy: MaskStrategy,
     selection: CosineSelection,
 ) -> tuple[object, object | None, int, int] | None:
-    variant_sample_cache: dict[str, object] = {}
-
-    def _load_variant(variant: str):
-        if variant not in variant_sample_cache:
-            samples = _load_variant_vectors(
-                store,
-                [variant],
-                mask_strategy,
-                persona_ids=selection.persona_ids,
-            )
-            variant_sample_cache[variant] = samples[variant]
-        return variant_sample_cache[variant]
-
     try:
-        samples_a = _load_variant(selection.variant_a)
-        samples_b = _load_variant(selection.variant_b)
+        by_variant = _load_variant_vectors(
+            store,
+            selection.variants,
+            mask_strategy,
+            persona_ids=selection.persona_ids,
+        )
+        samples_a = by_variant[selection.variant_a]
+        samples_b = by_variant[selection.variant_b]
     except Exception as exc:
         st.error(f"Could not load vectors: {exc}")
         return None
@@ -120,8 +113,8 @@ def _build_cosine_figures(
     pair_errors = []
     for left, right in combinations(selection.variants, 2):
         try:
-            left_samples = _load_variant(left)
-            right_samples = _load_variant(right)
+            left_samples = by_variant[left]
+            right_samples = by_variant[right]
             pair_traces.append(
                 (
                     f"{prompt_variant_label(left)} vs {prompt_variant_label(right)}",
@@ -207,7 +200,7 @@ def _render_cosine_similarity(
             _store_figure_state(cosine_fig_key, figures)
             progress.progress(100, text="Done.")
         finally:
-            _release_vector_memory(store, selection.variants)
+            _release_vector_memory()
             progress.empty()
 
     if cosine_fig_key in st.session_state:

tabs/analysis/dendrogram.py

@@ -13,7 +13,7 @@ from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
 
 from tabs.analysis._shared import (
     _load_persona_options,
-    _load_persona_vectors,
+    _load_variant_vectors,
     _plotly_chart,
     _release_vector_memory,
     _render_layer_frame_controls,
@@ -204,13 +204,14 @@ def _render_dendrogram_analysis(
     ):
         progress = st.progress(0, text="Loading first variant vectors…")
         try:
-            progress.progress(15, text="Loading first variant vectors…")
-            samples_a = _load_persona_vectors(
+            progress.progress(15, text="Loading variant vectors…")
+            by_variant = _load_variant_vectors(
                 store,
-                variant_a,
+                shared_variants,
                 mask_strategy,
                 persona_ids,
             )
+            samples_a = by_variant[variant_a]
             progress.progress(40, text="Building first dendrogram…")
             fig_a = plot_persona_dendrogram(
                 samples_a,
@@ -223,13 +224,8 @@ def _render_dendrogram_analysis(
             del samples_a
             fig_b = None
             if variant_a != variant_b:
-                progress.progress(60, text="Loading second variant vectors…")
-                samples_b = _load_persona_vectors(
-                    store,
-                    variant_b,
-                    mask_strategy,
-                    persona_ids,
-                )
+                progress.progress(60, text="Building second dendrogram…")
+                samples_b = by_variant[variant_b]
                 progress.progress(75, text="Building second dendrogram…")
                 fig_b = plot_persona_dendrogram(
                     samples_b,
@@ -250,7 +246,7 @@ def _render_dendrogram_analysis(
             st.error(f"Could not build dendrogram: {exc}")
             st.session_state.pop(fig_key, None)
         finally:
-            _release_vector_memory(store, shared_variants)
+            _release_vector_memory()
             progress.empty()
 
     if fig_key in st.session_state:

tabs/analysis/layered.py

@@ -11,14 +11,19 @@ from persona_vectors.plots import (
     build_layered_figure,
     build_pair_similarity_figure,
     build_similarity_figures,
-    prepare_layered_projection_data,
 )
 
 from utils.analysis_metadata import (
     synth_persona_attribute_names,
     synth_persona_dataset_cached,
 )
-from utils.analysis_sources import Store, store_id
+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 (
@@ -48,7 +53,7 @@ from tabs.analysis._state import (
     LayeredFigureStateKeys,
     ProjectionColorConfig,
     _clear_old_figure_states,
-    _clear_old_projection_states,
+    _clear_old_prepared_states,
     _highlight_persona_groups,
     _persona_display_label,
     _persona_names_state_key,
@@ -116,7 +121,7 @@ def _render_projection_color_config(
         key=color_mode_key,
         remember_key=_LAST_PROJECTION_COLOR_MODE_KEY,
         options=_PROJECTION_COLOR_MODES,
-        default="Persona",
+        default="Persona attribute",
     )
     if color_mode == "K-means clusters":
         max_clusters = min(10, len(persona_ids))
@@ -265,36 +270,34 @@ def _layered_figure_state_keys(
     )
     if figure_kind not in _PROJECTION_KINDS:
         return LayeredFigureStateKeys(figure=figure_key)
-
-    graph_overlay = figure_kind == "isomap"
-    projection_key = widget_key(
+    prepared_key = widget_key(
         "load",
-        f"{scope}_projection_state",
+        f"{scope}_projection_ready",
         store_id(store),
         store.model_name,
         mask_strategy.value,
         figure_kind,
         str(n_components),
-        str(graph_overlay),
+        str(figure_kind == "isomap"),
         str(_DEFAULT_GRAPH_NEIGHBORS),
         variant,
-        "persona_vector",
         persona_key,
         layer_key,
     )
-    return LayeredFigureStateKeys(figure=figure_key, projection=projection_key)
+    return LayeredFigureStateKeys(figure=figure_key, prepared=prepared_key)
 
 
 def _projection_build_kwargs(
-    samples,
     *,
+    store: Store,
+    mask_strategy: MaskStrategy,
+    variant: str,
     figure_kind: str,
     selected_layers: list[int],
     n_components: int,
     color_config: ProjectionColorConfig,
     persona_ids: list[str],
     persona_names: dict[str, str],
-    projection_key: str | None,
 ) -> dict:
     if figure_kind not in _PROJECTION_KINDS:
         return {}
@@ -305,22 +308,29 @@ def _projection_build_kwargs(
         "graph_overlay": graph_overlay,
         "graph_n_neighbors": _DEFAULT_GRAPH_NEIGHBORS,
     }
+    source, location, model_name = store_cache_parts(store)
+    cache_args = (
+        source,
+        location,
+        model_name,
+        mask_strategy.value,
+        variant,
+        tuple(persona_ids),
+        tuple(selected_layers),
+    )
+    build_kwargs["projection_data"] = projection_data_cached(
+        *cache_args,
+        figure_kind,
+        n_components,
+        graph_overlay,
+        _DEFAULT_GRAPH_NEIGHBORS,
+    )
     if color_config.n_clusters is not None:
-        build_kwargs["n_clusters"] = color_config.n_clusters
-        build_kwargs["cluster_mode"] = color_config.cluster_mode
-    if projection_key is not None:
-        projection_data = st.session_state.get(projection_key)
-        if projection_data is None:
-            projection_data = prepare_layered_projection_data(
-                samples,
-                figure_kind,
-                layers=selected_layers,
-                n_components=n_components,
-                graph_overlay=graph_overlay,
-                graph_n_neighbors=_DEFAULT_GRAPH_NEIGHBORS,
-            )
-            st.session_state[projection_key] = projection_data
-        build_kwargs["projection_data"] = projection_data
+        build_kwargs["groups"] = kmeans_groups_cached(
+            *cache_args,
+            color_config.n_clusters,
+            color_config.cluster_mode or "mean_across_layers",
+        )
     if color_config.attribute_name is not None:
         build_kwargs.update(
             attribute_color_kwargs(
@@ -487,8 +497,6 @@ def _render_layered_figure_analysis(
         selected_layers=selected_layers,
         pair_trajectories=pair_trajectories,
     )
-    if state_keys.projection is not None:
-        _clear_old_projection_states(state_keys.projection)
     filename = scope
     _clear_old_figure_states(state_keys.figure)
     persona_names = st.session_state.get(
@@ -496,7 +504,13 @@ def _render_layered_figure_analysis(
         {},
     )
 
-    if st.button(button_label, type="primary"):
+    build_clicked = st.button(button_label, type="primary")
+    recolor_from_warm_projection = (
+        state_keys.prepared is not None
+        and bool(st.session_state.get(state_keys.prepared))
+        and state_keys.figure not in st.session_state
+    )
+    if build_clicked or recolor_from_warm_projection:
         build_label = {
             "umap": "Computing UMAP projections…",
             "pca": "Computing PCA projections…",
@@ -514,14 +528,15 @@ def _render_layered_figure_analysis(
             )
             progress.progress(55, text=build_label)
             build_kwargs = _projection_build_kwargs(
-                samples,
+                store=store,
+                mask_strategy=mask_strategy,
+                variant=variant,
                 figure_kind=figure_kind,
                 selected_layers=selected_layers,
                 n_components=n_components,
                 color_config=color_config,
                 persona_ids=persona_ids,
                 persona_names=persona_names,
-                projection_key=state_keys.projection,
             )
             main_fig, extra_fig = _build_layered_analysis_figures(
                 samples,
@@ -541,12 +556,15 @@ def _render_layered_figure_analysis(
             n_samples = samples.vectors.shape[0]
             del samples
             _store_figure_state(state_keys.figure, (main_fig, extra_fig, n_samples))
+            if state_keys.prepared is not None:
+                _clear_old_prepared_states(state_keys.prepared)
+                st.session_state[state_keys.prepared] = True
             progress.progress(100, text="Done.")
         except Exception as exc:
             st.error(f"Could not build figure: {exc}")
             st.session_state.pop(state_keys.figure, None)
         finally:
-            _release_vector_memory(store, [variant])
+            _release_vector_memory()
             progress.empty()
 
     if state_keys.figure in st.session_state:

tabs/probe.py

@@ -23,6 +23,7 @@ from persona_vectors.plots import plot_metric_comparison, plot_metric_over_layer
 from persona_vectors.probes import (
     AttributeLabels,
     attribute_probe_labels,
+    default_probe_kinds,
     filter_attribute_samples_min_count,
     infer_probe_task,
     layer_matrix,
@@ -85,8 +86,9 @@ class _SweepInputs:
     mask_value: str
     variant: str
     persona_ids: tuple[str, ...]
-    attribute: str
+    attributes: tuple[str, ...]
     task: str
+    probe_kinds: tuple[str, ...]
     n_pca_components: int | None
     layers: tuple[int, ...]
     min_class_count: int
@@ -234,22 +236,62 @@ def _select_personas(
 # ---------------------------------------------------------------------------
 
 
-def _select_attribute() -> str:
+@st.cache_data(show_spinner=False)
+def _attribute_tasks() -> dict[str, str]:
     dataset = synth_persona_dataset_cached()
-    options = list(synth_persona_attribute_names())
-    if "sex" in options:
-        default_index = options.index("sex")
+    return {
+        name: infer_probe_task(dataset, name)
+        for name in synth_persona_attribute_names()
+    }
+
+
+def _select_attributes() -> list[str]:
+    """Multi-select locked to one task type.
+
+    Picking the first attribute fixes the task; only same-task attributes stay
+    selectable. Clearing the selection reopens every attribute again.
+    """
+    dataset = synth_persona_dataset_cached()
+    tasks = _attribute_tasks()
+    all_names = list(synth_persona_attribute_names())
+
+    key = "probe:attributes"
+    if key not in st.session_state:
+        st.session_state[key] = ["sex"] if "sex" in all_names else all_names[:1]
+
+    selected = st.session_state[key]
+    if selected:
+        locked = tasks[selected[0]]
+        options = [name for name in all_names if tasks[name] == locked]
     else:
-        default_index = 0
-    return st.selectbox(
-        "Attribute to probe",
+        options = all_names
+
+    return st.multiselect(
+        "Attributes to probe",
         options=options,
-        index=default_index,
         format_func=lambda name: attribute_display_label(dataset, name),
-        key="probe:attribute",
+        key=key,
+        help="Pick one or more attributes of the same task type. They are "
+        "overlaid in one figure. Remove all to switch to a different task type.",
     )
 
 
+def _select_probe_kinds(task: str) -> list[str]:
+    """Pick which probe families to fit. Only shown when the task has >1."""
+    available = list(default_probe_kinds(task))  # type: ignore[arg-type]
+    if len(available) < 2:
+        return available
+    selected = st.multiselect(
+        "Probe kinds to fit",
+        options=available,
+        default=available,
+        key=f"probe:kinds:{task}",
+        help="Which probe families to fit at each layer. Defaults to all "
+        "available for this task.",
+    )
+    return selected or available
+
+
 def _select_pca_components() -> int | None:
     use_pca = st.toggle(
         "Add PCA-compressed comparison",
@@ -298,61 +340,78 @@ 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]]], AttributeLabels, LayeredSamples]:
+) -> 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()
-    labels = attribute_probe_labels(
-        dataset, inputs.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
-    )
+    # 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(n_pca: int | None) -> list[dict[str, object]]:
+    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(None),
-            f"pca{inputs.n_pca_components}": _sweep(inputs.n_pca_components),
+            "full": _sweep_all(None),
+            f"pca{inputs.n_pca_components}": _sweep_all(inputs.n_pca_components),
         }
     else:
-        rows_by_label = {"full": _sweep(None)}
-    return rows_by_label, labels, probe_samples
+        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]]],
-    labels: AttributeLabels,
-    samples: LayeredSamples,
-    attribute: str,
+    per_attr: dict[str, tuple[AttributeLabels, LayeredSamples]],
+    attributes: tuple[str, ...],
     task: str,
     inputs: _SweepInputs,
 ) -> None:
     primary = _PRIMARY_METRIC[task]
     secondary = _SECONDARY_METRIC.get(task)
 
-    # Tolerate stale session state from a previous code version (bare rows).
-    if isinstance(rows_by_label, list):
-        rows_by_label = {"full": rows_by_label}
     primary_label = (
         f"pca{inputs.n_pca_components}" if inputs.n_pca_components else "full"
     )
     rows = rows_by_label.get(primary_label) or next(iter(rows_by_label.values()))
 
     def _plot(metric: str):
-        if len(rows_by_label) > 1:
-            return plot_metric_comparison(rows_by_label, attribute, metric=metric)
-        return plot_metric_over_layers(rows, attribute, metric=metric)
+        if len(rows_by_label) > 1 or len(attributes) > 1:
+            return plot_metric_comparison(
+                rows_by_label, list(attributes), metric=metric
+            )
+        return plot_metric_over_layers(rows, attributes[0], metric=metric)
 
     st.plotly_chart(_plot(primary), width="stretch")
     if secondary is not None:
@@ -377,21 +436,31 @@ def _show_sweep(
     if best is None:
         return
 
-    if len(rows_by_label) > 1:
+    multi_attr = len(attributes) > 1
+    if len(rows_by_label) > 1 or multi_attr:
         summary_rows = []
         for label, label_rows in rows_by_label.items():
-            label_best = _best_row(label_rows)
-            if label_best is None:
-                continue
-            summary_rows.append({
-                "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
-                ),
-            })
+            for attribute in attributes:
+                attr_rows = [
+                    row for row in label_rows
+                    if row.get("attribute") == attribute
+                ]
+                label_best = _best_row(attr_rows)
+                if label_best is None:
+                    continue
+                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_rows.append(summary_row)
         if summary_rows:
             st.dataframe(summary_rows, width="stretch", hide_index=True)
 
@@ -399,18 +468,26 @@ def _show_sweep(
         f" · pca{inputs.n_pca_components}" if inputs.n_pca_components else ""
     )
 
-    cols = st.columns([1, 1.2, 1.8])
-    cols[0].metric("Best layer", best["layer"])
-    cols[1].metric(
-        f"Best {primary}",
-        f"{best[primary]:.3f}",
-        delta=f"baseline {best.get(f'baseline_{primary}', float('nan')):.3f}",
-        delta_color="off",
-    )
-    cols[2].metric("Probe", f"{best['probe_kind']}{feature_desc}")
+    best_attr = str(best["attribute"])
+    labels, samples = per_attr[best_attr]
+    if multi_attr:
+        # The per-attribute summary table above already covers every result;
+        # a single "best" card would only show one attribute, so skip it and
+        # just say which one the controls below operate on.
+        st.caption(f"Controls below use the best result: **{best_attr}**.")
+    else:
+        cols = st.columns([1, 1.2, 1.8])
+        cols[0].metric("Best layer", best["layer"])
+        cols[1].metric(
+            f"Best {primary}",
+            f"{best[primary]:.3f}",
+            delta=f"baseline {best.get(f'baseline_{primary}', float('nan')):.3f}",
+            delta_color="off",
+        )
+        cols[2].metric("Probe", f"{best['probe_kind']}{feature_desc}")
 
     _render_selectivity_control(best, labels, samples, task, inputs)
-    _render_save_artifact(best, labels, samples, attribute, task, inputs)
+    _render_save_artifact(best, labels, samples, task, inputs)
 
 
 def _render_selectivity_control(
@@ -461,7 +538,6 @@ def _render_save_artifact(
     best: dict[str, object],
     labels: AttributeLabels,
     samples: LayeredSamples,
-    attribute: str,
     task: str,
     inputs: _SweepInputs,
 ) -> None:
@@ -540,12 +616,15 @@ def render_probing_tab() -> None:
         if not persona_ids:
             return
 
-    dataset = synth_persona_dataset_cached()
     with st.expander("Probe configuration", expanded=True):
-        attribute = _select_attribute()
-        task = infer_probe_task(dataset, attribute)
+        attributes = _select_attributes()
+        if not attributes:
+            st.info("Select at least one attribute to probe.")
+            return
+        task = _attribute_tasks()[attributes[0]]
         st.caption(f"Inferred task: **{task}**")
 
+        probe_kinds = _select_probe_kinds(task)
         n_pca_components = _select_pca_components()
 
         source, location, model_name = store_cache_parts(store)
@@ -563,17 +642,13 @@ def render_probing_tab() -> None:
         num_layers = max(available_layers) + 1
         layers = _select_layers(num_layers)
         min_class_count = _MIN_CLASS_COUNT
-        seed = st.number_input(
-            "Seed", min_value=0, max_value=10_000, value=0, step=1,
-            key="probe:seed",
-            help="Seeds the probe/PCA fit. The 80/20 split itself is fixed "
-            "(random_state=0).",
-        )
+        seed = 0
 
     inputs = _SweepInputs(
         source=source, location=location, model_name=model_name,
         mask_value=mask_strategy.value, variant=variant,
-        persona_ids=tuple(persona_ids), attribute=attribute, task=task,
+        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,
         seed=int(seed),
@@ -584,25 +659,21 @@ def render_probing_tab() -> None:
     if run:
         with st.spinner("Evaluating probes across layers..."):
             try:
-                sweep, labels, probe_samples = _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)
                 return
-        st.session_state[state_key] = (
-            sweep,
-            labels,
-            probe_samples,
-            attribute,
-            task,
-            inputs,
-        )
+        st.session_state[state_key] = (sweep, per_attr, inputs)
 
     if state_key in st.session_state:
         saved_result = st.session_state[state_key]
-        if len(saved_result) == 5:
-            sweep, labels, probe_samples, last_attr, last_task = saved_result
-            result_inputs = inputs
+        if len(saved_result) != 3:
+            # Stale shape from a previous code version — drop it.
+            st.session_state.pop(state_key, None)
         else:
-            sweep, labels, probe_samples, last_attr, last_task, result_inputs = saved_result
-        _show_sweep(sweep, labels, probe_samples, last_attr, last_task, result_inputs)
+            sweep, per_attr, result_inputs = saved_result
+            _show_sweep(
+                sweep, per_attr, result_inputs.attributes,
+                result_inputs.task, result_inputs,
+            )

tests/test_probes.py

@@ -12,9 +12,11 @@ two correctness fixes:
 import pytest
 import torch
 
+from persona_vectors.probes import ProbeArtifact
 from utils.probes import (
     LoadedProbe,
     _LinearProbe,
+    _loaded_probe_from_artifact,
     _normalize_labels,
     parse_probe_filename,
 )
@@ -196,3 +198,33 @@ def test_run_single_output_predicts_negative_when_score_low():
     result = probe.run(torch.tensor([1.0, 1.0]))
     assert result.predicted_index == 0
     assert result.predicted_label == "neg"
+
+
+# --------------------------------------------------------------------------- #
+# canonical persona-vectors artifacts
+# --------------------------------------------------------------------------- #
+
+
+def test_loaded_probe_from_canonical_artifact():
+    artifact = ProbeArtifact(
+        metadata={
+            "schema_version": 2,
+            "input_dim": 2,
+            "artifact_feature_dim": 2,
+            "class_names": ["neg", "pos"],
+            "task": "binary",
+            "probe_kind": "logistic_regression",
+            "layer": 3,
+        },
+        tensors={
+            "weight": torch.tensor([[-1.0, 0.0], [1.0, 0.0]]),
+            "bias": torch.zeros(2),
+        },
+    )
+    probe = _loaded_probe_from_artifact(
+        filename="m/answer_mean/templated/sex/logistic_regression_layer3/probe.json",
+        artifact=artifact,
+    )
+    assert probe.labels == ["neg", "pos"]
+    assert probe.layer == 3
+    assert probe.run(torch.tensor([1.0, 0.0])).predicted_label == "pos"

utils/analysis_sources.py

@@ -1,7 +1,11 @@
 import os
 
 import streamlit as st
-from persona_vectors.analysis import LayeredSamples, load_persona_vectors
+from persona_vectors.analysis import (
+    AnalysisDataset,
+    LayeredSamples,
+    load_analysis_dataset,
+)
 from persona_vectors.artifacts import (
     PersonaVectorStore,
     HFPersonaVectorStore,
@@ -10,6 +14,11 @@ from persona_vectors.artifacts import (
 )
 from persona_vectors.extraction import MaskStrategy
 from persona_vectors.hub import list_hub_vector_models
+from persona_vectors.plots import (
+    LayeredProjectionData,
+    prepare_kmeans_groups,
+    prepare_layered_projection_data,
+)
 
 from utils.helpers import env_int
 
@@ -26,7 +35,8 @@ SOURCES = (SOURCE_HUB, SOURCE_LOCAL)
 
 
 _STORE_CACHE_ENTRIES = env_int("PERSONA_UI_STORE_CACHE_ENTRIES", 4)
-_VECTOR_CACHE_ENTRIES = env_int("PERSONA_UI_VECTOR_CACHE_ENTRIES", 2)
+_VECTOR_CACHE_ENTRIES = env_int("PERSONA_UI_VECTOR_CACHE_ENTRIES", 4)
+_PREPARED_CACHE_ENTRIES = env_int("PERSONA_UI_PREPARED_CACHE_ENTRIES", 8)
 
 
 @st.cache_resource(show_spinner=False, max_entries=_STORE_CACHE_ENTRIES)
@@ -137,23 +147,41 @@ def local_model_matches(left: str, right: str) -> bool:
 
 
 @st.cache_resource(show_spinner=False, max_entries=_VECTOR_CACHE_ENTRIES)
-def load_persona_vectors_cached(
+def load_analysis_dataset_cached(
     source: str,
     location: str,
     model_name: str,
     mask_strategy_value: str,
-    variant: str,
+    variants: tuple[str, ...],
     persona_ids: tuple[str, ...],
-) -> LayeredSamples:
+) -> AnalysisDataset:
     store = activation_store_cached(source, location, model_name, mask_strategy_value)
-    return load_persona_vectors(
+    return load_analysis_dataset(
         store,
-        variant,
+        variants,
         mask_strategy=MaskStrategy(mask_strategy_value),
-        persona_ids=list(persona_ids),
+        persona_ids=persona_ids,
     )
 
 
+def load_persona_vectors_cached(
+    source: str,
+    location: str,
+    model_name: str,
+    mask_strategy_value: str,
+    variant: str,
+    persona_ids: tuple[str, ...],
+) -> LayeredSamples:
+    return load_analysis_dataset_cached(
+        source,
+        location,
+        model_name,
+        mask_strategy_value,
+        (variant,),
+        persona_ids,
+    ).samples(variant)
+
+
 def load_variant_vectors_cached(
     source: str,
     location: str,
@@ -162,12 +190,64 @@ def load_variant_vectors_cached(
     variants: tuple[str, ...],
     persona_ids: tuple[str, ...],
 ) -> dict[str, LayeredSamples]:
-    return {
-        variant: load_persona_vectors_cached(
-            source, location, model_name, mask_strategy_value, variant, persona_ids
-        )
-        for variant in variants
-    }
+    return load_analysis_dataset_cached(
+        source,
+        location,
+        model_name,
+        mask_strategy_value,
+        variants,
+        persona_ids,
+    ).samples_by_variant
+
+
+@st.cache_resource(show_spinner=False, max_entries=_PREPARED_CACHE_ENTRIES)
+def projection_data_cached(
+    source: str,
+    location: str,
+    model_name: str,
+    mask_strategy_value: str,
+    variant: str,
+    persona_ids: tuple[str, ...],
+    layers: tuple[int, ...],
+    kind: str,
+    n_components: int,
+    graph_overlay: bool,
+    graph_n_neighbors: int,
+) -> LayeredProjectionData:
+    samples = load_persona_vectors_cached(
+        source, location, model_name, mask_strategy_value, variant, persona_ids
+    )
+    return prepare_layered_projection_data(
+        samples,
+        kind,
+        layers=list(layers),
+        n_components=n_components,
+        graph_overlay=graph_overlay,
+        graph_n_neighbors=graph_n_neighbors,
+    )
+
+
+@st.cache_resource(show_spinner=False, max_entries=_PREPARED_CACHE_ENTRIES)
+def kmeans_groups_cached(
+    source: str,
+    location: str,
+    model_name: str,
+    mask_strategy_value: str,
+    variant: str,
+    persona_ids: tuple[str, ...],
+    layers: tuple[int, ...],
+    n_clusters: int,
+    cluster_mode: str,
+) -> list[str] | dict[int, list[str]]:
+    samples = load_persona_vectors_cached(
+        source, location, model_name, mask_strategy_value, variant, persona_ids
+    )
+    return prepare_kmeans_groups(
+        samples,
+        layers=list(layers),
+        n_clusters=n_clusters,
+        cluster_mode=cluster_mode,
+    )
 
 
 def prefetch_hub_metadata(
@@ -194,13 +274,3 @@ def prefetch_hub_metadata(
             mask_strategy_value,
             (variant,),
         )
-
-
-def release_hf_store_cache(
-    store: Store,
-    variants: list[str] | tuple[str, ...] | None = None,
-) -> None:
-    """Drop cached HF data for ``variants`` (or all) on Hub stores."""
-    release_cache = getattr(store, "release_cache", None)
-    if isinstance(store, HFPersonaVectorStore) and callable(release_cache):
-        release_cache(variants)

utils/probes.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import io
-import json
 import os
 import re
 from dataclasses import dataclass
@@ -12,6 +11,7 @@ import streamlit as st
 import torch
 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]+)_"
@@ -457,14 +457,19 @@ def _load_persona_probe_artifact(
     metadata_path: Path,
     weights_path: Path,
 ) -> LoadedProbe:
-    if not metadata_path.is_file():
-        raise FileNotFoundError(f"Missing probe metadata file: {metadata_path}")
-    if not weights_path.is_file():
-        raise FileNotFoundError(f"Missing probe weights file: {weights_path}")
-    from safetensors.torch import load_file
-
-    metadata = json.loads(metadata_path.read_text())
-    tensors = load_file(str(weights_path), device="cpu")
+    if metadata_path.parent != weights_path.parent:
+        raise ValueError("Canonical probe files must share one artifact directory.")
+    artifact = load_probe_artifact(metadata_path)
+    return _loaded_probe_from_artifact(filename=filename, artifact=artifact)
+
+
+def _loaded_probe_from_artifact(
+    *,
+    filename: str,
+    artifact: ProbeArtifact,
+) -> LoadedProbe:
+    metadata = artifact.metadata
+    tensors = artifact.tensors
     payload = {
         **metadata,
         "model_type": "linear",

uv.lock

@@ -1608,7 +1608,7 @@ requires-dist = [
     { name = "catppuccin", specifier = ">=2.5.0" },
     { name = "datasets", specifier = ">=4.8.5" },
     { name = "huggingface-hub", specifier = ">=1.14.0" },
-    { name = "persona-vectors", specifier = ">=0.8.2" },
+    { name = "persona-vectors", specifier = ">=0.8.3" },
     { name = "plotly", specifier = ">=6.6.0" },
     { name = "python-dotenv", specifier = ">=1.2.2" },
     { name = "safetensors", specifier = ">=0.7.0" },
@@ -1620,7 +1620,7 @@ dev = [{ name = "pytest", specifier = ">=9.0.3" }]
 
 [[package]]
 name = "persona-vectors"
-version = "0.8.2"
+version = "0.8.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "datasets" },
@@ -1639,9 +1639,9 @@ dependencies = [
     { name = "transformers" },
     { name = "umap-learn" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/ef/f4/66d2a1e30ed814a1ea945e27e2f9241cd7374872575e4d4c9e602a92a1cc/persona_vectors-0.8.2.tar.gz", hash = "sha256:f5b0776f8adbdfd38b9ad0f097daf88abb4c5dc504b3d3620af3f392e4a4621d", size = 42138, upload-time = "2026-05-16T22:11:53.019Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/c0/1d/472284f43e2a276a035e9e3de08a92654945193699598def6d6a2aa74c96/persona_vectors-0.8.3.tar.gz", hash = "sha256:f0519846b3712865bd2562cd239df05ddd006ac3d1e73e5ec5a6c860aaed5b2e", size = 43146, upload-time = "2026-05-17T12:43:13.601Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/15/02/7af86ed4040c4f866705a7ec28b50ebaa570502c3b74465fd9282856b2b7/persona_vectors-0.8.2-py3-none-any.whl", hash = "sha256:6bfa374e86d5cefc009cea07a8b43cc98d710e508d8f3e3394c24483d342799b", size = 52033, upload-time = "2026-05-16T22:11:54.128Z" },
+    { url = "https://files.pythonhosted.org/packages/60/d1/a38dc354718310122cd5d3de63e3aa9060490c8db4c2eadb1d4985684796/persona_vectors-0.8.3-py3-none-any.whl", hash = "sha256:2feeaf45b071ed417d88add48a1012455c8027e4f839e99658a9808c26786b8a", size = 53129, upload-time = "2026-05-17T12:43:12.693Z" },
 ]
 
 [[package]]