Adding color theme option

Support light and dark color themes

.streamlit/config.toml

@@ -1,9 +1,20 @@
-# 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.
+# Catppuccin theme. `base` is the default flavor; the built-in Streamlit
+# Settings menu (Light / Dark / System) switches between the two variants
+# below. Latte = light, Mocha = dark. Keep in sync with utils/theme.py.
 [theme]
 base = "dark"
-primaryColor = "#89b4fa"          # Mocha blue
-backgroundColor = "#1e1e2e"       # base
-secondaryBackgroundColor = "#313244"  # surface0
-textColor = "#cdd6f4"             # text
 font = "sans serif"
+
+[theme.light]                              # Catppuccin Latte (style-guide mapping)
+primaryColor = "#8839ef"                   # mauve   - accent (Catppuccin default)
+backgroundColor = "#eff1f5"                # base    - background pane
+secondaryBackgroundColor = "#dce0e8"       # crust   - secondary panes
+textColor = "#4c4f69"                      # text    - body copy
+linkColor = "#1e66f5"                      # blue    - links/URLs
+borderColor = "#9ca0b0"                    # overlay0 - inactive border
+
+[theme.dark]                               # Catppuccin Mocha (unchanged)
+primaryColor = "#89b4fa"                   # blue
+backgroundColor = "#1e1e2e"                # base
+secondaryBackgroundColor = "#313244"       # surface0
+textColor = "#cdd6f4"                      # text

app.py

@@ -12,7 +12,7 @@ from utils.analysis_sources import (
 from utils.helpers import DATASET_SOURCES, session_key, widget_key
 from utils.preload import preload_once
 from utils.runtime import list_remote_models
-from utils.theme import install_catppuccin_theme
+from utils.theme import active_base, install_catppuccin_theme
 
 load_dotenv()
 DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL", "google/gemma-2-2b-it")
@@ -233,7 +233,7 @@ def main() -> None:
     """Run the Streamlit app."""
 
     st.set_page_config(page_title="Persona UI", layout="wide")
-    install_catppuccin_theme(st.get_option("theme.base"))
+    install_catppuccin_theme(active_base())
 
     sidebar = _sidebar_controls()
 

tabs/analysis_core.py

@@ -58,7 +58,7 @@ from utils.helpers import (
     slugify,
     widget_key,
 )
-from utils.theme import style_plotly_layer_controls
+from utils.theme import active_base, style_plotly_layer_controls
 
 
 def _filename(*parts: str) -> str:
@@ -586,7 +586,7 @@ def _render_save_buttons(
 
 
 def _style_plotly_figures(figs: list[object]) -> None:
-    base = st.get_option("theme.base")
+    base = active_base()
     for fig in figs:
         if isinstance(fig, go.Figure):
             style_plotly_layer_controls(fig, base)

tabs/extract.py

@@ -14,7 +14,11 @@ from persona_vectors.extraction import (
 from persona_vectors.preview import TokenSegment, preview_token_segments
 
 from utils.controls import render_mask_strategy_select
-from utils.datasets import load_dataset, load_persona_list_from_dataset
+from utils.datasets import (
+    load_dataset,
+    load_persona_list_from_dataset,
+    warm_qa_in_background,
+)
 from utils.helpers import (
     NDIF_STATUS_ICONS,
     persona_label,
@@ -23,6 +27,7 @@ from utils.helpers import (
     widget_key,
 )
 from utils.runtime import cached_model
+from utils.theme import active_base
 
 _LAST_VARIANTS_KEY = "extract:last_variants"
 _LAST_BASELINE_KEY = "extract:last_include_baseline"
@@ -138,6 +143,10 @@ def _load_qa_dataset_personas(
             "Try another dataset source or check that the personas file is not empty."
         )
         return None
+
+    # Extract is the only tab that needs QA; warm it now so the parse overlaps
+    # with the user configuring the run instead of blocking the first extract.
+    warm_qa_in_background(dataset)
     return dataset, personas
 
 
@@ -171,7 +180,7 @@ _MAX_PREVIEW_SAMPLES = 3
 
 
 def _preview_palette():
-    flavor = PALETTE.latte if st.get_option("theme.base") == "light" else PALETTE.mocha
+    flavor = PALETTE.latte if active_base() == "light" else PALETTE.mocha
     return flavor.colors
 
 

utils/contrast.py

@@ -243,10 +243,24 @@ 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.
     """
+    # The model often opens a response with newline tokens; under pre-wrap
+    # those render as blank lines before the first word. Drop leading
+    # whitespace-only tokens (and left-trim the first visible one) so the
+    # contrast starts at real content. Display-only — weights stay aligned.
+    items = list(
+        zip(result.tokens, result.weights, result.raw_diffs, strict=True)
+    )
+    start = 0
+    while start < len(items) and not items[start][0].strip():
+        start += 1
+    if start >= len(items):
+        start = 0  # all-whitespace response: render as-is, not blank
+    items = items[start:]
+
     spans: list[str] = []
-    for token, weight, raw in zip(
-        result.tokens, result.weights, result.raw_diffs, strict=True
-    ):
+    for idx, (token, weight, raw) in enumerate(items):
+        if idx == 0:
+            token = token.lstrip()
         bg = _weight_to_bg(weight)
         tip = escape(f"Δlog P(A−B): {raw:+.3f}")
         text = escape(token)

utils/datasets.py

@@ -1,6 +1,7 @@
 import atexit
 import hashlib
 import shutil
+import threading
 from pathlib import Path
 from tempfile import mkdtemp
 from typing import Any
@@ -23,6 +24,30 @@ def _cached_dataset(cls: type) -> Any:
     return cls()
 
 
+_qa_warm_lock = threading.Lock()
+
+
+def warm_qa_in_background(dataset: Any) -> None:
+    """Trigger the dataset's lazy QA parse on a daemon thread, once.
+
+    QA loading is deferred in persona-data (large, unused outside Extract).
+    Kicking it off when the Extract tab opens means the parse overlaps with
+    the user picking personas/options instead of blocking the first run.
+    Idempotent across Streamlit reruns: guarded per cached dataset instance.
+    """
+
+    warm = getattr(dataset, "_load_qa", None)
+    if warm is None:
+        return  # persona-only dataset (e.g. Nemotron) has no QA
+    with _qa_warm_lock:
+        if getattr(dataset, "_qa_warm_started", False):
+            return
+        dataset._qa_warm_started = True
+    threading.Thread(
+        target=warm, name="persona-ui-warm-qa", daemon=True
+    ).start()
+
+
 @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."""

utils/theme.py

@@ -2,9 +2,19 @@
 
 import plotly.graph_objects as go
 import plotly.io as pio
+import streamlit as st
 from catppuccin import PALETTE
 
 
+def active_base() -> str:
+    """Return the live theme flavor ("light"/"dark").
+
+    Reflects the user's choice in Streamlit's built-in Settings-menu theme
+    toggle, falling back to the configured ``theme.base`` default.
+    """
+    return st.context.theme.type or st.get_option("theme.base")
+
+
 def _flavor(base: str | None):
     return PALETTE.latte if base == "light" else PALETTE.mocha
 
@@ -17,8 +27,14 @@ def install_catppuccin_theme(base: str | None = None) -> None:
     per-figure code.
     """
     c = _flavor(base).colors
-    bg, surface, line = c.base.hex, c.surface0.hex, c.surface1.hex
-    text, subtext = c.text.hex, c.subtext1.hex
+    bg, text, subtext = c.base.hex, c.text.hex, c.subtext1.hex
+    if base == "light":
+        # Latte's surface0/1 are heavy grays; use lighter slots so the grid
+        # stays subtle, the legend isn't a gray block, and diverging plots
+        # fade through near-white instead of mud.
+        surface, line, mid = c.mantle.hex, c.surface0.hex, c.base.hex
+    else:
+        surface, line, mid = c.surface0.hex, c.surface1.hex, c.surface0.hex
 
     axis = dict(
         gridcolor=line,
@@ -61,7 +77,7 @@ def install_catppuccin_theme(base: str | None = None) -> None:
             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]],
+                diverging=[[0.0, c.blue.hex], [0.5, mid], [1.0, c.red.hex]],
             ),
         )
     )