Imrpoved for performances and better support

pyproject.toml

@@ -5,7 +5,7 @@ description = "Streamlit UI for persona-vectors"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
-    "persona-vectors>=0.7.0",
+    "persona-vectors>=0.7.2",
     "persona-data>=0.4.2",
     "datasets>=4.8.5",
     "huggingface-hub>=1.14.0",

tabs/compare.py

@@ -11,6 +11,7 @@ from persona_vectors.extraction import MaskStrategy
 from persona_vectors.plots import (
     build_layered_figure,
     build_pair_similarity_figure,
+    build_similarity_figures,
     plot_layer_similarity,
     plot_persona_dendrogram,
     save_plot_html,
@@ -38,6 +39,7 @@ from utils.controls import render_mask_strategy_select
 from utils.helpers import (
     ANALYSIS_HELP_TEXT,
     ANALYSIS_MODES,
+    personas_fingerprint,
     prompt_variant_label,
     slugify,
     widget_key,
@@ -171,9 +173,7 @@ def _seed_persona_memory(
     )
     remembered_count = int(st.session_state.get(remembered_count_key, default_count))
     persona_count = min(max(remembered_count, 0), len(options.regular_ids))
-    include_assistant = bool(
-        st.session_state.get(remembered_assistant_key, options.assistant_id is not None)
-    )
+    include_assistant = bool(st.session_state.get(remembered_assistant_key, False))
     return persona_count, include_assistant
 
 
@@ -355,7 +355,7 @@ def _render_cosine_selection(
         variant_a=variant_a,
         variant_b=variant_b,
         persona_ids=persona_ids,
-        persona_key="_".join(sorted(persona_ids)),
+        persona_key=personas_fingerprint(persona_ids),
     )
 
 
@@ -363,30 +363,31 @@ def _build_cosine_figures(
     store: Store,
     selection: CosineSelection,
 ) -> tuple[object, object | None, int, int] | None:
-    variant_sample_cache = {}
+    variant_sample_cache: dict[str, object] = {}
 
-    def _load_pair(left: str, right: str):
-        key = tuple(sorted((left, right)))
-        if key not in variant_sample_cache:
-            variant_sample_cache[key] = load_variant_vectors(
+    def _load_variant(variant: str):
+        if variant not in variant_sample_cache:
+            samples = load_variant_vectors(
                 store,
-                [left, right],
+                [variant],
                 persona_ids=selection.persona_ids,
             )
-        return variant_sample_cache[key]
+            variant_sample_cache[variant] = samples[variant]
+        return variant_sample_cache[variant]
 
     try:
-        variant_samples = _load_pair(selection.variant_a, selection.variant_b)
+        samples_a = _load_variant(selection.variant_a)
+        samples_b = _load_variant(selection.variant_b)
     except Exception as exc:
         st.error(f"Could not load vectors: {exc}")
         return None
 
-    labels = variant_samples[selection.variant_a].labels
+    labels = samples_a.labels
     display_traces = [
         (
             label,
-            variant_samples[selection.variant_a].vectors[index],
-            variant_samples[selection.variant_b].vectors[index],
+            samples_a.vectors[index],
+            samples_b.vectors[index],
         )
         for index, label in enumerate(labels)
     ]
@@ -403,12 +404,13 @@ def _build_cosine_figures(
     pair_errors = []
     for left, right in combinations(selection.variants, 2):
         try:
-            pair_samples = _load_pair(left, right)
+            left_samples = _load_variant(left)
+            right_samples = _load_variant(right)
             pair_traces.append(
                 (
                     f"{prompt_variant_label(left)} vs {prompt_variant_label(right)}",
-                    pair_samples[left].vectors.mean(dim=0),
-                    pair_samples[right].vectors.mean(dim=0),
+                    left_samples.vectors.mean(dim=0),
+                    right_samples.vectors.mean(dim=0),
                 )
             )
         except Exception as exc:
@@ -477,11 +479,18 @@ def _render_cosine_similarity(
             selection.persona_key,
         ),
     ):
-        figures = _build_cosine_figures(store, selection)
-        if figures is None:
-            st.session_state.pop(cosine_fig_key, None)
-            return
-        st.session_state[cosine_fig_key] = figures
+        progress = st.progress(0, text="Loading activation vectors…")
+        try:
+            progress.progress(15, text="Loading activation vectors…")
+            figures = _build_cosine_figures(store, selection)
+            if figures is None:
+                st.session_state.pop(cosine_fig_key, None)
+                return
+            progress.progress(90, text="Storing figure state…")
+            st.session_state[cosine_fig_key] = figures
+            progress.progress(100, text="Done.")
+        finally:
+            progress.empty()
 
     if cosine_fig_key in st.session_state:
         fig, pair_fig, n_traces, n_pair_traces = st.session_state[cosine_fig_key]
@@ -526,7 +535,7 @@ def _select_single_variant_samples(
     if not persona_ids:
         return None
 
-    persona_key = "_".join(sorted(persona_ids))
+    persona_key = personas_fingerprint(persona_ids)
     layer_options = _layers_for_variant(store, variant, persona_ids, mask_strategy)
     if not layer_options:
         st.info("No shared layers are available for the selected personas.")
@@ -590,43 +599,66 @@ def _render_layered_figure_analysis(
     filename = scope
 
     if st.button(button_label, type="primary"):
+        build_label = {
+            "umap": "Computing UMAP projections…",
+            "pca": "Computing PCA projections…",
+            "similarity": "Computing similarity matrices…",
+        }.get(figure_kind, "Building figure…")
+        progress = st.progress(0, text="Loading activation vectors…")
         try:
+            progress.progress(15, text="Loading activation vectors…")
             samples = load_persona_vectors(
                 store,
                 variant,
                 mask_strategy=mask_strategy,
                 persona_ids=persona_ids,
             )
+            progress.progress(55, text=build_label)
             build_kwargs = {}
             if figure_kind in {"umap", "pca"}:
                 build_kwargs["n_components"] = n_components
                 if n_clusters is not None:
                     build_kwargs["n_clusters"] = n_clusters
-            main_fig = build_layered_figure(
-                samples,
-                figure_kind,
-                layers=selected_layers,
-                title=title_fn(variant),
-                **build_kwargs,
-            )
-            if figure_kind in {"umap", "pca"}:
-                main_fig.update_layout(height=700)
-            extra_fig = (
-                build_pair_similarity_figure(
+            if figure_kind == "similarity" and include_pair_trajectories:
+                main_fig, extra_fig = build_similarity_figures(
                     samples,
                     layers=selected_layers,
-                    title=(
+                    title=title_fn(variant),
+                    pair_title=(
                         "Pair similarity trajectories - "
                         f"{prompt_variant_label(variant)} - persona vectors"
                     ),
                 )
-                if include_pair_trajectories
-                else None
-            )
+            else:
+                main_fig = build_layered_figure(
+                    samples,
+                    figure_kind,
+                    layers=selected_layers,
+                    title=title_fn(variant),
+                    **build_kwargs,
+                )
+                if figure_kind in {"umap", "pca"}:
+                    main_fig.update_layout(height=700)
+                extra_fig = (
+                    build_pair_similarity_figure(
+                        samples,
+                        layers=selected_layers,
+                        title=(
+                            "Pair similarity trajectories - "
+                            f"{prompt_variant_label(variant)} - persona vectors"
+                        ),
+                    )
+                    if include_pair_trajectories
+                    else None
+                )
+            progress.progress(90, text="Storing figure state…")
             st.session_state[fig_key] = (main_fig, extra_fig, samples.vectors.shape[0])
+            progress.progress(100, text="Done.")
         except Exception as exc:
             st.error(f"Could not build figure: {exc}")
             st.session_state.pop(fig_key, None)
+        finally:
+            progress.empty()
 
     if fig_key in st.session_state:
         main_fig, extra_fig, n_samples = st.session_state[fig_key]
@@ -658,7 +690,11 @@ def _render_dendrogram_analysis(
     with st.expander("Variant selection", expanded=True):
         col1, col2 = st.columns(2)
         default_a = "biography" if "biography" in variants else variants[0]
-        default_b_idx = variants.index("templated") if "templated" in variants else min(1, len(variants) - 1)
+        default_b_idx = (
+            variants.index("templated")
+            if "templated" in variants
+            else min(1, len(variants) - 1)
+        )
         with col1:
             variant_a = st.selectbox(
                 "Variant A",
@@ -704,26 +740,37 @@ def _render_dendrogram_analysis(
             key=widget_key("load", "dendro_linkage", store_id(store)),
         )
 
-    persona_key = "_".join(sorted(persona_ids))
+    persona_key = personas_fingerprint(persona_ids)
     fig_key = widget_key(
-        "load", "dendro_fig_state",
+        "load",
+        "dendro_fig_state",
         store_id(store),
         store.model_name,
         mask_strategy.value,
-        variant_a, variant_b,
+        variant_a,
+        variant_b,
         persona_key,
-        str(layered_mode), linkage,
+        str(layered_mode),
+        linkage,
     )
 
     if st.button(
         "Generate dendrograms",
         type="primary",
-        key=widget_key("load", "dendro_btn", store_id(store), variant_a, variant_b, persona_key),
+        key=widget_key(
+            "load", "dendro_btn", store_id(store), variant_a, variant_b, persona_key
+        ),
     ):
+        progress = st.progress(0, text="Loading first variant vectors…")
         try:
+            progress.progress(15, text="Loading first variant vectors…")
             samples_a = load_persona_vectors(
-                store, variant_a, mask_strategy=mask_strategy, persona_ids=persona_ids,
+                store,
+                variant_a,
+                mask_strategy=mask_strategy,
+                persona_ids=persona_ids,
             )
+            progress.progress(40, text="Building first dendrogram…")
             fig_a = plot_persona_dendrogram(
                 samples_a,
                 layered=layered_mode,
@@ -733,9 +780,14 @@ def _render_dendrogram_analysis(
             fig_a.update_layout(height=750)
             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=mask_strategy, persona_ids=persona_ids,
+                    store,
+                    variant_b,
+                    mask_strategy=mask_strategy,
+                    persona_ids=persona_ids,
                 )
+                progress.progress(75, text="Building second dendrogram…")
                 fig_b = plot_persona_dendrogram(
                     samples_b,
                     layered=layered_mode,
@@ -743,10 +795,20 @@ def _render_dendrogram_analysis(
                     title=f"Dendrogram — {prompt_variant_label(variant_b)}",
                 )
                 fig_b.update_layout(height=750)
-            st.session_state[fig_key] = (fig_a, fig_b, len(persona_ids), variant_a, variant_b)
+            progress.progress(90, text="Storing figure state…")
+            st.session_state[fig_key] = (
+                fig_a,
+                fig_b,
+                len(persona_ids),
+                variant_a,
+                variant_b,
+            )
+            progress.progress(100, text="Done.")
         except Exception as exc:
             st.error(f"Could not build dendrogram: {exc}")
             st.session_state.pop(fig_key, None)
+        finally:
+            progress.empty()
 
     if fig_key in st.session_state:
         fig_a, fig_b, n_personas, va, vb = st.session_state[fig_key]
@@ -764,7 +826,11 @@ def _render_dendrogram_analysis(
         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 []),
+            *(
+                [_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).")
@@ -917,7 +983,9 @@ def render_compare_tab() -> None:
     """Render the analysis tab."""
 
     st.title("Analysis")
-    st.caption("Analyse persona vectors by cosine similarity, PCA, UMAP, or hierarchical clustering.")
+    st.caption(
+        "Analyse persona vectors by cosine similarity, PCA, UMAP, or hierarchical clustering."
+    )
 
     source = _render_source_select()
 

tabs/extract.py

@@ -92,7 +92,7 @@ def _render_variant_controls(
     )
     include_baseline = st.checkbox(
         "Extract Assistant baseline",
-        value=st.session_state.get(_LAST_BASELINE_KEY, True),
+        value=st.session_state.get(_LAST_BASELINE_KEY, False),
         key=_extract_widget_key(model_name, remote, dataset_source, "baseline"),
         help="Also extract the Assistant baseline persona using the first persona's QA set.",
     )

utils/compare_sources.py

@@ -22,7 +22,7 @@ SOURCE_LOCAL = "Local activations"
 SOURCES = (SOURCE_HUB, SOURCE_LOCAL)
 
 
-@st.cache_resource(show_spinner=False)
+@st.cache_resource(show_spinner=False, max_entries=1)
 def activation_store_cached(
     source: str,
     location: str,
@@ -35,7 +35,7 @@ def activation_store_cached(
     return ActivationStore(model_name, location, mask_strategy=mask_strategy)
 
 
-@st.cache_data(show_spinner=False, ttl=10)
+@st.cache_data(show_spinner=False)
 def available_variants_cached(
     source: str,
     location: str,
@@ -46,7 +46,7 @@ def available_variants_cached(
     return store.available_variants()
 
 
-@st.cache_data(show_spinner=False, ttl=10)
+@st.cache_data(show_spinner=False)
 def personas_cached(
     source: str,
     location: str,
@@ -61,7 +61,7 @@ def personas_cached(
     )
 
 
-@st.cache_data(show_spinner=False, ttl=10)
+@st.cache_data(show_spinner=False)
 def persona_names_cached(
     source: str,
     location: str,
@@ -78,7 +78,7 @@ def persona_names_cached(
     )
 
 
-@st.cache_data(show_spinner=False, ttl=10)
+@st.cache_data(show_spinner=False)
 def local_model_options_cached(
     artifacts_root: str, mask_strategy_value: str
 ) -> list[str]:

utils/helpers.py

@@ -1,4 +1,6 @@
+import hashlib
 import re
+from collections.abc import Iterable
 
 from persona_data.synth_persona import PersonaData
 
@@ -54,6 +56,18 @@ def widget_key(*parts: str) -> str:
     return "::".join(parts)
 
 
+def personas_fingerprint(persona_ids: Iterable[str]) -> str:
+    """Stable short fingerprint for a set of persona ids.
+
+    Used as a discriminator in widget keys and session-state keys. At ~1k
+    personas, joining ids would produce ~20 KB strings; the sha1 prefix is
+    fixed-length and keeps tracebacks readable.
+    """
+
+    joined = "|".join(sorted(persona_ids))
+    return hashlib.sha1(joined.encode()).hexdigest()[:16]
+
+
 def prompt_variant_label(variant: str) -> str:
     """Return a human-friendly prompt-variant label."""
 

uv.lock

@@ -1594,7 +1594,7 @@ requires-dist = [
     { name = "datasets", specifier = ">=4.8.5" },
     { name = "huggingface-hub", specifier = ">=1.14.0" },
     { name = "persona-data", specifier = ">=0.4.2" },
-    { name = "persona-vectors", specifier = ">=0.7.0" },
+    { name = "persona-vectors", specifier = ">=0.7.2" },
     { name = "plotly", specifier = ">=6.6.0" },
     { name = "python-dotenv", specifier = ">=1.2.2" },
     { name = "streamlit", specifier = ">=1.44.0" },
@@ -1602,7 +1602,7 @@ requires-dist = [
 
 [[package]]
 name = "persona-vectors"
-version = "0.7.0"
+version = "0.7.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "datasets" },
@@ -1621,9 +1621,9 @@ dependencies = [
     { name = "transformers" },
     { name = "umap-learn" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/21/03/c419bd2809447ae3de0dfbe06de8c7462a141feebaea7157f2e64a119416/persona_vectors-0.7.0.tar.gz", hash = "sha256:b3bb10d61fb58e7a151a2f715a3bf3afecc6a2e30da85b73ccb5e5485109909b", size = 27499, upload-time = "2026-05-10T08:59:29.558Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/03/ee/3550a015546ad9e7842e4d94346e2268c1aa65a7278f39ab574dd2b702bd/persona_vectors-0.7.2.tar.gz", hash = "sha256:22f3e4349bd020fcc287ed5e00ff9d0b752dfe95b3f145b8778efc485cf7f67f", size = 27777, upload-time = "2026-05-10T10:29:07.67Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/3f/cf/5603631767d29cff21a7b02fb21f4ca53b85bc4369bb0a5060bdee01ee4e/persona_vectors-0.7.0-py3-none-any.whl", hash = "sha256:89c3ddd1f1a0ffc0dfce59319ce9417e0764011f8680fd2e7844fbea2ee75afe", size = 32116, upload-time = "2026-05-10T08:59:28.674Z" },
+    { url = "https://files.pythonhosted.org/packages/c7/c2/334c66944f1669c7cc19e4e3c2f49ab65b0a186e6cd6cb709afd253b4268/persona_vectors-0.7.2-py3-none-any.whl", hash = "sha256:cbd5b6de2afb17fccef9414eba1b0096ffa3407a1496b253bb2d30ed548d123e", size = 32387, upload-time = "2026-05-10T10:29:06.808Z" },
 ]
 
 [[package]]