Cleaned up code with the new updates

README.md

@@ -24,13 +24,10 @@ persona-ui/
 │   ├── compare.py           # Activation comparison tab
 │   └── extract.py           # Extraction tab
 └── utils/
-    ├── artifacts.py         # Load saved activations metadata
     ├── chat.py              # Chat generation logic
     ├── chat_export.py       # Export chat logs to JSON
     ├── datasets.py          # Dataset loader wrapper
-    ├── extraction.py        # Extraction orchestration
     ├── helpers.py           # UI labels and slug helpers
-    ├── local_dataset.py     # Local JSONL dataset parsing
     └── runtime.py           # Model caching and NDIF queries
 ```
 
@@ -81,7 +78,7 @@ HF_HOME=...            # Optional: HuggingFace cache directory
 ARTIFACTS_DIR=...      # Optional: where activations are read from (default: ./artifacts)
 ```
 
-The app picks up this file automatically via `load_env()` on startup.
+The app picks up this file automatically via `load_dotenv()` on startup.
 
 ## Saved Artifacts
 

app.py

@@ -26,7 +26,7 @@ def _sidebar_controls() -> tuple[bool, str, str, str]:
             if st.button(
                 tab_name,
                 key=f"sidebar__tab__{tab_name.lower()}",
-                use_container_width=True,
+                width="stretch",
                 type="primary" if is_selected else "secondary",
                 icon=icon,
             ):

tabs/chat.py

@@ -1,10 +1,15 @@
-import threading
 from concurrent.futures import ThreadPoolExecutor
-from contextlib import nullcontext
+from typing import Any
 
 import streamlit as st
+from persona_data.synth_persona import PersonaData
 
-from state import chat_session_key, get_chat_state, reset_chat_state
+from state import (
+    _default_chat_state,
+    chat_session_key,
+    get_chat_state,
+    reset_chat_state,
+)
 from utils.chat import ChatReply, generate_chat_reply, resolve_system_prompt
 from utils.chat_export import save_chat_export
 from utils.datasets import load_dataset
@@ -12,14 +17,12 @@ from utils.helpers import (
     MODE_LABEL_TO_KEY,
     MODE_LABELS,
     VARIANT_LABELS,
+    VISIBLE_MESSAGE_COUNT,
     persona_label,
     widget_key,
 )
 from utils.runtime import cached_model
 
-_VISIBLE_MESSAGE_COUNT = 5
-_model_lock = threading.Lock()
-
 
 def _render_chat_message(message: dict[str, str]) -> None:
     if not message.get("content"):
@@ -33,6 +36,21 @@ def _clear_chat_ui_state(*keys: str) -> None:
         st.session_state.pop(key, None)
 
 
+def _reset_single_chat_context(
+    model_name: str,
+    remote: bool,
+    dataset_source: str,
+    chat_state: dict[str, object],
+    persona_id: str,
+    prompt_mode: str,
+    *ui_keys: str,
+) -> None:
+    reset_chat_state(model_name, remote, dataset_source)
+    chat_state["persona_id"] = persona_id
+    chat_state["prompt_mode"] = prompt_mode
+    _clear_chat_ui_state(*ui_keys)
+
+
 def _generation_dict(gen_kwargs: dict, advanced_generation: bool) -> dict[str, object]:
     return {
         "max_new_tokens": int(gen_kwargs["max_new_tokens"]),
@@ -46,186 +64,146 @@ def _generation_dict(gen_kwargs: dict, advanced_generation: bool) -> dict[str, o
     }
 
 
-# ── Compare mode helpers ───────────────────────────────────────────────────────
+def _render_persona_prompt_controls(
+    personas: list[PersonaData],
+    current_persona_id: str | None,
+    current_prompt_mode: str,
+    persona_key: str,
+    prompt_key: str,
+    column_widths: tuple[int, int] = (3, 2),
+) -> tuple[PersonaData, str, bool]:
+    """Render persona and prompt selectors, returning the selected values."""
 
-
-def _panel_state(panel_key: str) -> dict:
-    """Get or initialise compare-panel chat state stored in session_state."""
-    if panel_key not in st.session_state:
-        st.session_state[panel_key] = {
-            "messages": [],
-            "persona_id": None,
-            "prompt_mode": "templated",
-            "past_key_values": None,
-        }
-    return st.session_state[panel_key]
-
-
-def _render_compare_panel(
-    side: str,
-    context_key: str,
-    personas: list,
-    remote: bool,
-    model_name: str,
-    dataset_source: str,
-    gen_kwargs: dict,
-    advanced_generation: bool,
-) -> dict:
-    """Render persona/prompt controls + chat log for one compare panel.
-
-    Returns a dict with keys needed by the generation step:
-      panel_key, state, active_system_prompt, selected_persona, chat_log
-    """
-    panel_key = widget_key(context_key, f"cmp_{side}")
-    state = _panel_state(panel_key)
-
-    # ── Per-panel selectors ──────────────────────────────────────────────────
-    p_col, m_col = st.columns([3, 2])
+    p_col, m_col = st.columns(list(column_widths))
     with p_col:
         selected_index = next(
-            (i for i, p in enumerate(personas) if p.id == state["persona_id"]), 0
+            (i for i, p in enumerate(personas) if p.id == current_persona_id), 0
         )
         selected_persona = st.selectbox(
             "Persona",
             options=personas,
             index=selected_index,
             format_func=persona_label,
-            key=widget_key(panel_key, "persona"),
+            key=persona_key,
         )
     with m_col:
-        current_label = VARIANT_LABELS.get(state["prompt_mode"], "None")
+        current_label = VARIANT_LABELS.get(current_prompt_mode, "None")
         prompt_mode_label = st.selectbox(
             "Prompt",
             options=MODE_LABELS,
             index=MODE_LABELS.index(current_label),
-            key=widget_key(panel_key, "prompt_mode"),
+            key=prompt_key,
         )
     prompt_mode = MODE_LABEL_TO_KEY[prompt_mode_label]
-
-    # Reset state when persona or mode changes.
     changed = (
-        state["persona_id"] != selected_persona.id
-        or state["prompt_mode"] != prompt_mode
+        current_persona_id != selected_persona.id or current_prompt_mode != prompt_mode
     )
-    if changed:
-        state["messages"] = []
-        state["past_key_values"] = None
-        state["persona_id"] = selected_persona.id
-        state["prompt_mode"] = prompt_mode
-        _clear_chat_ui_state(
-            widget_key(panel_key, "custom_prompt"),
-            widget_key(panel_key, "show_all"),
+    return selected_persona, prompt_mode, changed
+
+
+def _render_system_prompt_editor(
+    prompt_key: str,
+    prompt_mode: str,
+    active_system_prompt: str | None,
+    *,
+    height: int,
+    label: str = "Prompt",
+) -> str | None:
+    """Render the editable system prompt area for a chat panel."""
+
+    if prompt_mode == "empty":
+        return active_system_prompt
+
+    if prompt_key not in st.session_state:
+        st.session_state[prompt_key] = active_system_prompt or ""
+
+    with st.expander("Edit prompt", expanded=False):
+        edited_prompt = (
+            st.text_area(
+                label,
+                key=prompt_key,
+                height=height,
+                label_visibility="collapsed",
+            )
+            or None
         )
+    return edited_prompt
 
-    # ── System prompt ────────────────────────────────────────────────────────
-    active_system_prompt = resolve_system_prompt(
-        persona=selected_persona, mode=prompt_mode
-    )
-    custom_prompt_key = widget_key(panel_key, "custom_prompt")
-    if prompt_mode != "empty":
-        if custom_prompt_key not in st.session_state:
-            st.session_state[custom_prompt_key] = active_system_prompt
-        with st.expander("Edit prompt", expanded=False):
-            active_system_prompt = (
-                st.text_area(
-                    "prompt",
-                    key=custom_prompt_key,
-                    height=150,
-                    label_visibility="collapsed",
-                )
-                or None
-            )
 
-    export_success_message: str | None = None
-    action_col1, action_col2 = st.columns(2)
-    with action_col1:
-        if st.button(
-            "Export chat",
-            key=widget_key(panel_key, "export_chat"),
-            use_container_width=True,
-        ):
-            export_path = save_chat_export(
-                model_name=model_name,
-                dataset_source=dataset_source,
-                persona_id=selected_persona.id,
-                persona_name=getattr(selected_persona, "name", None),
-                panel_label=side,
-                prompt_mode=prompt_mode,
-                system_prompt=active_system_prompt,
-                messages=state["messages"],
-                generation=_generation_dict(gen_kwargs, advanced_generation),
-            )
-            export_success_message = f"Saved chat export to {export_path}"
-    with action_col2:
-        if st.button(
-            "Reset chat",
-            key=widget_key(panel_key, "reset"),
-            use_container_width=True,
-            type="secondary",
-        ):
-            state["messages"] = []
-            state["past_key_values"] = None
-            _clear_chat_ui_state(
-                widget_key(panel_key, "custom_prompt"),
-                widget_key(panel_key, "show_all"),
-            )
-            st.rerun()
+def _render_chat_window(
+    *,
+    chat_log: Any,
+    messages: list[dict[str, str]],
+    show_all_key: str,
+    show_all_btn_key: str,
+    show_earlier_label: str,
+) -> Any:
+    """Render the visible chat history inside one container."""
 
-    if export_success_message:
-        st.success(export_success_message)
-
-    # ── Message history ──────────────────────────────────────────────────────
-    show_all_key = widget_key(panel_key, "show_all")
-    messages = state["messages"]
-    if len(messages) > _VISIBLE_MESSAGE_COUNT and not st.session_state.get(
-        show_all_key, False
-    ):
-        hidden_count = len(messages) - _VISIBLE_MESSAGE_COUNT
-        if st.button(
-            f"Show earlier ({hidden_count} hidden)",
-            key=widget_key(panel_key, "show_all_btn"),
+    with chat_log:
+        if len(messages) > VISIBLE_MESSAGE_COUNT and not st.session_state.get(
+            show_all_key, False
         ):
-            st.session_state[show_all_key] = True
-            st.rerun()
-        visible = messages[-_VISIBLE_MESSAGE_COUNT:]
-    else:
-        visible = messages
+            hidden_count = len(messages) - VISIBLE_MESSAGE_COUNT
+            if st.button(
+                f"{show_earlier_label} ({hidden_count} hidden)",
+                key=show_all_btn_key,
+            ):
+                st.session_state[show_all_key] = True
+                st.rerun()
+            visible_messages = messages[-VISIBLE_MESSAGE_COUNT:]
+        else:
+            visible_messages = messages
 
-    chat_log = st.container()
-    with chat_log:
-        for msg in visible:
-            _render_chat_message(msg)
+        for message in visible_messages:
+            _render_chat_message(message)
 
-    return {
-        "panel_key": panel_key,
-        "state": state,
-        "active_system_prompt": active_system_prompt,
-        "selected_persona": selected_persona,
-        "chat_log": chat_log,
-    }
+    return chat_log
 
 
-def _generate_for_panel(
-    panel: dict,
-    model,
-    remote: bool,
-    gen_kwargs: dict,
-) -> ChatReply:
-    """Run generate_chat_reply for one compare panel. Thread-safe."""
-    messages = []
-    if panel["active_system_prompt"]:
-        messages.append({"role": "system", "content": panel["active_system_prompt"]})
-    messages.extend(panel["state"]["messages"])
-
-    ctx = nullcontext() if remote else _model_lock
-    with ctx:
-        return generate_chat_reply(
-            model=model,
-            messages=messages,
-            remote=remote,
-            past_key_values=panel["state"]["past_key_values"],
-            **gen_kwargs,
-        )
+def _build_chat_messages(
+    system_prompt: str | None,
+    messages: list[dict[str, str]],
+) -> list[dict[str, str]]:
+    return (
+        [{"role": "system", "content": system_prompt}] if system_prompt else []
+    ) + messages
+
+
+def _save_chat_export_message(
+    *,
+    model_name: str,
+    dataset_source: str,
+    persona_id: str,
+    persona_name: str | None,
+    prompt_mode: str,
+    system_prompt: str | None,
+    messages: list[dict[str, str]],
+    generation: dict[str, object],
+    panel_label: str | None = None,
+) -> str:
+    export_path = save_chat_export(
+        model_name=model_name,
+        dataset_source=dataset_source,
+        persona_id=persona_id,
+        persona_name=persona_name,
+        panel_label=panel_label,
+        prompt_mode=prompt_mode,
+        system_prompt=system_prompt,
+        messages=messages,
+        generation=generation,
+    )
+    return f"Saved chat export to {export_path}"
+
+
+# ── Compare mode helpers ───────────────────────────────────────────────────────
+
+
+def _panel_state(panel_key: str) -> dict:
+    """Get or initialise compare-panel chat state stored in session_state."""
+    if panel_key not in st.session_state:
+        st.session_state[panel_key] = _default_chat_state()
+    return st.session_state[panel_key]
 
 
 def _render_compare_mode(
@@ -233,35 +211,90 @@ def _render_compare_mode(
     model_name: str,
     context_key: str,
     dataset_source: str,
-    personas: list,
+    personas: list[PersonaData],
     gen_kwargs: dict,
     advanced_generation: bool,
 ) -> None:
     """Render the full side-by-side comparison UI."""
     left_col, right_col = st.columns(2)
 
-    with left_col:
-        left = _render_compare_panel(
-            "left",
-            context_key,
+    def render_panel(side: str, column) -> tuple[dict[str, object], Any, str | None]:
+        panel_key = widget_key(context_key, f"cmp_{side}")
+        state = st.session_state.get(panel_key)
+        if state is None:
+            state = _default_chat_state()
+            st.session_state[panel_key] = state
+        prompt_key = widget_key(panel_key, "custom_prompt")
+        show_all_key = widget_key(panel_key, "show_all")
+
+        selected_persona, prompt_mode, changed = _render_persona_prompt_controls(
             personas,
-            remote,
-            model_name,
-            dataset_source,
-            gen_kwargs,
-            advanced_generation,
+            state["persona_id"],
+            state["prompt_mode"],
+            widget_key(panel_key, "persona"),
+            widget_key(panel_key, "prompt_mode"),
         )
-    with right_col:
-        right = _render_compare_panel(
-            "right",
-            context_key,
-            personas,
-            remote,
-            model_name,
-            dataset_source,
-            gen_kwargs,
-            advanced_generation,
+        if changed:
+            state["messages"] = []
+            state["past_key_values"] = None
+            state["persona_id"] = selected_persona.id
+            state["prompt_mode"] = prompt_mode
+            _clear_chat_ui_state(prompt_key, show_all_key)
+
+        active_system_prompt = resolve_system_prompt(
+            persona=selected_persona, mode=prompt_mode
+        )
+        active_system_prompt = _render_system_prompt_editor(
+            prompt_key,
+            prompt_mode,
+            active_system_prompt,
+            height=150,
+        )
+
+        btn_col1, btn_col2 = st.columns(2)
+        with btn_col1:
+            if st.button(
+                "Export chat", key=widget_key(panel_key, "export_chat"), width="stretch"
+            ):
+                st.success(
+                    _save_chat_export_message(
+                        model_name=model_name,
+                        dataset_source=dataset_source,
+                        persona_id=selected_persona.id,
+                        persona_name=getattr(selected_persona, "name", None),
+                        prompt_mode=prompt_mode,
+                        system_prompt=active_system_prompt,
+                        messages=state["messages"],
+                        generation=_generation_dict(gen_kwargs, advanced_generation),
+                        panel_label=side,
+                    )
+                )
+        with btn_col2:
+            if st.button(
+                "Reset chat",
+                key=widget_key(panel_key, "reset"),
+                width="stretch",
+                type="secondary",
+            ):
+                state["messages"] = []
+                state["past_key_values"] = None
+                _clear_chat_ui_state(prompt_key, show_all_key)
+                st.rerun()
+
+        chat_log = st.container()
+        _render_chat_window(
+            chat_log=chat_log,
+            messages=state["messages"],
+            show_all_key=show_all_key,
+            show_all_btn_key=widget_key(panel_key, "show_all_btn"),
+            show_earlier_label="Show earlier",
         )
+        return state, chat_log, active_system_prompt
+
+    with left_col:
+        left_state, left_log, left_prompt = render_panel("left", left_col)
+    with right_col:
+        right_state, right_log, right_prompt = render_panel("right", right_col)
 
     user_prompt = st.chat_input(
         "Ask both...",
@@ -271,43 +304,73 @@ def _render_compare_mode(
         return
 
     model = cached_model(model_name=model_name, remote=remote)
-    panels = [(left, left_col), (right, right_col)]
+    panels = [
+        (left_state, left_log, left_prompt),
+        (right_state, right_log, right_prompt),
+    ]
 
-    for panel, col in panels:
-        panel["state"]["messages"].append({"role": "user", "content": user_prompt})
-        with col:
-            with panel["chat_log"]:
-                _render_chat_message({"role": "user", "content": user_prompt})
+    for panel_state, panel_log, _panel_prompt in panels:
+        panel_state["messages"].append({"role": "user", "content": user_prompt})
+        with panel_log:
+            _render_chat_message({"role": "user", "content": user_prompt})
 
-    # Generate both responses in parallel (remote: truly concurrent; local: serialised via lock).
     with st.spinner("Generating..."):
-        with ThreadPoolExecutor(max_workers=2) as executor:
-            futures = [
-                executor.submit(_generate_for_panel, panel, model, remote, gen_kwargs)
-                for panel, col in panels
-            ]
+        if remote:
+            with ThreadPoolExecutor(max_workers=2) as executor:
+                futures = [
+                    executor.submit(
+                        generate_chat_reply,
+                        model=model,
+                        messages=(
+                            [{"role": "system", "content": panel_prompt}]
+                            if panel_prompt
+                            else []
+                        )
+                        + panel_state["messages"],
+                        remote=remote,
+                        past_key_values=panel_state["past_key_values"],
+                        **gen_kwargs,
+                    )
+                    for panel_state, _panel_log, panel_prompt in panels
+                ]
+                results: list[ChatReply | Exception] = []
+                for future in futures:
+                    try:
+                        results.append(future.result())
+                    except Exception as exc:
+                        results.append(exc)
+        else:
             results = []
-            for future in futures:
+            for panel_state, _panel_log, panel_prompt in panels:
                 try:
-                    results.append(future.result())
+                    results.append(
+                        generate_chat_reply(
+                            model=model,
+                            messages=(
+                                [{"role": "system", "content": panel_prompt}]
+                                if panel_prompt
+                                else []
+                            )
+                            + panel_state["messages"],
+                            remote=remote,
+                            past_key_values=panel_state["past_key_values"],
+                            **gen_kwargs,
+                        )
+                    )
                 except Exception as exc:
                     results.append(exc)
 
-    for (panel, col), result in zip(panels, results):
+    for (panel_state, panel_log, _panel_prompt), result in zip(panels, results):
         if isinstance(result, Exception):
-            with col:
-                with panel["chat_log"]:
-                    st.error(f"Generation failed: {result}")
-            panel["state"]["messages"].pop()
+            with panel_log:
+                st.error(f"Generation failed: {result}")
+            panel_state["messages"].pop()
             continue
 
-        panel["state"]["messages"].append({"role": "assistant", "content": result.text})
-        panel["state"]["past_key_values"] = (
-            result.past_key_values if not remote else None
-        )
-        with col:
-            with panel["chat_log"]:
-                _render_chat_message({"role": "assistant", "content": result.text})
+        panel_state["messages"].append({"role": "assistant", "content": result.text})
+        panel_state["past_key_values"] = result.past_key_values if not remote else None
+        with panel_log:
+            _render_chat_message({"role": "assistant", "content": result.text})
 
 
 # ── Main tab entry point ───────────────────────────────────────────────────────
@@ -465,6 +528,12 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     # ── Single-chat mode ──────────────────────────────────────────────────────
     persona_select_key = widget_key(context_key, "persona_select")
     prompt_mode_select_key = widget_key(context_key, "system_prompt_select")
+    prompt_key = widget_key(context_key, "custom_system_prompt")
+    show_all_key = widget_key(context_key, "show_all_messages")
+    chat_input_key = widget_key(context_key, "chat_input")
+    pending_key = widget_key(context_key, "pending_prompt")
+    export_key = widget_key(context_key, "export_chat")
+    reset_key = widget_key(context_key, "reset")
 
     col1, col2 = st.columns([2, 1])
     with col1:
@@ -481,66 +550,35 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
         )
     with col2:
         current_mode_label = VARIANT_LABELS.get(chat_state["prompt_mode"], "None")
-        prompt_mode_label = st.selectbox(
+        st.selectbox(
             "Prompt",
             options=MODE_LABELS,
             index=MODE_LABELS.index(current_mode_label),
             key=prompt_mode_select_key,
         )
-        prompt_mode = MODE_LABEL_TO_KEY[prompt_mode_label]
+        prompt_mode = MODE_LABEL_TO_KEY[st.session_state[prompt_mode_select_key]]
 
     active_system_prompt = resolve_system_prompt(
         persona=selected_persona,
         mode=prompt_mode,
     )
 
-    chat_input_key = widget_key(context_key, "chat_input")
-    show_all_key = widget_key(context_key, "show_all_messages")
-    custom_prompt_key = widget_key(context_key, "custom_system_prompt")
-    pending_key = widget_key(context_key, "pending_prompt")
-    export_success_message: str | None = None
-
-    action_col1, action_col2 = st.columns(2)
-    with action_col1:
-        if st.button("Reset chat", use_container_width=True, type="secondary"):
-            reset_chat_state(model_name, remote, dataset_source)
-            _clear_chat_ui_state(
-                chat_input_key,
-                show_all_key,
-                custom_prompt_key,
-                pending_key,
-            )
-            st.rerun()
-    with action_col2:
-        if st.button("Export chat", use_container_width=True):
-            export_path = save_chat_export(
-                model_name=model_name,
-                dataset_source=dataset_source,
-                persona_id=selected_persona.id,
-                persona_name=getattr(selected_persona, "name", None),
-                prompt_mode=prompt_mode,
-                system_prompt=active_system_prompt,
-                messages=chat_state["messages"],
-                generation=_generation_dict(gen_kwargs, advanced_generation),
-            )
-            export_success_message = f"Saved chat export to {export_path}"
-
-    if export_success_message:
-        st.success(export_success_message)
-
     changed_context = (
         chat_state["persona_id"] != selected_persona.id
         or chat_state["prompt_mode"] != prompt_mode
     )
     if changed_context:
         had_history = bool(chat_state["messages"])
-        chat_state["persona_id"] = selected_persona.id
-        chat_state["prompt_mode"] = prompt_mode
-        reset_chat_state(model_name, remote, dataset_source)
-        _clear_chat_ui_state(
+        _reset_single_chat_context(
+            model_name,
+            remote,
+            dataset_source,
+            chat_state,
+            selected_persona.id,
+            prompt_mode,
             chat_input_key,
             show_all_key,
-            custom_prompt_key,
+            prompt_key,
             pending_key,
         )
         if had_history:
@@ -548,40 +586,51 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
 
     chat_log = st.container()
 
-    with chat_log:
-        # System prompt as first item in conversation — collapsed by default, editable.
-        if prompt_mode != "empty":
-            if custom_prompt_key not in st.session_state:
-                st.session_state[custom_prompt_key] = active_system_prompt
-            with st.expander("Edit prompt", expanded=False):
-                active_system_prompt = (
-                    st.text_area(
-                        "Prompt",
-                        key=custom_prompt_key,
-                        height=200,
-                        label_visibility="collapsed",
-                    )
-                    or None
-                )
+    active_system_prompt = _render_system_prompt_editor(
+        prompt_key,
+        prompt_mode,
+        active_system_prompt,
+        height=200,
+    )
 
-        # Collapse older messages, show only the most recent ones.
-        messages = chat_state["messages"]
-        if len(messages) > _VISIBLE_MESSAGE_COUNT and not st.session_state.get(
-            show_all_key, False
-        ):
-            hidden_count = len(messages) - _VISIBLE_MESSAGE_COUNT
-            if st.button(
-                f"Show earlier messages ({hidden_count} hidden)",
-                key=widget_key(context_key, "show_all_btn"),
-            ):
-                st.session_state[show_all_key] = True
-                st.rerun()
-            visible_messages = messages[-_VISIBLE_MESSAGE_COUNT:]
-        else:
-            visible_messages = messages
+    action_col1, action_col2 = st.columns(2)
+    with action_col1:
+        if st.button("Export chat", key=export_key, width="stretch"):
+            st.success(
+                _save_chat_export_message(
+                    model_name=model_name,
+                    dataset_source=dataset_source,
+                    persona_id=selected_persona.id,
+                    persona_name=getattr(selected_persona, "name", None),
+                    prompt_mode=prompt_mode,
+                    system_prompt=active_system_prompt,
+                    messages=chat_state["messages"],
+                    generation=_generation_dict(gen_kwargs, advanced_generation),
+                )
+            )
+    with action_col2:
+        if st.button("Reset chat", key=reset_key, width="stretch", type="secondary"):
+            _reset_single_chat_context(
+                model_name,
+                remote,
+                dataset_source,
+                chat_state,
+                selected_persona.id,
+                prompt_mode,
+                chat_input_key,
+                show_all_key,
+                prompt_key,
+                pending_key,
+            )
+            st.rerun()
 
-        for message in visible_messages:
-            _render_chat_message(message)
+    _render_chat_window(
+        chat_log=chat_log,
+        messages=chat_state["messages"],
+        show_all_key=show_all_key,
+        show_all_btn_key=widget_key(context_key, "show_all_btn"),
+        show_earlier_label="Show earlier messages",
+    )
 
     user_prompt = st.chat_input(
         "Ask something...",
@@ -598,10 +647,7 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     if not st.session_state.pop(pending_key, False):
         return
 
-    messages = []
-    if active_system_prompt:
-        messages.append({"role": "system", "content": active_system_prompt})
-    messages.extend(chat_state["messages"])
+    messages = _build_chat_messages(active_system_prompt, chat_state["messages"])
 
     with st.spinner("Generating reply..."):
         model = cached_model(model_name=model_name, remote=remote)

tabs/compare.py

@@ -1,21 +1,18 @@
+from collections.abc import Callable
+from dataclasses import dataclass
+
 import streamlit as st
+import torch
 from persona_data.environment import get_artifacts_dir
 from persona_vectors.analysis import build_embedding_figure, project_pca, project_umap
-from persona_vectors.plots import (
-    plot_multiple_layer_similarities,
-    save_plot_html,
-    save_plot_png,
-)
+from persona_vectors.artifacts import ActivationStore
+from persona_vectors.artifacts import list_layers as list_available_layers
+from persona_vectors.artifacts import list_personas as list_available_personas
+from persona_vectors.artifacts import load_mean_activations, load_persona_names
+from persona_vectors.plots import plot_layer_similarity, save_plot_html, save_plot_png
 
-from utils.artifacts import (
-    artifact_persona_options,
-    list_available_layers,
-    load_cosine_traces,
-    load_embedding_samples,
-)
 from utils.helpers import (
     ANALYSIS_HELP_TEXT,
-    ANALYSIS_LABELS,
     ANALYSIS_MODES,
     PROMPT_VARIANTS,
     persona_display_label,
@@ -29,15 +26,151 @@ def _filename(*parts: str) -> str:
     return "__".join(slugify(part) for part in parts if part)
 
 
-def _select_artifact_personas(
-    artifacts_root: str,
+@dataclass(frozen=True)
+class ProjectionConfig:
+    title_prefix: str
+    x_label: str
+    y_label: str
+    project_fn: Callable[[torch.Tensor], torch.Tensor]
+
+
+_PROJECTION_CONFIGS: dict[str, ProjectionConfig] = {
+    "PCA": ProjectionConfig("PCA", "PC1", "PC2", project_pca),
+    "UMAP": ProjectionConfig("UMAP", "UMAP 1", "UMAP 2", project_umap),
+}
+
+
+@st.cache_data(show_spinner=False)
+def _list_layers(
+    root_dir: str,
     model_name: str,
     variants: list[str],
+    persona_ids: list[str],
+) -> list[int]:
+    return list_available_layers(root_dir, model_name, variants, persona_ids)
+
+
+def _load_embedding_samples(
+    store: ActivationStore,
+    persona_ids: list[str],
+    variant: str,
+    selected_layers: list[int],
+    project_fn: Callable[[torch.Tensor], torch.Tensor],
+    persona_names: dict[str, str],
+    progress_fn: Callable[[int, int, int], None] | None = None,
+) -> tuple[list[tuple[int, torch.Tensor, list[str], list[str]]], list[str]]:
+    """Load samples for 2D projections without re-reading each layer from disk."""
+
+    plots: list[tuple[int, torch.Tensor, list[str], list[str]]] = []
+    errors: list[str] = []
+    vectors_by_persona: dict[str, torch.Tensor] = {}
+
+    for persona_id in persona_ids:
+        try:
+            vectors, _ = store.load(variant, persona_id)
+        except (FileNotFoundError, KeyError, OSError, ValueError) as exc:
+            errors.append(f"{persona_id} / {variant}: {exc}")
+            continue
+
+        vectors_by_persona[persona_id] = vectors
+
+    total_layers = len(selected_layers)
+    for idx, layer_idx in enumerate(selected_layers, start=1):
+        samples: list[torch.Tensor] = []
+        labels: list[str] = []
+        hover_text: list[str] = []
+
+        for persona_id, vectors in vectors_by_persona.items():
+            if layer_idx >= vectors.shape[1]:
+                errors.append(f"{persona_id} / {variant}: missing layer {layer_idx}")
+                continue
+
+            layer_vectors = vectors[:, layer_idx, :]
+            samples.append(layer_vectors)
+            labels.extend([persona_id] * layer_vectors.shape[0])
+            display_name = persona_names.get(persona_id) or persona_id
+            hover_text.extend(
+                [f"<b>{display_name}</b><br>{variant}"] * layer_vectors.shape[0]
+            )
+
+        if not samples:
+            errors.append(f"Layer {layer_idx}: no selected personas have this layer")
+        else:
+            all_samples = torch.cat(samples, dim=0)
+            if all_samples.shape[0] < 2:
+                errors.append(
+                    f"Layer {layer_idx}: need at least 2 samples after filtering selected personas"
+                )
+            else:
+                try:
+                    coords = project_fn(all_samples)
+                    plots.append((layer_idx, coords, labels, hover_text))
+                except Exception as exc:
+                    errors.append(f"Layer {layer_idx}: {exc}")
+
+        if progress_fn is not None:
+            progress_fn(idx, total_layers, len(plots))
+
+    return plots, errors
+
+
+def _build_embedding_figures(
+    plots: list[tuple[int, torch.Tensor, list[str], list[str]]],
+    config: ProjectionConfig,
+) -> list[tuple[int, object]]:
+    return [
+        (
+            layer_idx,
+            build_embedding_figure(
+                coords=coords,
+                labels=labels,
+                title=f"{config.title_prefix}, layer {layer_idx}",
+                x_label=config.x_label,
+                y_label=config.y_label,
+                hover_text=hover_text,
+            ),
+        )
+        for layer_idx, coords, labels, hover_text in plots
+    ]
+
+
+def _render_embedding_results(
+    store: ActivationStore,
+    analysis_mode: str,
+    rendered_figures: list[tuple[int, object]],
+    saved_variant: str,
+    saved_persona_key: str,
+    total_samples: int,
+) -> None:
+    cols = st.columns(2)
+    for idx, (_, fig) in enumerate(rendered_figures):
+        with cols[idx % 2]:
+            st.plotly_chart(fig, width="stretch")
+
+    st.success(f"Loaded {total_samples} samples across {len(rendered_figures)} layers.")
+    filenames = [
+        _filename(
+            "compare",
+            analysis_mode,
+            store.model_name,
+            saved_variant,
+            saved_persona_key,
+            str(layer_idx),
+        )
+        for layer_idx, _ in rendered_figures
+    ]
+    _render_save_buttons([fig for _, fig in rendered_figures], filenames, analysis_mode)
+
+
+def _select_artifact_personas(
+    store: ActivationStore,
+    variants: list[str],
 ) -> tuple[list[str], dict[str, str]]:
-    persona_options, persona_names = artifact_persona_options(
-        artifacts_root,
-        model_name,
-        variants,
+    persona_options = list_available_personas(
+        store.root_dir, store.model_name, variants
+    )
+    persona_names = load_persona_names(
+        store.root_dir, store.model_name, variants, persona_options
     )
     if not persona_options:
         if len(variants) > 1:
@@ -55,15 +188,81 @@ def _select_artifact_personas(
         format_func=lambda persona_id: persona_display_label(
             persona_id, persona_names.get(persona_id)
         ),
-        key=widget_key("load", "personas", model_name, *variants),
+        key=widget_key("load", "personas", store.model_name, *variants),
     )
     return persona_ids, persona_names
 
 
-def _render_cosine_similarity(
-    artifacts_root: str,
-    model_name: str,
+def _render_save_buttons(
+    figs: list[object],
+    filenames: list[str],
+    key_suffix: str,
 ) -> None:
+    """Render Save HTML / Save PNG column buttons for one or more figures."""
+    col1, col2 = st.columns(2)
+    with col1:
+        if st.button("Save HTML", key=widget_key("load", "save_html", key_suffix)):
+            paths = [save_plot_html(fig, fn) for fig, fn in zip(figs, filenames)]
+            st.success(f"Saved {len(paths)} HTML file(s) to `artifacts/plots`.")
+    with col2:
+        if st.button("Save PNG", key=widget_key("load", "save_png", key_suffix)):
+            try:
+                paths = [save_plot_png(fig, fn) for fig, fn in zip(figs, filenames)]
+                st.success(f"Saved {len(paths)} PNG file(s) to `artifacts/plots`.")
+            except Exception as exc:
+                st.error(f"Could not save PNG: {exc}")
+
+
+def _select_embedding_config(
+    store: ActivationStore,
+) -> tuple[str, list[str], dict[str, str], list[int]] | None:
+    """Render variant / persona / layer selectors and return the selection, or None on early exit."""
+    selected_variant = st.selectbox(
+        "Variant",
+        options=PROMPT_VARIANTS,
+        format_func=prompt_variant_label,
+        key=widget_key("load", "variant"),
+    )
+
+    persona_ids, persona_names = _select_artifact_personas(store, [selected_variant])
+    if not persona_ids:
+        return None
+
+    layer_options = _list_layers(
+        str(store.root_dir),
+        store.model_name,
+        [selected_variant],
+        persona_ids,
+    )
+    if not layer_options:
+        st.info(
+            "No shared layers are available for the selected personas. Try fewer personas or a different variant."
+        )
+        return None
+
+    persona_key = "_".join(sorted(persona_ids))
+    layer_key = widget_key(
+        "load", "layers", store.model_name, selected_variant, persona_key
+    )
+    default_layers = [
+        layer
+        for layer in st.session_state.get(layer_key, layer_options[:3])
+        if layer in layer_options
+    ] or layer_options[:3]
+    selected_layers = st.multiselect(
+        "Layers",
+        options=layer_options,
+        default=default_layers,
+        key=layer_key,
+    )
+    if not selected_layers:
+        st.info("Select at least one layer.")
+        return None
+
+    return selected_variant, persona_ids, persona_names, selected_layers
+
+
+def _render_cosine_similarity(store: ActivationStore) -> None:
     col1, col2 = st.columns(2)
     with col1:
         variant_a = st.selectbox(
@@ -86,24 +285,16 @@ def _render_cosine_similarity(
         st.warning("Choose two different variants to compare.")
         return
 
-    persona_ids, _ = _select_artifact_personas(
-        artifacts_root,
-        model_name,
-        [variant_a, variant_b],
-    )
+    persona_ids, _ = _select_artifact_personas(store, [variant_a, variant_b])
     if not persona_ids:
         return
 
-    cosine_fig_key = widget_key("load", "cosine_fig_state", model_name)
-    filename = _filename("compare", "cosine", model_name, variant_a, variant_b)
+    cosine_fig_key = widget_key("load", "cosine_fig_state", store.model_name)
+    filename = _filename("compare", "cosine", store.model_name, variant_a, variant_b)
 
     if st.button("Compare vectors", type="primary"):
-        traces, loaded_names, errors = load_cosine_traces(
-            artifacts_root,
-            model_name,
-            persona_ids,
-            variant_a,
-            variant_b,
+        traces, loaded_names, errors = load_mean_activations(
+            store.root_dir, store.model_name, persona_ids, variant_a, variant_b
         )
 
         if errors:
@@ -125,7 +316,7 @@ def _render_cosine_similarity(
             )
             for persona_id, short, long in traces
         ]
-        fig = plot_multiple_layer_similarities(
+        fig = plot_layer_similarity(
             display_traces,
             title=f"{prompt_variant_label(variant_a)} vs {prompt_variant_label(variant_b)}",
             show=False,
@@ -134,82 +325,27 @@ def _render_cosine_similarity(
 
     if cosine_fig_key in st.session_state:
         fig, n_traces = st.session_state[cosine_fig_key]
-        st.plotly_chart(fig, use_container_width=True)
-        save_col1, save_col2 = st.columns(2)
-        with save_col1:
-            if st.button("Save HTML", key=widget_key("load", "save_cosine_html")):
-                output_path = save_plot_html(fig, filename)
-                st.success(f"Saved HTML to `{output_path}`")
-        with save_col2:
-            if st.button("Save PNG", key=widget_key("load", "save_cosine_png")):
-                try:
-                    output_path = save_plot_png(fig, filename)
-                    st.success(f"Saved PNG to `{output_path}`")
-                except Exception as exc:
-                    st.error(f"Could not save PNG: {exc}")
+        st.plotly_chart(fig, width="stretch")
+        _render_save_buttons([fig], [filename], "cosine")
         st.success(f"Loaded {n_traces} personas for cosine comparison.")
 
 
-def _render_embedding_analysis(
-    artifacts_root: str,
-    model_name: str,
-    analysis_mode: str,
-) -> None:
-    selected_variant = st.selectbox(
-        "Variant",
-        options=PROMPT_VARIANTS,
-        format_func=prompt_variant_label,
-        key=widget_key("load", "variant"),
-    )
-
-    persona_ids, persona_names = _select_artifact_personas(
-        artifacts_root,
-        model_name,
-        [selected_variant],
-    )
-    if not persona_ids:
-        return
-
-    layer_options = list_available_layers(
-        artifacts_root,
-        model_name,
-        [selected_variant],
-        persona_ids,
-    )
-    if not layer_options:
-        st.info(
-            "No shared layers are available for the selected personas. Try fewer personas or a different variant."
-        )
+def _render_embedding_analysis(store: ActivationStore, analysis_mode: str) -> None:
+    config = _select_embedding_config(store)
+    if config is None:
         return
-
+    selected_variant, persona_ids, persona_names, selected_layers = config
     persona_key = "_".join(sorted(persona_ids))
-    layer_key = widget_key("load", "layers", model_name, selected_variant, persona_key)
-    default_layers = [
-        layer
-        for layer in st.session_state.get(layer_key, layer_options[:3])
-        if layer in layer_options
-    ] or layer_options[:3]
-    selected_layers = st.multiselect(
-        "Layers",
-        options=layer_options,
-        default=default_layers,
-        key=layer_key,
-    )
-    if not selected_layers:
-        st.info("Select at least one layer.")
+    projection_config = _PROJECTION_CONFIGS.get(analysis_mode)
+    if projection_config is None:
+        st.error(f"Unsupported analysis mode: {analysis_mode}")
         return
 
-    button_label = (
-        "Generate PCA projection"
-        if analysis_mode == "PCA"
-        else "Generate UMAP projection"
-    )
-
     embedding_fig_key = widget_key(
-        "load", "embedding_fig_state", model_name, analysis_mode
+        "load", "embedding_fig_state", store.model_name, analysis_mode
     )
 
-    if st.button(button_label, type="primary"):
+    if st.button(f"Generate {analysis_mode} projection", type="primary"):
         progress = st.progress(0, text="Preparing projections...")
 
         def update_progress(current: int, total: int, loaded: int) -> None:
@@ -219,15 +355,13 @@ def _render_embedding_analysis(
                 text=f"Processing layer {current}/{total} ({loaded} plot(s) ready)",
             )
 
-        project_fn = project_pca if analysis_mode == "PCA" else project_umap
         try:
-            plots, errors = load_embedding_samples(
-                artifacts_root,
-                model_name,
+            plots, errors = _load_embedding_samples(
+                store,
                 persona_ids,
                 selected_variant,
                 selected_layers,
-                project_fn,
+                projection_config.project_fn,
                 persona_names,
                 progress_fn=update_progress,
             )
@@ -248,18 +382,7 @@ def _render_embedding_analysis(
                 st.info("Try fewer personas, fewer layers, or a different variant.")
                 st.session_state.pop(embedding_fig_key, None)
             else:
-                title_prefix, x_label, y_label = ANALYSIS_LABELS[analysis_mode]
-                rendered_figures: list[tuple[int, object]] = []
-                for layer_idx, coords, labels, hover_text in plots:
-                    fig = build_embedding_figure(
-                        coords=coords,
-                        labels=labels,
-                        title=f"{title_prefix}, layer {layer_idx}",
-                        x_label=x_label,
-                        y_label=y_label,
-                        hover_text=hover_text,
-                    )
-                    rendered_figures.append((layer_idx, fig))
+                rendered_figures = _build_embedding_figures(plots, projection_config)
                 total_samples = sum(coords.shape[0] for _, coords, _, _ in plots)
                 st.session_state[embedding_fig_key] = (
                     rendered_figures,
@@ -274,52 +397,14 @@ def _render_embedding_analysis(
         rendered_figures, saved_persona_key, saved_variant, total_samples = (
             st.session_state[embedding_fig_key]
         )
-        cols = st.columns(2)
-        for idx, (layer_idx, fig) in enumerate(rendered_figures):
-            with cols[idx % 2]:
-                st.plotly_chart(fig, use_container_width=True)
-        st.success(
-            f"Loaded {total_samples} samples across {len(rendered_figures)} layers."
+        _render_embedding_results(
+            store,
+            analysis_mode,
+            rendered_figures,
+            saved_variant,
+            saved_persona_key,
+            total_samples,
         )
-        filenames = [
-            _filename(
-                "compare",
-                analysis_mode,
-                model_name,
-                saved_variant,
-                saved_persona_key,
-                str(layer_idx),
-            )
-            for layer_idx, _ in rendered_figures
-        ]
-        save_col1, save_col2 = st.columns(2)
-        with save_col1:
-            if st.button(
-                "Save HTML",
-                key=widget_key("load", "save_embedding_html", analysis_mode),
-            ):
-                saved_paths = [
-                    save_plot_html(fig, fn)
-                    for (_, fig), fn in zip(rendered_figures, filenames)
-                ]
-                st.success(
-                    f"Saved {len(saved_paths)} HTML plot(s) to `artifacts/plots`."
-                )
-        with save_col2:
-            if st.button(
-                "Save PNG",
-                key=widget_key("load", "save_embedding_png", analysis_mode),
-            ):
-                try:
-                    saved_paths = [
-                        save_plot_png(fig, fn)
-                        for (_, fig), fn in zip(rendered_figures, filenames)
-                    ]
-                    st.success(
-                        f"Saved {len(saved_paths)} PNG plot(s) to `artifacts/plots`."
-                    )
-                except Exception as exc:
-                    st.error(f"Could not save PNGs: {exc}")
 
 
 def render_compare_tab(model_name: str) -> None:
@@ -336,6 +421,8 @@ def render_compare_tab(model_name: str) -> None:
             value=str(get_artifacts_dir() / "activations"),
         )
 
+    store = ActivationStore(model_name, artifacts_root)
+
     analysis_mode = st.segmented_control(
         "Analysis mode",
         options=ANALYSIS_MODES,
@@ -348,7 +435,7 @@ def render_compare_tab(model_name: str) -> None:
     st.caption(ANALYSIS_HELP_TEXT[analysis_mode])
 
     if analysis_mode == "Cosine similarity":
-        _render_cosine_similarity(artifacts_root, model_name)
+        _render_cosine_similarity(store)
         return
 
-    _render_embedding_analysis(artifacts_root, model_name, analysis_mode)
+    _render_embedding_analysis(store, analysis_mode)

tabs/extract.py

@@ -3,6 +3,7 @@ from persona_vectors.extraction import run_extraction
 
 from utils.datasets import load_dataset
 from utils.helpers import (
+    NDIF_STATUS_ICONS,
     PROMPT_VARIANTS,
     persona_label,
     prompt_variant_label,
@@ -84,8 +85,8 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         st.info("Select at least one persona.")
         return
 
-    qa_filter_type: str | None
-    qa_filter_difficulty: list[int] | None
+    runs = None
+    max_questions = 0
 
     with st.expander("Advanced", expanded=False):
         st.caption("Filters")
@@ -114,35 +115,38 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
             )
             qa_filter_difficulty = difficulty_values if difficulty_values else None
 
-        # Pre-load QA pairs for all selected personas to validate filters and set slider range.
-        qa_by_persona = {
-            p.id: dataset.get_qa(
-                p.id, type=qa_filter_type, difficulty=qa_filter_difficulty
+        runs, skipped = [], []
+        for persona in selected_personas:
+            qa = list(
+                dataset.get_qa(
+                    persona.id, type=qa_filter_type, difficulty=qa_filter_difficulty
+                )
             )
-            for p in selected_personas
-        }
-        personas_without_qa = [p for p in selected_personas if not qa_by_persona[p.id]]
-        if personas_without_qa:
-            names = ", ".join(p.name for p in personas_without_qa)
+            if qa:
+                runs.append((persona, qa))
+            else:
+                skipped.append(persona)
+        if skipped:
+            names = ", ".join(p.name for p in skipped)
             st.warning(f"No QA pairs match filters for: {names}. They will be skipped.")
 
-        personas_to_run = [p for p in selected_personas if qa_by_persona[p.id]]
-        if not personas_to_run:
+        if not runs:
             st.info("No personas have matching QA pairs. Widen the filters.")
             return
 
-        min_qa_count = min(len(qa_by_persona[p.id]) for p in personas_to_run)
+        max_q = min(len(qa_pairs) for _, qa_pairs in runs)
+        max_questions = st.slider(
+            "Max questions",
+            min_value=1,
+            max_value=max_q,
+            value=max_q,
+            key=_extract_widget_key(
+                model_name, remote, dataset_source, "max_questions"
+            ),
+        )
 
-        with col3:
-            max_questions = st.slider(
-                "Max questions",
-                min_value=1,
-                max_value=min_qa_count,
-                value=min_qa_count,
-                key=_extract_widget_key(
-                    model_name, remote, dataset_source, "max_questions"
-                ),
-            )
+    if runs is None:
+        return
 
     run_clicked = st.button("Run extraction", type="primary")
     if not run_clicked:
@@ -153,25 +157,19 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
     progress = st.progress(0, text="Preparing extraction...")
     ndif_status_box = st.empty()  # shows live NDIF job status when remote=True
 
-    _STATUS_ICONS = {
-        "RECEIVED": "◉", "QUEUED": "◎", "DISPATCHED": "◈",
-        "RUNNING": "●", "COMPLETED": "✓", "ERROR": "✗",
-    }
-
     def _on_ndif_status(job_id: str, status_name: str, description: str) -> None:
-        icon = _STATUS_ICONS.get(status_name, "•")
+        icon = NDIF_STATUS_ICONS.get(status_name, "•")
         ndif_status_box.caption(f"{icon} `{job_id}` **{status_name}** — {description}")
 
     with st.spinner("Loading model..."):
         model = cached_model(model_name=model_name, remote=remote)
 
     try:
-        total_steps = len(personas_to_run) * len(selected_variants)
+        total_steps = len(runs) * len(selected_variants)
         step = 0
         results = []
 
-        for persona in personas_to_run:
-            qa_pairs = qa_by_persona[persona.id][:max_questions]
+        for persona, qa_pairs in runs:
             for variant in selected_variants:
                 progress.progress(
                     step / total_steps if total_steps else 1.0,
@@ -181,7 +179,7 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
                     model=model,
                     model_name=model_name,
                     persona=persona,
-                    qa_pairs=qa_pairs,
+                    qa_pairs=qa_pairs[:max_questions],
                     variants=[variant],
                     remote=remote,
                     on_status=_on_ndif_status if remote else None,

utils/artifacts.py

@@ -1,244 +0,0 @@
-import logging
-from collections.abc import Callable
-from pathlib import Path
-
-import streamlit as st
-import torch
-from persona_vectors.activation_io import (
-    load_activation_metadata,
-    load_per_question_vectors,
-    model_dir_name,
-)
-
-logger = logging.getLogger(__name__)
-
-
-def list_available_personas(
-    artifacts_root: str | Path,
-    model_name: str,
-    variants: list[str],
-) -> list[str]:
-    """List persona ids available for every requested variant."""
-
-    shared_personas: set[str] | None = None
-    root = Path(artifacts_root)
-    for variant in variants:
-        model_dir = root / model_dir_name(model_name) / variant
-        if not model_dir.exists():
-            return []
-
-        variant_personas = {d.name for d in model_dir.iterdir() if d.is_dir()}
-        if shared_personas is None:
-            shared_personas = variant_personas
-        else:
-            shared_personas &= variant_personas
-
-        if not shared_personas:
-            return []
-
-    return sorted(shared_personas or set())
-
-
-def load_persona_names(
-    artifacts_root: str | Path,
-    model_name: str,
-    variants: list[str],
-    persona_ids: list[str],
-) -> dict[str, str]:
-    """Load display names from saved activation metadata."""
-
-    names: dict[str, str] = {}
-    for persona_id in persona_ids:
-        for variant in variants:
-            try:
-                metadata = load_activation_metadata(
-                    root_dir=artifacts_root,
-                    model_name=model_name,
-                    prompt_variant=variant,
-                    persona_id=persona_id,
-                )
-            except Exception:
-                logger.debug(
-                    "Failed to load metadata for persona %s variant %s",
-                    persona_id,
-                    variant,
-                    exc_info=True,
-                )
-                continue
-
-            persona_name = metadata.get("persona_name")
-            if isinstance(persona_name, str) and persona_name:
-                names[persona_id] = persona_name
-                break
-
-    return names
-
-
-def artifact_persona_options(
-    artifacts_root: str | Path,
-    model_name: str,
-    variants: list[str],
-) -> tuple[list[str], dict[str, str]]:
-    """Return persona ids and names for the selected artifacts."""
-
-    persona_options = list_available_personas(artifacts_root, model_name, variants)
-    persona_names = load_persona_names(
-        artifacts_root,
-        model_name,
-        variants,
-        persona_options,
-    )
-    return persona_options, persona_names
-
-
-@st.cache_data(show_spinner=False)
-def list_available_layers(
-    artifacts_root: str,
-    model_name: str,
-    variants: list[str],
-    persona_ids: list[str],
-) -> list[int]:
-    """List layer indices shared by all matching saved activation files."""
-
-    shared_layers: set[int] | None = None
-    for variant in variants:
-        for persona_id in persona_ids:
-            try:
-                vectors, _ = load_per_question_vectors(
-                    root_dir=artifacts_root,
-                    model_name=model_name,
-                    prompt_variant=variant,
-                    persona_id=persona_id,
-                )
-            except Exception:
-                logger.debug(
-                    "Failed to load vectors for persona %s variant %s",
-                    persona_id,
-                    variant,
-                    exc_info=True,
-                )
-                continue
-
-            layers = set(range(vectors.shape[1]))
-            if shared_layers is None:
-                shared_layers = layers
-            else:
-                shared_layers &= layers
-
-    return sorted(shared_layers or set())
-
-
-def load_cosine_traces(
-    artifacts_root: str | Path,
-    model_name: str,
-    persona_ids: list[str],
-    variant_a: str,
-    variant_b: str,
-) -> tuple[list[tuple[str, torch.Tensor, torch.Tensor]], dict[str, str], list[str]]:
-    """Load mean activation traces for pairwise cosine-similarity plots."""
-
-    persona_names = load_persona_names(
-        artifacts_root,
-        model_name,
-        [variant_a, variant_b],
-        persona_ids,
-    )
-    traces: list[tuple[str, torch.Tensor, torch.Tensor]] = []
-    errors: list[str] = []
-
-    for persona_id in persona_ids:
-        try:
-            vectors_a, _ = load_per_question_vectors(
-                root_dir=artifacts_root,
-                model_name=model_name,
-                prompt_variant=variant_a,
-                persona_id=persona_id,
-            )
-            vectors_b, _ = load_per_question_vectors(
-                root_dir=artifacts_root,
-                model_name=model_name,
-                prompt_variant=variant_b,
-                persona_id=persona_id,
-            )
-        except Exception as exc:
-            errors.append(f"{persona_id}: {exc}")
-            continue
-
-        traces.append(
-            (persona_id, vectors_a.float().mean(dim=0), vectors_b.float().mean(dim=0))
-        )
-
-    return traces, persona_names, errors
-
-
-def load_embedding_samples(
-    artifacts_root: str | Path,
-    model_name: str,
-    persona_ids: list[str],
-    variant: str,
-    selected_layers: list[int],
-    project_fn: Callable[[torch.Tensor], torch.Tensor],
-    persona_names: dict[str, str],
-    progress_fn: Callable[[int, int, int], None] | None = None,
-) -> tuple[list[tuple[int, torch.Tensor, list[str], list[str]]], list[str]]:
-    """Load samples for 2D projections without re-reading each layer from disk."""
-
-    plots: list[tuple[int, torch.Tensor, list[str], list[str]]] = []
-    errors: list[str] = []
-    vectors_by_persona: dict[str, torch.Tensor] = {}
-
-    for persona_id in persona_ids:
-        try:
-            vectors, _ = load_per_question_vectors(
-                root_dir=artifacts_root,
-                model_name=model_name,
-                prompt_variant=variant,
-                persona_id=persona_id,
-            )
-        except Exception as exc:
-            errors.append(f"{persona_id} / {variant}: {exc}")
-            continue
-
-        vectors_by_persona[persona_id] = vectors
-
-    total_layers = len(selected_layers)
-    for idx, layer_idx in enumerate(selected_layers, start=1):
-        samples: list[torch.Tensor] = []
-        labels: list[str] = []
-        hover_text: list[str] = []
-
-        for persona_id, vectors in vectors_by_persona.items():
-            if layer_idx >= vectors.shape[1]:
-                errors.append(f"{persona_id} / {variant}: missing layer {layer_idx}")
-                continue
-
-            layer_vectors = vectors[:, layer_idx, :]
-            samples.append(layer_vectors)
-            labels.extend([persona_id] * layer_vectors.shape[0])
-            display_name = persona_names.get(persona_id) or persona_id
-            hover_text.extend(
-                [
-                    f"<b>{display_name}</b><br>{variant}",
-                ]
-                * layer_vectors.shape[0]
-            )
-
-        if not samples:
-            errors.append(f"Layer {layer_idx}: no selected personas have this layer")
-        else:
-            all_samples = torch.cat(samples, dim=0)
-            if all_samples.shape[0] < 2:
-                errors.append(
-                    f"Layer {layer_idx}: need at least 2 samples after filtering selected personas"
-                )
-            else:
-                try:
-                    coords = project_fn(all_samples)
-                    plots.append((layer_idx, coords, labels, hover_text))
-                except Exception as exc:
-                    errors.append(f"Layer {layer_idx}: {exc}")
-
-        if progress_fn is not None:
-            progress_fn(idx, total_layers, len(plots))
-
-    return plots, errors

utils/chat.py

@@ -52,7 +52,6 @@ def resolve_system_prompt(
         return format_biography_prompt(persona.biography_md)
     if mode == "custom":
         return format_empty_persona_prompt()
-    return ""
 
 
 def _format_plain_messages(

utils/chat_export.py

@@ -3,24 +3,23 @@ from datetime import datetime, timezone
 from pathlib import Path
 
 from persona_data.environment import get_artifacts_dir
-from persona_vectors.activation_io import model_dir_name
 
 from utils.helpers import slugify
 
 
-def build_chat_export_payload(
+def save_chat_export(
     *,
     model_name: str,
     dataset_source: str,
     persona_id: str,
     persona_name: str | None,
-    panel_label: str | None,
     prompt_mode: str,
     system_prompt: str | None,
     messages: list[dict[str, str]],
     generation: dict[str, object],
-) -> dict[str, object]:
-    """Build a JSON-serializable snapshot of the current chat session.
+    panel_label: str | None = None,
+) -> Path:
+    """Save the current chat session to ``artifacts/chats`` as JSON.
 
     Args:
         model_name: Model identifier used for the chat.
@@ -28,14 +27,15 @@ def build_chat_export_payload(
         persona_id: Selected persona id.
         persona_name: Selected persona display name, if available.
         prompt_mode: Active system prompt mode.
+        system_prompt: Current system prompt text, if any.
         messages: Conversation messages without the system prompt.
         generation: Generation settings used for the chat.
 
     Returns:
-        A JSON-serializable dictionary.
+        The path the export was written to.
     """
 
-    return {
+    payload = {
         "model_name": model_name,
         "dataset_source": dataset_source,
         "persona": {
@@ -51,50 +51,10 @@ def build_chat_export_payload(
         + messages,
     }
 
-
-def save_chat_export(
-    *,
-    model_name: str,
-    dataset_source: str,
-    persona_id: str,
-    persona_name: str | None,
-    prompt_mode: str,
-    system_prompt: str | None,
-    messages: list[dict[str, str]],
-    generation: dict[str, object],
-    panel_label: str | None = None,
-) -> Path:
-    """Save the current chat session to ``artifacts/chats`` as JSON.
-
-    Args:
-        model_name: Model identifier used for the chat.
-        dataset_source: Human-readable dataset source label.
-        persona_id: Selected persona id.
-        persona_name: Selected persona display name, if available.
-        prompt_mode: Active system prompt mode.
-        system_prompt: Current system prompt text, if any.
-        messages: Conversation messages without the system prompt.
-        generation: Generation settings used for the chat.
-
-    Returns:
-        The path the export was written to.
-    """
-
-    payload = build_chat_export_payload(
-        model_name=model_name,
-        dataset_source=dataset_source,
-        persona_id=persona_id,
-        persona_name=persona_name,
-        panel_label=panel_label,
-        prompt_mode=prompt_mode,
-        system_prompt=system_prompt,
-        messages=messages,
-        generation=generation,
-    )
     export_dir = (
         get_artifacts_dir()
         / "chats"
-        / model_dir_name(model_name)
+        / model_name.replace("/", "__")
         / slugify(dataset_source)
         / slugify(persona_id)
     )

utils/datasets.py

@@ -1,4 +1,5 @@
 import atexit
+import hashlib
 import shutil
 from pathlib import Path
 from tempfile import mkdtemp
@@ -31,10 +32,13 @@ def _upload_cache_dir() -> Path:
 def _uploaded_file_to_temp_path(uploaded_file: Any, stem: str) -> Path:
     suffix = Path(uploaded_file.name).suffix or ".jsonl"
     temp_path = _upload_cache_dir() / f"{stem}{suffix}"
+    hash_path = temp_path.with_suffix(temp_path.suffix + ".sha256")
     data = uploaded_file.getvalue()
-    if temp_path.exists() and temp_path.stat().st_size == len(data):
+    digest = hashlib.sha256(data).hexdigest()
+    if temp_path.exists() and hash_path.exists() and hash_path.read_text() == digest:
         return temp_path
     temp_path.write_bytes(data)
+    hash_path.write_text(digest)
     return temp_path
 
 

utils/helpers.py

@@ -1,5 +1,7 @@
+import re
+
 from persona_data.synth_persona import PersonaData
-from persona_vectors.extraction import SUPPORTED_VARIANTS
+from persona_vectors.artifacts import SUPPORTED_VARIANTS
 
 # Variant key -> human-readable label mapping
 VARIANT_LABELS = {
@@ -18,25 +20,29 @@ MODE_LABELS = list(VARIANT_LABELS.values())
 # Reverse lookup: label -> key
 MODE_LABEL_TO_KEY = {v: k for k, v in VARIANT_LABELS.items()}
 
+VISIBLE_MESSAGE_COUNT = 5
+
 DATASET_SOURCES = ["HuggingFace: synth-persona", "Local JSONL upload"]
 ANALYSIS_MODES = ["Cosine similarity", "PCA", "UMAP"]
 
-ANALYSIS_LABELS = {
-    "PCA": ("PCA", "PC1", "PC2"),
-    "UMAP": ("UMAP", "UMAP 1", "UMAP 2"),
-}
-
 ANALYSIS_HELP_TEXT = {
     "Cosine similarity": "Compare layer-wise alignment between variants.",
     "PCA": "Project the selected layers into a global 2D view.",
     "UMAP": "Project the selected layers into a local-neighborhood 2D view.",
 }
 
+NDIF_STATUS_ICONS = {
+    "RECEIVED": "◉",
+    "QUEUED": "◎",
+    "DISPATCHED": "◈",
+    "RUNNING": "●",
+    "COMPLETED": "✓",
+    "ERROR": "✗",
+}
 
-def slugify(value: str) -> str:
-    """Convert a string to a slug safe for filenames and URLs."""
 
-    import re
+def slugify(value: str) -> str:
+    """Convert a string to a filesystem-safe slug."""
 
     return re.sub(r"[^a-z0-9]+", "_", value.lower()).strip("_") or "unknown"