Switching to uniform colortheme catpuccin

.streamlit/config.toml

@@ -0,0 +1,9 @@
+# Catppuccin Mocha theme. Switch base to "light" and swap the four colors
+# below to the Latte equivalents (see utils/theme.py) for the light flavor.
+[theme]
+base = "dark"
+primaryColor = "#89b4fa"          # Mocha blue
+backgroundColor = "#1e1e2e"       # base
+secondaryBackgroundColor = "#313244"  # surface0
+textColor = "#cdd6f4"             # text
+font = "sans serif"

app.py

@@ -1,5 +1,4 @@
 import os
-import threading
 
 import streamlit as st
 from dotenv import load_dotenv
@@ -17,22 +16,6 @@ _TABS = ["Chat", "Compare", "Extract"]
 _TAB_ICONS = [":material/chat:", ":material/search:", ":material/tune:"]
 
 
-def _preload_default_model() -> None:
-    """Background-warm the default local model so the first chat is instant."""
-    try:
-        import torch
-
-        torch.set_grad_enabled(False)
-        from utils.runtime import cached_model
-
-        cached_model(DEFAULT_MODEL)
-    except Exception:
-        pass
-
-
-threading.Thread(target=_preload_default_model, daemon=True).start()
-
-
 def _remote_model_input(remote_models: list[str]) -> str:
     """Return the active remote model id, picking from running NDIF deployments or a custom value."""
 
@@ -142,13 +125,17 @@ def _sidebar_controls() -> tuple[bool, str, str, str]:
 def main() -> None:
     """Run the Streamlit app."""
 
+    st.set_page_config(page_title="Persona UI", layout="wide")
+    from utils.theme import install_catppuccin_theme
+
+    install_catppuccin_theme(st.get_option("theme.base"))
+
     # Deferred: importing torch is slow; keep it after dotenv load (done at
     # module level above) so the Streamlit page config renders immediately.
     import torch
 
     torch.set_grad_enabled(False)
 
-    st.set_page_config(page_title="Persona UI", layout="wide")
     remote, model_name, dataset_source, active_tab = _sidebar_controls()
 
     if active_tab == "Extract":

pyproject.toml

@@ -10,6 +10,7 @@ dependencies = [
     "streamlit>=1.44.0",
     "plotly>=6.6.0",
     "python-dotenv>=1.2.2",
+    "catppuccin>=2.5.0",
 ]
 
 # Local development:

state.py

@@ -17,7 +17,6 @@ class ChatState(TypedDict):
     messages: list[ChatMessage]
     persona_id: str | None
     prompt_mode: str
-    past_key_values: object | None
 
 
 def chat_session_key(model_name: str, dataset_source: str) -> str:
@@ -31,7 +30,6 @@ def default_chat_state() -> ChatState:
         "messages": [],
         "persona_id": None,
         "prompt_mode": "templated",
-        "past_key_values": None,
     }
 
 
@@ -44,18 +42,15 @@ def reset_chat_context_state(
     """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 get_chat_state(model_name: str, remote: bool, dataset_source: str) -> ChatState:
+def get_chat_state(model_name: str, _remote: bool, dataset_source: str) -> ChatState:
     """Return the mutable chat state for the active context."""
 
     key = chat_session_key(model_name, dataset_source)
     state = st.session_state.setdefault(key, default_chat_state())
-    if remote and state.get("past_key_values") is not None:
-        state["past_key_values"] = None
     return state

tabs/chat.py

@@ -113,7 +113,6 @@ def _handle_single_chat_generation(
                 model=model,
                 messages=messages,
                 remote=remote,
-                past_key_values=chat_state["past_key_values"],
                 **generation.to_generate_kwargs(),
             )
         except Exception as exc:
@@ -125,7 +124,6 @@ def _handle_single_chat_generation(
             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()
 
 
@@ -218,6 +216,13 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
             prompt_key,
             prompt_mode,
             active_system_prompt,
+            on_save=lambda: reset_chat_context_state(
+                chat_state,
+                selected_persona.id,
+                prompt_mode,
+                chat_input_key,
+                pending_key,
+            ),
         )
 
     render_probe_inspector(
@@ -232,7 +237,6 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     render_chat_window(
         chat_log=chat_log,
         messages=chat_state["messages"],
-        chat_state=chat_state,
         edit_key=edit_key,
         pending_key=pending_key,
     )

tabs/chat_ui.py

@@ -1,3 +1,4 @@
+from collections.abc import Callable
 from dataclasses import asdict, dataclass
 from typing import Any
 
@@ -74,7 +75,6 @@ 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]
@@ -102,7 +102,6 @@ def _open_edit_dialog(
             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] = "regenerate_after_edit"
             st.rerun()
@@ -112,7 +111,12 @@ def _open_edit_dialog(
 
 
 @st.dialog("Edit system prompt", width="large")
-def _open_system_prompt_dialog(*, prompt_key: str, current_value: str) -> None:
+def _open_system_prompt_dialog(
+    *,
+    prompt_key: str,
+    current_value: str,
+    on_save: Callable[[], None] | None = None,
+) -> None:
     new_value = st.text_area(
         "System prompt",
         value=current_value,
@@ -123,6 +127,8 @@ def _open_system_prompt_dialog(*, prompt_key: str, current_value: str) -> None:
     with save_col:
         if st.button("Save", type="primary", use_container_width=True):
             st.session_state[prompt_key] = new_value
+            if on_save is not None:
+                on_save()
             st.rerun()
     with cancel_col:
         if st.button("Cancel", use_container_width=True):
@@ -329,7 +335,6 @@ def render_chat_window(
     *,
     chat_log: Any,
     messages: list[dict[str, str]],
-    chat_state: dict[str, object],
     edit_key: str,
     pending_key: str,
     show_contrast: bool = False,
@@ -354,7 +359,6 @@ def render_chat_window(
                     _open_edit_dialog(
                         msg_index=i,
                         messages=messages,
-                        chat_state=chat_state,
                         pending_key=pending_key,
                     )
 
@@ -372,6 +376,8 @@ def render_system_prompt(
     prompt_key: str,
     prompt_mode: str,
     active_system_prompt: str | None,
+    *,
+    on_save: Callable[[], None] | None = None,
 ) -> str | None:
     if prompt_key not in st.session_state:
         st.session_state[prompt_key] = active_system_prompt or ""
@@ -381,7 +387,11 @@ def render_system_prompt(
         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)
+            _open_system_prompt_dialog(
+                prompt_key=prompt_key,
+                current_value=current,
+                on_save=on_save,
+            )
     return st.session_state.get(prompt_key) or None
 
 

tabs/compare.py

@@ -42,6 +42,24 @@ def _filename(*parts: str) -> str:
 
 _list_layers_cached = st.cache_data(show_spinner=False)(list_local_layers)
 
+
+@st.cache_data(show_spinner=False)
+def _hub_layers_cached(
+    repo_id: str,
+    model_name: str,
+    mask_strategy_value: str,
+    variant: str,
+    persona_id: str,
+) -> list[int]:
+    store = HFActivationStore(
+        repo_id,
+        model_name,
+        mask_strategy=MaskStrategy(mask_strategy_value),
+    )
+    sample = store.load(variant, persona_id)
+    return list(range(int(sample.shape[0])))
+
+
 # Keep compare-tab selection state separate so projection defaults do not
 # overwrite cosine similarity defaults.
 _LAST_COSINE_PERSONAS_KEY = "compare:last_personas:cosine"
@@ -75,8 +93,13 @@ def _layers_for_variant(
     if isinstance(store, HFActivationStore):
         if not persona_ids:
             return []
-        sample = store.load(variant, persona_ids[0])
-        return list(range(int(sample.shape[0])))
+        return _hub_layers_cached(
+            store.repo_id,
+            store.model_name,
+            mask_strategy.value,
+            variant,
+            persona_ids[0],
+        )
     return _list_layers_cached(
         str(store.root_dir),
         store.model_name,
@@ -234,12 +257,20 @@ def _build_cosine_figures(
     store: Store,
     selection: CosineSelection,
 ) -> tuple[object, object | None, int, int] | None:
+    variant_sample_cache = {}
+
+    def _load_pair(left: str, right: str):
+        key = tuple(sorted((left, right)))
+        if key not in variant_sample_cache:
+            variant_sample_cache[key] = load_variant_vectors(
+                store,
+                [left, right],
+                persona_ids=selection.persona_ids,
+            )
+        return variant_sample_cache[key]
+
     try:
-        variant_samples = load_variant_vectors(
-            store,
-            [selection.variant_a, selection.variant_b],
-            persona_ids=selection.persona_ids,
-        )
+        variant_samples = _load_pair(selection.variant_a, selection.variant_b)
     except Exception as exc:
         st.error(f"Could not load vectors: {exc}")
         return None
@@ -266,15 +297,7 @@ def _build_cosine_figures(
     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_vectors(
-                    store,
-                    [left, right],
-                    persona_ids=selection.persona_ids,
-                )
-            )
+            pair_samples = _load_pair(left, right)
             pair_traces.append(
                 (
                     f"{prompt_variant_label(left)} vs {prompt_variant_label(right)}",

tabs/compare_chat.py

@@ -105,6 +105,12 @@ def _render_compare_panel(
             prompt_key,
             prompt_mode,
             active_system_prompt,
+            on_save=lambda: reset_chat_context_state(
+                state,
+                selected_persona.id,
+                prompt_mode,
+                pending_key,
+            ),
         )
 
     return ComparePanel(
@@ -138,7 +144,6 @@ def _generate_panels(
                             panel.prompt, panel.state["messages"]
                         ),
                         remote=remote,
-                        past_key_values=panel.state["past_key_values"],
                         **generation.to_generate_kwargs(),
                     )
                 )
@@ -151,7 +156,6 @@ 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] = []
@@ -165,7 +169,6 @@ def _apply_panel_results(
             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
 
@@ -242,7 +245,6 @@ def _render_compare_history(
         render_chat_window(
             chat_log=panel.log,
             messages=panel.state["messages"],
-            chat_state=panel.state,
             edit_key=panel.edit_key,
             pending_key=panel.pending_key,
             show_contrast=contrast_enabled,
@@ -409,7 +411,6 @@ def render_compare_mode(
         _apply_panel_results(
             panels=regen_panels,
             results=results,
-            remote=remote,
             rollback_user_on_error=False,
         )
         st.rerun()
@@ -452,7 +453,6 @@ def render_compare_mode(
     valid_results = _apply_panel_results(
         panels=panels,
         results=results,
-        remote=remote,
         rollback_user_on_error=True,
     )
     if contrast_enabled:

tabs/extract.py

@@ -12,7 +12,7 @@ from persona_vectors.extraction import (
 )
 from persona_vectors.preview import TokenSegment, preview_token_segments
 
-from utils.datasets import load_dataset
+from utils.datasets import load_dataset, load_persona_list
 from utils.helpers import (
     NDIF_STATUS_ICONS,
     persona_label,
@@ -111,6 +111,11 @@ def _load_qa_dataset_personas(
             personas_file=st.session_state.get("extract__personas_file"),
             qa_file=st.session_state.get("extract__qa_file"),
         )
+        personas, _ = load_persona_list(
+            dataset_source,
+            personas_file=st.session_state.get("extract__personas_file"),
+            qa_file=st.session_state.get("extract__qa_file"),
+        )
         st.caption(dataset_status)
     except Exception as exc:
         st.error(f"Could not load data: {exc}")
@@ -123,7 +128,6 @@ def _load_qa_dataset_personas(
         st.info("This dataset is persona-only for now. Use Chat to browse personas.")
         return None
 
-    personas = list(dataset)
     if not personas:
         st.warning("No personas found in the selected dataset.")
         st.info(

tabs/probe_ui.py

@@ -1,3 +1,5 @@
+import html
+
 import streamlit as st
 import torch
 
@@ -103,7 +105,10 @@ def _render_token_picker(trace: ConversationTrace, context_key: str) -> int:
     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)
+        token_repr = html.escape(token_repr)
+        parts.append(
+            f"<strong>[{token_repr}]</strong>" 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;"

utils/chat.py

@@ -15,7 +15,6 @@ SystemPromptMode = Literal["empty", "templated", "biography", "custom"]
 @dataclass
 class ChatReply:
     text: str
-    past_key_values: object | None
     generated_ids: torch.Tensor | None = None
 
 
@@ -171,7 +170,6 @@ def generate_chat_reply(
     model: StandardizedTransformer,
     messages: list[dict[str, str]],
     remote: bool,
-    past_key_values: object | None = None,
     max_new_tokens: int = 256,
     do_sample: bool = False,
     temperature: float = 1.0,
@@ -183,14 +181,12 @@ def generate_chat_reply(
     """Generate one assistant reply from a full chat history.
 
     The helper uses ``model.generate`` so it works with both local and NDIF-backed
-    nnsight models. The full conversation is re-rendered each turn and the cache from
-    the previous turn is reused when available.
+    nnsight models. The full conversation is re-rendered each turn.
 
     Args:
         model: Loaded standardized nnterp model.
         messages: Full chat history, including any system prompt as the first message.
         remote: Whether to execute the generation on NDIF.
-        past_key_values: Cache returned by the previous generation step.
         max_new_tokens: Maximum number of assistant tokens to generate.
         do_sample: Whether to sample from the model distribution.
         temperature: Sampling temperature, used only when sampling is enabled.
@@ -200,7 +196,7 @@ def generate_chat_reply(
         seed: Optional local RNG seed for sampled generation.
 
     Returns:
-        ChatReply with generated text and the updated cache.
+        ChatReply with generated text and token ids.
     """
 
     tokenizer = model.tokenizer
@@ -220,9 +216,6 @@ def generate_chat_reply(
         generation_kwargs["top_k"] = top_k
     if repetition_penalty != 1.0:
         generation_kwargs["repetition_penalty"] = repetition_penalty
-    if past_key_values is not None and not remote:
-        generation_kwargs["past_key_values"] = past_key_values
-
     # `remote` is captured by nnsight's RemoteableMixin.trace() and is NOT
     # forwarded to the underlying model's generate
     with _seeded_rng(seed if do_sample and not remote else None):
@@ -240,8 +233,5 @@ def generate_chat_reply(
     text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
     return ChatReply(
         text=text,
-        past_key_values=(
-            getattr(generated, "past_key_values", None) if not remote else None
-        ),
         generated_ids=generated_ids.detach().cpu(),
     )

utils/datasets.py

@@ -23,6 +23,13 @@ def _cached_dataset(cls: type) -> Any:
     return cls()
 
 
+@st.cache_resource(show_spinner=False)
+def _cached_local_dataset(personas_path: str, qa_path: str) -> LocalPersonaDataset:
+    """Instantiate and cache a local upload dataset for stable temp paths."""
+
+    return LocalPersonaDataset(personas_path=Path(personas_path), qa_path=Path(qa_path))
+
+
 def _upload_cache_dir() -> Path:
     cache_dir = st.session_state.get("_upload_cache_dir")
     if cache_dir is None:
@@ -35,14 +42,12 @@ def _upload_cache_dir() -> Path:
 
 def _uploaded_file_to_temp_path(uploaded_file: Any, stem: str) -> Path:
     suffix = Path(uploaded_file.name).suffix or ".jsonl"
-    temp_path = _upload_cache_dir() / f"{stem}{suffix}"
-    hash_path = temp_path.with_suffix(temp_path.suffix + ".sha256")
     data = uploaded_file.getvalue()
     digest = hashlib.sha256(data).hexdigest()
-    if temp_path.exists() and hash_path.exists() and hash_path.read_text() == digest:
+    temp_path = _upload_cache_dir() / f"{stem}_{digest[:16]}{suffix}"
+    if temp_path.exists():
         return temp_path
     temp_path.write_bytes(data)
-    hash_path.write_text(digest)
     return temp_path
 
 
@@ -95,7 +100,4 @@ def load_dataset(
 
     personas_path = _uploaded_file_to_temp_path(personas_file, stem="personas")
     qa_path = _uploaded_file_to_temp_path(qa_file, stem="qa")
-    return (
-        LocalPersonaDataset(personas_path=personas_path, qa_path=qa_path),
-        "Local upload",
-    )
+    return _cached_local_dataset(str(personas_path), str(qa_path)), "Local upload"

utils/theme.py

@@ -0,0 +1,67 @@
+"""Catppuccin Plotly template installer."""
+
+import plotly.graph_objects as go
+import plotly.io as pio
+from catppuccin import PALETTE
+
+
+def install_catppuccin_theme(base: str | None = None) -> None:
+    """Register a Catppuccin template and alias it as ``plotly_white``.
+
+    Call once at startup. Persona-vectors pins ``template="plotly_white"`` on
+    every figure, so replacing that entry themes all plots without any
+    per-figure code.
+    """
+    flavor = PALETTE.latte if base == "light" else PALETTE.mocha
+    c = flavor.colors
+    bg, surface, line = c.base.hex, c.surface0.hex, c.surface1.hex
+    text, subtext = c.text.hex, c.subtext1.hex
+
+    axis = dict(
+        gridcolor=line,
+        zerolinecolor=line,
+        linecolor=line,
+        tickcolor=line,
+        tickfont=dict(color=subtext),
+        title=dict(font=dict(color=text)),
+    )
+    scene_axis = dict(
+        backgroundcolor=bg,
+        gridcolor=line,
+        zerolinecolor=line,
+        showbackground=True,
+        color=text,
+        tickfont=dict(color=subtext),
+    )
+
+    template = go.layout.Template(
+        layout=dict(
+            paper_bgcolor=bg,
+            plot_bgcolor=bg,
+            font=dict(color=text),
+            colorway=[
+                c.blue.hex,
+                c.mauve.hex,
+                c.green.hex,
+                c.peach.hex,
+                c.teal.hex,
+                c.pink.hex,
+                c.yellow.hex,
+                c.sapphire.hex,
+                c.lavender.hex,
+                c.red.hex,
+                c.sky.hex,
+                c.maroon.hex,
+            ],
+            xaxis=axis,
+            yaxis=axis,
+            scene=dict(xaxis=scene_axis, yaxis=scene_axis, zaxis=scene_axis),
+            legend=dict(bgcolor=surface, bordercolor=line, font=dict(color=text)),
+            colorscale=dict(
+                diverging=[[0.0, c.blue.hex], [0.5, surface], [1.0, c.red.hex]],
+            ),
+        )
+    )
+    pio.templates["catppuccin"] = template
+    pio.templates["plotly_white"] = template
+    pio.templates.default = "catppuccin"

uv.lock

@@ -236,6 +236,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/bf/0f/f897abe4ea0a8c408ae65c8c83bffab4936ad65d6032d4fb4cd35bbdc3ee/cachetools-7.1.1-py3-none-any.whl", hash = "sha256:0335cd7a0952d2b22327441fb0628139e234c565559eeb91a8a4ac7551c5353d", size = 16775, upload-time = "2026-05-03T20:00:27.857Z" },
 ]
 
+[[package]]
+name = "catppuccin"
+version = "2.5.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/b6/31/87c3ca25d29678e076e1f0d151ef8792f0557b26a8dc865035b5f1fa96c4/catppuccin-2.5.0.tar.gz", hash = "sha256:3035f3bf35bc2369d1cb4c754272a494708592a1e62d42655a371863ac7c6834", size = 1918017, upload-time = "2025-08-03T20:06:40.089Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c8/9c/7908f34009eec72884fe9448a62188423a3621f04ee58e2abaec379c9cc0/catppuccin-2.5.0-py3-none-any.whl", hash = "sha256:74a1f1db79d527905225f1afce1c8858d06e2dd1231e4300a72a812797cdb572", size = 19600, upload-time = "2025-08-03T20:06:38.912Z" },
+]
+
 [[package]]
 name = "certifi"
 version = "2026.4.22"
@@ -1569,6 +1578,7 @@ name = "persona-ui"
 version = "0.3.0"
 source = { virtual = "." }
 dependencies = [
+    { name = "catppuccin" },
     { name = "persona-data" },
     { name = "persona-vectors" },
     { name = "plotly" },
@@ -1578,6 +1588,7 @@ dependencies = [
 
 [package.metadata]
 requires-dist = [
+    { name = "catppuccin", specifier = ">=2.5.0" },
     { name = "persona-data", specifier = ">=0.4.2" },
     { name = "persona-vectors", specifier = ">=0.6.4" },
     { name = "plotly", specifier = ">=6.6.0" },