Small cleanups

Small reset bottom cleanup

state.py

@@ -1,6 +1,7 @@
 import streamlit as st
 
 _CHAT_STATE_PREFIX = "chat_state::"
+_CHAT_KEYS_REGISTRY = "chat_state::_registered_keys"
 
 
 def chat_session_key(model_name: str, dataset_source: str) -> str:
@@ -37,13 +38,9 @@ def reset_chat_context_state(
 def _evict_inactive_kv_caches(active_key: str) -> None:
     """Drop past_key_values from every chat context except the active one."""
 
-    for key in st.session_state:
-        if (
-            isinstance(key, str)
-            and key.startswith(_CHAT_STATE_PREFIX)
-            and key != active_key
-        ):
-            state = st.session_state[key]
+    for key in st.session_state.get(_CHAT_KEYS_REGISTRY, ()):
+        if key != active_key:
+            state = st.session_state.get(key)
             if isinstance(state, dict) and state.get("past_key_values") is not None:
                 state["past_key_values"] = None
 
@@ -54,13 +51,21 @@ def get_chat_state(
     """Return the mutable chat state for the active context."""
 
     key = chat_session_key(model_name, dataset_source)
+    registry = st.session_state.get(_CHAT_KEYS_REGISTRY)
+    if registry is None:
+        registry = set()
+        st.session_state[_CHAT_KEYS_REGISTRY] = registry
+    registry.add(key)
+
     state = st.session_state.get(key)
     if state is None:
         state = default_chat_state()
         st.session_state[key] = state
     else:
-        for default_key, default_value in default_chat_state().items():
-            state.setdefault(default_key, default_value)
+        state.setdefault("messages", [])
+        state.setdefault("persona_id", None)
+        state.setdefault("prompt_mode", "templated")
+        state.setdefault("past_key_values", None)
     _evict_inactive_kv_caches(key)
     if remote and state.get("past_key_values") is not None:
         state["past_key_values"] = None

tabs/chat.py

@@ -17,10 +17,19 @@ from utils.helpers import (
 )
 from utils.runtime import cached_model
 
-
-def _render_collapsible_markdown(content: str) -> None:
-    st.markdown(content)
-
+# ── Persistence keys for surviving model / remote switches ────────────────────
+_LAST_PERSONA_ID_KEY = "chat:last_persona_id"
+_LAST_PROMPT_MODE_KEY = "chat:last_prompt_mode"
+_LAST_COMPARE_MODE_KEY = "chat:last_compare_mode"
+
+# ── Generation defaults (single source of truth) ─────────────────────────────
+_GEN_DEFAULTS = {
+    "max_new_tokens": 256,
+    "temperature": 1.0,
+    "top_p": 1.0,
+    "top_k": 50,
+    "repetition_penalty": 1.0,
+}
 
 # ── Dialogs ───────────────────────────────────────────────────────────────────
 
@@ -91,7 +100,7 @@ def _open_system_prompt_dialog(*, prompt_key: str, current_value: str) -> None:
 # ── Message renderers ─────────────────────────────────────────────────────────
 
 
-def _render_chat_message(
+def render_chat_message(
     message: dict[str, str],
     show_contrast: bool = False,
 ) -> None:
@@ -103,7 +112,7 @@ def _render_chat_message(
         if tc is not None:
             st.html(render_contrast_html(tc))
         else:
-            _render_collapsible_markdown(message["content"])
+            st.markdown(message["content"])
 
 
 def _render_editable_message(
@@ -129,7 +138,7 @@ def _render_editable_message(
             if tc is not None:
                 st.html(render_contrast_html(tc))
             else:
-                _render_collapsible_markdown(message["content"])
+                st.markdown(message["content"])
     with edit_col:
         if st.button(
             "", icon=":material/edit:", key=f"{edit_key}_edit_{msg_index}", help="Edit"
@@ -142,7 +151,7 @@ def _render_editable_message(
             )
 
 
-def _render_system_prompt(
+def render_system_prompt(
     prompt_key: str,
     prompt_mode: str,
     active_system_prompt: str | None,
@@ -159,7 +168,7 @@ def _render_system_prompt(
     return st.session_state.get(prompt_key) or None
 
 
-def _generation_dict(gen_kwargs: dict, advanced_generation: bool) -> dict[str, object]:
+def generation_dict(gen_kwargs: dict, advanced_generation: bool) -> dict[str, object]:
     return {
         "max_new_tokens": int(gen_kwargs["max_new_tokens"]),
         "advanced_generation": bool(advanced_generation),
@@ -172,7 +181,7 @@ def _generation_dict(gen_kwargs: dict, advanced_generation: bool) -> dict[str, o
     }
 
 
-def _render_persona_prompt_controls(
+def render_persona_prompt_controls(
     personas: list[PersonaData],
     current_persona_id: str | None,
     current_prompt_mode: str,
@@ -209,7 +218,7 @@ def _render_persona_prompt_controls(
     return selected_persona, prompt_mode, changed
 
 
-def _render_chat_window(
+def render_chat_window(
     *,
     chat_log: Any,
     messages: list[dict[str, str]],
@@ -233,10 +242,10 @@ def _render_chat_window(
                     column_ratio=edit_column_ratio,
                 )
             else:
-                _render_chat_message(message, show_contrast=show_contrast)
+                render_chat_message(message, show_contrast=show_contrast)
 
 
-def _build_chat_messages(
+def build_chat_messages(
     system_prompt: str | None,
     messages: list[dict[str, str]],
 ) -> list[dict[str, str]]:
@@ -245,31 +254,6 @@ def _build_chat_messages(
     ) + messages
 
 
-def _save_chat_export_message(
-    *,
-    model_name: str,
-    dataset_source: str,
-    persona_id: str,
-    persona_name: str | None,
-    prompt_mode: str,
-    system_prompt: str | None,
-    messages: list[dict[str, str]],
-    generation: dict[str, object],
-    panel_label: str | None = None,
-) -> None:
-    save_chat_export(
-        model_name=model_name,
-        dataset_source=dataset_source,
-        persona_id=persona_id,
-        persona_name=persona_name,
-        panel_label=panel_label,
-        prompt_mode=prompt_mode,
-        system_prompt=system_prompt,
-        messages=messages,
-        generation=generation,
-    )
-
-
 # ── Main tab entry point ───────────────────────────────────────────────────────
 
 
@@ -286,7 +270,7 @@ def _render_generation_settings(context_key: str, remote: bool) -> tuple[dict, b
                 "Max new tokens",
                 min_value=16,
                 max_value=512,
-                value=256,
+                value=_GEN_DEFAULTS["max_new_tokens"],
                 step=16,
                 key=widget_key(context_key, "max_new_tokens"),
             )
@@ -295,7 +279,7 @@ def _render_generation_settings(context_key: str, remote: bool) -> tuple[dict, b
                 "Repetition penalty",
                 min_value=0.5,
                 max_value=2.0,
-                value=1.0,
+                value=_GEN_DEFAULTS["repetition_penalty"],
                 step=0.05,
                 key=widget_key(context_key, "repetition_penalty"),
             )
@@ -313,7 +297,7 @@ def _render_generation_settings(context_key: str, remote: bool) -> tuple[dict, b
                 "Temperature",
                 min_value=0.01,
                 max_value=2.0,
-                value=1.0,
+                value=_GEN_DEFAULTS["temperature"],
                 step=0.01,
                 disabled=sampling_disabled,
                 key=widget_key(context_key, "temperature"),
@@ -323,7 +307,7 @@ def _render_generation_settings(context_key: str, remote: bool) -> tuple[dict, b
                 "Top-p",
                 min_value=0.01,
                 max_value=1.0,
-                value=1.0,
+                value=_GEN_DEFAULTS["top_p"],
                 step=0.01,
                 disabled=sampling_disabled,
                 key=widget_key(context_key, "top_p"),
@@ -333,7 +317,7 @@ def _render_generation_settings(context_key: str, remote: bool) -> tuple[dict, b
                 "Top-k (0 = off)",
                 min_value=0,
                 max_value=100,
-                value=50,
+                value=_GEN_DEFAULTS["top_k"],
                 step=1,
                 disabled=sampling_disabled,
                 key=widget_key(context_key, "top_k"),
@@ -365,12 +349,12 @@ def _render_generation_settings(context_key: str, remote: bool) -> tuple[dict, b
             st.caption("Seed is local-only and disabled for remote runs.")
 
     advanced_generation = (
-        max_new_tokens != 256
+        max_new_tokens != _GEN_DEFAULTS["max_new_tokens"]
         or use_sampling
-        or temperature != 1.0
-        or top_p != 1.0
-        or top_k != 50
-        or repetition_penalty != 1.0
+        or temperature != _GEN_DEFAULTS["temperature"]
+        or top_p != _GEN_DEFAULTS["top_p"]
+        or top_k != _GEN_DEFAULTS["top_k"]
+        or repetition_penalty != _GEN_DEFAULTS["repetition_penalty"]
         or seed is not None
     )
 
@@ -395,6 +379,14 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
 
     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"
+        )
+
     try:
         dataset, dataset_status = load_dataset(
             dataset_source,
@@ -416,12 +408,17 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     gen_kwargs, advanced_generation = _render_generation_settings(context_key, remote)
 
     # ── 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
+        )
     compare_mode = st.toggle(
         "Compare mode",
-        value=False,
-        key=widget_key(context_key, "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:
         from tabs.compare_chat import render_compare_mode
@@ -458,7 +455,7 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
         )
         st.session_state.pop(edit_key, None)
 
-    selected_persona, prompt_mode, changed_context = _render_persona_prompt_controls(
+    selected_persona, prompt_mode, changed_context = render_persona_prompt_controls(
         personas,
         chat_state["persona_id"],
         chat_state["prompt_mode"],
@@ -466,6 +463,8 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
         prompt_mode_select_key,
         column_widths=(2, 1),
     )
+    st.session_state[_LAST_PERSONA_ID_KEY] = selected_persona.id
+    st.session_state[_LAST_PROMPT_MODE_KEY] = prompt_mode
 
     active_system_prompt = resolve_system_prompt(
         persona=selected_persona,
@@ -481,13 +480,13 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     chat_log = st.container()
 
     with chat_log:
-        active_system_prompt = _render_system_prompt(
+        active_system_prompt = render_system_prompt(
             prompt_key,
             prompt_mode,
             active_system_prompt,
         )
 
-    _render_chat_window(
+    render_chat_window(
         chat_log=chat_log,
         messages=chat_state["messages"],
         chat_state=chat_state,
@@ -505,7 +504,7 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
                 key=export_key,
                 help="Export chat",
             ):
-                _save_chat_export_message(
+                save_chat_export(
                     model_name=model_name,
                     dataset_source=dataset_source,
                     persona_id=selected_persona.id,
@@ -513,7 +512,7 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
                     prompt_mode=prompt_mode,
                     system_prompt=active_system_prompt,
                     messages=chat_state["messages"],
-                    generation=_generation_dict(gen_kwargs, advanced_generation),
+                    generation=generation_dict(gen_kwargs, advanced_generation),
                 )
                 st.toast("Exported", icon=":material/check:")
         with rst_col:
@@ -538,7 +537,7 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     if not st.session_state.pop(pending_key, False):
         return
 
-    messages = _build_chat_messages(active_system_prompt, chat_state["messages"])
+    messages = build_chat_messages(active_system_prompt, chat_state["messages"])
 
     with st.spinner("Generating reply..."):
         model = cached_model(model_name=model_name, remote=remote)
@@ -559,15 +558,4 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
 
     chat_state["messages"].append({"role": "assistant", "content": reply.text})
     chat_state["past_key_values"] = reply.past_key_values if not remote else None
-
-    save_chat_export(
-        model_name=model_name,
-        dataset_source=dataset_source,
-        persona_id=selected_persona.id,
-        persona_name=getattr(selected_persona, "name", None),
-        prompt_mode=prompt_mode,
-        system_prompt=active_system_prompt,
-        messages=chat_state["messages"],
-        generation=_generation_dict(gen_kwargs, advanced_generation),
-    )
     st.rerun()

tabs/compare_chat.py

@@ -4,18 +4,18 @@ 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.chat_export import save_chat_export
 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,
+    build_chat_messages,
+    generation_dict,
+    render_chat_message,
+    render_chat_window,
+    render_persona_prompt_controls,
+    render_system_prompt,
 )
 
 
@@ -47,7 +47,7 @@ def _generate_panel_reply(
 ) -> ChatReply:
     return generate_chat_reply(
         model=model,
-        messages=_build_chat_messages(panel_prompt, panel_state["messages"]),
+        messages=build_chat_messages(panel_prompt, panel_state["messages"]),
         remote=remote,
         past_key_values=panel_state["past_key_values"],
         **gen_kwargs,
@@ -90,17 +90,28 @@ def render_compare_mode(
     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)
+
+        # 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")
+
         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(
+        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,
@@ -117,7 +128,7 @@ def render_compare_mode(
 
         chat_log = st.container()
         with chat_log:
-            active_system_prompt = _render_system_prompt(
+            active_system_prompt = render_system_prompt(
                 prompt_key,
                 prompt_mode,
                 active_system_prompt,
@@ -220,10 +231,10 @@ def render_compare_mode(
                     ):
                         msg.pop("_needs_contrast", None)
                         continue
-                    context_a = _build_chat_messages(
+                    context_a = build_chat_messages(
                         left_prompt, left_state["messages"][:msg_idx]
                     )
-                    context_b = _build_chat_messages(
+                    context_b = build_chat_messages(
                         right_prompt, right_state["messages"][:msg_idx]
                     )
                     try:
@@ -256,7 +267,7 @@ def render_compare_mode(
         panel_edit_key,
         _,
     ) in panels:
-        _render_chat_window(
+        render_chat_window(
             chat_log=panel_log,
             messages=panel_state["messages"],
             chat_state=panel_state,
@@ -267,6 +278,9 @@ def render_compare_mode(
         )
 
     footer = st.container()
+    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
     with footer:
         exp_col, rst_col, _spacer = st.columns([0.5, 0.5, 10], gap="xsmall")
         with exp_col:
@@ -280,7 +294,7 @@ def render_compare_mode(
                     ("left", left_state, left_prompt, left_persona),
                     ("right", right_state, right_prompt, right_persona),
                 ):
-                    _save_chat_export_message(
+                    save_chat_export(
                         model_name=model_name,
                         dataset_source=dataset_source,
                         persona_id=panel_persona.id,
@@ -288,15 +302,21 @@ def render_compare_mode(
                         prompt_mode=panel_state["prompt_mode"],
                         system_prompt=panel_prompt,
                         messages=panel_state["messages"],
-                        generation=_generation_dict(gen_kwargs, advanced_generation),
+                        generation=generation_dict(gen_kwargs, advanced_generation),
                         panel_label=side,
                     )
                 st.toast("Exported", icon=":material/check:")
         with rst_col:
+            popover_key = widget_key(
+                context_key,
+                "cmp_reset_menu",
+                str(st.session_state[reset_menu_nonce_key]),
+            )
             with st.popover(
                 "",
                 icon=":material/delete_sweep:",
                 help="Reset chat",
+                key=popover_key,
             ):
                 if st.button(
                     "Reset left",
@@ -310,6 +330,7 @@ def render_compare_mode(
                         left_prompt_key,
                         left_pending_key,
                     )
+                    st.session_state[reset_menu_nonce_key] += 1
                     st.rerun()
                 if st.button(
                     "Reset right",
@@ -323,6 +344,7 @@ def render_compare_mode(
                         right_prompt_key,
                         right_pending_key,
                     )
+                    st.session_state[reset_menu_nonce_key] += 1
                     st.rerun()
                 if st.button(
                     "Reset both",
@@ -345,6 +367,7 @@ def render_compare_mode(
                         right_prompt_key,
                         right_pending_key,
                     )
+                    st.session_state[reset_menu_nonce_key] += 1
                     st.rerun()
 
     user_prompt = st.chat_input(
@@ -360,11 +383,11 @@ def render_compare_mode(
     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})
+            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"])
+        build_chat_messages(panel_prompt, panel_state["messages"])
         for panel_state, _panel_log, panel_prompt, _p_pending, _panel_edit_key, _ in panels
     ]
 

utils/chat.py

@@ -73,7 +73,7 @@ def _format_plain_messages(
     return "\n\n".join(lines)
 
 
-def _format_generation_prompt(
+def format_generation_prompt(
     messages: list[dict[str, str]], tokenizer: object
 ) -> tuple[str, int]:
     """Render messages into a single prompt string and count prompt tokens.
@@ -169,7 +169,7 @@ def generate_chat_reply(
     """
 
     tokenizer = model.tokenizer
-    prompt, prompt_token_count = _format_generation_prompt(messages, tokenizer)
+    prompt, prompt_token_count = format_generation_prompt(messages, tokenizer)
 
     generation_kwargs: dict[str, object] = {
         "max_new_tokens": max_new_tokens,

utils/contrast.py

@@ -1,7 +1,3 @@
-# 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.
 
@@ -15,13 +11,16 @@ 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
 
 import torch
 from nnterp import StandardizedTransformer
 
-from utils.chat import _format_generation_prompt
+from utils.chat import format_generation_prompt
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -48,6 +47,7 @@ def _normalise_diffs(diffs: torch.Tensor) -> list[float]:
 
 
 def _decode_ids(tokenizer: object, ids: list[int]) -> str:
+    """Decode token IDs, falling back when clean_up_tokenization_spaces is unsupported."""
     try:
         return tokenizer.decode(
             ids,
@@ -79,15 +79,18 @@ def _prepare_trace_text(
     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_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"
+        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),
         )
     n_ctx = len(context_ids)
     n_resp = len(response_ids)