Updated to latest persona-vector and loading from HF

README.md

@@ -20,7 +20,7 @@ Streamlit interface for persona vector extraction, analysis, and chat.
 A web app built on top of [persona-vectors](../persona-vectors) that provides three tabs:
 
 - **Chat** — interactive conversations with a model using persona-based system prompts (templated or biography)
-- **Compare** — load saved activations and explore layer-wise cosine similarity, persona-mean PCA, UMAP, and similarity projections
+- **Compare** — load local or Hub persona vectors and explore cosine similarity, PCA, UMAP, and similarity views
 - **Extract** — run activation extraction from HuggingFace persona datasets or a local JSONL dataset directly from the browser
 
 ## Repository Layout
@@ -29,19 +29,20 @@ A web app built on top of [persona-vectors](../persona-vectors) that provides th
 persona-ui/
 ├── app.py                   # Main entry point (Streamlit)
 ├── state.py                 # Session state management (chat history, KV cache)
-├── scripts/
-│   └── oracle_probe.py      # Notebook-style activation oracle script
 ├── tabs/
 │   ├── chat.py              # Chat tab
 │   ├── compare.py           # Activation comparison tab
 │   ├── compare_chat.py      # Side-by-side chat comparison mode
-│   └── extract.py           # Extraction tab
+│   ├── extract.py           # Extraction tab
+│   └── probe_ui.py          # Probe upload and tracing controls
 └── utils/
     ├── chat.py              # Chat generation logic
     ├── chat_export.py       # Export chat logs to JSON
     ├── contrast.py          # Contrastive token log-prob coloring
     ├── datasets.py          # Dataset loader wrapper
     ├── helpers.py           # UI labels and slug helpers
+    ├── probe_trace.py       # Chat-token activation tracing
+    ├── probes.py            # Probe loading and scoring
     └── runtime.py           # Model caching and NDIF queries
 ```
 
@@ -58,13 +59,11 @@ cp .env.example .env
 
 ## Local Development
 
-The committed dependency graph uses git sources so `persona-ui` can install cleanly in a Hugging Face Space or any isolated environment.
+This checkout is configured to use the sibling `../persona-vectors` package as
+an editable dependency. For deployment, switch `persona-vectors` back to the
+published package or another installable source.
 
-For local sibling checkouts, uncomment the `path` sources in `persona-ui/pyproject.toml` and `persona-vectors/pyproject.toml`, then comment out the git sources.
-
-## Local Setup Note
-
-For local development, `persona-data` and `persona-vectors` can still be checked out in the parent directory of `persona-ui`.
+`persona-data` can also be checked out next to this repo for local package work.
 
 Example:
 
@@ -112,13 +111,16 @@ Copy `.env.example` to `.env` and fill in:
 NDIF_API_KEY=...       # Required for remote (NDIF) model execution
 HF_HOME=...            # Optional: HuggingFace cache directory
 ARTIFACTS_DIR=...      # Optional: where activations are read from (default: ./artifacts)
+PERSONA_VECTORS_HUB_REPO=...  # Optional: default Compare-tab Hub dataset repo
 ```
 
 The app picks up this file automatically via `load_dotenv()` on startup.
 
-## Saved Artifacts
+## Persona Vectors
 
-The Compare and Extract tabs read from / write to:
+The Compare tab reads persona vectors from either a Hugging Face dataset created
+by `persona-vectors/scripts/push_to_hf.py` or from local artifacts. The Extract
+tab writes local artifacts to:
 
 ```
 artifacts/

pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "persona-ui"
-version = "0.2.1"
+version = "0.3.0"
 description = "Streamlit UI for persona-vectors"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
-    "persona-vectors>=0.5.3",
+    "persona-vectors>=0.6.1",
     "persona-data>=0.4.1",
     "streamlit>=1.44.0",
     "plotly>=6.6.0",

tabs/chat_ui.py

@@ -6,14 +6,13 @@ from persona_data.synth_persona import PersonaData
 
 from utils.contrast import TokenContrast, render_contrast_html
 from utils.helpers import (
-    CHAT_PROMPT_MODE_LABELS,
     CHAT_PROMPT_MODE_LABEL_TO_KEY,
+    CHAT_PROMPT_MODE_LABELS,
     VARIANT_LABELS,
     persona_label,
     widget_key,
 )
 
-
 GENERATION_DEFAULTS = {
     "max_new_tokens": 256,
     "temperature": 1.0,

tabs/compare.py

@@ -1,15 +1,16 @@
+import os
 from collections.abc import Callable
-from itertools import combinations
 from dataclasses import dataclass
+from itertools import combinations
 
 import streamlit as st
 from persona_data.environment import get_artifacts_dir
 from persona_vectors.analysis import (
-    load_persona_mean_samples,
-    load_variant_mean_samples,
+    load_persona_vectors,
+    load_variant_vectors,
 )
-from persona_vectors.artifacts import ActivationStore
-from persona_vectors.artifacts import list_layers as list_available_layers
+from persona_vectors.artifacts import ActivationStore, HFActivationStore
+from persona_vectors.artifacts import list_layers as list_local_layers
 from persona_vectors.extraction import MaskStrategy
 from persona_vectors.plots import (
     build_layered_figure,
@@ -28,18 +29,29 @@ from utils.helpers import (
     widget_key,
 )
 
+Store = ActivationStore | HFActivationStore
+
+DEFAULT_HUB_REPO = os.environ.get(
+    "PERSONA_VECTORS_HUB_REPO",
+    "implicit-personalization/synth-persona-vectors",
+)
+SOURCE_HUB = "Hugging Face Hub"
+SOURCE_LOCAL = "Local activations"
+SOURCES = (SOURCE_HUB, SOURCE_LOCAL)
+
 
 def _filename(*parts: str) -> str:
     return "__".join(slugify(part) for part in parts if part)
 
 
-_list_layers_cached = st.cache_data(show_spinner=False)(list_available_layers)
+_list_layers_cached = st.cache_data(show_spinner=False)(list_local_layers)
 
 # Keep compare-tab selection state separate so projection defaults do not
 # overwrite cosine similarity defaults.
 _LAST_COSINE_PERSONAS_KEY = "compare:last_personas:cosine"
 _LAST_PROJECTION_PERSONAS_KEY = "compare:last_personas:projection"
 _LAST_MASK_STRATEGY_KEY = "compare:last_mask_strategy"
+_LAST_SOURCE_KEY = "compare:last_source"
 
 
 @dataclass(frozen=True)
@@ -51,8 +63,35 @@ class CosineSelection:
     persona_key: str
 
 
+def _store_id(store: Store) -> str:
+    """Stable identifier for cache/widget keys that distinguishes Hub vs local."""
+    if isinstance(store, HFActivationStore):
+        return f"hub:{store.repo_id}"
+    return f"local:{store.root_dir}"
+
+
+def _layers_for_variant(
+    store: Store,
+    variant: str,
+    persona_ids: list[str],
+    mask_strategy: MaskStrategy,
+) -> list[int]:
+    if isinstance(store, HFActivationStore):
+        if not persona_ids:
+            return []
+        sample = store.load(variant, persona_ids[0])
+        return list(range(int(sample.shape[0])))
+    return _list_layers_cached(
+        str(store.root_dir),
+        store.model_name,
+        [variant],
+        persona_ids,
+        mask_strategy=mask_strategy,
+    )
+
+
 def _select_artifact_personas(
-    store: ActivationStore,
+    store: Store,
     variants: list[str],
     mask_strategy: MaskStrategy,
     *,
@@ -61,17 +100,15 @@ def _select_artifact_personas(
     default_all: bool = False,
 ) -> tuple[list[str], dict[str, str]]:
     persona_options = store.list_personas(variants)
-    persona_names = store.persona_names(
-        persona_options,
-        variants=variants,
-    )
+    persona_names = store.persona_names(persona_options, variants=variants)
     if not persona_options:
         if len(variants) > 1:
             st.info(
-                "No personas have saved activations for all selected variants. Run extraction for both variants first."
+                "No personas have vectors for all selected variants. "
+                "Pick a single variant or change the source."
             )
         else:
-            st.info("No personas found for this model yet. Run extraction first.")
+            st.info("No personas found for this model and variant.")
         return [], persona_names
 
     last_personas: list[str] = st.session_state.get(remember_key, [])
@@ -147,19 +184,19 @@ def _render_mask_strategy_select(scope: str) -> MaskStrategy:
         ),
         format_func=lambda strategy: strategy.value.replace("_", " ").title(),
         key=widget_key("load", "mask_strategy", scope),
-        help="Which extracted activation artifact set to load.",
+        help="Which extracted activation set to load.",
     )
     st.session_state[_LAST_MASK_STRATEGY_KEY] = selected.value
     return selected
 
 
 def _render_cosine_selection(
-    store: ActivationStore,
+    store: Store,
     mask_strategy: MaskStrategy,
 ) -> CosineSelection | None:
-    variants = list(store.variants)
+    variants = store.available_variants()
     if len(variants) < 2:
-        st.info("Need at least two non-baseline variants for cosine comparison.")
+        st.info("Need at least two variants with saved vectors for cosine comparison.")
         return None
 
     with st.expander("Vector selection", expanded=True):
@@ -170,7 +207,7 @@ def _render_cosine_selection(
                 options=variants,
                 index=0,
                 format_func=prompt_variant_label,
-                key=widget_key("load", "variant_a"),
+                key=widget_key("load", "variant_a", _store_id(store)),
             )
         with col2:
             variant_b = st.selectbox(
@@ -178,7 +215,7 @@ def _render_cosine_selection(
                 options=variants,
                 index=min(1, len(variants) - 1),
                 format_func=prompt_variant_label,
-                key=widget_key("load", "variant_b"),
+                key=widget_key("load", "variant_b", _store_id(store)),
             )
 
         if variant_a == variant_b:
@@ -189,7 +226,7 @@ def _render_cosine_selection(
             store,
             [variant_a, variant_b],
             mask_strategy,
-            widget_scope="cosine",
+            widget_scope=f"cosine:{_store_id(store)}",
             remember_key=_LAST_COSINE_PERSONAS_KEY,
         )
     if not persona_ids:
@@ -204,11 +241,11 @@ def _render_cosine_selection(
 
 
 def _build_cosine_figures(
-    store: ActivationStore,
+    store: Store,
     selection: CosineSelection,
 ) -> tuple[object, object | None, int, int] | None:
     try:
-        variant_samples = load_variant_mean_samples(
+        variant_samples = load_variant_vectors(
             store,
             [selection.variant_a, selection.variant_b],
             persona_ids=selection.persona_ids,
@@ -242,7 +279,7 @@ def _build_cosine_figures(
             pair_samples = (
                 variant_samples
                 if {left, right} == {selection.variant_a, selection.variant_b}
-                else load_variant_mean_samples(
+                else load_variant_vectors(
                     store,
                     [left, right],
                     persona_ids=selection.persona_ids,
@@ -274,7 +311,7 @@ def _build_cosine_figures(
 
 
 def _render_cosine_similarity(
-    store: ActivationStore,
+    store: Store,
     mask_strategy: MaskStrategy,
 ) -> None:
     selection = _render_cosine_selection(store, mask_strategy)
@@ -284,6 +321,7 @@ def _render_cosine_similarity(
     cosine_fig_key = widget_key(
         "load",
         "cosine_fig_state",
+        _store_id(store),
         store.model_name,
         mask_strategy.value,
         selection.variant_a,
@@ -312,6 +350,7 @@ def _render_cosine_similarity(
         key=widget_key(
             "load",
             "compare_vectors",
+            _store_id(store),
             store.model_name,
             mask_strategy.value,
             selection.variant_a,
@@ -342,27 +381,26 @@ def _render_cosine_similarity(
 
 
 def _select_single_variant_samples(
-    store: ActivationStore,
+    store: Store,
     mask_strategy: MaskStrategy,
     scope: str,
 ) -> tuple[str, list[str], str, list[int]] | None:
-    variants = list(store.variants)
+    variants = store.available_variants()
+    if not variants:
+        st.info("No variants with saved vectors for this model.")
+        return None
     variant = st.selectbox(
         "Variant",
         options=variants,
-        index=(
-            variants.index("biography")
-            if "biography" in variants
-            else 0
-        ),
+        index=variants.index("biography") if "biography" in variants else 0,
         format_func=prompt_variant_label,
-        key=widget_key("load", "variant", scope),
+        key=widget_key("load", "variant", scope, _store_id(store)),
     )
     persona_ids, _ = _select_artifact_personas(
         store,
         [variant],
         mask_strategy,
-        widget_scope=scope,
+        widget_scope=f"{scope}:{_store_id(store)}",
         remember_key=_LAST_PROJECTION_PERSONAS_KEY,
         default_all=True,
     )
@@ -370,13 +408,7 @@ def _select_single_variant_samples(
         return None
 
     persona_key = "_".join(sorted(persona_ids))
-    layer_options = _list_layers_cached(
-        str(store.root_dir),
-        store.model_name,
-        [variant],
-        persona_ids,
-        mask_strategy=mask_strategy,
-    )
+    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.")
         return None
@@ -389,6 +421,7 @@ def _select_single_variant_samples(
             "load",
             "layers",
             scope,
+            _store_id(store),
             store.model_name,
             mask_strategy.value,
             variant,
@@ -403,7 +436,7 @@ def _select_single_variant_samples(
 
 
 def _render_layered_figure_analysis(
-    store: ActivationStore,
+    store: Store,
     mask_strategy: MaskStrategy,
     *,
     scope: str,
@@ -425,11 +458,12 @@ def _render_layered_figure_analysis(
     fig_key = widget_key(
         "load",
         f"{scope}_fig_state",
+        _store_id(store),
         store.model_name,
         mask_strategy.value,
         figure_kind,
         variant,
-        "persona_mean",
+        "persona_vector",
         persona_key,
     )
     filename = _filename(
@@ -438,13 +472,13 @@ def _render_layered_figure_analysis(
         store.model_name,
         mask_strategy.value,
         variant,
-        "persona_mean",
+        "persona_vector",
         persona_key,
     )
 
     if st.button(button_label, type="primary"):
         try:
-            samples = load_persona_mean_samples(
+            samples = load_persona_vectors(
                 store,
                 variant,
                 mask_strategy=mask_strategy,
@@ -462,8 +496,7 @@ def _render_layered_figure_analysis(
                     layers=selected_layers,
                     title=(
                         "Pair similarity trajectories - "
-                        f"{prompt_variant_label(variant)} - "
-                        "persona mean activations"
+                        f"{prompt_variant_label(variant)} - persona vectors"
                     ),
                 )
                 if include_pair_trajectories
@@ -488,17 +521,45 @@ def _render_layered_figure_analysis(
         st.success(f"Loaded {n_samples} samples.")
 
 
+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 _build_store(source: str, model_name: str, mask_strategy: MaskStrategy) -> Store:
+    if source == SOURCE_HUB:
+        repo = st.text_input(
+            "Hub repo",
+            value=st.session_state.get("compare:hub_repo", DEFAULT_HUB_REPO),
+            key="compare:hub_repo",
+            help="Hugging Face dataset published by `scripts/push_to_hf.py`.",
+        )
+        return HFActivationStore(repo, model_name, mask_strategy=mask_strategy)
+    artifacts_root = st.text_input(
+        "Artifacts root",
+        value=str(get_artifacts_dir() / "activations"),
+        key="compare:artifacts_root",
+    )
+    return ActivationStore(model_name, artifacts_root, mask_strategy=mask_strategy)
+
+
 def render_compare_tab(model_name: str) -> None:
     """Render the compare tab."""
 
     st.title("Compare")
-    st.caption("Compare saved activations by cosine similarity, PCA, or UMAP.")
+    st.caption("Compare persona vectors by cosine similarity, PCA, or UMAP.")
 
-    with st.expander("Artifact settings", expanded=False):
-        artifacts_root = st.text_input(
-            "Artifacts root",
-            value=str(get_artifacts_dir() / "activations"),
-        )
+    source = _render_source_select()
 
     analysis_mode = st.segmented_control(
         "Analysis mode",
@@ -510,9 +571,10 @@ def render_compare_tab(model_name: str) -> None:
     if analysis_mode is None:
         analysis_mode = ANALYSIS_MODES[0]
     st.caption(ANALYSIS_HELP_TEXT[analysis_mode])
-    with st.expander("Activation settings", expanded=False):
+
+    with st.expander("Source settings", expanded=False):
         mask_strategy = _render_mask_strategy_select(analysis_mode)
-    store = ActivationStore(model_name, artifacts_root, mask_strategy=mask_strategy)
+        store = _build_store(source, model_name, mask_strategy)
 
     if analysis_mode == "Cosine similarity":
         _render_cosine_similarity(store, mask_strategy)
@@ -525,8 +587,7 @@ def render_compare_tab(model_name: str) -> None:
             figure_kind="similarity",
             button_label="Generate similarity matrix",
             title_fn=lambda v: (
-                "Centered similarity - "
-                f"{prompt_variant_label(v)} - persona mean activations"
+                f"Centered similarity - {prompt_variant_label(v)} - persona vectors"
             ),
             include_pair_trajectories=True,
         )
@@ -539,6 +600,6 @@ def render_compare_tab(model_name: str) -> None:
         figure_kind=analysis_mode.lower(),
         button_label=f"Generate {analysis_mode} projection",
         title_fn=lambda v: (
-            f"{analysis_mode} - {prompt_variant_label(v)} - Persona means"
+            f"{analysis_mode} - {prompt_variant_label(v)} - persona vectors"
         ),
     )

tabs/extract.py

@@ -7,11 +7,10 @@ from persona_data.synth_persona import BASELINE_PERSONA_ID, PersonaData, QAPair
 from persona_vectors.artifacts import PERSONA_VARIANTS
 from persona_vectors.extraction import (
     MaskStrategy,
-    TokenSegment,
     prepare_inputs_for_strategy,
-    preview_token_segments,
     run_extraction,
 )
+from persona_vectors.preview import TokenSegment, preview_token_segments
 
 from utils.datasets import load_dataset
 from utils.helpers import (
@@ -33,7 +32,6 @@ _DEFAULT_MAX_QUESTIONS = 50
 
 @dataclass(frozen=True)
 class ExtractSettings:
-    runs: list[tuple[PersonaData, list[QAPair]]]
     mask_strategy: MaskStrategy
     max_questions: int
 
@@ -307,7 +305,6 @@ def _render_extract_actions() -> tuple[bool, bool]:
 
 def _render_token_preview(
     *,
-    remote: bool,
     model_name: str,
     run_plan: list[tuple[PersonaData, list[QAPair], str]],
     settings: ExtractSettings,
@@ -387,7 +384,7 @@ def _run_extraction_plan(
         progress.empty()
         ndif_status_box.empty()
 
-    status_box.success("Extraction complete")
+    status_box.empty()
     st.success(f"Saved {len(results)} artifact set(s)")
 
     for result in results:
@@ -448,7 +445,6 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         dataset_source=dataset_source,
     )
     settings = ExtractSettings(
-        runs=runs,
         mask_strategy=mask_strategy,
         max_questions=max_questions,
     )
@@ -458,7 +454,6 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
 
     if preview_clicked:
         _render_token_preview(
-            remote=remote,
             model_name=model_name,
             run_plan=run_plan,
             settings=settings,

utils/helpers.py

@@ -28,9 +28,9 @@ ANALYSIS_MODES = ["Cosine similarity", "Similarity matrix", "PCA", "UMAP"]
 
 ANALYSIS_HELP_TEXT = {
     "Cosine similarity": "Compare layer-wise alignment between variants.",
-    "Similarity matrix": "Compare centered pairwise similarity between persona means by layer, with pair trajectories across layers.",
-    "PCA": "Project per-persona mean activations into a 2D global view.",
-    "UMAP": "Project per-persona mean activations into a 2D local-neighborhood view.",
+    "Similarity matrix": "Compare centered pairwise similarity between persona vectors by layer, with pair trajectories across layers.",
+    "PCA": "Project per-persona vectors into a 2D global view.",
+    "UMAP": "Project per-persona vectors into a 2D local-neighborhood view.",
 }
 
 NDIF_STATUS_ICONS = {

uv.lock

@@ -1566,7 +1566,7 @@ wheels = [
 
 [[package]]
 name = "persona-ui"
-version = "0.2.1"
+version = "0.3.0"
 source = { virtual = "." }
 dependencies = [
     { name = "persona-data" },
@@ -1579,7 +1579,7 @@ dependencies = [
 [package.metadata]
 requires-dist = [
     { name = "persona-data", specifier = ">=0.4.1" },
-    { name = "persona-vectors", specifier = ">=0.5.3" },
+    { name = "persona-vectors", specifier = ">=0.6.1" },
     { name = "plotly", specifier = ">=6.6.0" },
     { name = "python-dotenv", specifier = ">=1.2.2" },
     { name = "streamlit", specifier = ">=1.44.0" },
@@ -1587,7 +1587,7 @@ requires-dist = [
 
 [[package]]
 name = "persona-vectors"
-version = "0.5.3"
+version = "0.6.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "datasets" },
@@ -1606,9 +1606,9 @@ dependencies = [
     { name = "transformers" },
     { name = "umap-learn" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/c7/53/20c77e298eb864ab917d58312b679007b936af64ede5fbb72d409268d62e/persona_vectors-0.5.3.tar.gz", hash = "sha256:d8bcb088a1814702401d22e21c39662ab840a4fb4b4f57dfd79999c9debfc1b8", size = 22791, upload-time = "2026-05-07T12:16:49.969Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/69/f3/6da35af90c8ea5333db1763ece04a3230353ac5a76c0dc8fea705a6e86cf/persona_vectors-0.6.1.tar.gz", hash = "sha256:552ac9a0d739a453c5d9eb612cb0d0d2820a1b53ce84f490295a84105a71f7cc", size = 24311, upload-time = "2026-05-07T15:07:29.951Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/32/26/a0197928e5202403094883331ff87799b54803bd3fd749b7d9c11f7332b3/persona_vectors-0.5.3-py3-none-any.whl", hash = "sha256:e44af7a6846d6d9249da12707dfae57f89d2a50a4ced05cdb4a844d39a9f03e8", size = 26805, upload-time = "2026-05-07T12:16:51.035Z" },
+    { url = "https://files.pythonhosted.org/packages/86/66/91df378258e2c0cbc7860652b07b5e65ee1949ba14be2efdb6c646a933f1/persona_vectors-0.6.1-py3-none-any.whl", hash = "sha256:593977ad19c9f23df7d86e302fe4bcf49159425da67d83281a11858026c5e85e", size = 28683, upload-time = "2026-05-07T15:07:30.791Z" },
 ]
 
 [[package]]