Performance improvements

.env.example

@@ -18,3 +18,8 @@ ARTIFACTS_DIR=artifacts
 # Default model IDs shown in the sidebar (optional — change to override the built-in defaults)
 # DEFAULT_MODEL=google/gemma-2-2b-it
 # REMOTE_DEFAULT_MODEL=google/gemma-2-9b-it
+
+# Cache sizing knobs (optional)
+# Keep model cache at 1 unless you have enough RAM for multiple loaded models.
+# PERSONA_UI_MODEL_CACHE_ENTRIES=1
+# PERSONA_UI_STORE_CACHE_ENTRIES=4

app.py

@@ -5,6 +5,7 @@ import streamlit as st
 from dotenv import load_dotenv
 
 from utils.helpers import DATASET_SOURCES, session_key
+from utils.preload import preload_once
 from utils.runtime import list_remote_models
 from utils.theme import install_catppuccin_theme
 
@@ -28,6 +29,15 @@ _SIDEBAR_DATASET_SOURCE_KEY = session_key("sidebar", "dataset_source")
 
 _TABS = ["Chat", "Analysis", "Extract"]
 _TAB_ICONS = [":material/chat:", ":material/search:", ":material/tune:"]
+_TAB_PRELOAD_MODULES = {
+    "Chat": ("tabs.analysis_core", "tabs.extract", "tabs.compare_chat"),
+    "Analysis": ("tabs.chat", "tabs.extract"),
+    "Extract": ("tabs.chat", "tabs.analysis_core"),
+}
+_TAB_PRELOAD_FUNCTIONS = {
+    "Chat": ("utils.analysis_metadata:synth_persona_attribute_names",),
+    "Extract": ("utils.analysis_metadata:synth_persona_attribute_names",),
+}
 
 
 @dataclass(frozen=True)
@@ -181,6 +191,12 @@ def main() -> None:
 
         render_chat_tab(sidebar.remote, sidebar.model_name, sidebar.dataset_source)
 
+    preload_once(
+        f"after-{sidebar.active_tab.lower()}",
+        modules=_TAB_PRELOAD_MODULES.get(sidebar.active_tab, ()),
+        functions=_TAB_PRELOAD_FUNCTIONS.get(sidebar.active_tab, ()),
+    )
+
 
 if __name__ == "__main__":
     main()

tabs/analysis_core.py

@@ -7,7 +7,7 @@ from pathlib import Path
 import plotly.graph_objects as go
 import streamlit as st
 from persona_data.environment import get_artifacts_dir
-from persona_data.synth_persona import BASELINE_PERSONA_ID, SynthPersonaDataset
+from persona_data.synth_persona import BASELINE_PERSONA_ID
 from persona_vectors.attributes import (
     DEFAULT_MAX_ATTRIBUTE_CATEGORIES,
     attribute_color_kwargs,
@@ -45,6 +45,10 @@ from utils.analysis_sources import (
     store_id,
     store_layers_cached,
 )
+from utils.analysis_metadata import (
+    synth_persona_attribute_names,
+    synth_persona_dataset_cached,
+)
 from utils.controls import render_mask_strategy_select
 from utils.helpers import (
     ANALYSIS_HELP_TEXT,
@@ -99,9 +103,8 @@ _PROJECTION_COLOR_MODES = ["Persona", "K-means clusters", "Persona attribute"]
 _MAX_ATTRIBUTE_CATEGORIES = DEFAULT_MAX_ATTRIBUTE_CATEGORIES
 
 
-@st.cache_resource(show_spinner=False)
-def _synth_persona_dataset() -> SynthPersonaDataset:
-    return SynthPersonaDataset()
+def _synth_persona_dataset():
+    return synth_persona_dataset_cached()
 
 
 def _is_assistant_persona(persona_id: str, persona_name: str | None = None) -> bool:
@@ -983,7 +986,7 @@ def _render_projection_color_config(
 
     if color_mode == "Persona attribute":
         persona_dataset = _synth_persona_dataset()
-        attribute_options = list(persona_dataset.attribute_names)
+        attribute_options = list(synth_persona_attribute_names())
         if not attribute_options:
             st.info("No persona attributes are available for this dataset.")
             return None

tabs/chat.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import streamlit as st
-from persona_data.synth_persona import PersonaData
 
 from state import (
     ChatState,
@@ -29,6 +28,9 @@ from utils.chat_export import save_chat_export
 from utils.helpers import session_key, widget_key
 from utils.runtime import cached_model
 
+if TYPE_CHECKING:
+    from persona_data.synth_persona import PersonaData
+
 _LAST_PERSONA_ID_KEY = session_key("chat", "last_persona_id")
 _LAST_PROMPT_MODE_KEY = session_key("chat", "last_prompt_mode")
 _LAST_COMPARE_MODE_KEY = session_key("chat", "last_compare_mode")

tabs/chat_shared.py

@@ -2,9 +2,9 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import streamlit as st
-from persona_data.synth_persona import PersonaData
 
 from state import ChatState
 from tabs.chat_ui import GenerationConfig, render_persona_prompt_controls
@@ -12,6 +12,9 @@ from utils.chat import ChatReply, generate_chat_reply
 from utils.datasets import load_persona_list
 from utils.helpers import session_key
 
+if TYPE_CHECKING:
+    from persona_data.synth_persona import PersonaData
+
 
 @dataclass(frozen=True)
 class ChatSelection:

tabs/chat_ui.py

@@ -5,7 +5,6 @@ from dataclasses import asdict, dataclass
 from typing import TYPE_CHECKING, Any
 
 import streamlit as st
-from persona_data.synth_persona import PersonaData
 
 from utils.helpers import (
     CHAT_PROMPT_MODE_LABEL_TO_KEY,
@@ -16,6 +15,7 @@ from utils.helpers import (
 )
 
 if TYPE_CHECKING:
+    from persona_data.synth_persona import PersonaData
     from utils.contrast import TokenContrast
 
 GENERATION_DEFAULTS = {

tabs/compare_chat.py

@@ -1,9 +1,9 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import streamlit as st
-from nnterp import StandardizedTransformer
-from persona_data.synth_persona import PersonaData
 
 from state import ChatState, default_chat_state, reset_chat_context_state
 from tabs.chat_shared import (
@@ -24,6 +24,10 @@ from .chat_ui import (
     render_system_prompt,
 )
 
+if TYPE_CHECKING:
+    from nnterp import StandardizedTransformer
+    from persona_data.synth_persona import PersonaData
+
 
 @dataclass(frozen=True)
 class ComparePanel:

utils/analysis_metadata.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Any
+
+
+@lru_cache(maxsize=1)
+def synth_persona_dataset_cached() -> Any:
+    from persona_data.synth_persona import SynthPersonaDataset
+
+    return SynthPersonaDataset()
+
+
+@lru_cache(maxsize=1)
+def synth_persona_attribute_names() -> tuple[str, ...]:
+    return tuple(synth_persona_dataset_cached().attribute_names)

utils/analysis_sources.py

@@ -11,6 +11,8 @@ from persona_vectors.artifacts import (
 from persona_vectors.extraction import MaskStrategy
 from persona_vectors.hub import list_hub_vector_models
 
+from utils.helpers import env_int
+
 Store = ActivationStore | HFActivationStore
 
 DEFAULT_HUB_REPO = os.environ.get(
@@ -23,7 +25,10 @@ SOURCE_LOCAL = "Local activations"
 SOURCES = (SOURCE_HUB, SOURCE_LOCAL)
 
 
-@st.cache_resource(show_spinner=False, max_entries=1)
+_STORE_CACHE_ENTRIES = env_int("PERSONA_UI_STORE_CACHE_ENTRIES", 4)
+
+
+@st.cache_resource(show_spinner=False, max_entries=_STORE_CACHE_ENTRIES)
 def activation_store_cached(
     source: str,
     location: str,

utils/chat.py

@@ -3,14 +3,14 @@ from __future__ import annotations
 import logging
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
-import torch
 from persona_data.prompts import format_messages, format_prompt, normalize_messages
-from persona_data.synth_persona import PersonaData
 
 if TYPE_CHECKING:
+    import torch
     from nnterp import StandardizedTransformer
+    from persona_data.synth_persona import PersonaData
 
 logger = logging.getLogger(__name__)
 SystemPromptMode = Literal["empty", "templated", "biography", "custom"]
@@ -19,7 +19,7 @@ SystemPromptMode = Literal["empty", "templated", "biography", "custom"]
 @dataclass
 class ChatReply:
     text: str
-    generated_ids: torch.Tensor | None = None
+    generated_ids: Any | None = None
 
 
 def build_chat_messages(
@@ -133,6 +133,8 @@ def format_generation_prompt(
 
 def resolve_saved_tensor(value: object) -> torch.Tensor:
     """Resolve an nnsight ``.save()`` proxy (or raw tensor) to a CPU tensor."""
+    import torch
+
     resolved = value.value if getattr(value, "value", None) is not None else value
     if not isinstance(resolved, torch.Tensor):
         raise TypeError(f"Trace result did not resolve to a tensor: {type(resolved)!r}")
@@ -158,6 +160,8 @@ def _seeded_rng(seed: int | None):
         yield
         return
 
+    import torch
+
     cuda_ctx = torch.random.fork_rng(devices=range(torch.cuda.device_count()))
     mps_ctx = (
         torch.random.fork_rng(devices=range(1), device_type="mps")
@@ -203,6 +207,8 @@ def generate_chat_reply(
         ChatReply with generated text and token ids.
     """
 
+    import torch
+
     tokenizer = model.tokenizer
     prompt, prompt_token_count = format_generation_prompt(messages, tokenizer)
 

utils/helpers.py

@@ -1,9 +1,17 @@
+from __future__ import annotations
+
 import hashlib
+import logging
+import os
 import re
 from collections.abc import Iterable
 from enum import Enum
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from persona_data.synth_persona import PersonaData
 
-from persona_data.synth_persona import PersonaData
+logger = logging.getLogger(__name__)
 
 
 class DatasetSource(str, Enum):
@@ -74,6 +82,16 @@ def session_key(*parts: str) -> str:
     return ":".join(parts)
 
 
+def env_int(name: str, default: int, *, minimum: int = 1) -> int:
+    """Read a bounded integer from the environment."""
+
+    try:
+        return max(minimum, int(os.environ.get(name, str(default))))
+    except ValueError:
+        logger.warning("Ignoring invalid integer for %s", name)
+        return default
+
+
 def personas_fingerprint(persona_ids: Iterable[str]) -> str:
     """Stable short fingerprint for a set of persona ids.
 

utils/preload.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import importlib
+import logging
+import threading
+import time
+from collections.abc import Iterable
+
+logger = logging.getLogger(__name__)
+
+_started: set[tuple[str, ...]] = set()
+_lock = threading.Lock()
+
+
+def _warm_imports(
+    modules: tuple[str, ...],
+    functions: tuple[str, ...],
+    delay_seconds: float,
+) -> None:
+    if delay_seconds > 0:
+        time.sleep(delay_seconds)
+    for module in modules:
+        try:
+            importlib.import_module(module)
+        except Exception:
+            logger.debug("Background preload failed for %s", module, exc_info=True)
+    for function_path in functions:
+        try:
+            module_name, function_name = function_path.split(":", 1)
+            function = getattr(importlib.import_module(module_name), function_name)
+            function()
+        except Exception:
+            logger.debug(
+                "Background preload failed for %s", function_path, exc_info=True
+            )
+
+
+def preload_once(
+    name: str,
+    *,
+    modules: Iterable[str] = (),
+    functions: Iterable[str] = (),
+    delay_seconds: float = 0.25,
+) -> None:
+    """Warm small predictable costs on a daemon thread after the visible render.
+
+    Keep this limited to imports and tiny local metadata. Avoid model
+    construction, Hub requests, and Streamlit cache population because those can
+    steal enough CPU or I/O to make the visible page feel slower.
+    """
+
+    module_tuple = tuple(dict.fromkeys(modules))
+    function_tuple = tuple(dict.fromkeys(functions))
+    if not module_tuple and not function_tuple:
+        return
+
+    key = (name, *module_tuple, *function_tuple)
+    with _lock:
+        if key in _started:
+            return
+        _started.add(key)
+
+    thread = threading.Thread(
+        target=_warm_imports,
+        args=(module_tuple, function_tuple, delay_seconds),
+        name=f"persona-ui-preload-{name}",
+        daemon=True,
+    )
+    thread.start()

utils/runtime.py

@@ -4,9 +4,12 @@ from collections.abc import Iterable
 
 import streamlit as st
 
+from utils.helpers import env_int
+
 logger = logging.getLogger(__name__)
 _LANGUAGE_MODEL_CLASSES = {"LanguageModel", "StandardizedTransformer"}
 _EXPECTED_NDIF_STATES = {"RUNNING", "NOT DEPLOYED", "DEPLOYING", "DELETING"}
+_MODEL_CACHE_ENTRIES = env_int("PERSONA_UI_MODEL_CACHE_ENTRIES", 1)
 
 
 def _iter_deployments(raw: object) -> Iterable[dict]:
@@ -91,16 +94,17 @@ def list_remote_models() -> list[str]:
     return sorted(set(model_names))
 
 
-@st.cache_resource(show_spinner=False, max_entries=1)
+@st.cache_resource(show_spinner=False, max_entries=_MODEL_CACHE_ENTRIES)
 def cached_model(model_name: str):
     """Load and cache a standardized nnterp model.
 
     Streamlit reruns this app on every interaction, so caching keeps one loaded
-    model instance per model name instead of reloading weights on every widget
-    change. ``remote`` is intentionally not part of the cache key: it matters
-    at generation/trace time, but the current ``StandardizedTransformer``
+    model instance instead of reloading weights on every widget change.
+    ``remote`` is intentionally not part of the cache key: it matters at
+    generation/trace time, but the current ``StandardizedTransformer``
     constructor ignores it, and excluding it avoids loading duplicate local
-    model objects when toggling NDIF.
+    model objects when toggling NDIF. The cache defaults to one model to avoid
+    keeping multiple large models in RAM.
     """
 
     import torch