Updated code supporting latest version of persona-vector and data

pyproject.toml

@@ -5,12 +5,11 @@ description = "Streamlit UI for persona-vectors"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
-    "persona-vectors>=0.4.4",
-    "persona-data>=0.3.4",
+    "persona-vectors>=0.5.1",
+    "persona-data>=0.4.0",
     "streamlit>=1.44.0",
     "plotly>=6.6.0",
     "python-dotenv>=1.2.2",
-    "transformers>=5.5.0",
 ]
 
 # Local development:

tabs/chat.py

@@ -1,11 +1,11 @@
 import streamlit as st
+from persona_data.synth_persona import PersonaData
 
-from state import chat_session_key, get_chat_state, reset_chat_context_state
+from state import ChatState, chat_session_key, get_chat_state, reset_chat_context_state
 from tabs.chat_ui import (
     GenerationConfig,
-    generation_dict,
+    render_advanced_settings,
     render_chat_window,
-    render_generation_settings,
     render_persona_prompt_controls,
     render_system_prompt,
 )
@@ -27,21 +27,7 @@ _LAST_PROMPT_MODE_KEY = "chat:last_prompt_mode"
 _LAST_COMPARE_MODE_KEY = "chat:last_compare_mode"
 
 
-def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
-    """Render the chat tab."""
-
-    st.title("Chat")
-
-    context_key = chat_session_key(model_name, dataset_source)
-    chat_state = get_chat_state(model_name, remote, dataset_source)
-
-    # Carry over persona / prompt selections across model or remote switches.
-    if chat_state["persona_id"] is None:
-        chat_state["persona_id"] = st.session_state.get(_LAST_PERSONA_ID_KEY)
-        chat_state["prompt_mode"] = st.session_state.get(
-            _LAST_PROMPT_MODE_KEY, "templated"
-        )
-
+def _load_personas(dataset_source: str) -> list[PersonaData] | None:
     try:
         dataset, dataset_status = load_dataset(
             dataset_source,
@@ -52,36 +38,124 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     except Exception as exc:
         st.error(f"Could not load data: {exc}")
         st.info("Check the selected dataset source or upload both JSONL files.")
-        return
+        return None
 
     personas = list(dataset)
     if not personas:
         st.warning("No personas found in the selected dataset.")
         st.info("Try a different dataset source or upload a non-empty personas file.")
-        return
+        return None
+    return personas
 
-    generation: GenerationConfig = render_generation_settings(context_key, remote)
-    probe_enabled = st.toggle(
-        "Probe tools",
-        value=False,
-        key=widget_key(context_key, "probe_enabled"),
-        help="Trace chat activations and run compatible `.pt` probes on tapped tokens.",
-    )
 
-    # ── Mode toggle ───────────────────────────────────────────────────────────
-    compare_key = widget_key(context_key, "compare_mode")
-    if compare_key not in st.session_state:
-        st.session_state[compare_key] = st.session_state.get(
-            _LAST_COMPARE_MODE_KEY, False
+def _render_single_chat_footer(
+    *,
+    model_name: str,
+    dataset_source: str,
+    persona: PersonaData,
+    prompt_mode: str,
+    system_prompt: str | None,
+    chat_state: ChatState,
+    generation: GenerationConfig,
+    export_key: str,
+    reset_key: str,
+    on_reset,
+) -> None:
+    footer = st.container()
+    with footer:
+        exp_col, rst_col, _spacer = st.columns([0.5, 0.5, 10], gap="xsmall")
+        with exp_col:
+            if st.button(
+                "",
+                icon=":material/download:",
+                key=export_key,
+                help="Export chat",
+            ):
+                save_chat_export(
+                    model_name=model_name,
+                    dataset_source=dataset_source,
+                    persona_id=persona.id,
+                    persona_name=getattr(persona, "name", None),
+                    prompt_mode=prompt_mode,
+                    system_prompt=system_prompt,
+                    messages=chat_state["messages"],
+                    generation=generation.to_export_dict(),
+                )
+                st.toast("Exported", icon=":material/check:")
+        with rst_col:
+            if st.button(
+                "",
+                icon=":material/delete_sweep:",
+                key=reset_key,
+                help="Reset chat",
+            ):
+                on_reset()
+                st.rerun()
+
+
+def _handle_single_chat_generation(
+    *,
+    remote: bool,
+    model_name: str,
+    chat_state: ChatState,
+    active_system_prompt: str | None,
+    generation: GenerationConfig,
+    pending_action: object,
+    chat_log,
+) -> None:
+    messages = build_chat_messages(active_system_prompt, chat_state["messages"])
+
+    with st.spinner("Generating reply..."):
+        model = cached_model(model_name=model_name)
+        try:
+            reply: ChatReply = generate_chat_reply(
+                model=model,
+                messages=messages,
+                remote=remote,
+                past_key_values=chat_state["past_key_values"],
+                **generation.to_generate_kwargs(),
+            )
+        except Exception as exc:
+            with chat_log:
+                st.error(f"Could not generate a reply: {exc}")
+                st.info("Try a shorter prompt, reset the chat, or switch personas.")
+            if pending_action == "new_user_prompt" and chat_state["messages"]:
+                chat_state["messages"].pop()
+            return
+
+    chat_state["messages"].append({"role": "assistant", "content": reply.text})
+    chat_state["past_key_values"] = reply.past_key_values if not remote else None
+    st.rerun()
+
+
+
+
+def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
+    """Render the chat tab."""
+
+    st.title("Chat")
+    st.caption("Chat with a persona, optionally side-by-side or with token contrast.")
+
+    context_key = chat_session_key(model_name, dataset_source)
+    chat_state = get_chat_state(model_name, remote, dataset_source)
+
+    # Carry over persona / prompt selections across model or remote switches.
+    if chat_state["persona_id"] is None:
+        chat_state["persona_id"] = st.session_state.get(_LAST_PERSONA_ID_KEY)
+        chat_state["prompt_mode"] = st.session_state.get(
+            _LAST_PROMPT_MODE_KEY, "templated"
         )
-    compare_mode = st.toggle(
-        "Compare mode",
-        key=compare_key,
-        help="Side-by-side: send one message to two independent persona/prompt configurations.",
-    )
-    st.session_state[_LAST_COMPARE_MODE_KEY] = compare_mode
 
-    if compare_mode:
+    personas = _load_personas(dataset_source)
+    if personas is None:
+        return
+
+    generation, tools = render_advanced_settings(
+        context_key,
+        remote,
+        last_compare_mode_key=_LAST_COMPARE_MODE_KEY,
+    )
+    if tools.compare_mode:
         render_compare_mode(
             remote,
             model_name,
@@ -89,6 +163,7 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
             dataset_source,
             personas,
             generation,
+            contrast_enabled=tools.token_contrast,
         )
         return
 
@@ -150,7 +225,7 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
         remote=remote,
         active_system_prompt=active_system_prompt,
         chat_state=chat_state,
-        enabled=probe_enabled,
+        enabled=tools.probe_enabled,
     )
 
     render_chat_window(
@@ -161,36 +236,18 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
         pending_key=pending_key,
     )
 
-    footer = st.container()
-    with footer:
-        exp_col, rst_col, _spacer = st.columns([0.5, 0.5, 10], gap="xsmall")
-        with exp_col:
-            if st.button(
-                "",
-                icon=":material/download:",
-                key=export_key,
-                help="Export chat",
-            ):
-                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(generation),
-                )
-                st.toast("Exported", icon=":material/check:")
-        with rst_col:
-            if st.button(
-                "",
-                icon=":material/delete_sweep:",
-                key=reset_key,
-                help="Reset chat",
-            ):
-                _reset_active_chat_context()
-                st.rerun()
+    _render_single_chat_footer(
+        model_name=model_name,
+        dataset_source=dataset_source,
+        persona=selected_persona,
+        prompt_mode=prompt_mode,
+        system_prompt=active_system_prompt,
+        chat_state=chat_state,
+        generation=generation,
+        export_key=export_key,
+        reset_key=reset_key,
+        on_reset=_reset_active_chat_context,
+    )
 
     user_prompt = st.chat_input("Ask something...", key=chat_input_key)
 
@@ -205,26 +262,12 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     if not pending_action:
         return
 
-    messages = build_chat_messages(active_system_prompt, chat_state["messages"])
-
-    with st.spinner("Generating reply..."):
-        model = cached_model(model_name=model_name, remote=remote)
-        try:
-            reply: ChatReply = generate_chat_reply(
-                model=model,
-                messages=messages,
-                remote=remote,
-                past_key_values=chat_state["past_key_values"],
-                **generation.to_generate_kwargs(),
-            )
-        except Exception as exc:
-            with chat_log:
-                st.error(f"Could not generate a reply: {exc}")
-                st.info("Try a shorter prompt, reset the chat, or switch personas.")
-            if pending_action == "new_user_prompt" and chat_state["messages"]:
-                chat_state["messages"].pop()
-            return
-
-    chat_state["messages"].append({"role": "assistant", "content": reply.text})
-    chat_state["past_key_values"] = reply.past_key_values if not remote else None
-    st.rerun()
+    _handle_single_chat_generation(
+        remote=remote,
+        model_name=model_name,
+        chat_state=chat_state,
+        active_system_prompt=active_system_prompt,
+        generation=generation,
+        pending_action=pending_action,
+        chat_log=chat_log,
+    )

tabs/chat_ui.py

@@ -48,6 +48,13 @@ class GenerationConfig:
         }
 
 
+@dataclass(frozen=True)
+class ChatTools:
+    probe_enabled: bool
+    compare_mode: bool
+    token_contrast: bool
+
+
 @st.dialog("Edit", width="medium")
 def _open_edit_dialog(
     *,
@@ -108,13 +115,54 @@ def _open_system_prompt_dialog(*, prompt_key: str, current_value: str) -> None:
             st.rerun()
 
 
-def generation_dict(config: GenerationConfig) -> dict[str, object]:
-    return config.to_export_dict()
+def render_advanced_settings(
+    context_key: str,
+    remote: bool,
+    *,
+    last_compare_mode_key: str,
+) -> tuple[GenerationConfig, ChatTools]:
+    """Render the Advanced expander: tool toggles + generation settings."""
+    with st.expander("Advanced", expanded=False):
+        st.caption("Tools")
 
+        compare_key = widget_key(context_key, "compare_mode")
+        if compare_key not in st.session_state:
+            st.session_state[compare_key] = st.session_state.get(
+                last_compare_mode_key, False
+            )
+
+        tools_col1, tools_col2, tools_col3 = st.columns(3)
+        with tools_col1:
+            probe_enabled = st.toggle(
+                "Probe tools",
+                value=False,
+                key=widget_key(context_key, "probe_enabled"),
+                help="Trace chat activations and run compatible `.pt` probes on tapped tokens.",
+            )
+        with tools_col2:
+            compare_mode = st.toggle(
+                "Compare mode",
+                key=compare_key,
+                help="Side-by-side: send one message to two independent persona/prompt configurations.",
+            )
+        with tools_col3:
+            token_contrast = st.toggle(
+                "Token contrast",
+                value=False,
+                key=widget_key(context_key, "token_contrast"),
+                disabled=not compare_mode,
+                help=(
+                    "Color each generated token by how characteristic it is of each persona. "
+                    "Red = more likely under the left persona, blue = more likely under the "
+                    "right. Requires up to four extra scoring passes after each turn. "
+                    "Available only in Compare mode."
+                ),
+            )
+        st.session_state[last_compare_mode_key] = compare_mode
+
+        st.divider()
+        st.caption("Generation")
 
-def render_generation_settings(context_key: str, remote: bool) -> GenerationConfig:
-    """Render the Advanced generation settings expander."""
-    with st.expander("Advanced", expanded=False):
         config_col1, config_col2 = st.columns([2, 1])
         with config_col1:
             max_new_tokens = st.slider(
@@ -199,7 +247,7 @@ def render_generation_settings(context_key: str, remote: bool) -> GenerationConf
             st.caption("Seed is local-only and disabled for remote runs.")
 
     do_sample = bool(use_sampling)
-    return GenerationConfig(
+    generation = GenerationConfig(
         max_new_tokens=int(max_new_tokens),
         do_sample=do_sample,
         temperature=float(temperature),
@@ -208,6 +256,12 @@ def render_generation_settings(context_key: str, remote: bool) -> GenerationConf
         repetition_penalty=float(repetition_penalty),
         seed=seed if do_sample and seed is not None and not remote else None,
     )
+    tools = ChatTools(
+        probe_enabled=probe_enabled,
+        compare_mode=compare_mode,
+        token_contrast=token_contrast and compare_mode,
+    )
+    return generation, tools
 
 
 def render_chat_message(

tabs/compare.py

@@ -1,8 +1,9 @@
+from collections.abc import Callable
 from itertools import combinations
+from dataclasses import dataclass
 
 import streamlit as st
 from persona_data.environment import get_artifacts_dir
-from persona_data.prompts import BASELINE_PERSONA_ID
 from persona_vectors.analysis import (
     load_persona_mean_samples,
     load_variant_mean_samples,
@@ -41,6 +42,15 @@ _LAST_PROJECTION_PERSONAS_KEY = "compare:last_personas:projection"
 _LAST_MASK_STRATEGY_KEY = "compare:last_mask_strategy"
 
 
+@dataclass(frozen=True)
+class CosineSelection:
+    variants: list[str]
+    variant_a: str
+    variant_b: str
+    persona_ids: list[str]
+    persona_key: str
+
+
 def _select_artifact_personas(
     store: ActivationStore,
     variants: list[str],
@@ -143,71 +153,157 @@ def _render_mask_strategy_select(scope: str) -> MaskStrategy:
     return selected
 
 
-def _render_cosine_similarity(
+def _render_cosine_selection(
     store: ActivationStore,
     mask_strategy: MaskStrategy,
-) -> None:
+) -> CosineSelection | None:
     variants = list(store.variants)
     if len(variants) < 2:
         st.info("Need at least two non-baseline variants for cosine comparison.")
-        return
+        return None
 
-    col1, col2 = st.columns(2)
-    with col1:
-        variant_a = st.selectbox(
-            "Variant A",
-            options=variants,
-            index=0,
-            format_func=prompt_variant_label,
-            key=widget_key("load", "variant_a"),
+    with st.expander("Vector selection", expanded=True):
+        col1, col2 = st.columns(2)
+        with col1:
+            variant_a = st.selectbox(
+                "Variant A",
+                options=variants,
+                index=0,
+                format_func=prompt_variant_label,
+                key=widget_key("load", "variant_a"),
+            )
+        with col2:
+            variant_b = st.selectbox(
+                "Variant B",
+                options=variants,
+                index=min(1, len(variants) - 1),
+                format_func=prompt_variant_label,
+                key=widget_key("load", "variant_b"),
+            )
+
+        if variant_a == variant_b:
+            st.warning("Choose two different variants to compare.")
+            return None
+
+        persona_ids, _ = _select_artifact_personas(
+            store,
+            [variant_a, variant_b],
+            mask_strategy,
+            widget_scope="cosine",
+            remember_key=_LAST_COSINE_PERSONAS_KEY,
         )
-    with col2:
-        variant_b = st.selectbox(
-            "Variant B",
-            options=variants,
-            index=min(1, len(variants) - 1),
-            format_func=prompt_variant_label,
-            key=widget_key("load", "variant_b"),
+    if not persona_ids:
+        return None
+    return CosineSelection(
+        variants=variants,
+        variant_a=variant_a,
+        variant_b=variant_b,
+        persona_ids=persona_ids,
+        persona_key="_".join(sorted(persona_ids)),
+    )
+
+
+def _build_cosine_figures(
+    store: ActivationStore,
+    selection: CosineSelection,
+) -> tuple[object, object | None, int, int] | None:
+    try:
+        variant_samples = load_variant_mean_samples(
+            store,
+            [selection.variant_a, selection.variant_b],
+            persona_ids=selection.persona_ids,
         )
+    except Exception as exc:
+        st.error(f"Could not load vectors: {exc}")
+        return None
 
-    if variant_a == variant_b:
-        st.warning("Choose two different variants to compare.")
-        return
+    labels = variant_samples[selection.variant_a].labels
+    display_traces = [
+        (
+            label,
+            variant_samples[selection.variant_a].vectors[index],
+            variant_samples[selection.variant_b].vectors[index],
+        )
+        for index, label in enumerate(labels)
+    ]
+    fig = plot_layer_similarity(
+        display_traces,
+        title=(
+            f"{prompt_variant_label(selection.variant_a)} vs "
+            f"{prompt_variant_label(selection.variant_b)}"
+        ),
+        show=False,
+    )
 
-    persona_ids, _ = _select_artifact_personas(
-        store,
-        [variant_a, variant_b],
-        mask_strategy,
-        widget_scope="cosine",
-        remember_key=_LAST_COSINE_PERSONAS_KEY,
+    pair_traces = []
+    pair_errors = []
+    for left, right in combinations(selection.variants, 2):
+        try:
+            pair_samples = (
+                variant_samples
+                if {left, right} == {selection.variant_a, selection.variant_b}
+                else load_variant_mean_samples(
+                    store,
+                    [left, right],
+                    persona_ids=selection.persona_ids,
+                )
+            )
+            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),
+                )
+            )
+        except Exception as exc:
+            pair_errors.append(f"{left} vs {right}: {exc}")
+            continue
+
+    for err in pair_errors:
+        st.warning(f"Skipped pair trace: `{err}`")
+    pair_fig = (
+        plot_layer_similarity(
+            pair_traces,
+            title="Variant-pair cosine similarity averaged over selected personas",
+            show=False,
+        )
+        if pair_traces
+        else None
     )
-    if not persona_ids:
+    return fig, pair_fig, len(display_traces), len(pair_traces)
+
+
+def _render_cosine_similarity(
+    store: ActivationStore,
+    mask_strategy: MaskStrategy,
+) -> None:
+    selection = _render_cosine_selection(store, mask_strategy)
+    if selection is None:
         return
-    persona_key = "_".join(sorted(persona_ids))
 
     cosine_fig_key = widget_key(
         "load",
         "cosine_fig_state",
         store.model_name,
         mask_strategy.value,
-        variant_a,
-        variant_b,
-        persona_key,
+        selection.variant_a,
+        selection.variant_b,
+        selection.persona_key,
     )
     filename = _filename(
         "compare",
         "cosine",
         store.model_name,
         mask_strategy.value,
-        variant_a,
-        variant_b,
+        selection.variant_a,
+        selection.variant_b,
     )
     pairs_filename = _filename(
         "compare",
         "cosine_pairs",
         store.model_name,
         mask_strategy.value,
-        "_".join(variants),
+        "_".join(selection.variants),
     )
 
     if st.button(
@@ -218,79 +314,16 @@ def _render_cosine_similarity(
             "compare_vectors",
             store.model_name,
             mask_strategy.value,
-            variant_a,
-            variant_b,
-            persona_key,
+            selection.variant_a,
+            selection.variant_b,
+            selection.persona_key,
         ),
     ):
-        try:
-            variant_samples = load_variant_mean_samples(
-                store,
-                [variant_a, variant_b],
-                persona_ids=persona_ids,
-            )
-        except Exception as exc:
-            st.error(f"Could not load vectors: {exc}")
+        figures = _build_cosine_figures(store, selection)
+        if figures is None:
             st.session_state.pop(cosine_fig_key, None)
             return
-
-        labels = variant_samples[variant_a].labels
-        display_traces = [
-            (
-                label,
-                variant_samples[variant_a].vectors[index],
-                variant_samples[variant_b].vectors[index],
-            )
-            for index, label in enumerate(labels)
-        ]
-        fig = plot_layer_similarity(
-            display_traces,
-            title=f"{prompt_variant_label(variant_a)} vs {prompt_variant_label(variant_b)}",
-            show=False,
-        )
-
-        pair_traces = []
-        pair_errors = []
-        for left, right in combinations(variants, 2):
-            try:
-                pair_samples = (
-                    variant_samples
-                    if {left, right} == {variant_a, variant_b}
-                    else load_variant_mean_samples(
-                        store,
-                        [left, right],
-                        persona_ids=persona_ids,
-                    )
-                )
-            except Exception as exc:
-                pair_errors.append(f"{left} vs {right}: {exc}")
-                continue
-            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),
-                )
-            )
-
-        if pair_errors:
-            for err in pair_errors:
-                st.warning(f"Skipped pair trace: `{err}`")
-        pair_fig = (
-            plot_layer_similarity(
-                pair_traces,
-                title="Variant-pair cosine similarity averaged over selected personas",
-                show=False,
-            )
-            if pair_traces
-            else None
-        )
-        st.session_state[cosine_fig_key] = (
-            fig,
-            pair_fig,
-            len(display_traces),
-            len(pair_traces),
-        )
+        st.session_state[cosine_fig_key] = figures
 
     if cosine_fig_key in st.session_state:
         fig, pair_fig, n_traces, n_pair_traces = st.session_state[cosine_fig_key]
@@ -369,190 +402,89 @@ def _select_single_variant_samples(
     return variant, persona_ids, persona_key, selected_layers
 
 
-def _baseline_available(
-    store: ActivationStore,
-) -> bool:
-    return BASELINE_PERSONA_ID in store.list_personas(
-        [BASELINE_PERSONA_ID],
-        warn_missing=False,
-    )
-
-
-def _render_baseline_reference_toggle(
+def _render_layered_figure_analysis(
     store: ActivationStore,
     mask_strategy: MaskStrategy,
+    *,
     scope: str,
-) -> bool:
-    available = _baseline_available(store)
-    return st.checkbox(
-        "Include Assistant baseline reference",
-        value=available,
-        disabled=not available,
-        key=widget_key("load", "include_baseline", scope, mask_strategy.value),
-        help=(
-            "Adds the single saved baseline artifact as one reference sample."
-            if available
-            else "Run Assistant baseline extraction first."
-        ),
-    )
-
-
-def _render_similarity_matrix(
-    store: ActivationStore,
-    mask_strategy: MaskStrategy,
+    figure_kind: str,
+    button_label: str,
+    title_fn: Callable[[str], str],
+    include_pair_trajectories: bool = False,
 ) -> None:
-    selected = _select_single_variant_samples(
-        store,
-        mask_strategy,
-        "similarity_matrix",
-    )
-    if selected is None:
-        return
-    variant, persona_ids, persona_key, selected_layers = selected
-    include_baseline = _render_baseline_reference_toggle(
-        store,
-        mask_strategy,
-        "similarity_matrix",
-    )
-
-    fig_key = widget_key(
-        "load",
-        "similarity_matrix_fig_state",
-        store.model_name,
-        mask_strategy.value,
-        variant,
-        "persona_mean",
-        persona_key,
-        BASELINE_PERSONA_ID if include_baseline else "no_baseline",
-    )
-    filename = _filename(
-        "compare",
-        "similarity_matrix",
-        store.model_name,
-        mask_strategy.value,
-        variant,
-        "persona_mean",
-        persona_key,
-        BASELINE_PERSONA_ID if include_baseline else "",
-    )
+    """Render a single-variant layered analysis: select → button → figure(s).
 
-    if st.button("Generate similarity matrix", type="primary"):
-        try:
-            samples = load_persona_mean_samples(
-                store,
-                variant,
-                mask_strategy=mask_strategy,
-                persona_ids=persona_ids,
-                include_baseline=include_baseline,
-            )
-            matrix_fig = build_layered_figure(
-                samples,
-                "similarity",
-                layers=selected_layers,
-                title=(
-                    "Centered similarity - "
-                    f"{prompt_variant_label(variant)} - personas averaged over questions"
-                ),
-            )
-            trajectory_fig = build_pair_similarity_figure(
-                samples,
-                layers=selected_layers,
-                title=(
-                    "Pair similarity trajectories - "
-                    f"{prompt_variant_label(variant)} - personas averaged over questions"
-                ),
-            )
-            st.session_state[fig_key] = (
-                matrix_fig,
-                trajectory_fig,
-                samples.vectors.shape[0],
-            )
-        except Exception as exc:
-            st.error(f"Could not build similarity matrix: {exc}")
-            st.session_state.pop(fig_key, None)
-
-    if fig_key in st.session_state:
-        matrix_fig, trajectory_fig, n_samples = st.session_state[fig_key]
-        st.plotly_chart(matrix_fig, width="stretch")
-        st.subheader("Pair trajectories")
-        st.plotly_chart(trajectory_fig, width="stretch")
-        _render_save_buttons(
-            [matrix_fig, trajectory_fig],
-            [filename, f"{filename}__pair_trajectories"],
-            "similarity_matrix",
-        )
-        st.success(f"Loaded {n_samples} samples.")
-
-
-def _render_embedding_analysis(
-    store: ActivationStore,
-    analysis_mode: str,
-    mask_strategy: MaskStrategy,
-) -> None:
-    selected = _select_single_variant_samples(
-        store,
-        mask_strategy,
-        analysis_mode.lower(),
-    )
+    Used for similarity matrix, PCA, and UMAP. Set ``include_pair_trajectories``
+    to add the pair-similarity-trajectory figure (similarity matrix only).
+    """
+    selected = _select_single_variant_samples(store, mask_strategy, scope)
     if selected is None:
         return
     variant, persona_ids, persona_key, selected_layers = selected
 
-    figure_kind = analysis_mode.lower()
-    include_baseline = _render_baseline_reference_toggle(
-        store,
-        mask_strategy,
-        analysis_mode.lower(),
-    )
-
     fig_key = widget_key(
         "load",
-        "embedding_fig_state",
+        f"{scope}_fig_state",
         store.model_name,
         mask_strategy.value,
         figure_kind,
         variant,
         "persona_mean",
         persona_key,
-        BASELINE_PERSONA_ID if include_baseline else "no_baseline",
     )
     filename = _filename(
         "compare",
-        figure_kind,
+        scope,
         store.model_name,
         mask_strategy.value,
         variant,
         "persona_mean",
         persona_key,
-        BASELINE_PERSONA_ID if include_baseline else "",
     )
 
-    if st.button(f"Generate {analysis_mode} projection", type="primary"):
+    if st.button(button_label, type="primary"):
         try:
             samples = load_persona_mean_samples(
                 store,
                 variant,
                 mask_strategy=mask_strategy,
                 persona_ids=persona_ids,
-                include_baseline=include_baseline,
             )
-            fig = build_layered_figure(
+            main_fig = build_layered_figure(
                 samples,
                 figure_kind,
                 layers=selected_layers,
-                title=(
-                    f"{analysis_mode} - {prompt_variant_label(variant)} - Persona means"
-                ),
+                title=title_fn(variant),
+            )
+            extra_fig = (
+                build_pair_similarity_figure(
+                    samples,
+                    layers=selected_layers,
+                    title=(
+                        "Pair similarity trajectories - "
+                        f"{prompt_variant_label(variant)} - "
+                        "persona mean activations"
+                    ),
+                )
+                if include_pair_trajectories
+                else None
             )
-            st.session_state[fig_key] = (fig, samples.vectors.shape[0])
+            st.session_state[fig_key] = (main_fig, extra_fig, samples.vectors.shape[0])
         except Exception as exc:
-            st.error(f"Could not build {analysis_mode}: {exc}")
+            st.error(f"Could not build figure: {exc}")
             st.session_state.pop(fig_key, None)
 
     if fig_key in st.session_state:
-        fig, n_samples = st.session_state[fig_key]
-        st.plotly_chart(fig, width="stretch")
-        _render_save_buttons([fig], [filename], figure_kind)
+        main_fig, extra_fig, n_samples = st.session_state[fig_key]
+        st.plotly_chart(main_fig, width="stretch")
+        figs = [main_fig]
+        filenames = [filename]
+        if extra_fig is not None:
+            st.subheader("Pair trajectories")
+            st.plotly_chart(extra_fig, width="stretch")
+            figs.append(extra_fig)
+            filenames.append(f"{filename}__pair_trajectories")
+        _render_save_buttons(figs, filenames, scope)
         st.success(f"Loaded {n_samples} samples.")
 
 
@@ -562,9 +494,7 @@ def render_compare_tab(model_name: str) -> None:
     st.title("Compare")
     st.caption("Compare saved activations by cosine similarity, PCA, or UMAP.")
 
-    st.subheader("Analysis")
-
-    with st.expander("Advanced", expanded=False):
+    with st.expander("Artifact settings", expanded=False):
         artifacts_root = st.text_input(
             "Artifacts root",
             value=str(get_artifacts_dir() / "activations"),
@@ -580,14 +510,35 @@ 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])
-    mask_strategy = _render_mask_strategy_select(analysis_mode)
+    with st.expander("Activation settings", expanded=False):
+        mask_strategy = _render_mask_strategy_select(analysis_mode)
     store = ActivationStore(model_name, artifacts_root, mask_strategy=mask_strategy)
 
     if analysis_mode == "Cosine similarity":
         _render_cosine_similarity(store, mask_strategy)
         return
     if analysis_mode == "Similarity matrix":
-        _render_similarity_matrix(store, mask_strategy)
+        _render_layered_figure_analysis(
+            store,
+            mask_strategy,
+            scope="similarity_matrix",
+            figure_kind="similarity",
+            button_label="Generate similarity matrix",
+            title_fn=lambda v: (
+                "Centered similarity - "
+                f"{prompt_variant_label(v)} - persona mean activations"
+            ),
+            include_pair_trajectories=True,
+        )
         return
 
-    _render_embedding_analysis(store, analysis_mode, mask_strategy)
+    _render_layered_figure_analysis(
+        store,
+        mask_strategy,
+        scope=analysis_mode.lower(),
+        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"
+        ),
+    )

tabs/compare_chat.py

@@ -1,10 +1,11 @@
-from typing import Any, NamedTuple
+from dataclasses import dataclass
+from typing import Any
 
 import streamlit as st
 from nnterp import StandardizedTransformer
 from persona_data.synth_persona import PersonaData
 
-from state import default_chat_state, reset_chat_context_state
+from state import ChatState, default_chat_state, reset_chat_context_state
 from utils.chat import (
     ChatReply,
     build_chat_messages,
@@ -18,7 +19,6 @@ from utils.runtime import cached_model
 
 from .chat_ui import (
     GenerationConfig,
-    generation_dict,
     render_chat_message,
     render_chat_window,
     render_persona_prompt_controls,
@@ -26,9 +26,10 @@ from .chat_ui import (
 )
 
 
-class ComparePanel(NamedTuple):
+@dataclass(frozen=True)
+class ComparePanel:
     side: str
-    state: dict[str, object]
+    state: ChatState
     log: Any
     prompt: str | None
     persona: PersonaData
@@ -37,6 +38,13 @@ class ComparePanel(NamedTuple):
     pending_key: str
 
 
+def _get_compare_state(context_key: str, side: str) -> tuple[str, ChatState]:
+    panel_key = widget_key(context_key, f"cmp_{side}")
+    if panel_key not in st.session_state:
+        st.session_state[panel_key] = default_chat_state()
+    return panel_key, st.session_state[panel_key]
+
+
 def _reset_compare_panel(panel: ComparePanel) -> None:
     reset_chat_context_state(
         panel.state,
@@ -48,195 +56,188 @@ def _reset_compare_panel(panel: ComparePanel) -> None:
     st.session_state.pop(panel.edit_key, None)
 
 
-def _generate_panel_reply(
+def _render_compare_panel(
     *,
-    model: StandardizedTransformer,
-    remote: bool,
-    panel_state: dict[str, object],
-    panel_prompt: str | None,
-    generation: GenerationConfig,
-) -> ChatReply:
-    return generate_chat_reply(
-        model=model,
-        messages=build_chat_messages(panel_prompt, panel_state["messages"]),
-        remote=remote,
-        past_key_values=panel_state["past_key_values"],
-        **generation.to_generate_kwargs(),
-    )
-
-
-def render_compare_mode(
-    remote: bool,
-    model_name: str,
     context_key: str,
-    dataset_source: str,
+    side: str,
     personas: list[PersonaData],
-    generation: GenerationConfig,
-) -> None:
-    """Render the full side-by-side comparison UI."""
-    model: StandardizedTransformer | None = None
-
-    def _get_model() -> StandardizedTransformer:
-        nonlocal model
-        if model is None:
-            model = cached_model(model_name=model_name, remote=remote)
-        return model
-
-    contrast_key = widget_key(context_key, "token_contrast")
-    contrast_enabled = st.toggle(
-        "Token contrast",
-        value=False,
-        key=contrast_key,
-        help=(
-            "Color each generated token by how characteristic it is of each persona. "
-            "Red = more likely under the left persona, blue = more likely under the right. "
-            "Requires up to four extra scoring passes after each turn."
-        ),
+) -> ComparePanel:
+    panel_key, state = _get_compare_state(context_key, side)
+
+    prompt_key = widget_key(panel_key, "custom_prompt")
+    edit_key = widget_key(panel_key, "edit_idx")
+    pending_key = widget_key(panel_key, "pending_regen")
+
+    persist_persona_key = f"chat:last_cmp_{side}_persona"
+    persist_prompt_key = f"chat:last_cmp_{side}_prompt"
+    if state["persona_id"] is None:
+        state["persona_id"] = st.session_state.get(persist_persona_key)
+        state["prompt_mode"] = st.session_state.get(persist_prompt_key, "templated")
+
+    selected_persona, prompt_mode, changed = render_persona_prompt_controls(
+        personas,
+        state["persona_id"],
+        state["prompt_mode"],
+        widget_key(panel_key, "persona"),
+        widget_key(panel_key, "prompt_mode"),
     )
-
-    def render_panel(side: str) -> ComparePanel:
-        panel_key = widget_key(context_key, f"cmp_{side}")
-        if panel_key not in st.session_state:
-            st.session_state[panel_key] = default_chat_state()
-        state = st.session_state[panel_key]
-
-        prompt_key = widget_key(panel_key, "custom_prompt")
-        edit_key = widget_key(panel_key, "edit_idx")
-        pending_key = widget_key(panel_key, "pending_regen")
-
-        # Carry over persona / prompt selections across model or remote switches.
-        persist_persona_key = f"chat:last_cmp_{side}_persona"
-        persist_prompt_key = f"chat:last_cmp_{side}_prompt"
-        if state["persona_id"] is None:
-            state["persona_id"] = st.session_state.get(persist_persona_key)
-            state["prompt_mode"] = st.session_state.get(persist_prompt_key, "templated")
-
-        selected_persona, prompt_mode, changed = render_persona_prompt_controls(
-            personas,
-            state["persona_id"],
-            state["prompt_mode"],
-            widget_key(panel_key, "persona"),
-            widget_key(panel_key, "prompt_mode"),
+    st.session_state[persist_persona_key] = selected_persona.id
+    st.session_state[persist_prompt_key] = prompt_mode
+
+    if changed:
+        reset_chat_context_state(
+            state,
+            selected_persona.id,
+            prompt_mode,
+            prompt_key,
+            pending_key,
         )
-        st.session_state[persist_persona_key] = selected_persona.id
-        st.session_state[persist_prompt_key] = prompt_mode
+        st.session_state.pop(edit_key, None)
 
-        if changed:
-            reset_chat_context_state(
-                state, selected_persona.id, prompt_mode, prompt_key, pending_key
-            )
-            st.session_state.pop(edit_key, None)
+    active_system_prompt = resolve_system_prompt(
+        persona=selected_persona,
+        mode=prompt_mode,
+    )
 
-        active_system_prompt = resolve_system_prompt(
-            persona=selected_persona, mode=prompt_mode
+    chat_log = st.container()
+    with chat_log:
+        active_system_prompt = render_system_prompt(
+            prompt_key,
+            prompt_mode,
+            active_system_prompt,
         )
 
-        chat_log = st.container()
-        with chat_log:
-            active_system_prompt = render_system_prompt(
-                prompt_key, prompt_mode, active_system_prompt
-            )
-        return ComparePanel(
-            side=side,
-            state=state,
-            log=chat_log,
-            prompt=active_system_prompt,
-            persona=selected_persona,
-            prompt_key=prompt_key,
-            edit_key=edit_key,
-            pending_key=pending_key,
-        )
+    return ComparePanel(
+        side=side,
+        state=state,
+        log=chat_log,
+        prompt=active_system_prompt,
+        persona=selected_persona,
+        prompt_key=prompt_key,
+        edit_key=edit_key,
+        pending_key=pending_key,
+    )
 
-    left_col, right_col = st.columns(2)
-    with left_col:
-        left = render_panel("left")
-    with right_col:
-        right = render_panel("right")
-    panels: list[ComparePanel] = [left, right]
 
-    # Handle per-panel regeneration triggered by message edits
-    regen_panels = [p for p in panels if st.session_state.pop(p.pending_key, False)]
-    if regen_panels:
-        model = _get_model()
-
-        results: list[ChatReply | Exception] = []
-        with st.spinner("Regenerating..."):
-            for panel in regen_panels:
-                try:
-                    results.append(
-                        _generate_panel_reply(
-                            model=model,
-                            remote=remote,
-                            panel_state=panel.state,
-                            panel_prompt=panel.prompt,
-                            generation=generation,
-                        )
+def _generate_panels(
+    *,
+    model: StandardizedTransformer,
+    remote: bool,
+    panels: list[ComparePanel],
+    generation: GenerationConfig,
+    spinner_label: str,
+) -> list[ChatReply | Exception]:
+    results: list[ChatReply | Exception] = []
+    with st.spinner(spinner_label):
+        for panel in panels:
+            try:
+                results.append(
+                    generate_chat_reply(
+                        model=model,
+                        messages=build_chat_messages(
+                            panel.prompt, panel.state["messages"]
+                        ),
+                        remote=remote,
+                        past_key_values=panel.state["past_key_values"],
+                        **generation.to_generate_kwargs(),
                     )
-                except Exception as exc:
-                    results.append(exc)
+                )
+            except Exception as exc:
+                results.append(exc)
+    return results
 
-        for panel, result in zip(regen_panels, results):
-            if isinstance(result, Exception):
-                with panel.log:
-                    st.error(f"Generation failed: {result}")
+
+def _apply_panel_results(
+    *,
+    panels: list[ComparePanel],
+    results: list[ChatReply | Exception],
+    remote: bool,
+    rollback_user_on_error: bool,
+) -> list[ChatReply | None]:
+    valid_results: list[ChatReply | None] = []
+    for panel, result in zip(panels, results, strict=True):
+        if isinstance(result, Exception):
+            with panel.log:
+                st.error(f"Generation failed: {result}")
+            if rollback_user_on_error and panel.state["messages"]:
+                panel.state["messages"].pop()
+            valid_results.append(None)
+            continue
+
+        panel.state["messages"].append({"role": "assistant", "content": result.text})
+        panel.state["past_key_values"] = result.past_key_values if not remote else None
+        valid_results.append(result)
+    return valid_results
+
+
+def _pending_contrast_edits(panels: list[ComparePanel]) -> list[tuple[int, int]]:
+    return [
+        (panel_idx, msg_idx)
+        for panel_idx, panel in enumerate(panels)
+        for msg_idx, msg in enumerate(panel.state["messages"])
+        if msg.get("_needs_contrast") and msg.get("role") == "assistant"
+    ]
+
+
+def _recompute_pending_contrast(
+    *,
+    model: StandardizedTransformer,
+    remote: bool,
+    panels: list[ComparePanel],
+) -> bool:
+    pending_edits = _pending_contrast_edits(panels)
+    if not pending_edits:
+        return False
+
+    left, right = panels
+    label_a = persona_label(left.persona)
+    label_b = persona_label(right.persona)
+    with st.spinner("Recomputing token contrast..."):
+        for panel_idx, msg_idx in pending_edits:
+            panel = panels[panel_idx]
+            msg = panel.state["messages"][msg_idx]
+            if msg_idx >= len(left.state["messages"]) or msg_idx >= len(
+                right.state["messages"]
+            ):
+                msg.pop("_needs_contrast", None)
                 continue
-            panel.state["messages"].append(
-                {"role": "assistant", "content": result.text}
+
+            context_a = build_chat_messages(
+                left.prompt,
+                left.state["messages"][:msg_idx],
             )
-            panel.state["past_key_values"] = (
-                result.past_key_values if not remote else None
+            context_b = build_chat_messages(
+                right.prompt,
+                right.state["messages"][:msg_idx],
             )
-        st.rerun()
+            try:
+                response_ids = model.tokenizer(
+                    msg["content"],
+                    add_special_tokens=False,
+                    return_tensors="pt",
+                ).input_ids[0]
+                contrast = compute_contrast(
+                    model=model,
+                    context_a=context_a,
+                    context_b=context_b,
+                    response_ids=response_ids,
+                    label_a=label_a,
+                    label_b=label_b,
+                    remote=remote,
+                )
+                if contrast is not None:
+                    msg["_contrast"] = contrast
+            except Exception as exc:
+                st.warning(f"Token contrast recompute failed: {exc}")
+            msg.pop("_needs_contrast", None)
+    return True
 
-    # Recompute contrast for assistant messages that were edited in place.
-    if contrast_enabled:
-        pending_edits: list[tuple[int, int]] = [
-            (panel_idx, msg_idx)
-            for panel_idx, panel in enumerate(panels)
-            for msg_idx, msg in enumerate(panel.state["messages"])
-            if msg.get("_needs_contrast") and msg.get("role") == "assistant"
-        ]
-        if pending_edits:
-            model = _get_model()
-            label_a = persona_label(left.persona)
-            label_b = persona_label(right.persona)
-            with st.spinner("Recomputing token contrast…"):
-                for panel_idx, msg_idx in pending_edits:
-                    panel = panels[panel_idx]
-                    msg = panel.state["messages"][msg_idx]
-                    if msg_idx >= len(left.state["messages"]) or msg_idx >= len(
-                        right.state["messages"]
-                    ):
-                        msg.pop("_needs_contrast", None)
-                        continue
-                    context_a = build_chat_messages(
-                        left.prompt, left.state["messages"][:msg_idx]
-                    )
-                    context_b = build_chat_messages(
-                        right.prompt, right.state["messages"][:msg_idx]
-                    )
-                    try:
-                        response_ids = model.tokenizer(
-                            msg["content"],
-                            add_special_tokens=False,
-                            return_tensors="pt",
-                        ).input_ids[0]
-                        tc = compute_contrast(
-                            model=model,
-                            context_a=context_a,
-                            context_b=context_b,
-                            response_ids=response_ids,
-                            label_a=label_a,
-                            label_b=label_b,
-                            remote=remote,
-                        )
-                        if tc is not None:
-                            msg["_contrast"] = tc
-                    except Exception as exc:
-                        st.warning(f"Token contrast recompute failed: {exc}")
-                    msg.pop("_needs_contrast", None)
-            st.rerun()
 
+def _render_compare_history(
+    *,
+    panels: list[ComparePanel],
+    contrast_enabled: bool,
+) -> None:
     for panel in panels:
         render_chat_window(
             chat_log=panel.log,
@@ -248,10 +249,23 @@ def render_compare_mode(
             edit_column_ratio=(10, 1),
         )
 
-    footer = st.container()
+
+def _render_compare_footer(
+    *,
+    context_key: str,
+    model_name: str,
+    dataset_source: str,
+    panels: list[ComparePanel],
+    generation: GenerationConfig,
+) -> None:
+    # Bumping this nonce after a reset gives the popover a fresh widget key,
+    # which forces Streamlit to re-mount it closed (popovers don't auto-close
+    # on click).
     reset_menu_nonce_key = widget_key(context_key, "cmp_reset_menu_nonce")
     if reset_menu_nonce_key not in st.session_state:
         st.session_state[reset_menu_nonce_key] = 0
+
+    footer = st.container()
     with footer:
         exp_col, rst_col, _spacer = st.columns([0.5, 0.5, 10], gap="xsmall")
         with exp_col:
@@ -270,7 +284,7 @@ def render_compare_mode(
                         prompt_mode=panel.state["prompt_mode"],
                         system_prompt=panel.prompt,
                         messages=panel.state["messages"],
-                        generation=generation_dict(generation),
+                        generation=generation.to_export_dict(),
                         panel_label=panel.side,
                     )
                 st.toast("Exported", icon=":material/check:")
@@ -304,81 +318,150 @@ def render_compare_mode(
                     st.session_state[reset_menu_nonce_key] += 1
                     st.rerun()
 
-    user_prompt = st.chat_input(
-        "Ask both...",
-        key=widget_key(context_key, "cmp_input"),
-    )
-
-    if not user_prompt:
-        return
-
-    model = cached_model(model_name=model_name, remote=remote)
 
+def _append_user_prompt(panels: list[ComparePanel], user_prompt: str) -> None:
     for panel in panels:
         panel.state["messages"].append({"role": "user", "content": user_prompt})
         with panel.log:
             render_chat_message({"role": "user", "content": user_prompt})
 
-    # Snapshot contexts before the new assistant turn is appended (needed for contrast).
-    pre_gen_contexts = [
-        build_chat_messages(panel.prompt, panel.state["messages"]) for panel in panels
-    ]
 
-    results: list[ChatReply | Exception] = []
-    with st.spinner("Generating..."):
-        # Sequential generation keeps both panels using model/session state safely.
-        for panel in panels:
-            try:
-                results.append(
-                    _generate_panel_reply(
-                        model=model,
-                        remote=remote,
-                        panel_state=panel.state,
-                        panel_prompt=panel.prompt,
-                        generation=generation,
-                    )
-                )
-            except Exception as exc:
-                results.append(exc)
+def _compute_new_reply_contrast(
+    *,
+    model: StandardizedTransformer,
+    remote: bool,
+    panels: list[ComparePanel],
+    pre_gen_contexts: list[list[dict[str, str]]],
+    results: list[ChatReply | None],
+) -> None:
+    if len(results) != 2 or any(
+        result is None or result.generated_ids is None for result in results
+    ):
+        return
 
-    valid_results: list[ChatReply | None] = []
-    for panel, result in zip(panels, results):
-        if isinstance(result, Exception):
-            with panel.log:
-                st.error(f"Generation failed: {result}")
-            panel.state["messages"].pop()
-            valid_results.append(None)
-            continue
+    left, right = panels
+    with st.spinner("Computing token contrast..."):
+        try:
+            left_contrast, right_contrast = compute_contrast_pair(
+                model=model,
+                context_a=pre_gen_contexts[0],
+                context_b=pre_gen_contexts[1],
+                response_ids_a=results[0].generated_ids,
+                response_ids_b=results[1].generated_ids,
+                label_a=persona_label(left.persona),
+                label_b=persona_label(right.persona),
+                remote=remote,
+            )
+            if left_contrast is not None:
+                left.state["messages"][-1]["_contrast"] = left_contrast
+            if right_contrast is not None:
+                right.state["messages"][-1]["_contrast"] = right_contrast
+        except Exception as exc:
+            st.warning(f"Token contrast failed: {exc}")
 
-        panel.state["messages"].append({"role": "assistant", "content": result.text})
-        panel.state["past_key_values"] = result.past_key_values if not remote else None
-        valid_results.append(result)
 
-    # Compute contrastive token coloring when both panels succeeded.
-    if (
-        contrast_enabled
-        and len(valid_results) == 2
-        and all(r is not None and r.generated_ids is not None for r in valid_results)
+def _render_compare_panels(
+    *,
+    context_key: str,
+    personas: list[PersonaData],
+) -> list[ComparePanel]:
+    left_col, right_col = st.columns(2)
+    with left_col:
+        left = _render_compare_panel(
+            context_key=context_key,
+            side="left",
+            personas=personas,
+        )
+    with right_col:
+        right = _render_compare_panel(
+            context_key=context_key,
+            side="right",
+            personas=personas,
+        )
+    return [left, right]
+
+
+def render_compare_mode(
+    remote: bool,
+    model_name: str,
+    context_key: str,
+    dataset_source: str,
+    personas: list[PersonaData],
+    generation: GenerationConfig,
+    *,
+    contrast_enabled: bool,
+) -> None:
+    """Render the full side-by-side comparison UI."""
+
+    panels = _render_compare_panels(context_key=context_key, personas=personas)
+
+    regen_panels = [
+        panel for panel in panels if st.session_state.pop(panel.pending_key, False)
+    ]
+    if regen_panels:
+        results = _generate_panels(
+            model=cached_model(model_name=model_name),
+            remote=remote,
+            panels=regen_panels,
+            generation=generation,
+            spinner_label="Regenerating...",
+        )
+        _apply_panel_results(
+            panels=regen_panels,
+            results=results,
+            remote=remote,
+            rollback_user_on_error=False,
+        )
+        st.rerun()
+
+    if contrast_enabled and _recompute_pending_contrast(
+        model=cached_model(model_name=model_name),
+        remote=remote,
+        panels=panels,
     ):
-        with st.spinner("Computing token contrast…"):
-            try:
-                tc_a, tc_b = compute_contrast_pair(
-                    model=model,
-                    context_a=pre_gen_contexts[0],
-                    context_b=pre_gen_contexts[1],
-                    response_ids_a=valid_results[0].generated_ids,
-                    response_ids_b=valid_results[1].generated_ids,
-                    label_a=persona_label(left.persona),
-                    label_b=persona_label(right.persona),
-                    remote=remote,
-                )
-                if tc_a is not None:
-                    left.state["messages"][-1]["_contrast"] = tc_a
-                if tc_b is not None:
-                    right.state["messages"][-1]["_contrast"] = tc_b
-            except Exception as exc:
-                st.warning(f"Token contrast failed: {exc}")
+        st.rerun()
+
+    _render_compare_history(panels=panels, contrast_enabled=contrast_enabled)
+    _render_compare_footer(
+        context_key=context_key,
+        model_name=model_name,
+        dataset_source=dataset_source,
+        panels=panels,
+        generation=generation,
+    )
+
+    user_prompt = st.chat_input(
+        "Ask both...",
+        key=widget_key(context_key, "cmp_input"),
+    )
+    if not user_prompt:
+        return
+
+    _append_user_prompt(panels, user_prompt)
+    pre_gen_contexts = [
+        build_chat_messages(panel.prompt, panel.state["messages"]) for panel in panels
+    ]
+    model = cached_model(model_name=model_name)
+    results = _generate_panels(
+        model=model,
+        remote=remote,
+        panels=panels,
+        generation=generation,
+        spinner_label="Generating...",
+    )
+    valid_results = _apply_panel_results(
+        panels=panels,
+        results=results,
+        remote=remote,
+        rollback_user_on_error=True,
+    )
+    if contrast_enabled:
+        _compute_new_reply_contrast(
+            model=model,
+            remote=remote,
+            panels=panels,
+            pre_gen_contexts=pre_gen_contexts,
+            results=valid_results,
+        )
 
-    # Rerun so the newly appended turns are redrawn through the editable history
-    # renderer instead of only appearing in the one-off generation pass.
     st.rerun()

tabs/extract.py

@@ -1,12 +1,9 @@
 import html
+from dataclasses import dataclass
 from typing import Literal, cast
 
 import streamlit as st
-from persona_data.prompts import (
-    BASELINE_PERSONA_ID,
-    BASELINE_PERSONA_NAME,
-    format_prompt,
-)
+from persona_data.prompts import format_prompt
 from persona_data.synth_persona import PersonaData, QAPair
 from persona_vectors.artifacts import PERSONA_VARIANTS
 from persona_vectors.extraction import (
@@ -42,8 +39,11 @@ _ITEM_TYPE_OPTIONS = ["all", "mcq", "frq"]
 _SCOPE_OPTIONS = ["all", "individual", "shared"]
 
 
-def _option_index(options: list[str], value: str) -> int:
-    return options.index(value) if value in options else 0
+@dataclass(frozen=True)
+class ExtractSettings:
+    runs: list[tuple[PersonaData, list[QAPair]]]
+    mask_strategy: MaskStrategy
+    max_questions: int
 
 
 def _remembered_select(
@@ -53,10 +53,11 @@ def _remembered_select(
     key: str,
     default: str = "all",
 ) -> str:
+    current = st.session_state.get(state_key, default)
     selected = st.selectbox(
         label,
         options=options,
-        index=_option_index(options, st.session_state.get(state_key, default)),
+        index=options.index(current) if current in options else 0,
         key=key,
     )
     st.session_state[state_key] = selected
@@ -66,21 +67,13 @@ def _remembered_select(
 def _build_run_plan(
     selected_variants: list[str],
     runs: list[tuple[PersonaData, list[QAPair]]],
-) -> list[tuple[PersonaData | None, list[QAPair], str]]:
-    """Expand selected variants × personas into one (persona, qa, variant) per call.
-
-    The baseline variant is run once across the first persona's QA pairs and
-    has no associated persona.
-    """
-    plan: list[tuple[PersonaData | None, list[QAPair], str]] = []
-    for variant in selected_variants:
-        if variant == BASELINE_PERSONA_ID:
-            _, qa_pairs = runs[0]
-            plan.append((None, qa_pairs, variant))
-        else:
-            for persona, qa_pairs in runs:
-                plan.append((persona, qa_pairs, variant))
-    return plan
+) -> list[tuple[PersonaData, list[QAPair], str]]:
+    """Cartesian product of personas × variants."""
+    return [(p, qa, v) for v in selected_variants for p, qa in runs]
+
+
+def _row_label(persona: PersonaData, variant: str) -> str:
+    return f"{persona.name} · {prompt_variant_label(variant)}"
 
 
 def _extract_widget_key(
@@ -89,103 +82,55 @@ def _extract_widget_key(
     return widget_key("extract", str(remote), model_name, dataset_source, suffix)
 
 
-_TOKEN_LEGEND = (
-    '<div style="display:flex;gap:12px;flex-wrap:wrap;font-size:0.8em;margin-bottom:8px">'
-    '<span style="background:#86efac;color:black;padding:1px 6px;border-radius:3px">masked</span>'
-    '<span style="color:#fde047;padding:1px 6px">question</span>'
-    '<span style="color:#22d3ee;padding:1px 6px">response</span>'
-    '<span style="color:#d946ef;font-weight:bold;padding:1px 6px">special</span>'
-    '<span style="color:#9ca3af;padding:1px 6px">template</span>'
-    "</div>"
-)
-
-_MAX_PREVIEW_SAMPLES = 3
-
-
-def _token_style(segment: TokenSegment) -> str:
-    style = {
-        "response": "color:#22d3ee",
-        "question": "color:#fde047",
-    }.get(segment.role, "color:#9ca3af")
-
-    if segment.is_special:
-        style = "color:#d946ef;font-weight:bold"
-    if segment.is_masked:
-        style = f"{style};background:#86efac;border-radius:2px;padding:0 1px"
-    return style
-
-
-def _render_sample_tokens_html(p, tokenizer, *, max_tokens: int = 200) -> str:
-    spans: list[str] = []
-    for segment in preview_token_segments(p, tokenizer, max_tokens=max_tokens):
-        spans.append(
-            f'<span style="{_token_style(segment)}">{html.escape(segment.text)}</span>'
+def _render_local_dataset_upload(dataset_source: str) -> None:
+    if dataset_source != "Local JSONL upload":
+        return
+    with st.expander("Local dataset upload", expanded=True):
+        st.file_uploader(
+            "personas.jsonl",
+            type=["jsonl"],
+            key="extract__personas_file",
+            help="Expected fields: id, persona, templated_view, biography_view",
+        )
+        st.file_uploader(
+            "qa.jsonl",
+            type=["jsonl"],
+            key="extract__qa_file",
+            help="Expected fields: id, qid, type, item_type, scope, question, answer",
         )
 
-    return (
-        '<pre style="white-space:pre-wrap;font-size:0.82em;line-height:1.5;'
-        "background:#0e1117;padding:8px 10px;border-radius:6px;"
-        'border:1px solid #333;margin:0">'
-        f"{''.join(spans)}</pre>"
-    )
-
-
-def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> None:
-    """Render the extraction tab."""
-
-    st.title("Extract")
-
-    if dataset_source == "Local JSONL upload":
-        with st.expander("Local dataset upload", expanded=True):
-            st.file_uploader(
-                "personas.jsonl",
-                type=["jsonl"],
-                key="extract__personas_file",
-                help="Expected fields: id, persona, templated_view, biography_view",
-            )
-            st.file_uploader(
-                "qa.jsonl",
-                type=["jsonl"],
-                key="extract__qa_file",
-                help="Expected fields: id, qid, type, item_type, scope, question, answer",
-            )
 
-    last_variants = st.session_state.get(
-        _LAST_VARIANTS_KEY, [*PERSONA_VARIANTS, BASELINE_PERSONA_ID]
-    )
-    default_persona_variants = [
-        v for v in last_variants if v in PERSONA_VARIANTS
-    ] or list(PERSONA_VARIANTS)
-    selected_persona_variants = st.multiselect(
+def _render_variant_controls(
+    *,
+    model_name: str,
+    remote: bool,
+    dataset_source: str,
+) -> tuple[list[str], bool] | None:
+    default_variants = st.session_state.get(_LAST_VARIANTS_KEY, list(PERSONA_VARIANTS))
+    selected_variants = st.multiselect(
         "Persona variants",
         options=PERSONA_VARIANTS,
-        default=default_persona_variants,
+        default=[v for v in default_variants if v in PERSONA_VARIANTS]
+        or list(PERSONA_VARIANTS),
         format_func=prompt_variant_label,
         key=_extract_widget_key(model_name, remote, dataset_source, "persona_variants"),
         help="Extract these variants for each selected persona.",
     )
     include_baseline = st.checkbox(
         "Extract Assistant baseline",
-        value=st.session_state.get(
-            _LAST_BASELINE_KEY,
-            BASELINE_PERSONA_ID in last_variants,
-        ),
+        value=st.session_state.get(_LAST_BASELINE_KEY, True),
         key=_extract_widget_key(model_name, remote, dataset_source, "baseline"),
-        help=(
-            "Extracts the persona-less Assistant prompt once using the first "
-            "selected persona's QA set."
-        ),
+        help="Also extract the Assistant baseline persona using the first persona's QA set.",
     )
-    selected_variants = [
-        *selected_persona_variants,
-        *([BASELINE_PERSONA_ID] if include_baseline else []),
-    ]
     st.session_state[_LAST_VARIANTS_KEY] = selected_variants
     st.session_state[_LAST_BASELINE_KEY] = include_baseline
     if not selected_variants:
-        st.info("Select at least one persona variant or enable the baseline.")
-        return
+        st.info("Select at least one persona variant.")
+        return None
+    return selected_variants, include_baseline
+
 
+def _load_qa_dataset_personas(dataset_source: str) -> tuple[object, list[PersonaData]] | None:
     try:
         dataset, dataset_status = load_dataset(
             dataset_source,
@@ -198,11 +143,11 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         st.info(
             "Upload both JSONL files or switch to the built-in SynthPersona source."
         )
-        return
+        return None
 
     if not getattr(dataset, "supports_qa", True):
         st.info("This dataset is persona-only for now. Use Chat to browse personas.")
-        return
+        return None
 
     personas = list(dataset)
     if not personas:
@@ -210,8 +155,17 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         st.info(
             "Try another dataset source or check that the personas file is not empty."
         )
-        return
+        return None
+    return dataset, personas
+
 
+def _render_persona_select(
+    *,
+    personas: list[PersonaData],
+    model_name: str,
+    remote: bool,
+    dataset_source: str,
+) -> list[PersonaData] | None:
     last_persona_ids: set[str] = set(st.session_state.get(_LAST_PERSONA_IDS_KEY, []))
     default_personas = [p for p in personas if p.id in last_persona_ids] or [
         personas[0]
@@ -227,185 +181,295 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
 
     if not selected_personas:
         st.info("Select at least one persona.")
-        return
+        return None
+    return selected_personas
 
-    with st.expander("Advanced", expanded=False):
-        st.caption("Filters")
 
-        col1, col2, col3 = st.columns(3)
-        with col1:
-            qa_type_select = _remembered_select(
-                "QA type",
-                _QA_TYPE_OPTIONS,
-                _LAST_QA_TYPE_KEY,
-                key=_extract_widget_key(
-                    model_name, remote, dataset_source, "qa_type_select"
-                ),
-            )
-            qa_filter_type: Literal["explicit", "implicit"] | None = (
-                cast(Literal["explicit", "implicit"], qa_type_select)
-                if qa_type_select in ("explicit", "implicit")
-                else None
-            )
-        with col2:
-            item_type_select = _remembered_select(
-                "Item type",
-                _ITEM_TYPE_OPTIONS,
-                _LAST_ITEM_TYPE_KEY,
-                key=_extract_widget_key(
-                    model_name, remote, dataset_source, "item_type_select"
-                ),
-            )
-            qa_filter_item_type: Literal["mcq", "frq"] | None = (
-                cast(Literal["mcq", "frq"], item_type_select)
-                if item_type_select in ("mcq", "frq")
-                else None
-            )
-        with col3:
-            scope_select = _remembered_select(
-                "Scope",
-                _SCOPE_OPTIONS,
-                _LAST_SCOPE_KEY,
-                key=_extract_widget_key(
-                    model_name,
-                    remote,
-                    dataset_source,
-                    "scope_select",
-                ),
-            )
-            qa_filter_scope: Literal["individual", "shared"] | None = (
-                cast(Literal["individual", "shared"], scope_select)
-                if scope_select in ("individual", "shared")
-                else None
-            )
+_TOKEN_LEGEND = (
+    '<div style="display:flex;gap:12px;flex-wrap:wrap;font-size:0.8em;margin-bottom:8px">'
+    '<span style="background:#86efac;color:black;padding:1px 6px;border-radius:3px">masked</span>'
+    '<span style="color:#fde047;padding:1px 6px">question</span>'
+    '<span style="color:#22d3ee;padding:1px 6px">response</span>'
+    '<span style="color:#d946ef;font-weight:bold;padding:1px 6px">special</span>'
+    '<span style="color:#9ca3af;padding:1px 6px">template</span>'
+    "</div>"
+)
 
-        st.caption("Extraction settings")
-        last_strategy = st.session_state.get(
-            _LAST_MASK_STRATEGY_KEY,
-            MaskStrategy.ANSWER_MEAN.value,
+_MAX_PREVIEW_SAMPLES = 3
+
+
+def _token_style(segment: TokenSegment) -> str:
+    style = {
+        "response": "color:#22d3ee",
+        "question": "color:#fde047",
+    }.get(segment.role, "color:#9ca3af")
+
+    if segment.is_special:
+        style = "color:#d946ef;font-weight:bold"
+    if segment.is_masked:
+        style = f"{style};background:#86efac;border-radius:2px;padding:0 1px"
+    return style
+
+
+def _render_sample_tokens_html(p, tokenizer, *, max_tokens: int = 200) -> str:
+    spans: list[str] = []
+    for segment in preview_token_segments(p, tokenizer, max_tokens=max_tokens):
+        spans.append(
+            f'<span style="{_token_style(segment)}">{html.escape(segment.text)}</span>'
         )
-        strategy_options = list(MaskStrategy)
-        mask_strategy = st.selectbox(
-            "Mask strategy",
-            options=strategy_options,
-            index=next(
-                (
-                    idx
-                    for idx, strategy in enumerate(strategy_options)
-                    if strategy.value == last_strategy
-                ),
-                0,
-            ),
-            format_func=lambda s: s.value.replace("_", " ").title(),
+
+    return (
+        '<pre style="white-space:pre-wrap;font-size:0.82em;line-height:1.5;'
+        "background:var(--secondary-background-color,rgba(127,127,127,0.08));"
+        "padding:8px 10px;border-radius:6px;"
+        'border:1px solid rgba(127,127,127,0.25);margin:0">'
+        f"{''.join(spans)}</pre>"
+    )
+
+
+def _render_filter_controls(
+    *,
+    model_name: str,
+    remote: bool,
+    dataset_source: str,
+) -> tuple[
+    Literal["explicit", "implicit"] | None,
+    Literal["mcq", "frq"] | None,
+    Literal["individual", "shared"] | None,
+]:
+    col1, col2, col3 = st.columns(3)
+    with col1:
+        qa_type_select = _remembered_select(
+            "QA type",
+            _QA_TYPE_OPTIONS,
+            _LAST_QA_TYPE_KEY,
+            key=_extract_widget_key(model_name, remote, dataset_source, "qa_type_select"),
+        )
+    with col2:
+        item_type_select = _remembered_select(
+            "Item type",
+            _ITEM_TYPE_OPTIONS,
+            _LAST_ITEM_TYPE_KEY,
             key=_extract_widget_key(
                 model_name,
                 remote,
                 dataset_source,
-                "mask_strategy",
+                "item_type_select",
             ),
-            help="Which tokens contribute to the averaged hidden state.",
         )
-        st.session_state[_LAST_MASK_STRATEGY_KEY] = mask_strategy.value
-
-        runs, skipped = [], []
-        for persona in selected_personas:
-            qa = list(
-                dataset.get_qa(
-                    persona.id,
-                    type=qa_filter_type,
-                    item_type=qa_filter_item_type,
-                    scope=qa_filter_scope,
-                )
-            )
-            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.")
-
-        if not runs:
-            st.info("No personas have matching QA pairs. Widen the filters.")
-            return
-
-        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=min(
-                max(st.session_state.get(_LAST_MAX_QUESTIONS_KEY, max_q), 1),
-                max_q,
-            ),
-            key=_extract_widget_key(
-                model_name, remote, dataset_source, "max_questions"
+    with col3:
+        scope_select = _remembered_select(
+            "Scope",
+            _SCOPE_OPTIONS,
+            _LAST_SCOPE_KEY,
+            key=_extract_widget_key(model_name, remote, dataset_source, "scope_select"),
+        )
+
+    return (
+        cast(Literal["explicit", "implicit"], qa_type_select)
+        if qa_type_select in ("explicit", "implicit")
+        else None,
+        cast(Literal["mcq", "frq"], item_type_select)
+        if item_type_select in ("mcq", "frq")
+        else None,
+        cast(Literal["individual", "shared"], scope_select)
+        if scope_select in ("individual", "shared")
+        else None,
+    )
+
+
+def _render_mask_strategy_select(
+    *,
+    model_name: str,
+    remote: bool,
+    dataset_source: str,
+) -> MaskStrategy:
+    last_strategy = st.session_state.get(
+        _LAST_MASK_STRATEGY_KEY,
+        MaskStrategy.ANSWER_MEAN.value,
+    )
+    strategy_options = list(MaskStrategy)
+    mask_strategy = st.selectbox(
+        "Mask strategy",
+        options=strategy_options,
+        index=next(
+            (
+                idx
+                for idx, strategy in enumerate(strategy_options)
+                if strategy.value == last_strategy
             ),
+            0,
+        ),
+        format_func=lambda s: s.value.replace("_", " ").title(),
+        key=_extract_widget_key(model_name, remote, dataset_source, "mask_strategy"),
+        help="Which tokens contribute to the averaged hidden state.",
+    )
+    st.session_state[_LAST_MASK_STRATEGY_KEY] = mask_strategy.value
+    return mask_strategy
+
+
+def _collect_runs(
+    *,
+    dataset,
+    selected_personas: list[PersonaData],
+    qa_filter_type: Literal["explicit", "implicit"] | None,
+    qa_filter_item_type: Literal["mcq", "frq"] | None,
+    qa_filter_scope: Literal["individual", "shared"] | None,
+) -> list[tuple[PersonaData, list[QAPair]]] | None:
+    runs, skipped = [], []
+    for persona in selected_personas:
+        qa = list(
+            dataset.get_qa(
+                persona.id,
+                type=qa_filter_type,
+                item_type=qa_filter_item_type,
+                scope=qa_filter_scope,
+            )
+        )
+        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.")
+
+    if not runs:
+        st.info("No personas have matching QA pairs. Widen the filters.")
+        return None
+    return runs
+
+
+def _render_max_questions(
+    *,
+    model_name: str,
+    remote: bool,
+    dataset_source: str,
+    runs: list[tuple[PersonaData, list[QAPair]]],
+) -> int:
+    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=min(max(st.session_state.get(_LAST_MAX_QUESTIONS_KEY, max_q), 1), max_q),
+        key=_extract_widget_key(model_name, remote, dataset_source, "max_questions"),
+    )
+    st.session_state[_LAST_MAX_QUESTIONS_KEY] = max_questions
+    return max_questions
+
+
+def _render_advanced_settings(
+    *,
+    dataset,
+    selected_personas: list[PersonaData],
+    model_name: str,
+    remote: bool,
+    dataset_source: str,
+) -> ExtractSettings | None:
+    with st.expander("Advanced", expanded=False):
+        st.caption("Filters")
+        qa_filter_type, qa_filter_item_type, qa_filter_scope = _render_filter_controls(
+            model_name=model_name,
+            remote=remote,
+            dataset_source=dataset_source,
         )
-        st.session_state[_LAST_MAX_QUESTIONS_KEY] = max_questions
 
+        st.caption("Extraction settings")
+        mask_strategy = _render_mask_strategy_select(
+            model_name=model_name,
+            remote=remote,
+            dataset_source=dataset_source,
+        )
+        runs = _collect_runs(
+            dataset=dataset,
+            selected_personas=selected_personas,
+            qa_filter_type=qa_filter_type,
+            qa_filter_item_type=qa_filter_item_type,
+            qa_filter_scope=qa_filter_scope,
+        )
+        if runs is None:
+            return None
+
+        max_questions = _render_max_questions(
+            model_name=model_name,
+            remote=remote,
+            dataset_source=dataset_source,
+            runs=runs,
+        )
+
+    return ExtractSettings(
+        runs=runs,
+        mask_strategy=mask_strategy,
+        max_questions=max_questions,
+    )
+
+
+def _render_extract_actions() -> tuple[bool, bool]:
     run_col, preview_col, _spacer = st.columns([1, 1, 4], gap="small")
     with run_col:
         run_clicked = st.button(
-            "Run extraction", type="primary", use_container_width=True
+            "Run extraction",
+            type="primary",
+            use_container_width=True,
         )
     with preview_col:
         preview_clicked = st.button("Preview tokens", use_container_width=True)
-
-    run_plan = _build_run_plan(selected_variants, runs)
-
-    def _row_label(persona: PersonaData | None, variant: str) -> str:
-        name = persona.name if persona is not None else BASELINE_PERSONA_NAME
-        return f"{name} · {prompt_variant_label(variant)}"
-
-    if preview_clicked:
-        with st.spinner("Loading tokenizer..."):
-            model = cached_model(model_name=model_name, remote=remote)
-        st.markdown(_TOKEN_LEGEND, unsafe_allow_html=True)
-        for persona, qa_pairs, variant in run_plan:
-            system_prompt = (
-                format_prompt()
-                if persona is None
-                else format_prompt(persona, variant)  # type: ignore[arg-type]
-            )
-            prepared = prepare_inputs_for_strategy(
-                tokenizer=model.tokenizer,
-                system_prompt=system_prompt,
-                qa_pairs=qa_pairs[:max_questions],
-                mask_strategy=mask_strategy,
-            )
-            st.caption(_row_label(persona, variant))
-            for i, p in enumerate(prepared[:_MAX_PREVIEW_SAMPLES]):
-                question = (
-                    p.question if len(p.question) <= 60 else p.question[:57] + "..."
+    return run_clicked, preview_clicked
+
+
+def _render_token_preview(
+    *,
+    remote: bool,
+    model_name: str,
+    run_plan: list[tuple[PersonaData, list[QAPair], str]],
+    settings: ExtractSettings,
+) -> None:
+    with st.spinner("Loading tokenizer..."):
+        model = cached_model(model_name=model_name)
+    st.markdown(_TOKEN_LEGEND, unsafe_allow_html=True)
+    for persona, qa_pairs, variant in run_plan:
+        system_prompt = format_prompt(persona, variant)  # type: ignore[arg-type]
+        prepared = prepare_inputs_for_strategy(
+            tokenizer=model.tokenizer,
+            system_prompt=system_prompt,
+            qa_pairs=qa_pairs[: settings.max_questions],
+            mask_strategy=settings.mask_strategy,
+        )
+        st.caption(_row_label(persona, variant))
+        for i, p in enumerate(prepared[:_MAX_PREVIEW_SAMPLES]):
+            question = p.question if len(p.question) <= 60 else p.question[:57] + "..."
+            seq_len = int(p.input_ids.shape[0])
+            masked = int(p.token_mask.sum())
+            label = f"sample {i} — {question}  (len={seq_len}, masked={masked})"
+            with st.expander(label):
+                st.markdown(
+                    _render_sample_tokens_html(p, model.tokenizer),
+                    unsafe_allow_html=True,
                 )
-                seq_len = int(p.input_ids.shape[0])
-                masked = int(p.token_mask.sum())
-                label = f"sample {i} — {question}  (len={seq_len}, masked={masked})"
-                with st.expander(label):
-                    st.markdown(
-                        _render_sample_tokens_html(p, model.tokenizer),
-                        unsafe_allow_html=True,
-                    )
-            if len(prepared) > _MAX_PREVIEW_SAMPLES:
-                remaining = len(prepared) - _MAX_PREVIEW_SAMPLES
-                st.caption(f"… and {remaining} more sample(s) not shown.")
-        return
-
-    if not run_clicked:
-        return
-
+        if len(prepared) > _MAX_PREVIEW_SAMPLES:
+            remaining = len(prepared) - _MAX_PREVIEW_SAMPLES
+            st.caption(f"… and {remaining} more sample(s) not shown.")
+
+
+def _run_extraction_plan(
+    *,
+    remote: bool,
+    model_name: str,
+    run_plan: list[tuple[PersonaData, list[QAPair], str]],
+    settings: ExtractSettings,
+) -> None:
     status_box = st.empty()
     status_box.info("Extraction in progress...")
     progress = st.progress(0, text="Preparing extraction...")
-    ndif_status_box = st.empty()  # shows live NDIF job status when remote=True
+    ndif_status_box = st.empty()
 
     def _on_ndif_status(job_id: str, status_name: str, description: str) -> None:
         icon = NDIF_STATUS_ICONS.get(status_name, "•")
         ndif_status_box.caption(f"{icon} `{job_id}` **{status_name}** — {description}")
 
     with st.spinner("Loading model..."):
-        model = cached_model(model_name=model_name, remote=remote)
+        model = cached_model(model_name=model_name)
 
     try:
         total_steps = len(run_plan)
@@ -419,10 +483,10 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
                 run_extraction(
                     model=model,
                     model_name=model_name,
-                    qa_pairs=qa_pairs[:max_questions],
+                    qa_pairs=qa_pairs[: settings.max_questions],
                     variants=(variant,),
                     persona=persona,
-                    mask_strategy=mask_strategy,
+                    mask_strategy=settings.mask_strategy,
                     remote=remote,
                     on_status=_on_ndif_status if remote else None,
                 )
@@ -444,3 +508,71 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
             f"- **{result.persona_name}** · {prompt_variant_label(result.variant)}: "
             f"{result.n_questions} questions"
         )
+
+
+def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> None:
+    """Render the extraction tab."""
+
+    st.title("Extract")
+    st.caption("Extract per-persona activation vectors from QA pairs.")
+
+    _render_local_dataset_upload(dataset_source)
+    variant_choice = _render_variant_controls(
+        model_name=model_name,
+        remote=remote,
+        dataset_source=dataset_source,
+    )
+    if variant_choice is None:
+        return
+    selected_variants, include_baseline = variant_choice
+
+    loaded = _load_qa_dataset_personas(dataset_source)
+    if loaded is None:
+        return
+    dataset, personas = loaded
+
+    selected_personas = _render_persona_select(
+        personas=personas,
+        model_name=model_name,
+        remote=remote,
+        dataset_source=dataset_source,
+    )
+    if selected_personas is None:
+        return
+
+    settings = _render_advanced_settings(
+        dataset=dataset,
+        selected_personas=selected_personas,
+        model_name=model_name,
+        remote=remote,
+        dataset_source=dataset_source,
+    )
+    if settings is None:
+        return
+
+    runs = list(settings.runs)
+    baseline = getattr(dataset, "baseline", None)
+    if include_baseline and baseline is not None and runs:
+        runs.append((baseline, runs[0][1]))
+
+    run_clicked, preview_clicked = _render_extract_actions()
+    run_plan = _build_run_plan(selected_variants, runs)
+
+    if preview_clicked:
+        _render_token_preview(
+            remote=remote,
+            model_name=model_name,
+            run_plan=run_plan,
+            settings=settings,
+        )
+        return
+
+    if not run_clicked:
+        return
+
+    _run_extraction_plan(
+        remote=remote,
+        model_name=model_name,
+        run_plan=run_plan,
+        settings=settings,
+    )

tabs/probe_ui.py

@@ -16,13 +16,6 @@ from utils.probes import (
 from utils.runtime import cached_model
 
 
-def _token_button_label(index: int, token: str) -> str:
-    display = token.encode("unicode_escape").decode("ascii") or "<empty>"
-    if len(display) > 18:
-        display = display[:15] + "..."
-    return f"{index}: {display}"
-
-
 def _render_probe_results(result: ProbeRunResult, probe: LoadedProbe) -> None:
     top_k = min(5, int(result.probabilities.numel()))
     if top_k == 0:
@@ -95,55 +88,27 @@ def _load_probe_from_controls(context_key: str) -> LoadedProbe | None:
     return load_probe(repo_id.strip(), selected_file)
 
 
-def _render_token_buttons(trace: ConversationTrace, context_key: str) -> int:
-    selected_key = widget_key(
-        context_key,
-        "probe_selected_token",
-        trace.prompt_hash[:12],
-    )
-    selected = int(st.session_state.get(selected_key, trace.n_tokens - 1))
-    selected = max(0, min(selected, trace.n_tokens - 1))
-
-    window_size = st.slider(
-        "Token window",
-        min_value=8,
-        max_value=min(96, max(8, trace.n_tokens)),
-        value=min(32, max(8, trace.n_tokens)),
-        step=8,
-        key=widget_key(context_key, "probe_token_window", trace.prompt_hash[:12]),
-    )
-    center = st.slider(
-        "Window center",
+def _render_token_picker(trace: ConversationTrace, context_key: str) -> int:
+    selected = st.slider(
+        "Token index",
         min_value=0,
         max_value=trace.n_tokens - 1,
-        value=selected,
-        key=widget_key(context_key, "probe_token_center", trace.prompt_hash[:12]),
+        value=trace.n_tokens - 1,
+        key=widget_key(context_key, "probe_selected_token", trace.prompt_hash[:12]),
     )
-    start = max(0, center - window_size // 2)
-    end = min(trace.n_tokens, start + window_size)
-    start = max(0, end - window_size)
-
-    cols = st.columns(8)
-    for offset, token_index in enumerate(range(start, end)):
-        col = cols[offset % len(cols)]
-        token = trace.tokens[token_index]
-        if col.button(
-            _token_button_label(token_index, token),
-            key=widget_key(
-                context_key,
-                "probe_token",
-                trace.prompt_hash[:12],
-                str(token_index),
-            ),
-            type="primary" if token_index == selected else "secondary",
-            help=token.encode("unicode_escape").decode("ascii"),
-        ):
-            selected = token_index
-            st.session_state[selected_key] = token_index
 
-    st.caption(
-        f"Selected token {selected}: "
-        f"`{trace.tokens[selected].encode('unicode_escape').decode('ascii')}`"
+    window = 8
+    start = max(0, selected - window)
+    end = min(trace.n_tokens, selected + window + 1)
+    parts: list[str] = []
+    for i in range(start, end):
+        token_repr = trace.tokens[i].encode("unicode_escape").decode("ascii") or "·"
+        parts.append(f"**[{token_repr}]**" if i == selected else token_repr)
+    st.markdown(
+        f"<div style='font-family:ui-monospace,monospace;font-size:0.85em;"
+        f"line-height:1.6;background:rgba(127,127,127,0.08);padding:6px 10px;"
+        f"border-radius:4px;'>{' '.join(parts)}</div>",
+        unsafe_allow_html=True,
     )
     return selected
 
@@ -163,6 +128,128 @@ def _model_dimensions(model: object) -> tuple[int, int]:
     return int(hidden_size), int(num_layers)
 
 
+def _load_model_with_dimensions(model_name: str) -> tuple[object, int, int] | None:
+    with st.spinner("Loading model metadata..."):
+        model = cached_model(model_name=model_name)
+    try:
+        hidden_size, num_layers = _model_dimensions(model)
+    except Exception as exc:
+        st.error(str(exc))
+        return None
+    return model, hidden_size, num_layers
+
+
+def _select_probe_target(
+    *,
+    probe: LoadedProbe,
+    context_key: str,
+    num_layers: int,
+) -> tuple[int, str]:
+    layer = probe.layer
+    if layer is None:
+        layer = int(
+            st.number_input(
+                "Layer",
+                min_value=0,
+                max_value=max(0, num_layers - 1),
+                value=min(15, max(0, num_layers - 1)),
+                step=1,
+                key=widget_key(context_key, "probe_layer"),
+            )
+        )
+
+    location = probe.location
+    if location is None:
+        location = st.selectbox(
+            "Activation location",
+            options=("post_reasoning", "pre_reasoning"),
+            key=widget_key(context_key, "probe_location"),
+        )
+    return layer, location
+
+
+def _probe_target_is_valid(
+    *,
+    probe: LoadedProbe,
+    layer: int,
+    num_layers: int,
+    hidden_size: int,
+) -> bool:
+    if not 0 <= layer < num_layers:
+        st.error(f"Probe layer {layer} is outside the model's {num_layers} layers.")
+        return False
+    if probe.input_dim != hidden_size:
+        st.warning(
+            "This probe input dim does not match a single-token activation "
+            "for the active model."
+        )
+        return False
+    return True
+
+
+def _trace_requested(context_key: str) -> bool:
+    trace_key = widget_key(context_key, "probe_trace_enabled")
+    if st.button(
+        "Trace conversation",
+        key=widget_key(context_key, "probe_trace"),
+        use_container_width=True,
+    ):
+        st.session_state[trace_key] = True
+    return bool(st.session_state.get(trace_key, False))
+
+
+def _trace_active_conversation(
+    *,
+    model: object,
+    model_name: str,
+    remote: bool,
+    active_system_prompt: str | None,
+    chat_state: dict[str, object],
+    layer: int,
+    location: str,
+) -> ConversationTrace | None:
+    messages = build_chat_messages(active_system_prompt, chat_state["messages"])
+    with st.spinner("Tracing conversation..."):
+        trace = trace_conversation(
+            model=model,
+            model_name=model_name,
+            messages=messages,
+            layer=layer,
+            location=location,
+            remote=remote,
+        )
+
+    st.caption(
+        f"Cached {trace.n_tokens} tokens from layer {trace.layer}; "
+        f"prompt hash `{trace.prompt_hash[:10]}`"
+    )
+    if trace.n_tokens == 0:
+        st.warning("The traced conversation produced no tokens.")
+        return None
+    return trace
+
+
+def _run_probe_on_selected_token(
+    *,
+    trace: ConversationTrace,
+    context_key: str,
+    probe: LoadedProbe,
+) -> None:
+    selected_token = _render_token_picker(trace, context_key)
+    try:
+        vector = vectorize_token(trace, token_index=selected_token)
+        result = probe.run(vector.vector)
+    except Exception as exc:
+        st.error(f"Probe execution failed: {exc}")
+        return
+
+    st.caption(
+        f"Vectorization {vector.mode}; token {vector.token_index}; "
+        f"vector dim {int(vector.vector.shape[0])}"
+    )
+    _render_probe_results(result, probe)
+
+
 def render_probe_inspector(
     *,
     context_key: str,
@@ -188,88 +275,40 @@ def render_probe_inspector(
         if probe is None:
             return
 
-        with st.spinner("Loading model metadata..."):
-            model = cached_model(model_name=model_name, remote=remote)
-        try:
-            hidden_size, num_layers = _model_dimensions(model)
-        except Exception as exc:
-            st.error(str(exc))
+        loaded = _load_model_with_dimensions(model_name)
+        if loaded is None:
             return
+        model, hidden_size, num_layers = loaded
 
-        layer = probe.layer
-        if layer is None:
-            layer = int(
-                st.number_input(
-                    "Layer",
-                    min_value=0,
-                    max_value=max(0, num_layers - 1),
-                    value=min(15, max(0, num_layers - 1)),
-                    step=1,
-                    key=widget_key(context_key, "probe_layer"),
-                )
-            )
-
-        location = probe.location
-        if location is None:
-            location = st.selectbox(
-                "Activation location",
-                options=("post_reasoning", "pre_reasoning"),
-                key=widget_key(context_key, "probe_location"),
-            )
-
+        layer, location = _select_probe_target(
+            probe=probe,
+            context_key=context_key,
+            num_layers=num_layers,
+        )
         st.caption(
             f"Probe layer {layer}; {location}; input dim {probe.input_dim}; "
             f"model hidden size {hidden_size}"
         )
-        if not 0 <= layer < num_layers:
-            st.error(f"Probe layer {layer} is outside the model's {num_layers} layers.")
-            return
-        if probe.input_dim != hidden_size:
-            st.warning(
-                "This probe input dim does not match a single-token activation "
-                "for the active model."
-            )
-            return
-
-        trace_key = widget_key(context_key, "probe_trace_enabled")
-        if st.button(
-            "Trace conversation",
-            key=widget_key(context_key, "probe_trace"),
-            use_container_width=True,
+        if not _probe_target_is_valid(
+            probe=probe,
+            layer=layer,
+            num_layers=num_layers,
+            hidden_size=hidden_size,
         ):
-            st.session_state[trace_key] = True
-        if not st.session_state.get(trace_key, False):
-            return
-
-        messages = build_chat_messages(active_system_prompt, chat_state["messages"])
-        with st.spinner("Tracing conversation..."):
-            trace = trace_conversation(
-                model=model,
-                model_name=model_name,
-                messages=messages,
-                layer=layer,
-                location=location,
-                remote=remote,
-            )
-
-        st.caption(
-            f"Cached {trace.n_tokens} tokens from layer {trace.layer}; "
-            f"prompt hash `{trace.prompt_hash[:10]}`"
-        )
-        if trace.n_tokens == 0:
-            st.warning("The traced conversation produced no tokens.")
             return
 
-        selected_token = _render_token_buttons(trace, context_key)
-        try:
-            vector = vectorize_token(trace, token_index=selected_token)
-            result = probe.run(vector.vector)
-        except Exception as exc:
-            st.error(f"Probe execution failed: {exc}")
+        if not _trace_requested(context_key):
             return
 
-        st.caption(
-            f"Vectorization {vector.mode}; token {vector.token_index}; "
-            f"vector dim {int(vector.vector.shape[0])}"
+        trace = _trace_active_conversation(
+            model=model,
+            model_name=model_name,
+            remote=remote,
+            active_system_prompt=active_system_prompt,
+            chat_state=chat_state,
+            layer=layer,
+            location=location,
         )
-        _render_probe_results(result, probe)
+        if trace is None:
+            return
+        _run_probe_on_selected_token(trace=trace, context_key=context_key, probe=probe)

utils/contrast.py

@@ -11,7 +11,6 @@ Negative (blue) → token is more characteristic of persona B.
 Near-zero (gray) → both personas would emit this token with similar likelihood.
 """
 
-import logging
 from dataclasses import dataclass
 from html import escape
 
@@ -20,8 +19,6 @@ from nnterp import StandardizedTransformer
 
 from utils.chat import format_generation_prompt
 
-logger = logging.getLogger(__name__)
-
 
 @dataclass
 class TokenContrast:
@@ -73,28 +70,18 @@ def _strip_special_ids(
     return ids[keep], keep
 
 
-def _prepare_trace_text(
+def _prepare_trace_input_ids(
     tokenizer: object,
     context_messages: list[dict[str, str]],
     response_ids: torch.Tensor,
-) -> tuple[str, int, int]:
-    """Build the trace text and return ``(full_text, n_ctx, n_resp)``."""
+) -> tuple[torch.Tensor, int, int]:
+    """Build exact trace input ids and return ``(input_ids, n_ctx, n_resp)``."""
     context_prompt, _ = format_generation_prompt(context_messages, tokenizer)
     context_ids = tokenizer(context_prompt, return_tensors="pt").input_ids[0]
-    response_text = _decode_ids(tokenizer, response_ids.tolist())
-    full_text = context_prompt + response_text
-    full_ids = tokenizer(full_text, return_tensors="pt").input_ids[0]
-    expected_ids = torch.cat([context_ids, response_ids.cpu()])
-    if full_ids.tolist() != expected_ids.tolist():
-        logger.warning(
-            "contrast trace text did not round-trip to the expected token ids "
-            "(expected %d tokens, got %d); contrast scores may be slightly misaligned",
-            len(expected_ids),
-            len(full_ids),
-        )
+    input_ids = torch.cat([context_ids.cpu(), response_ids.detach().cpu()])
     n_ctx = len(context_ids)
     n_resp = len(response_ids)
-    return full_text, n_ctx, n_resp
+    return input_ids, n_ctx, n_resp
 
 
 def _build_contrast(
@@ -122,8 +109,8 @@ def _token_display(tokenizer: object, token_id: int) -> str:
     return _decode_ids(tokenizer, [token_id])
 
 
-# Each spec: (key, full_text, n_ctx, n_resp, target_ids).
-PassSpec = tuple[str, str, int, int, torch.Tensor]
+# Each spec: (key, input_ids, n_ctx, n_resp, target_ids).
+PassSpec = tuple[str, torch.Tensor, int, int, torch.Tensor]
 
 
 def _score_passes(
@@ -140,12 +127,12 @@ def _score_passes(
     """
 
     def _score_pass(
-        full_text: str,
+        input_ids: torch.Tensor,
         n_ctx: int,
         n_resp: int,
         target_ids: torch.Tensor,
     ) -> torch.Tensor:
-        with torch.no_grad(), model.trace(full_text, remote=remote):
+        with torch.no_grad(), model.trace(input_ids, remote=remote):
             # logit at position i predicts token i+1, so response token j
             # (at full-text position n_ctx+j) uses logit at n_ctx+j-1.
             resp_logits = model.logits[0, n_ctx - 1 : n_ctx - 1 + n_resp].float()
@@ -163,8 +150,8 @@ def _score_passes(
         return out.detach().cpu()
 
     return {
-        key: _score_pass(full_text, n_ctx, n_resp, target_ids)
-        for key, full_text, n_ctx, n_resp, target_ids in specs
+        key: _score_pass(input_ids, n_ctx, n_resp, target_ids)
+        for key, input_ids, n_ctx, n_resp, target_ids in specs
     }
 
 
@@ -176,11 +163,13 @@ def _specs_for_response(
     prefix: str,
 ) -> list[PassSpec]:
     """Build the (under_a, under_b) pass specs for a single response."""
-    text_a, n_ctx_a, n_resp = _prepare_trace_text(tokenizer, context_a, response_ids)
-    text_b, n_ctx_b, _ = _prepare_trace_text(tokenizer, context_b, response_ids)
+    input_a, n_ctx_a, n_resp = _prepare_trace_input_ids(
+        tokenizer, context_a, response_ids
+    )
+    input_b, n_ctx_b, _ = _prepare_trace_input_ids(tokenizer, context_b, response_ids)
     return [
-        (f"{prefix}_under_a", text_a, n_ctx_a, n_resp, response_ids),
-        (f"{prefix}_under_b", text_b, n_ctx_b, n_resp, response_ids),
+        (f"{prefix}_under_a", input_a, n_ctx_a, n_resp, response_ids),
+        (f"{prefix}_under_b", input_b, n_ctx_b, n_resp, response_ids),
     ]
 
 

utils/runtime.py

@@ -1,8 +1,56 @@
+import json
 import logging
+from collections.abc import Iterable
 
 import streamlit as st
 
 logger = logging.getLogger(__name__)
+_LANGUAGE_MODEL_CLASSES = {"LanguageModel", "StandardizedTransformer"}
+_EXPECTED_NDIF_STATES = {"RUNNING", "NOT DEPLOYED", "DEPLOYING", "DELETING"}
+
+
+def _iter_deployments(raw: object) -> Iterable[dict]:
+    if not isinstance(raw, dict):
+        return ()
+    deployments = raw.get("deployments", {})
+    if not isinstance(deployments, dict):
+        return ()
+    return (value for value in deployments.values() if isinstance(value, dict))
+
+
+def _is_visible_deployment(deployment: dict) -> bool:
+    return deployment.get("deployment_level") in {"HOT", "WARM"} or (
+        "schedule" in deployment
+    )
+
+
+def _repo_id_from_model_key(model_key: str) -> str:
+    try:
+        repo_id = json.loads(model_key.split(":", 1)[-1]).get("repo_id")
+    except Exception:
+        return model_key
+    return repo_id if isinstance(repo_id, str) else model_key
+
+
+def _running_language_model(deployment: dict) -> str | None:
+    if not _is_visible_deployment(deployment):
+        return None
+
+    model_key = deployment.get("model_key", "")
+    model_class = model_key.split(":", 1)[0].split(".")[-1]
+    if model_class not in _LANGUAGE_MODEL_CLASSES:
+        return None
+    if deployment.get("application_state", "NOT DEPLOYED") != "RUNNING":
+        return None
+    return _repo_id_from_model_key(model_key)
+
+
+def _unexpected_state(deployment: dict) -> tuple[str, str] | None:
+    state = deployment.get("application_state", "NOT DEPLOYED")
+    if state in _EXPECTED_NDIF_STATES:
+        return None
+    model_key = deployment.get("model_key", "")
+    return _repo_id_from_model_key(model_key), state
 
 
 @st.cache_data(show_spinner=False, ttl=30)
@@ -16,8 +64,6 @@ def list_remote_models() -> list[str]:
     the whole response. See nnsight 0.6.3 ``ndif.py::status``.
     """
 
-    import json
-
     import nnsight
 
     try:
@@ -29,32 +75,11 @@ def list_remote_models() -> list[str]:
     model_names: list[str] = []
     bad_states: list[tuple[str, str]] = []  # (repo_id_or_key, application_state)
 
-    for value in (raw or {}).get("deployments", {}).values():
-        if not isinstance(value, dict):
-            continue
-        if (
-            value.get("deployment_level") not in {"HOT", "WARM"}
-            and "schedule" not in value
-        ):
-            continue
-
-        model_key = value.get("model_key", "")
-        model_class = model_key.split(":", 1)[0].split(".")[-1]
-        try:
-            repo_id = json.loads(model_key.split(":", 1)[-1]).get("repo_id")
-        except Exception:
-            repo_id = model_key
-
-        state = value.get("application_state", "NOT DEPLOYED")
-        if state not in {"RUNNING", "NOT DEPLOYED", "DEPLOYING", "DELETING"}:
-            bad_states.append((repo_id or model_key, state))
-
-        if model_class not in {"LanguageModel", "StandardizedTransformer"}:
-            continue
-        if state != "RUNNING":
-            continue
-        if isinstance(repo_id, str):
-            model_names.append(repo_id)
+    for deployment in _iter_deployments(raw):
+        if bad_state := _unexpected_state(deployment):
+            bad_states.append(bad_state)
+        if model_name := _running_language_model(deployment):
+            model_names.append(model_name)
 
     if bad_states:
         logger.warning(
@@ -67,27 +92,17 @@ def list_remote_models() -> list[str]:
 
 
 @st.cache_resource(show_spinner=False, max_entries=1)
-def _cached_model_by_name(model_name: str):
+def cached_model(model_name: str):
     """Load and cache a standardized nnterp model.
 
     Streamlit reruns this app on every interaction, so caching keeps one loaded
     model instance per model name instead of reloading weights on every widget
-    change.
+    change. ``remote`` is intentionally not part of the cache key: it matters
+    at generation/trace time, but the current ``StandardizedTransformer``
+    constructor ignores it, and excluding it avoids loading duplicate local
+    model objects when toggling NDIF.
     """
 
     from nnterp import StandardizedTransformer
 
-    # The remote constructor path is currently unstable for this model wrapper.
-    # return StandardizedTransformer(model_name, remote=remote, check_renaming=False)
     return StandardizedTransformer(model_name)
-
-
-def cached_model(model_name: str, remote: bool):
-    """Return the cached model for ``model_name``.
-
-    ``remote`` still matters at generation/trace time, but the current
-    ``StandardizedTransformer`` constructor ignores it. Keeping it out of the
-    cache key avoids loading duplicate local model objects when toggling NDIF.
-    """
-
-    return _cached_model_by_name(model_name)