Updated to new chat edit and comparison

- Cleaned up chat
- Improved flow and UI
- Added comparison
- Fixed remote to make it more robus

README.md

@@ -29,13 +29,17 @@ A web app built on top of [persona-vectors](../persona-vectors) that provides th
 persona-ui/
 ├── app.py                   # Main entry point (Streamlit)
 ├── state.py                 # Session state management (chat history, KV cache)
+├── scripts/
+│   └── oracle_probe.py      # Notebook-style activation oracle script
 ├── tabs/
 │   ├── chat.py              # Chat tab
 │   ├── compare.py           # Activation comparison tab
+│   ├── compare_chat.py      # Side-by-side chat comparison mode
 │   └── extract.py           # Extraction tab
 └── utils/
     ├── chat.py              # Chat generation logic
     ├── chat_export.py       # Export chat logs to JSON
+    ├── contrast.py          # Contrastive token log-prob coloring
     ├── datasets.py          # Dataset loader wrapper
     ├── helpers.py           # UI labels and slug helpers
     └── runtime.py           # Model caching and NDIF queries
@@ -121,8 +125,8 @@ artifacts/
 ├── activations/<model_dir>/<prompt_variant>/<persona_id>/
 │   ├── activations.safetensors
 │   └── metadata.json   # used for persona names and layer counts
-└── chats/<model_dir>/<prompt_variant>/
+└── chats/<model_dir>/<persona_id>/
     └── <export>.json
 ```
 
-`<model_dir>` is the model name with `/` replaced by `__` (e.g. `google__gemma-2-9b-it`).
+`<model_dir>` is the model name with `/` replaced by `__` (e.g. `google__gemma-2-9b-it`). Chat exports still store `dataset_source` in the JSON payload.

state.py

@@ -9,7 +9,7 @@ def chat_session_key(model_name: str, dataset_source: str) -> str:
     return f"{_CHAT_STATE_PREFIX}{model_name}::{dataset_source}"
 
 
-def _default_chat_state() -> dict[str, object]:
+def default_chat_state() -> dict[str, object]:
     return {
         "messages": [],
         "persona_id": None,
@@ -18,6 +18,22 @@ def _default_chat_state() -> dict[str, object]:
     }
 
 
+def reset_chat_context_state(
+    state: dict[str, object],
+    persona_id: str,
+    prompt_mode: str,
+    *ui_keys: str,
+) -> None:
+    """Reset one chat context and clear any related widget state."""
+
+    state["messages"] = []
+    state["past_key_values"] = None
+    state["persona_id"] = persona_id
+    state["prompt_mode"] = prompt_mode
+    for key in ui_keys:
+        st.session_state.pop(key, None)
+
+
 def _evict_inactive_kv_caches(active_key: str) -> None:
     """Drop past_key_values from every chat context except the active one."""
 
@@ -40,22 +56,12 @@ def get_chat_state(
     key = chat_session_key(model_name, dataset_source)
     state = st.session_state.get(key)
     if state is None:
-        state = _default_chat_state()
+        state = default_chat_state()
         st.session_state[key] = state
     else:
-        for default_key, default_value in _default_chat_state().items():
+        for default_key, default_value in default_chat_state().items():
             state.setdefault(default_key, default_value)
     _evict_inactive_kv_caches(key)
     if remote and state.get("past_key_values") is not None:
         state["past_key_values"] = None
     return state
-
-
-def reset_chat_state(model_name: str, dataset_source: str) -> None:
-    """Reset chat history and cache for the active context."""
-
-    key = chat_session_key(model_name, dataset_source)
-    if key in st.session_state:
-        state = st.session_state[key]
-        state["messages"] = []
-        state["past_key_values"] = None

tabs/chat.py

@@ -1,72 +1,109 @@
-from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
 import streamlit as st
 from persona_data.synth_persona import PersonaData
 
-from state import (
-    _default_chat_state,
-    chat_session_key,
-    get_chat_state,
-    reset_chat_state,
-)
+from state import chat_session_key, get_chat_state, reset_chat_context_state
 from utils.chat import ChatReply, generate_chat_reply, resolve_system_prompt
 from utils.chat_export import save_chat_export
+from utils.contrast import TokenContrast, render_contrast_html
 from utils.datasets import load_dataset
 from utils.helpers import (
     MODE_LABEL_TO_KEY,
     MODE_LABELS,
     VARIANT_LABELS,
-    VISIBLE_MESSAGE_COUNT,
     persona_label,
     widget_key,
 )
 from utils.runtime import cached_model
 
-COLLAPSED_MESSAGE_CHAR_LIMIT = 500
-
 
 def _render_collapsible_markdown(content: str) -> None:
-    if len(content) <= COLLAPSED_MESSAGE_CHAR_LIMIT:
-        st.markdown(content)
-        return
+    st.markdown(content)
 
-    with st.expander(f"Show full text ({len(content)} chars)", expanded=False):
-        st.markdown(content)
 
+# ── Dialogs ───────────────────────────────────────────────────────────────────
 
-def _render_chat_message(message: dict[str, str]) -> None:
-    if not message.get("content"):
-        return
-    with st.container(border=True):
-        st.caption(message["role"])
-        _render_collapsible_markdown(message["content"])
 
+@st.dialog("Edit", width="medium")
+def _open_edit_dialog(
+    *,
+    msg_index: int,
+    messages: list[dict[str, str]],
+    chat_state: dict[str, object],
+    pending_key: str,
+) -> None:
+    message = messages[msg_index]
+    role = message["role"]
+
+    n_after = len(messages) - msg_index - 1
+    st.caption(
+        f"**{role}**"
+        + (
+            f" — {n_after} subsequent {'message' if n_after == 1 else 'messages'} will be cleared"
+            if n_after > 0
+            else ""
+        )
+    )
 
-def _render_inline_system_prompt(
-    prompt_key: str,
-    prompt_mode: str,
-    active_system_prompt: str | None,
-    height: int = 200,
-) -> str | None:
-    """Render the system prompt as an always-editable text area at the top of the chat."""
-    if prompt_mode == "empty":
-        return active_system_prompt
+    new_content = st.text_area(
+        "Content",
+        value=message["content"],
+        height=320,
+        label_visibility="collapsed",
+    )
 
-    if prompt_key not in st.session_state:
-        st.session_state[prompt_key] = active_system_prompt or ""
+    save_col, cancel_col = st.columns(2)
+    with save_col:
+        if st.button("Save", type="primary", use_container_width=True):
+            messages[msg_index]["content"] = new_content
+            messages[msg_index].pop("_contrast", None)
+            if role == "assistant":
+                messages[msg_index]["_needs_contrast"] = True
+            del messages[msg_index + 1 :]
+            chat_state["past_key_values"] = None
+            if role == "user":
+                st.session_state[pending_key] = True
+            st.rerun()
+    with cancel_col:
+        if st.button("Cancel", use_container_width=True):
+            st.rerun()
 
-    with st.container(border=True):
-        st.caption("System prompt")
-        st.text_area(
-            "system_prompt_edit",
-            value=st.session_state[prompt_key],
-            height=height,
-            label_visibility="collapsed",
-            key=prompt_key,
-        )
 
-    return st.session_state.get(prompt_key) or None
+@st.dialog("Edit system prompt", width="large")
+def _open_system_prompt_dialog(*, prompt_key: str, current_value: str) -> None:
+    new_value = st.text_area(
+        "System prompt",
+        value=current_value,
+        height=320,
+        label_visibility="collapsed",
+    )
+    save_col, cancel_col = st.columns(2)
+    with save_col:
+        if st.button("Save", type="primary", use_container_width=True):
+            st.session_state[prompt_key] = new_value
+            st.rerun()
+    with cancel_col:
+        if st.button("Cancel", use_container_width=True):
+            st.rerun()
+
+
+# ── Message renderers ─────────────────────────────────────────────────────────
+
+
+def _render_chat_message(
+    message: dict[str, str],
+    show_contrast: bool = False,
+) -> None:
+    if not message.get("content"):
+        return
+    role = message["role"]
+    tc: TokenContrast | None = message.get("_contrast") if show_contrast else None
+    with st.chat_message(role):
+        if tc is not None:
+            st.html(render_contrast_html(tc))
+        else:
+            _render_collapsible_markdown(message["content"])
 
 
 def _render_editable_message(
@@ -76,63 +113,50 @@ def _render_editable_message(
     chat_state: dict[str, object],
     edit_key: str,
     pending_key: str,
+    show_contrast: bool = False,
+    column_ratio: tuple[int, int] = (25, 1),
 ) -> None:
-    """Render a single message with an inline edit button."""
     if not message.get("content"):
         return
+    role = message["role"]
+    tc: TokenContrast | None = message.get("_contrast") if show_contrast else None
 
-    is_editing = st.session_state.get(edit_key) == msg_index
-
-    with st.container(border=True):
-        st.caption(message["role"])
-        if is_editing:
-            new_content = st.text_area(
-                "Edit",
-                value=message["content"],
-                height=100,
-                label_visibility="collapsed",
-                key=f"{edit_key}_msg_{msg_index}",
+    msg_col, edit_col = st.columns(
+        list(column_ratio), gap="xsmall", vertical_alignment="center"
+    )
+    with msg_col:
+        with st.chat_message(role):
+            if tc is not None:
+                st.html(render_contrast_html(tc))
+            else:
+                _render_collapsible_markdown(message["content"])
+    with edit_col:
+        if st.button(
+            "", icon=":material/edit:", key=f"{edit_key}_edit_{msg_index}", help="Edit"
+        ):
+            _open_edit_dialog(
+                msg_index=msg_index,
+                messages=messages,
+                chat_state=chat_state,
+                pending_key=pending_key,
             )
-            c1, c2 = st.columns(2)
-            with c1:
-                if st.button(
-                    "Save", key=f"{edit_key}_msg_save_{msg_index}", type="primary"
-                ):
-                    messages[msg_index]["content"] = new_content
-                    del messages[msg_index + 1 :]
-                    chat_state["past_key_values"] = None
-                    st.session_state[edit_key] = None
-                    if message["role"] == "user":
-                        st.session_state[pending_key] = True
-                    st.rerun()
-            with c2:
-                if st.button("Cancel", key=f"{edit_key}_msg_cancel_{msg_index}"):
-                    st.session_state[edit_key] = None
-                    st.rerun()
-        else:
-            st.markdown(message["content"])
-            if st.button("Edit", key=f"{edit_key}_msg_edit_{msg_index}"):
-                st.session_state[edit_key] = msg_index
-                st.rerun()
-
 
-def _clear_chat_ui_state(*keys: str) -> None:
-    for key in keys:
-        st.session_state.pop(key, None)
 
-
-def _reset_single_chat_context(
-    model_name: str,
-    dataset_source: str,
-    chat_state: dict[str, object],
-    persona_id: str,
+def _render_system_prompt(
+    prompt_key: str,
     prompt_mode: str,
-    *ui_keys: str,
-) -> None:
-    reset_chat_state(model_name, dataset_source)
-    chat_state["persona_id"] = persona_id
-    chat_state["prompt_mode"] = prompt_mode
-    _clear_chat_ui_state(*ui_keys)
+    active_system_prompt: str | None,
+) -> str | None:
+    if prompt_key not in st.session_state:
+        st.session_state[prompt_key] = active_system_prompt or ""
+    current = st.session_state.get(prompt_key) or ""
+    with st.expander("System prompt"):
+        st.markdown(current or "*empty*")
+        if prompt_mode != "empty" and st.button(
+            "Edit", icon=":material/edit:", key=f"{prompt_key}_edit"
+        ):
+            _open_system_prompt_dialog(prompt_key=prompt_key, current_value=current)
+    return st.session_state.get(prompt_key) or None
 
 
 def _generation_dict(gen_kwargs: dict, advanced_generation: bool) -> dict[str, object]:
@@ -189,42 +213,27 @@ def _render_chat_window(
     *,
     chat_log: Any,
     messages: list[dict[str, str]],
-    show_all_key: str,
-    show_all_btn_key: str,
-    show_earlier_label: str,
     chat_state: dict[str, object] | None = None,
     edit_key: str | None = None,
     pending_key: str | None = None,
-) -> Any:
-    """Render the visible chat history inside one container."""
-
+    show_contrast: bool = False,
+    edit_column_ratio: tuple[int, int] = (25, 1),
+) -> None:
     with chat_log:
-        if len(messages) > VISIBLE_MESSAGE_COUNT and not st.session_state.get(
-            show_all_key, False
-        ):
-            hidden_count = len(messages) - VISIBLE_MESSAGE_COUNT
-            if st.button(
-                f"{show_earlier_label} ({hidden_count} hidden)",
-                key=show_all_btn_key,
-            ):
-                st.session_state[show_all_key] = True
-                st.rerun()
-            visible_messages = messages[-VISIBLE_MESSAGE_COUNT:]
-            index_offset = len(messages) - VISIBLE_MESSAGE_COUNT
-        else:
-            visible_messages = messages
-            index_offset = 0
-
-        for i, message in enumerate(visible_messages):
-            actual_index = index_offset + i
+        for i, message in enumerate(messages):
             if edit_key and pending_key:
                 _render_editable_message(
-                    message, actual_index, messages, chat_state, edit_key, pending_key
+                    message,
+                    i,
+                    messages,
+                    chat_state,
+                    edit_key,
+                    pending_key,
+                    show_contrast=show_contrast,
+                    column_ratio=edit_column_ratio,
                 )
             else:
-                _render_chat_message(message)
-
-    return chat_log
+                _render_chat_message(message, show_contrast=show_contrast)
 
 
 def _build_chat_messages(
@@ -247,8 +256,8 @@ def _save_chat_export_message(
     messages: list[dict[str, str]],
     generation: dict[str, object],
     panel_label: str | None = None,
-) -> str:
-    export_path = save_chat_export(
+) -> None:
+    save_chat_export(
         model_name=model_name,
         dataset_source=dataset_source,
         persona_id=persona_id,
@@ -259,230 +268,12 @@ def _save_chat_export_message(
         messages=messages,
         generation=generation,
     )
-    return f"Saved chat export to {export_path}"
-
-
-# ── Compare mode helpers ───────────────────────────────────────────────────────
-
-
-def _panel_state(panel_key: str) -> dict:
-    """Get or initialise compare-panel chat state stored in session_state."""
-    if panel_key not in st.session_state:
-        st.session_state[panel_key] = _default_chat_state()
-    return st.session_state[panel_key]
-
-
-def _render_compare_mode(
-    remote: bool,
-    model_name: str,
-    context_key: str,
-    dataset_source: str,
-    personas: list[PersonaData],
-    gen_kwargs: dict,
-    advanced_generation: bool,
-) -> None:
-    """Render the full side-by-side comparison UI."""
-    left_col, right_col = st.columns(2)
-
-    def render_panel(side: str) -> tuple[dict[str, object], Any, str | None, str]:
-        panel_key = widget_key(context_key, f"cmp_{side}")
-        state = _panel_state(panel_key)
-        prompt_key = widget_key(panel_key, "custom_prompt")
-        show_all_key = widget_key(panel_key, "show_all")
-        edit_key = widget_key(panel_key, "edit_idx")
-        pending_regen_key = widget_key(panel_key, "pending_regen")
-
-        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"),
-        )
-        if changed:
-            state["messages"] = []
-            state["past_key_values"] = None
-            state["persona_id"] = selected_persona.id
-            state["prompt_mode"] = prompt_mode
-            _clear_chat_ui_state(prompt_key, show_all_key)
-            st.session_state.pop(edit_key, None)
-
-        active_system_prompt = resolve_system_prompt(
-            persona=selected_persona, mode=prompt_mode
-        )
-
-        btn_col1, btn_col2 = st.columns(2)
-        with btn_col1:
-            if st.button(
-                "Export chat", key=widget_key(panel_key, "export_chat"), width="stretch"
-            ):
-                st.success(
-                    _save_chat_export_message(
-                        model_name=model_name,
-                        dataset_source=dataset_source,
-                        persona_id=selected_persona.id,
-                        persona_name=getattr(selected_persona, "name", None),
-                        prompt_mode=prompt_mode,
-                        system_prompt=active_system_prompt,
-                        messages=state["messages"],
-                        generation=_generation_dict(gen_kwargs, advanced_generation),
-                        panel_label=side,
-                    )
-                )
-        with btn_col2:
-            if st.button(
-                "Reset chat",
-                key=widget_key(panel_key, "reset"),
-                width="stretch",
-                type="secondary",
-            ):
-                state["messages"] = []
-                state["past_key_values"] = None
-                _clear_chat_ui_state(prompt_key, show_all_key)
-                st.session_state.pop(edit_key, None)
-                st.rerun()
-
-        chat_log = st.container()
-        with chat_log:
-            active_system_prompt = _render_inline_system_prompt(
-                prompt_key,
-                prompt_mode,
-                active_system_prompt,
-                height=150,
-            )
-        _render_chat_window(
-            chat_log=chat_log,
-            messages=state["messages"],
-            show_all_key=show_all_key,
-            show_all_btn_key=widget_key(panel_key, "show_all_btn"),
-            show_earlier_label="Show earlier",
-            chat_state=state,
-            edit_key=edit_key,
-            pending_key=pending_regen_key,
-        )
-        return state, chat_log, active_system_prompt, pending_regen_key
-
-    with left_col:
-        left_state, left_log, left_prompt, left_pending = render_panel("left")
-    with right_col:
-        right_state, right_log, right_prompt, right_pending = render_panel("right")
-
-    panels = [
-        (left_state, left_log, left_prompt, left_pending),
-        (right_state, right_log, right_prompt, right_pending),
-    ]
-
-    # Handle per-panel regeneration triggered by message edits
-    any_regen = any(st.session_state.get(p_pending) for _, _, _, p_pending in panels)
-    if any_regen:
-        model = cached_model(model_name=model_name, remote=remote)
-        for panel_state, panel_log, panel_prompt, p_pending in panels:
-            if not st.session_state.pop(p_pending, False):
-                continue
-            regen_messages = _build_chat_messages(panel_prompt, panel_state["messages"])
-            with st.spinner("Regenerating..."):
-                try:
-                    result = generate_chat_reply(
-                        model=model,
-                        messages=regen_messages,
-                        remote=remote,
-                        past_key_values=panel_state["past_key_values"],
-                        **gen_kwargs,
-                    )
-                except Exception as exc:
-                    with panel_log:
-                        st.error(f"Generation failed: {exc}")
-                    panel_state["messages"].pop()
-                    continue
-            panel_state["messages"].append(
-                {"role": "assistant", "content": result.text}
-            )
-            panel_state["past_key_values"] = (
-                result.past_key_values if not remote else None
-            )
-            with panel_log:
-                _render_chat_message({"role": "assistant", "content": result.text})
-        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)
-
-    for panel_state, panel_log, _panel_prompt, _p_pending in panels:
-        panel_state["messages"].append({"role": "user", "content": user_prompt})
-        with panel_log:
-            _render_chat_message({"role": "user", "content": user_prompt})
-
-    with st.spinner("Generating..."):
-        if remote:
-            with ThreadPoolExecutor(max_workers=2) as executor:
-                futures = [
-                    executor.submit(
-                        generate_chat_reply,
-                        model=model,
-                        messages=_build_chat_messages(
-                            panel_prompt, panel_state["messages"]
-                        ),
-                        remote=remote,
-                        past_key_values=panel_state["past_key_values"],
-                        **gen_kwargs,
-                    )
-                    for panel_state, _panel_log, panel_prompt, _p_pending in panels
-                ]
-                results: list[ChatReply | Exception] = []
-                for future in futures:
-                    try:
-                        results.append(future.result())
-                    except Exception as exc:
-                        results.append(exc)
-        else:
-            results = []
-            for panel_state, _panel_log, panel_prompt, _p_pending 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"],
-                            **gen_kwargs,
-                        )
-                    )
-                except Exception as exc:
-                    results.append(exc)
-
-    for (panel_state, panel_log, _panel_prompt, _p_pending), result in zip(
-        panels, results
-    ):
-        if isinstance(result, Exception):
-            with panel_log:
-                st.error(f"Generation failed: {result}")
-            panel_state["messages"].pop()
-            continue
-
-        panel_state["messages"].append({"role": "assistant", "content": result.text})
-        panel_state["past_key_values"] = result.past_key_values if not remote else None
-        with panel_log:
-            _render_chat_message({"role": "assistant", "content": result.text})
-
-    # 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()
 
 
 # ── Main tab entry point ───────────────────────────────────────────────────────
 
 
-def _render_generation_settings(
-    context_key: str, remote: bool
-) -> tuple[dict, bool]:
+def _render_generation_settings(context_key: str, remote: bool) -> tuple[dict, bool]:
     """Render the Advanced generation settings expander.
 
     Returns ``(gen_kwargs, advanced_generation)`` where ``advanced_generation``
@@ -633,7 +424,9 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     )
 
     if compare_mode:
-        _render_compare_mode(
+        from tabs.compare_chat import render_compare_mode
+
+        render_compare_mode(
             remote,
             model_name,
             context_key,
@@ -648,76 +441,70 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     persona_select_key = widget_key(context_key, "persona_select")
     prompt_mode_select_key = widget_key(context_key, "system_prompt_select")
     prompt_key = widget_key(context_key, "custom_system_prompt")
-    show_all_key = widget_key(context_key, "show_all_messages")
     chat_input_key = widget_key(context_key, "chat_input")
     pending_key = widget_key(context_key, "pending_prompt")
     export_key = widget_key(context_key, "export_chat")
     reset_key = widget_key(context_key, "reset")
     edit_key = widget_key(context_key, "edit_idx")
 
-    col1, col2 = st.columns([2, 1])
-    with col1:
-        selected_index = next(
-            (i for i, p in enumerate(personas) if p.id == chat_state["persona_id"]),
-            0,
-        )
-        selected_persona = st.selectbox(
-            "Persona",
-            options=personas,
-            index=selected_index,
-            format_func=persona_label,
-            key=persona_select_key,
-        )
-    with col2:
-        current_mode_label = VARIANT_LABELS.get(chat_state["prompt_mode"], "None")
-        st.selectbox(
-            "Prompt",
-            options=MODE_LABELS,
-            index=MODE_LABELS.index(current_mode_label),
-            key=prompt_mode_select_key,
+    def _reset_active_chat_context() -> None:
+        reset_chat_context_state(
+            chat_state,
+            selected_persona.id,
+            prompt_mode,
+            chat_input_key,
+            prompt_key,
+            pending_key,
         )
-        prompt_mode = MODE_LABEL_TO_KEY[st.session_state[prompt_mode_select_key]]
+        st.session_state.pop(edit_key, None)
+
+    selected_persona, prompt_mode, changed_context = _render_persona_prompt_controls(
+        personas,
+        chat_state["persona_id"],
+        chat_state["prompt_mode"],
+        persona_select_key,
+        prompt_mode_select_key,
+        column_widths=(2, 1),
+    )
 
     active_system_prompt = resolve_system_prompt(
         persona=selected_persona,
         mode=prompt_mode,
     )
 
-    changed_context = (
-        chat_state["persona_id"] != selected_persona.id
-        or chat_state["prompt_mode"] != prompt_mode
-    )
     if changed_context:
         had_history = bool(chat_state["messages"])
-        _reset_single_chat_context(
-            model_name,
-            dataset_source,
-            chat_state,
-            selected_persona.id,
-            prompt_mode,
-            chat_input_key,
-            show_all_key,
-            prompt_key,
-            pending_key,
-        )
-        st.session_state.pop(edit_key, None)
+        _reset_active_chat_context()
         if had_history:
             st.info("Chat history reset because the persona or system prompt changed.")
 
     chat_log = st.container()
 
     with chat_log:
-        active_system_prompt = _render_inline_system_prompt(
+        active_system_prompt = _render_system_prompt(
             prompt_key,
             prompt_mode,
             active_system_prompt,
-            height=200,
         )
 
-    action_col1, action_col2 = st.columns(2)
-    with action_col1:
-        if st.button("Export chat", key=export_key, width="stretch"):
-            st.success(
+    _render_chat_window(
+        chat_log=chat_log,
+        messages=chat_state["messages"],
+        chat_state=chat_state,
+        edit_key=edit_key,
+        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_message(
                     model_name=model_name,
                     dataset_source=dataset_source,
@@ -728,38 +515,18 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
                     messages=chat_state["messages"],
                     generation=_generation_dict(gen_kwargs, advanced_generation),
                 )
-            )
-    with action_col2:
-        if st.button("Reset chat", key=reset_key, width="stretch", type="secondary"):
-            _reset_single_chat_context(
-                model_name,
-                dataset_source,
-                chat_state,
-                selected_persona.id,
-                prompt_mode,
-                chat_input_key,
-                show_all_key,
-                prompt_key,
-                pending_key,
-            )
-            st.session_state.pop(edit_key, None)
-            st.rerun()
-
-    _render_chat_window(
-        chat_log=chat_log,
-        messages=chat_state["messages"],
-        show_all_key=show_all_key,
-        show_all_btn_key=widget_key(context_key, "show_all_btn"),
-        show_earlier_label="Show earlier messages",
-        chat_state=chat_state,
-        edit_key=edit_key,
-        pending_key=pending_key,
-    )
+                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()
 
-    user_prompt = st.chat_input(
-        "Ask something...",
-        key=chat_input_key,
-    )
+    user_prompt = st.chat_input("Ask something...", key=chat_input_key)
 
     # Pass 1: user submitted — append message and rerun so it renders before generation.
     if user_prompt:

tabs/compare_chat.py

@@ -0,0 +1,443 @@
+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 utils.chat import ChatReply, generate_chat_reply, resolve_system_prompt
+from utils.contrast import compute_contrast, compute_contrast_pair
+from utils.helpers import persona_label, widget_key
+from utils.runtime import cached_model
+
+from .chat import (
+    _build_chat_messages,
+    _generation_dict,
+    _render_chat_message,
+    _render_chat_window,
+    _render_persona_prompt_controls,
+    _render_system_prompt,
+    _save_chat_export_message,
+)
+
+
+def _panel_state(panel_key: str) -> dict[str, object]:
+    """Get or initialise compare-panel chat state stored in session_state."""
+    if panel_key not in st.session_state:
+        st.session_state[panel_key] = default_chat_state()
+    return st.session_state[panel_key]
+
+
+def _reset_compare_panel(
+    panel_state: dict,
+    edit_key: str,
+    persona_id: str,
+    prompt_mode: str,
+    *ui_keys: str,
+) -> None:
+    reset_chat_context_state(panel_state, persona_id, prompt_mode, *ui_keys)
+    st.session_state.pop(edit_key, None)
+
+
+def _generate_panel_reply(
+    *,
+    model: StandardizedTransformer,
+    remote: bool,
+    panel_state: dict[str, object],
+    panel_prompt: str | None,
+    gen_kwargs: dict,
+) -> 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"],
+        **gen_kwargs,
+    )
+
+
+def render_compare_mode(
+    remote: bool,
+    model_name: str,
+    context_key: str,
+    dataset_source: str,
+    personas: list[PersonaData],
+    gen_kwargs: dict,
+    advanced_generation: bool,
+) -> None:
+    """Render the full side-by-side comparison UI."""
+    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 four extra forward passes after each turn (batched into one "
+            "remote session when running on NDIF)."
+        ),
+    )
+
+    left_col, right_col = st.columns(2)
+    left_panel_key = widget_key(context_key, "cmp_left")
+    right_panel_key = widget_key(context_key, "cmp_right")
+    left_prompt_key = widget_key(left_panel_key, "custom_prompt")
+    right_prompt_key = widget_key(right_panel_key, "custom_prompt")
+    left_edit_key = widget_key(left_panel_key, "edit_idx")
+    right_edit_key = widget_key(right_panel_key, "edit_idx")
+    left_pending_key = widget_key(left_panel_key, "pending_regen")
+    right_pending_key = widget_key(right_panel_key, "pending_regen")
+
+    def render_panel(side: str) -> tuple[dict, object, str | None, str, PersonaData]:
+        panel_key = widget_key(context_key, f"cmp_{side}")
+        state = _panel_state(panel_key)
+        prompt_key = widget_key(panel_key, "custom_prompt")
+        edit_key = widget_key(panel_key, "edit_idx")
+        pending_regen_key = widget_key(panel_key, "pending_regen")
+
+        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"),
+        )
+        if changed:
+            reset_chat_context_state(
+                state,
+                selected_persona.id,
+                prompt_mode,
+                prompt_key,
+                pending_regen_key,
+            )
+            st.session_state.pop(edit_key, None)
+
+        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,
+            )
+        return (
+            state,
+            chat_log,
+            active_system_prompt,
+            pending_regen_key,
+            selected_persona,
+        )
+
+    with left_col:
+        left_state, left_log, left_prompt, left_pending, left_persona = render_panel(
+            "left"
+        )
+    with right_col:
+        right_state, right_log, right_prompt, right_pending, right_persona = (
+            render_panel("right")
+        )
+
+    panels = [
+        (
+            left_state,
+            left_log,
+            left_prompt,
+            left_pending,
+            left_edit_key,
+            left_persona,
+        ),
+        (
+            right_state,
+            right_log,
+            right_prompt,
+            right_pending,
+            right_edit_key,
+            right_persona,
+        ),
+    ]
+
+    # Handle per-panel regeneration triggered by message edits
+    regen_panels = [
+        (panel_state, panel_log, panel_prompt)
+        for panel_state, panel_log, panel_prompt, p_pending, _panel_edit_key, _ in panels
+        if st.session_state.pop(p_pending, False)
+    ]
+    if regen_panels:
+        model = cached_model(model_name=model_name, remote=remote)
+
+        results: list[ChatReply | Exception] = []
+        with st.spinner("Regenerating..."):
+            for panel_state, _panel_log, panel_prompt in regen_panels:
+                try:
+                    results.append(
+                        _generate_panel_reply(
+                            model=model,
+                            remote=remote,
+                            panel_state=panel_state,
+                            panel_prompt=panel_prompt,
+                            gen_kwargs=gen_kwargs,
+                        )
+                    )
+                except Exception as exc:
+                    results.append(exc)
+
+        for (panel_state, panel_log, _panel_prompt), result in zip(
+            regen_panels, results
+        ):
+            if isinstance(result, Exception):
+                with panel_log:
+                    st.error(f"Generation failed: {result}")
+                panel_state["messages"].pop()
+                continue
+            panel_state["messages"].append(
+                {"role": "assistant", "content": result.text}
+            )
+            panel_state["past_key_values"] = (
+                result.past_key_values if not remote else None
+            )
+        st.rerun()
+
+    # 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_state, *_rest) 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 = cached_model(model_name=model_name, remote=remote)
+            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_state = panels[panel_idx][0]
+                    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()
+
+    for (
+        panel_state,
+        panel_log,
+        _panel_prompt,
+        panel_pending,
+        panel_edit_key,
+        _,
+    ) in panels:
+        _render_chat_window(
+            chat_log=panel_log,
+            messages=panel_state["messages"],
+            chat_state=panel_state,
+            edit_key=panel_edit_key,
+            pending_key=panel_pending,
+            show_contrast=contrast_enabled,
+            edit_column_ratio=(10, 1),
+        )
+
+    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=widget_key(context_key, "cmp_export"),
+                help="Export both chats",
+            ):
+                for side, panel_state, panel_prompt, panel_persona in (
+                    ("left", left_state, left_prompt, left_persona),
+                    ("right", right_state, right_prompt, right_persona),
+                ):
+                    _save_chat_export_message(
+                        model_name=model_name,
+                        dataset_source=dataset_source,
+                        persona_id=panel_persona.id,
+                        persona_name=getattr(panel_persona, "name", None),
+                        prompt_mode=panel_state["prompt_mode"],
+                        system_prompt=panel_prompt,
+                        messages=panel_state["messages"],
+                        generation=_generation_dict(gen_kwargs, advanced_generation),
+                        panel_label=side,
+                    )
+                st.toast("Exported", icon=":material/check:")
+        with rst_col:
+            with st.popover(
+                "",
+                icon=":material/delete_sweep:",
+                help="Reset chat",
+            ):
+                if st.button(
+                    "Reset left",
+                    key=widget_key(context_key, "cmp_reset_left"),
+                ):
+                    _reset_compare_panel(
+                        left_state,
+                        left_edit_key,
+                        left_persona.id,
+                        left_state["prompt_mode"],
+                        left_prompt_key,
+                        left_pending_key,
+                    )
+                    st.rerun()
+                if st.button(
+                    "Reset right",
+                    key=widget_key(context_key, "cmp_reset_right"),
+                ):
+                    _reset_compare_panel(
+                        right_state,
+                        right_edit_key,
+                        right_persona.id,
+                        right_state["prompt_mode"],
+                        right_prompt_key,
+                        right_pending_key,
+                    )
+                    st.rerun()
+                if st.button(
+                    "Reset both",
+                    key=widget_key(context_key, "cmp_reset_both"),
+                    type="primary",
+                ):
+                    _reset_compare_panel(
+                        left_state,
+                        left_edit_key,
+                        left_persona.id,
+                        left_state["prompt_mode"],
+                        left_prompt_key,
+                        left_pending_key,
+                    )
+                    _reset_compare_panel(
+                        right_state,
+                        right_edit_key,
+                        right_persona.id,
+                        right_state["prompt_mode"],
+                        right_prompt_key,
+                        right_pending_key,
+                    )
+                    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)
+
+    for panel_state, panel_log, _panel_prompt, _p_pending, _panel_edit_key, _ 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_state, _panel_log, panel_prompt, _p_pending, _panel_edit_key, _ in panels
+    ]
+
+    results: list[ChatReply | Exception] = []
+    with st.spinner("Generating..."):
+        # Keep compare-mode generation sequential so both panels use the same
+        # model/session state safely.
+        for (
+            panel_state,
+            _panel_log,
+            panel_prompt,
+            _p_pending,
+            _panel_edit_key,
+            _,
+        ) in panels:
+            try:
+                results.append(
+                    _generate_panel_reply(
+                        model=model,
+                        remote=remote,
+                        panel_state=panel_state,
+                        panel_prompt=panel_prompt,
+                        gen_kwargs=gen_kwargs,
+                    )
+                )
+            except Exception as exc:
+                results.append(exc)
+
+    valid_results: list[ChatReply | None] = []
+    for (
+        panel_state,
+        panel_log,
+        _panel_prompt,
+        _p_pending,
+        _panel_edit_key,
+        _,
+    ), 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
+
+        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)
+    ):
+        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}")
+
+    # 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

@@ -111,8 +111,6 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         st.info("Select at least one persona.")
         return
 
-    max_questions = 0
-
     with st.expander("Advanced", expanded=False):
         st.caption("Filters")
 

utils/chat.py

@@ -1,5 +1,5 @@
 import logging
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from typing import Literal
 
@@ -15,9 +15,8 @@ SystemPromptMode = Literal["empty", "templated", "biography", "custom"]
 @dataclass
 class ChatReply:
     text: str
-    prompt_tokens: int
-    output_tokens: int
     past_key_values: object | None
+    generated_ids: torch.Tensor | None = None
 
 
 def resolve_system_prompt(
@@ -204,13 +203,10 @@ def generate_chat_reply(
 
     generated_ids = sequences[0, prompt_token_count:]
     text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
-    output_tokens = int(sequences.shape[1] - prompt_token_count)
-
     return ChatReply(
         text=text,
-        prompt_tokens=prompt_token_count,
-        output_tokens=max(0, output_tokens),
         past_key_values=(
             getattr(generated, "past_key_values", None) if not remote else None
         ),
+        generated_ids=generated_ids.detach().cpu(),
     )

utils/chat_export.py

@@ -30,6 +30,7 @@ def save_chat_export(
         system_prompt: Current system prompt text, if any.
         messages: Conversation messages without the system prompt.
         generation: Generation settings used for the chat.
+        panel_label: Optional side label (e.g. "left"/"right") for compare-mode exports.
 
     Returns:
         The path the export was written to.
@@ -55,7 +56,6 @@ def save_chat_export(
         get_artifacts_dir()
         / "chats"
         / "__".join(slugify(part) for part in model_name.split("/"))
-        / slugify(dataset_source)
         / slugify(persona_id)
     )
     export_dir.mkdir(parents=True, exist_ok=True)

utils/contrast.py

@@ -0,0 +1,311 @@
+# WARNING: This is mostly vibecoded and need reviews
+# - Check that the model is runned once with normally for gneration and things are beeing traced perphaps at the last step of generation with iter.last or somrething liek that from the docs
+# - Then the model is runned again with the entire context of the conversation from the other context on the rifht ? or on the left dependeing on which one we are doing at the moment. And this will then compute the prob diff and show them.
+
+"""
+Contrastive token-level log-probability comparison for compare mode.
+
+For a pair of responses generated under different persona contexts, each token
+gets a weight:
+
+    w(token) = log P(token | context_A) − log P(token | context_B)
+
+Positive (red)  → token is more characteristic of persona A.
+Negative (blue) → token is more characteristic of persona B.
+Near-zero (gray) → both personas would emit this token with similar likelihood.
+"""
+
+from dataclasses import dataclass
+from html import escape
+
+import torch
+from nnterp import StandardizedTransformer
+
+from utils.chat import _format_generation_prompt
+
+
+@dataclass
+class TokenContrast:
+    tokens: list[str]
+    weights: list[float]  # normalised to [-1, 1], used for coloring
+    raw_diffs: list[float]  # unclipped log P(A) - log P(B) per token
+    label_a: str
+    label_b: str
+
+
+# ── Weight computation ────────────────────────────────────────────────────────
+
+
+def _normalise_diffs(diffs: torch.Tensor) -> list[float]:
+    """
+    Clip at the 95th percentile of |diff| and scale to [-1, 1] so a few
+    high-magnitude tokens don't wash out everything else.
+    """
+    if len(diffs) < 2:
+        return diffs.tolist()
+    clip_val = max(torch.quantile(diffs.abs(), 0.95).item(), 0.3)
+    return (diffs.float().clamp(-clip_val, clip_val) / clip_val).tolist()
+
+
+def _decode_ids(tokenizer: object, ids: list[int]) -> str:
+    try:
+        return tokenizer.decode(
+            ids,
+            skip_special_tokens=False,
+            clean_up_tokenization_spaces=False,
+        )
+    except TypeError:
+        return tokenizer.decode(ids, skip_special_tokens=False)
+
+
+def _strip_special_ids(
+    ids: torch.Tensor,
+    tokenizer: object,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return display ids and a mask that excludes special tokens."""
+    ids = ids.cpu()
+    special_ids = set(getattr(tokenizer, "all_special_ids", []) or [])
+    if not special_ids or ids.numel() == 0:
+        return ids, torch.ones(ids.shape[0], dtype=torch.bool)
+    keep = torch.tensor(
+        [tid.item() not in special_ids for tid in ids], dtype=torch.bool
+    )
+    return ids[keep], keep
+
+
+def _prepare_trace_text(
+    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)``."""
+    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():
+        raise ValueError(
+            "contrast trace text did not round-trip to the expected token ids"
+        )
+    n_ctx = len(context_ids)
+    n_resp = len(response_ids)
+    return full_text, n_ctx, n_resp
+
+
+def _build_contrast(
+    tokenizer: object,
+    response_ids: torch.Tensor,
+    lp_a: torch.Tensor,
+    lp_b: torch.Tensor,
+    label_a: str,
+    label_b: str,
+) -> TokenContrast:
+    diffs = (lp_a - lp_b).cpu()
+    display_ids, keep_mask = _strip_special_ids(response_ids, tokenizer)
+    display_diffs = diffs[keep_mask]
+    return TokenContrast(
+        tokens=[_token_display(tokenizer, tid.item()) for tid in display_ids],
+        weights=_normalise_diffs(display_diffs),
+        raw_diffs=display_diffs.float().tolist(),
+        label_a=label_a,
+        label_b=label_b,
+    )
+
+
+def _token_display(tokenizer: object, token_id: int) -> str:
+    """Render a single token id as normal decoded text."""
+    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]
+
+
+def _score_passes(
+    model: StandardizedTransformer,
+    specs: list[PassSpec],
+    remote: bool,
+) -> dict[str, torch.Tensor]:
+    """
+    Run one forward pass per spec and return reduced per-token logprobs.
+
+    The log-softmax and target-pick happen *inside* the trace, so only the
+    reduced ``[n_resp]`` logprob vector per pass is shipped back — not the full
+    ``[1, seq, vocab]`` logits (which would be hundreds of MB per pass on NDIF).
+    """
+
+    def _score_pass(
+        full_text: str,
+        n_ctx: int,
+        n_resp: int,
+        target_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        with torch.no_grad(), model.trace(full_text, 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()
+            log_probs = torch.log_softmax(resp_logits, dim=-1)
+            targets = target_ids.to(log_probs.device).view(-1, 1)
+            picked = log_probs.gather(1, targets).view(-1)
+            out = picked.detach().cpu().save()
+
+        if hasattr(out, "value") and getattr(out, "value") is not None:
+            out = out.value
+        if not isinstance(out, torch.Tensor):
+            raise TypeError(
+                f"contrast score did not resolve to a tensor: {type(out)!r}"
+            )
+        return out.detach().cpu()
+
+    saved = [
+        _score_pass(full_text, n_ctx, n_resp, target_ids)
+        for _key, full_text, n_ctx, n_resp, target_ids in specs
+    ]
+
+    if len(saved) != len(specs):
+        raise RuntimeError(
+            f"contrast scoring returned {len(saved)} result(s) for {len(specs)} spec(s)"
+        )
+
+    return {spec[0]: tensor for spec, tensor in zip(specs, saved)}
+
+
+def _specs_for_response(
+    tokenizer: object,
+    response_ids: torch.Tensor,
+    context_a: list[dict[str, str]],
+    context_b: list[dict[str, str]],
+    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)
+    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),
+    ]
+
+
+def compute_contrast(
+    model: StandardizedTransformer,
+    context_a: list[dict[str, str]],
+    context_b: list[dict[str, str]],
+    response_ids: torch.Tensor,
+    label_a: str,
+    label_b: str,
+    remote: bool = False,
+) -> "TokenContrast | None":
+    """Compute per-token contrast weights for a single response (2 forward passes)."""
+    tokenizer = model.tokenizer
+    if response_ids.numel() == 0:
+        return None
+
+    specs = _specs_for_response(tokenizer, response_ids, context_a, context_b, "r")
+    out = _score_passes(model, specs, remote)
+    return _build_contrast(
+        tokenizer, response_ids, out["r_under_a"], out["r_under_b"], label_a, label_b
+    )
+
+
+def compute_contrast_pair(
+    model: StandardizedTransformer,
+    context_a: list[dict[str, str]],
+    context_b: list[dict[str, str]],
+    response_ids_a: torch.Tensor,
+    response_ids_b: torch.Tensor,
+    label_a: str,
+    label_b: str,
+    remote: bool = False,
+) -> tuple["TokenContrast | None", "TokenContrast | None"]:
+    """
+    Compute contrast weights for both panel responses (up to 4 remote passes).
+    """
+    tokenizer = model.tokenizer
+    if response_ids_a.numel() == 0 and response_ids_b.numel() == 0:
+        return None, None
+
+    specs: list[PassSpec] = []
+    if response_ids_a.numel() > 0:
+        specs += _specs_for_response(
+            tokenizer, response_ids_a, context_a, context_b, "a"
+        )
+    if response_ids_b.numel() > 0:
+        specs += _specs_for_response(
+            tokenizer, response_ids_b, context_a, context_b, "b"
+        )
+
+    out = _score_passes(model, specs, remote)
+
+    def _build(resp_ids: torch.Tensor, prefix: str) -> "TokenContrast | None":
+        k_a, k_b = f"{prefix}_under_a", f"{prefix}_under_b"
+        if resp_ids.numel() == 0 or k_a not in out or k_b not in out:
+            return None
+        return _build_contrast(
+            tokenizer, resp_ids, out[k_a], out[k_b], label_a, label_b
+        )
+
+    return _build(response_ids_a, "a"), _build(response_ids_b, "b")
+
+
+# ── HTML rendering ────────────────────────────────────────────────────────────
+
+
+def _weight_to_bg(w: float) -> str:
+    """Map a normalised weight in [-1, 1] to a CSS rgba background color."""
+    w = max(-1.0, min(1.0, w))
+    alpha = abs(w) * 0.5  # cap at 0.5 opacity so text stays readable
+    if w > 0.05:
+        return f"rgba(210,60,60,{alpha:.3f})"
+    if w < -0.05:
+        return f"rgba(50,110,210,{alpha:.3f})"
+    return "rgba(0,0,0,0)"
+
+
+_CONTRAST_CSS = (
+    "<style>"
+    ".contrast-tok{position:relative;border-radius:2px;padding:0 1px;"
+    "cursor:default;white-space:pre;}"
+    ".contrast-tok>.contrast-tip{display:none;position:absolute;bottom:100%;"
+    "left:50%;transform:translateX(-50%);margin-bottom:4px;padding:2px 6px;"
+    "border-radius:3px;background:#222;color:#eee;font-size:0.72em;"
+    "font-family:ui-monospace,monospace;white-space:nowrap;pointer-events:none;"
+    "z-index:10;box-shadow:0 2px 6px rgba(0,0,0,0.3);}"
+    ".contrast-tok:hover>.contrast-tip{display:block;}"
+    "</style>"
+)
+
+
+def render_contrast_html(result: TokenContrast) -> str:
+    """
+    Render each token with a colored background reflecting how A- or B-specific
+    it is, with a hover tooltip showing the raw Δlog P, plus a legend.
+    """
+    spans: list[str] = []
+    for token, weight, raw in zip(result.tokens, result.weights, result.raw_diffs):
+        bg = _weight_to_bg(weight)
+        tip = escape(f"Δlog P(A−B): {raw:+.3f}")
+        text = escape(token)
+        spans.append(
+            f'<span class="contrast-tok" style="background:{bg};">'
+            f'{text}<span class="contrast-tip">{tip}</span></span>'
+        )
+
+    la = escape(result.label_a)
+    lb = escape(result.label_b)
+
+    return (
+        _CONTRAST_CSS + '<div style="font-family:inherit;line-height:1.75;'
+        'white-space:pre-wrap;word-break:break-word;padding:2px 0 6px 0;">'
+        + "".join(spans)
+        + '<div style="margin-top:10px;font-size:0.72em;color:#888;'
+        + 'display:flex;gap:12px;flex-wrap:wrap;">'
+        + f'<span><span style="background:rgba(210,60,60,0.45);'
+        + f'padding:1px 6px;border-radius:2px;">&thinsp;</span>&nbsp;{la}</span>'
+        + f'<span><span style="background:rgba(50,110,210,0.45);'
+        + f'padding:1px 6px;border-radius:2px;">&thinsp;</span>&nbsp;{lb}</span>'
+        + '<span style="color:#aaa;">gray = shared by both</span>'
+        + "</div>"
+        + "</div>"
+    )

utils/helpers.py

@@ -16,8 +16,6 @@ MODE_LABELS = list(VARIANT_LABELS.values())
 # Reverse lookup: label -> key
 MODE_LABEL_TO_KEY = {v: k for k, v in VARIANT_LABELS.items()}
 
-VISIBLE_MESSAGE_COUNT = 5
-
 DATASET_SOURCES = ["HuggingFace: synth-persona", "Local JSONL upload"]
 ANALYSIS_MODES = ["Cosine similarity", "PCA", "UMAP"]
 

utils/runtime.py

@@ -7,33 +7,62 @@ logger = logging.getLogger(__name__)
 
 @st.cache_data(show_spinner=False, ttl=30)
 def list_remote_models() -> list[str]:
-    """Return the NDIF language models that are currently running."""
+    """Return the NDIF language models that are currently running.
+
+    Parses the raw NDIF response directly instead of going through
+    ``nnsight.ndif_status()`` because that call crashes whenever NDIF reports
+    any deployment with an ``application_state`` that isn't in nnsight's
+    ``ModelStatus`` enum (e.g. ``UNHEALTHY``) — one bad deployment poisons
+    the whole response. See nnsight 0.6.3 ``ndif.py::status``.
+    """
+
+    import json
 
     import nnsight
 
     try:
-        status = nnsight.ndif_status()
+        raw = nnsight.ndif_status(raw=True)
     except Exception:
         logger.warning("Failed to fetch NDIF status", exc_info=True)
         return []
 
     model_names: list[str] = []
+    bad_states: list[tuple[str, str]] = []  # (repo_id_or_key, application_state)
 
-    for entry in status.values():
-        if not isinstance(entry, dict):
+    for value in (raw or {}).get("deployments", {}).values():
+        if not isinstance(value, dict):
             continue
-        if entry.get("model_class") not in {"LanguageModel", "StandardizedTransformer"}:
+        if (
+            value.get("deployment_level") not in {"HOT", "WARM"}
+            and "schedule" not in value
+        ):
             continue
 
-        state = entry.get("state")
-        state_name = getattr(state, "name", None) or getattr(state, "value", None)
-        if state_name != "RUNNING":
-            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
 
-        repo_id = entry.get("repo_id")
+        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)
 
+    if bad_states:
+        logger.warning(
+            "NDIF reported deployments with unexpected application_state values "
+            "(nnsight's ModelStatus enum may not know about these): %s",
+            bad_states,
+        )
+
     return sorted(set(model_names))