Testing newer version implemetation

pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "persona-vectors>=0.4.3",
-    "persona-data>=0.2.7",
+    "persona-data>=0.3.4",
     "streamlit>=1.44.0",
     "plotly>=6.6.0",
     "python-dotenv>=1.2.2",
@@ -15,7 +15,7 @@ dependencies = [
 
 # Local development:
 [tool.uv.sources]
-# persona-vectors = { path = "../persona-vectors", editable = true }
+persona-vectors = { path = "../persona-vectors", editable = true }
 # persona-data = { path = "../persona-data", editable = true }
 
 # [build-system]

state.py

@@ -1,7 +1,6 @@
 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:
@@ -35,38 +34,13 @@ def reset_chat_context_state(
         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."""
-
-    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
-
-
 def get_chat_state(
     model_name: str, remote: bool, dataset_source: str
 ) -> dict[str, object]:
     """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:
-        state.setdefault("messages", [])
-        state.setdefault("persona_id", None)
-        state.setdefault("prompt_mode", "templated")
-        state.setdefault("past_key_values", None)
-    _evict_inactive_kv_caches(key)
+    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

@@ -169,10 +169,9 @@ 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) -> dict[str, object]:
     return {
         "max_new_tokens": int(gen_kwargs["max_new_tokens"]),
-        "advanced_generation": bool(advanced_generation),
         "use_sampling": bool(gen_kwargs["do_sample"]),
         "temperature": float(gen_kwargs["temperature"]),
         "top_p": float(gen_kwargs["top_p"]),
@@ -258,12 +257,8 @@ def build_chat_messages(
 # ── Main tab entry point ───────────────────────────────────────────────────────
 
 
-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``
-    is True when any generation setting differs from its default.
-    """
+def _render_generation_settings(context_key: str, remote: bool) -> dict:
+    """Render the Advanced generation settings expander and return generation kwargs."""
     with st.expander("Advanced", expanded=False):
         config_col1, config_col2 = st.columns([2, 1])
         with config_col1:
@@ -349,19 +344,9 @@ def _render_generation_settings(context_key: str, remote: bool) -> tuple[dict, b
         if remote:
             st.caption("Seed is local-only and disabled for remote runs.")
 
-    advanced_generation = (
-        max_new_tokens != _GEN_DEFAULTS["max_new_tokens"]
-        or use_sampling
-        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
-    )
-
     do_sample = bool(use_sampling)
     generation_seed = seed if do_sample and seed is not None and not remote else None
-    gen_kwargs = {
+    return {
         "max_new_tokens": int(max_new_tokens),
         "do_sample": do_sample,
         "temperature": temperature,
@@ -370,7 +355,6 @@ def _render_generation_settings(context_key: str, remote: bool) -> tuple[dict, b
         "repetition_penalty": repetition_penalty,
         "seed": generation_seed,
     }
-    return gen_kwargs, advanced_generation
 
 
 def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
@@ -406,7 +390,7 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
         st.info("Try a different dataset source or upload a non-empty personas file.")
         return
 
-    gen_kwargs, advanced_generation = _render_generation_settings(context_key, remote)
+    gen_kwargs = _render_generation_settings(context_key, remote)
 
     # ── Mode toggle ───────────────────────────────────────────────────────────
     compare_key = widget_key(context_key, "compare_mode")
@@ -431,7 +415,6 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
             dataset_source,
             personas,
             gen_kwargs,
-            advanced_generation,
         )
         return
 
@@ -513,7 +496,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),
                 )
                 st.toast("Exported", icon=":material/check:")
         with rst_col:

tabs/compare.py

@@ -124,18 +124,17 @@ def _render_mask_strategy_select(scope: str) -> MaskStrategy:
         MaskStrategy.ANSWER_MEAN.value,
     )
     strategies = list(MaskStrategy)
-    default_index = next(
-        (
-            idx
-            for idx, strategy in enumerate(strategies)
-            if strategy.value == last_strategy
-        ),
-        0,
-    )
     selected = st.selectbox(
         "Mask strategy",
         options=strategies,
-        index=default_index,
+        index=next(
+            (
+                idx
+                for idx, strategy in enumerate(strategies)
+                if strategy.value == last_strategy
+            ),
+            0,
+        ),
         format_func=lambda strategy: strategy.value.replace("_", " ").title(),
         key=widget_key("load", "mask_strategy", scope),
         help="Which extracted activation artifact set to load.",

tabs/compare_chat.py

@@ -1,3 +1,5 @@
+from typing import Any, NamedTuple
+
 import streamlit as st
 from nnterp import StandardizedTransformer
 from persona_data.synth_persona import PersonaData
@@ -19,22 +21,26 @@ from .chat import (
 )
 
 
-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)
+class ComparePanel(NamedTuple):
+    side: str
+    state: dict[str, object]
+    log: Any
+    prompt: str | None
+    persona: PersonaData
+    prompt_key: str
+    edit_key: str
+    pending_key: str
+
+
+def _reset_compare_panel(panel: ComparePanel) -> None:
+    reset_chat_context_state(
+        panel.state,
+        panel.persona.id,
+        panel.state["prompt_mode"],
+        panel.prompt_key,
+        panel.pending_key,
+    )
+    st.session_state.pop(panel.edit_key, None)
 
 
 def _generate_panel_reply(
@@ -61,7 +67,6 @@ def render_compare_mode(
     dataset_source: str,
     personas: list[PersonaData],
     gen_kwargs: dict,
-    advanced_generation: bool,
 ) -> None:
     """Render the full side-by-side comparison UI."""
     model: StandardizedTransformer | None = None
@@ -85,19 +90,15 @@ def render_compare_mode(
         ),
     )
 
-    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]:
+    def render_panel(side: str) -> ComparePanel:
         panel_key = widget_key(context_key, f"cmp_{side}")
-        state = _panel_state(panel_key)
+        if panel_key not in st.session_state:
+            st.session_state[panel_key] = default_chat_state()
+        state = st.session_state[panel_key]
+
+        prompt_key = widget_key(panel_key, "custom_prompt")
+        edit_key = widget_key(panel_key, "edit_idx")
+        pending_key = widget_key(panel_key, "pending_regen")
 
         # Carry over persona / prompt selections across model or remote switches.
         persist_persona_key = f"chat:last_cmp_{side}_persona"
@@ -106,10 +107,6 @@ def render_compare_mode(
             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(
             personas,
             state["persona_id"],
@@ -122,11 +119,7 @@ def render_compare_mode(
 
         if changed:
             reset_chat_context_state(
-                state,
-                selected_persona.id,
-                prompt_mode,
-                prompt_key,
-                pending_regen_key,
+                state, selected_persona.id, prompt_mode, prompt_key, pending_key
             )
             st.session_state.pop(edit_key, None)
 
@@ -137,83 +130,57 @@ def render_compare_mode(
         chat_log = st.container()
         with chat_log:
             active_system_prompt = render_system_prompt(
-                prompt_key,
-                prompt_mode,
-                active_system_prompt,
+                prompt_key, prompt_mode, active_system_prompt
             )
-        return (
-            state,
-            chat_log,
-            active_system_prompt,
-            pending_regen_key,
-            selected_persona,
+        return ComparePanel(
+            side=side,
+            state=state,
+            log=chat_log,
+            prompt=active_system_prompt,
+            persona=selected_persona,
+            prompt_key=prompt_key,
+            edit_key=edit_key,
+            pending_key=pending_key,
         )
 
+    left_col, right_col = st.columns(2)
     with left_col:
-        left_state, left_log, left_prompt, left_pending, left_persona = render_panel(
-            "left"
-        )
+        left = 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,
-        ),
-    ]
+        right = render_panel("right")
+    panels: list[ComparePanel] = [left, right]
 
     # 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)
-    ]
+    regen_panels = [p for p in panels if st.session_state.pop(p.pending_key, False)]
     if regen_panels:
         model = _get_model()
 
         results: list[ChatReply | Exception] = []
         with st.spinner("Regenerating..."):
-            for panel_state, _panel_log, panel_prompt in regen_panels:
+            for panel in regen_panels:
                 try:
                     results.append(
                         _generate_panel_reply(
                             model=model,
                             remote=remote,
-                            panel_state=panel_state,
-                            panel_prompt=panel_prompt,
+                            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
-        ):
+        for panel, result in zip(regen_panels, results):
             if isinstance(result, Exception):
-                with panel_log:
+                with panel.log:
                     st.error(f"Generation failed: {result}")
-                panel_state["messages"].pop()
+                panel.state["messages"].pop()
                 continue
-            panel_state["messages"].append(
+            panel.state["messages"].append(
                 {"role": "assistant", "content": result.text}
             )
-            panel_state["past_key_values"] = (
+            panel.state["past_key_values"] = (
                 result.past_key_values if not remote else None
             )
         st.rerun()
@@ -222,28 +189,28 @@ def render_compare_mode(
     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"])
+            for panel_idx, panel 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 = _get_model()
-            label_a = persona_label(left_persona)
-            label_b = persona_label(right_persona)
+            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"]
+                    panel = panels[panel_idx]
+                    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]
+                        left.prompt, left.state["messages"][:msg_idx]
                     )
                     context_b = build_chat_messages(
-                        right_prompt, right_state["messages"][:msg_idx]
+                        right.prompt, right.state["messages"][:msg_idx]
                     )
                     try:
                         response_ids = model.tokenizer(
@@ -267,20 +234,13 @@ def render_compare_mode(
                     msg.pop("_needs_contrast", None)
             st.rerun()
 
-    for (
-        panel_state,
-        panel_log,
-        _panel_prompt,
-        panel_pending,
-        panel_edit_key,
-        _,
-    ) in panels:
+    for panel 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,
+            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,
             edit_column_ratio=(10, 1),
         )
@@ -298,20 +258,17 @@ def render_compare_mode(
                 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),
-                ):
+                for panel in panels:
                     save_chat_export(
                         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,
+                        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),
+                        panel_label=panel.side,
                     )
                 st.toast("Exported", icon=":material/check:")
         with rst_col:
@@ -326,55 +283,21 @@ def render_compare_mode(
                 help="Reset chat",
                 key=popover_key,
             ):
-                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.session_state[reset_menu_nonce_key] += 1
-                    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.session_state[reset_menu_nonce_key] += 1
-                    st.rerun()
+                for panel in panels:
+                    if st.button(
+                        f"Reset {panel.side}",
+                        key=widget_key(context_key, f"cmp_reset_{panel.side}"),
+                    ):
+                        _reset_compare_panel(panel)
+                        st.session_state[reset_menu_nonce_key] += 1
+                        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,
-                    )
+                    for panel in panels:
+                        _reset_compare_panel(panel)
                     st.session_state[reset_menu_nonce_key] += 1
                     st.rerun()
 
@@ -388,36 +311,27 @@ def render_compare_mode(
 
     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:
+    for panel 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
+        build_chat_messages(panel.prompt, panel.state["messages"]) for panel 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:
+        # Sequential generation keeps both panels using model/session state safely.
+        for panel in panels:
             try:
                 results.append(
                     _generate_panel_reply(
                         model=model,
                         remote=remote,
-                        panel_state=panel_state,
-                        panel_prompt=panel_prompt,
+                        panel_state=panel.state,
+                        panel_prompt=panel.prompt,
                         gen_kwargs=gen_kwargs,
                     )
                 )
@@ -425,23 +339,16 @@ def render_compare_mode(
                 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):
+    for panel, result in zip(panels, results):
         if isinstance(result, Exception):
-            with panel_log:
+            with panel.log:
                 st.error(f"Generation failed: {result}")
-            panel_state["messages"].pop()
+            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
+        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.
@@ -458,14 +365,14 @@ def render_compare_mode(
                     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),
+                    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
+                    left.state["messages"][-1]["_contrast"] = tc_a
                 if tc_b is not None:
-                    right_state["messages"][-1]["_contrast"] = tc_b
+                    right.state["messages"][-1]["_contrast"] = tc_b
             except Exception as exc:
                 st.warning(f"Token contrast failed: {exc}")
 

tabs/extract.py

@@ -32,11 +32,18 @@ _LAST_VARIANTS_KEY = "extract:last_variants"
 _LAST_BASELINE_KEY = "extract:last_include_baseline"
 _LAST_PERSONA_IDS_KEY = "extract:last_persona_ids"
 _LAST_QA_TYPE_KEY = "extract:last_qa_type"
-_LAST_DIFFICULTY_KEY = "extract:last_difficulty"
+_LAST_ITEM_TYPE_KEY = "extract:last_item_type"
+_LAST_SCOPE_KEY = "extract:last_scope"
 _LAST_MAX_QUESTIONS_KEY = "extract:last_max_questions"
 _LAST_MASK_STRATEGY_KEY = "extract:last_mask_strategy"
 
 _QA_TYPE_OPTIONS = ["all", "explicit", "implicit"]
+_ITEM_TYPE_OPTIONS = ["all", "mcq", "frq"]
+_SCOPE_OPTIONS = ["all", "individual", "shared"]
+
+
+def _option_index(options: list[str], value: str) -> int:
+    return options.index(value) if value in options else 0
 
 
 def _build_run_plan(
@@ -106,31 +113,25 @@ def _render_sample_tokens_html(p, tokenizer, *, max_tokens: int = 200) -> str:
     )
 
 
-def _render_local_dataset_uploads() -> None:
-    """Render file inputs for local dataset uploads."""
-
-    with st.expander("Local dataset upload", expanded=True):
-        st.file_uploader(
-            "personas.jsonl",
-            type=["jsonl"],
-            key="extract__personas_file",
-            help="Expected fields: id, persona, templated_view, biography_view",
-        )
-        st.file_uploader(
-            "qa.jsonl",
-            type=["jsonl"],
-            key="extract__qa_file",
-            help="Expected fields: id, qid, type, question, answer, difficulty",
-        )
-
-
 def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> None:
     """Render the extraction tab."""
 
     st.title("Extract")
 
     if dataset_source == "Local JSONL upload":
-        _render_local_dataset_uploads()
+        with st.expander("Local dataset upload", expanded=True):
+            st.file_uploader(
+                "personas.jsonl",
+                type=["jsonl"],
+                key="extract__personas_file",
+                help="Expected fields: id, persona, templated_view, biography_view",
+            )
+            st.file_uploader(
+                "qa.jsonl",
+                type=["jsonl"],
+                key="extract__qa_file",
+                help="Expected fields: id, qid, type, item_type, scope, question, answer",
+            )
 
     last_variants = st.session_state.get(
         _LAST_VARIANTS_KEY, [*PERSONA_VARIANTS, BASELINE_PERSONA_ID]
@@ -146,12 +147,12 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         key=_extract_widget_key(model_name, remote, dataset_source, "persona_variants"),
         help="Extract these variants for each selected persona.",
     )
-    include_baseline_default = st.session_state.get(
-        _LAST_BASELINE_KEY, BASELINE_PERSONA_ID in last_variants
-    )
     include_baseline = st.checkbox(
         "Extract Assistant baseline",
-        value=include_baseline_default,
+        value=st.session_state.get(
+            _LAST_BASELINE_KEY,
+            BASELINE_PERSONA_ID in last_variants,
+        ),
         key=_extract_widget_key(model_name, remote, dataset_source, "baseline"),
         help=(
             "Extracts the persona-less Assistant prompt once using the first "
@@ -214,18 +215,15 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
     with st.expander("Advanced", expanded=False):
         st.caption("Filters")
 
-        col1, col2, _ = st.columns([2, 2, 1])
+        col1, col2, col3 = st.columns(3)
         with col1:
-            last_qa_type = st.session_state.get(_LAST_QA_TYPE_KEY, "all")
-            qa_type_index = (
-                _QA_TYPE_OPTIONS.index(last_qa_type)
-                if last_qa_type in _QA_TYPE_OPTIONS
-                else 0
-            )
             qa_type_select = st.selectbox(
                 "QA type",
                 options=_QA_TYPE_OPTIONS,
-                index=qa_type_index,
+                index=_option_index(
+                    _QA_TYPE_OPTIONS,
+                    st.session_state.get(_LAST_QA_TYPE_KEY, "all"),
+                ),
                 key=_extract_widget_key(
                     model_name, remote, dataset_source, "qa_type_select"
                 ),
@@ -237,39 +235,68 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
                 else None
             )
         with col2:
-            last_difficulty = st.session_state.get(_LAST_DIFFICULTY_KEY, [1, 2, 3])
-            default_difficulty = [d for d in last_difficulty if d in (1, 2, 3)] or [
-                1,
-                2,
-                3,
-            ]
-            difficulty_values = st.multiselect(
-                "Difficulty",
-                options=[1, 2, 3],
-                default=default_difficulty,
+            item_type_select = st.selectbox(
+                "Item type",
+                options=_ITEM_TYPE_OPTIONS,
+                index=_option_index(
+                    _ITEM_TYPE_OPTIONS,
+                    st.session_state.get(_LAST_ITEM_TYPE_KEY, "all"),
+                ),
+                key=_extract_widget_key(
+                    model_name, remote, dataset_source, "item_type_select"
+                ),
+            )
+            st.session_state[_LAST_ITEM_TYPE_KEY] = item_type_select
+            qa_filter_item_type: Literal["mcq", "frq"] | None = (
+                cast(Literal["mcq", "frq"], item_type_select)
+                if item_type_select in ("mcq", "frq")
+                else None
+            )
+        with col3:
+            scope_select = st.selectbox(
+                "Scope",
+                options=_SCOPE_OPTIONS,
+                index=_option_index(
+                    _SCOPE_OPTIONS,
+                    st.session_state.get(_LAST_SCOPE_KEY, "all"),
+                ),
                 key=_extract_widget_key(
-                    model_name, remote, dataset_source, "difficulty_select"
+                    model_name,
+                    remote,
+                    dataset_source,
+                    "scope_select",
                 ),
             )
-            st.session_state[_LAST_DIFFICULTY_KEY] = difficulty_values
-            qa_filter_difficulty = difficulty_values if difficulty_values else None
+            st.session_state[_LAST_SCOPE_KEY] = scope_select
+            qa_filter_scope: Literal["individual", "shared"] | None = (
+                cast(Literal["individual", "shared"], scope_select)
+                if scope_select in ("individual", "shared")
+                else None
+            )
 
         st.caption("Extraction settings")
         last_strategy = st.session_state.get(
-            _LAST_MASK_STRATEGY_KEY, MaskStrategy.ANSWER_MEAN.value
+            _LAST_MASK_STRATEGY_KEY,
+            MaskStrategy.ANSWER_MEAN.value,
         )
         strategy_options = list(MaskStrategy)
-        strategy_index = next(
-            (i for i, s in enumerate(strategy_options) if s.value == last_strategy),
-            0,
-        )
         mask_strategy = st.selectbox(
             "Mask strategy",
             options=strategy_options,
-            index=strategy_index,
+            index=next(
+                (
+                    idx
+                    for idx, strategy in enumerate(strategy_options)
+                    if strategy.value == last_strategy
+                ),
+                0,
+            ),
             format_func=lambda s: s.value.replace("_", " ").title(),
             key=_extract_widget_key(
-                model_name, remote, dataset_source, "mask_strategy"
+                model_name,
+                remote,
+                dataset_source,
+                "mask_strategy",
             ),
             help="Which tokens contribute to the averaged hidden state.",
         )
@@ -279,7 +306,10 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         for persona in selected_personas:
             qa = list(
                 dataset.get_qa(
-                    persona.id, type=qa_filter_type, difficulty=qa_filter_difficulty
+                    persona.id,
+                    type=qa_filter_type,
+                    item_type=qa_filter_item_type,
+                    scope=qa_filter_scope,
                 )
             )
             if qa:
@@ -295,13 +325,14 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
             return
 
         max_q = min(len(qa_pairs) for _, qa_pairs in runs)
-        last_max = st.session_state.get(_LAST_MAX_QUESTIONS_KEY, max_q)
-        default_max = min(max(last_max, 1), max_q)
         max_questions = st.slider(
             "Max questions",
             min_value=1,
             max_value=max_q,
-            value=default_max,
+            value=min(
+                max(st.session_state.get(_LAST_MAX_QUESTIONS_KEY, max_q), 1),
+                max_q,
+            ),
             key=_extract_widget_key(
                 model_name, remote, dataset_source, "max_questions"
             ),

uv.lock

@@ -122,11 +122,11 @@ wheels = [
 
 [[package]]
 name = "cachetools"
-version = "7.0.6"
+version = "7.1.1"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/76/7b/1755ed2c6bfabd1d98b37ae73152f8dcf94aa40fee119d163c19ed484704/cachetools-7.0.6.tar.gz", hash = "sha256:e5d524d36d65703a87243a26ff08ad84f73352adbeafb1cde81e207b456aaf24", size = 37526, upload-time = "2026-04-20T19:02:23.289Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/ff/e2/85f227594656000ff4d8adadae91a21f536d4a84c6c716a86bd6685874be/cachetools-7.1.1.tar.gz", hash = "sha256:27bdf856d68fd3c71c26c01b5edc312124ed427524d1ddb31aa2b7746fe20d4b", size = 40202, upload-time = "2026-05-03T20:00:29.391Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/fe/c4/cf76242a5da1410917107ff14551764aa405a5fd10cd10cf9a5ca8fa77f4/cachetools-7.0.6-py3-none-any.whl", hash = "sha256:4e94956cfdd3086f12042cdd29318f5ced3893014f7d0d059bf3ead3f85b7f8b", size = 13976, upload-time = "2026-04-20T19:02:21.187Z" },
+    { 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]]
@@ -511,15 +511,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/5d/13/ad7d7ca3808a898b4612b6fe93cde56b53f3034dcde235acb1f0e1df24c6/idna-3.13-py3-none-any.whl", hash = "sha256:892ea0cde124a99ce773decba204c5552b69c3c67ffd5f232eb7696135bc8bb3", size = 68629, upload-time = "2026-04-22T16:42:40.909Z" },
 ]
 
-[[package]]
-name = "iniconfig"
-version = "2.3.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
-]
-
 [[package]]
 name = "ipython"
 version = "9.13.0"
@@ -565,14 +556,14 @@ wheels = [
 
 [[package]]
 name = "jedi"
-version = "0.19.2"
+version = "0.20.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "parso" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/72/3a/79a912fbd4d8dd6fbb02bf69afd3bb72cf0c729bb3063c6f4498603db17a/jedi-0.19.2.tar.gz", hash = "sha256:4770dc3de41bde3966b02eb84fbcf557fb33cce26ad23da12c742fb50ecb11f0", size = 1231287, upload-time = "2024-11-11T01:41:42.873Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/46/b7/a3635f6a2d7cf5b5dd98064fc1d5fbbafcb25477bcea204a3a92145d158b/jedi-0.20.0.tar.gz", hash = "sha256:c3f4ccbd276696f4b19c54618d4fb18f9fc24b0aef02acf704b23f487daa1011", size = 3119416, upload-time = "2026-05-01T23:38:47.814Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/c0/5a/9cac0c82afec3d09ccd97c8b6502d48f165f9124db81b4bcb90b4af974ee/jedi-0.19.2-py2.py3-none-any.whl", hash = "sha256:a8ef22bde8490f57fe5c7681a3c83cb58874daf72b4784de3cce5b6ef6edb5b9", size = 1572278, upload-time = "2024-11-11T01:41:40.175Z" },
+    { url = "https://files.pythonhosted.org/packages/9a/93/242e2eab5fe682ffcb8b0084bde703a41d51e17ee0f3a31ff0d9d813620a/jedi-0.20.0-py2.py3-none-any.whl", hash = "sha256:7bdd9c2634f56713299976f4cbd59cb3fa92165cc5e05ea811fb253480728b67", size = 4884812, upload-time = "2026-05-01T23:38:43.919Z" },
 ]
 
 [[package]]
@@ -625,18 +616,17 @@ wheels = [
 
 [[package]]
 name = "kaleido"
-version = "1.2.0"
+version = "1.3.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "choreographer" },
     { name = "logistro" },
     { name = "orjson" },
     { name = "packaging" },
-    { name = "pytest-timeout" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/38/ad/76eec859b71eda803a88ea50ed3f270281254656bb23d19eb0a39aa706a0/kaleido-1.2.0.tar.gz", hash = "sha256:fa621a14423e8effa2895a2526be00af0cf21655be1b74b7e382c171d12e71ef", size = 64160, upload-time = "2025-11-04T21:24:23.833Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/e0/64/53eac73d31dbfc3310ee2e87bcac1ae7417427f0fbe3dd800eaf676db324/kaleido-1.3.0.tar.gz", hash = "sha256:5e0378a7475e98852773deeb6483dee91f8aa7b364dde7b5f2b3622cb468a3e6", size = 68938, upload-time = "2026-05-04T19:45:28.932Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/4b/97/f6de8d4af54d6401d6581a686cce3e3e2371a79ba459a449104e026c08bc/kaleido-1.2.0-py3-none-any.whl", hash = "sha256:c27ed82b51df6b923d0e656feac221343a0dbcd2fb9bc7e6b1db97f61e9a1513", size = 68997, upload-time = "2025-11-04T21:24:21.704Z" },
+    { url = "https://files.pythonhosted.org/packages/9e/b9/a6d8bb7d228940f01885bd9f327ab7f9d366a9be775c4bf366bf9d9477ae/kaleido-1.3.0-py3-none-any.whl", hash = "sha256:52714dfd38e8f2a114831826200c40bb10d0ca0c11d4272f3f48ad499cd8f8ea", size = 55580, upload-time = "2026-05-04T19:45:27.483Z" },
 ]
 
 [[package]]
@@ -797,7 +787,7 @@ wheels = [
 
 [[package]]
 name = "nnsight"
-version = "0.6.3"
+version = "0.7.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "accelerate" },
@@ -813,24 +803,24 @@ dependencies = [
     { name = "transformers" },
     { name = "zstandard" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/a6/75/e2f5c0913b078e08718e92cf51ddf963579dde59485b024052c7d11bd25d/nnsight-0.6.3.tar.gz", hash = "sha256:d9fa2aa0eed3b3335bbe5d7511a5127ccc3c79bafba6275369697b882e2b8249", size = 1263681, upload-time = "2026-03-20T16:46:33.329Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/a6/9e/76fd632deef926599d3d16c02fd736cf5f8465f49d708d5be13cd6638484/nnsight-0.7.0.tar.gz", hash = "sha256:5bc6678d567ecc5590b823b7bbab2c310c69f8dda6f4064684c6d488563eebee", size = 1912851, upload-time = "2026-05-05T05:40:50.084Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/e4/81/6dacc633467d4b91a7b46871f10be7fa83d1648986b0640fc9946d153c05/nnsight-0.6.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6f8ab8bb002da985bc0ca550f43b9f978d876269953c09709de32afb9c2378ad", size = 176203, upload-time = "2026-03-20T16:46:13.683Z" },
-    { url = "https://files.pythonhosted.org/packages/4d/51/e7d148ca030c7b1ff269268b82c918db88c95a7f378cd26b5d5885345fbf/nnsight-0.6.3-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4ea008fa818a6a0027a260399f8e921264b85d21cc3e97980a2c3e838068356a", size = 183668, upload-time = "2026-03-20T16:46:14.87Z" },
-    { url = "https://files.pythonhosted.org/packages/0c/0b/57b48373160978dbb37ff0ece06a6d740922deda15745e7afa20659e46a1/nnsight-0.6.3-cp312-cp312-win32.whl", hash = "sha256:2c6be467ac7c5d0ecbb196d251ae07cfc176da9b6b2dbe40f6db887f404d97d7", size = 177888, upload-time = "2026-03-20T16:46:16.138Z" },
-    { url = "https://files.pythonhosted.org/packages/87/8e/a6e739011fe6395c11b5570c2397c5a37c976ad541af820f690489864e7a/nnsight-0.6.3-cp312-cp312-win_amd64.whl", hash = "sha256:ac77ab7717c886412cf022bbd73c7da60e921200f6857255f1779fc40f3a7deb", size = 178305, upload-time = "2026-03-20T16:46:17.082Z" },
-    { url = "https://files.pythonhosted.org/packages/43/b3/62404d8019cd68339d3279c95bc993076445532b5879fea5893c9c03f0e6/nnsight-0.6.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2c70554b13996d33d46b05caa064de411e29440d356d61dbefecf8e5b8cd8987", size = 176185, upload-time = "2026-03-20T16:46:17.964Z" },
-    { url = "https://files.pythonhosted.org/packages/64/6d/8496e6957698b16cf9ad58d68efa8c5b1e24abdab3d306611d7935d560f8/nnsight-0.6.3-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:cf078a2bec36f07183b06fbc049ffd617d21770d2ff927fe722e6a374403c81f", size = 183717, upload-time = "2026-03-20T16:46:18.987Z" },
-    { url = "https://files.pythonhosted.org/packages/f7/8b/a79760b3af92a4110e1deabde3a5f39ddfeb585aaf7033323a120a4cc034/nnsight-0.6.3-cp313-cp313-win32.whl", hash = "sha256:96af51d141079b514488a36771713783cbc3b68aa2a818e33860b2bf2ba3bdc9", size = 177882, upload-time = "2026-03-20T16:46:19.966Z" },
-    { url = "https://files.pythonhosted.org/packages/c9/97/48309d97608a921e5450ea6fc568affca8793905772025d98a14d594a129/nnsight-0.6.3-cp313-cp313-win_amd64.whl", hash = "sha256:2c4b2d6b354c7e8552bd455d090460c38b98cfabee3fccbe079f10839df7b0c2", size = 178306, upload-time = "2026-03-20T16:46:21.237Z" },
-    { url = "https://files.pythonhosted.org/packages/65/0c/2713448586db6c594a5e4ccb84c137ac0ca83a188ee461a937880d763102/nnsight-0.6.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:76dce23d21641110efc6b6ec878c3cb21fd86548effe3ed1ba85d53d1a6794a2", size = 176196, upload-time = "2026-03-20T16:46:22.16Z" },
-    { url = "https://files.pythonhosted.org/packages/c7/0d/ff22ecc5a12075de4dda31ea32b88d945351da2669591485928bc3ab6a64/nnsight-0.6.3-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:cd1e4eb30640ca581731a6983d4f48504bc13b5cbe5ab3eb10fae232eb3741de", size = 183746, upload-time = "2026-03-20T16:46:23.104Z" },
-    { url = "https://files.pythonhosted.org/packages/ef/cb/473396dca499ea007e6f442015349ba89c4a023f72145af58c719c43abc5/nnsight-0.6.3-cp314-cp314-win32.whl", hash = "sha256:72c857c19b100feb598ca814bcec3436e091942c97dd6fa4d684ff6c797bbc2b", size = 177986, upload-time = "2026-03-20T16:46:24.391Z" },
-    { url = "https://files.pythonhosted.org/packages/3f/fb/3a02d1e8e9d0a40c9f7ca5468d7aca1a37a82684ad523f4b2c33bc6ff454/nnsight-0.6.3-cp314-cp314-win_amd64.whl", hash = "sha256:6bfafeb25841f4bbbcf8445154d3307a05876e7de1a6eaab995e520ebe9ac05f", size = 178413, upload-time = "2026-03-20T16:46:27.461Z" },
-    { url = "https://files.pythonhosted.org/packages/7c/4f/d02d99774ad8f779518c328604fe4414cf8e1f6cb3ae911e799280b172f1/nnsight-0.6.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:78b477920b3ecbfd5155126574321edf6b1d994e151a9bede6f9bc86794da8d6", size = 176308, upload-time = "2026-03-20T16:46:28.456Z" },
-    { url = "https://files.pythonhosted.org/packages/f4/84/6e7704256559e2eb54d79c7de290640d5fecd52d3e458e86866d78e8d5ad/nnsight-0.6.3-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:dfade5095d38dec52c439c738d1c7967abeef8394b242a973ac48a2a6b448217", size = 184544, upload-time = "2026-03-20T16:46:29.966Z" },
-    { url = "https://files.pythonhosted.org/packages/f9/0d/22a10ff3cbbc0f4a0036da287b303e5e5e789e28b9fe2defc0a14a2701bf/nnsight-0.6.3-cp314-cp314t-win32.whl", hash = "sha256:d477e2a0bb983b511f118a32d2180882f1345bd73c4e28ab060fb158e1e9b4e0", size = 178092, upload-time = "2026-03-20T16:46:31.216Z" },
-    { url = "https://files.pythonhosted.org/packages/df/4f/d66a76d2ecb7a309719da1e91becd81b8b7b801e60ec403f7405d21cc424/nnsight-0.6.3-cp314-cp314t-win_amd64.whl", hash = "sha256:b473c3b6cf50d057213b5d42bb0018116c24838a04c039e39745dd2a2b47ace5", size = 178527, upload-time = "2026-03-20T16:46:32.314Z" },
+    { url = "https://files.pythonhosted.org/packages/a8/4d/634c1f08e5d34d9d145f81a5b872c5d2a2fcae714c77450e4f14660670b8/nnsight-0.7.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d988fd477abcc15d413012aa1b6eaef667da81b2d61c86c528b2a154adf259c5", size = 265015, upload-time = "2026-05-05T05:40:30.453Z" },
+    { url = "https://files.pythonhosted.org/packages/f5/fa/aee0a17356528f6cf4e0deaf3af2f0a88f660f793fb1fbb10f363682ef6a/nnsight-0.7.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:000f5285a754d63fafc0a677e4518192f16c5fa3f2af347a46e2f27c793ef2b2", size = 272487, upload-time = "2026-05-05T05:40:31.545Z" },
+    { url = "https://files.pythonhosted.org/packages/85/77/c59d1e983936e77b2d9fcd4694461f78c66d0af95982b42ffc365fa0c224/nnsight-0.7.0-cp312-cp312-win32.whl", hash = "sha256:2601ccc8a352c09f6f17576ea07943ed29eb3c8d1bba2ec174fbb554942a3821", size = 267080, upload-time = "2026-05-05T05:40:32.909Z" },
+    { url = "https://files.pythonhosted.org/packages/f1/1b/d3d60dfbbd1167320ec7f02cf06f4c589a3d2d462c3748f45bba7eaa17eb/nnsight-0.7.0-cp312-cp312-win_amd64.whl", hash = "sha256:408e2ae8b0da3af0b71f3c8dbb2301b671801205f7a3f95ede7f3f6672c5a8a5", size = 267491, upload-time = "2026-05-05T05:40:34.469Z" },
+    { url = "https://files.pythonhosted.org/packages/d3/dd/2e2f800876f00a0f38e7ef5e536c53bacf9ac7775efc8b337ba117a4459c/nnsight-0.7.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:753acc23dcc97ed32bfc3f650400f049e94621dc56cd9a372dbbfef868f98753", size = 265005, upload-time = "2026-05-05T05:40:36.066Z" },
+    { url = "https://files.pythonhosted.org/packages/0a/e0/335c94482b3c739994cdabaef02a12458c806cbe7ca5883d03380f6fdf8a/nnsight-0.7.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3b59ff8a2d41ed482313fa21f62e5a035068bd68828201cccb928c6e0ded2b4f", size = 272540, upload-time = "2026-05-05T05:40:37.293Z" },
+    { url = "https://files.pythonhosted.org/packages/93/18/65f473ae3147156dec3e6d30b5fa386491474964db1af24f478ef467718d/nnsight-0.7.0-cp313-cp313-win32.whl", hash = "sha256:ed4c01b7cc882e3699de56f1ee77cedd7f720797bb3670a690ec1edeeb34f9bb", size = 267067, upload-time = "2026-05-05T05:40:38.35Z" },
+    { url = "https://files.pythonhosted.org/packages/67/0f/eb2a6cdff12abcd7264b6e9992c9fe907cfd4c445ea3df1b3aa4165061ff/nnsight-0.7.0-cp313-cp313-win_amd64.whl", hash = "sha256:8cd6b5aac59175a6c436fb4fca713a6a58fc85407f9b4f938a819a467c17108e", size = 267492, upload-time = "2026-05-05T05:40:39.66Z" },
+    { url = "https://files.pythonhosted.org/packages/ba/76/6148cd42a1323d218c64ed122f4da775233fb2e44edef0341fd8aa976c7e/nnsight-0.7.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:baaa6409717a8e95cadf7d71fbb376103524f9137a0cc0d8ea55c2656add4326", size = 265011, upload-time = "2026-05-05T05:40:40.693Z" },
+    { url = "https://files.pythonhosted.org/packages/29/a7/b1ba47acbe95fafc854170390d29ef3dd0bdaac44576f1554f25e53c8309/nnsight-0.7.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:22f172565dab55583fb2eb0d615697af76e2c22f12644309738bf7174090ccc6", size = 272569, upload-time = "2026-05-05T05:40:42.003Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/38/05845f3fc29631450efa5d9b8e17aaff37d2ae4e6eee8f8e38ce0ba5b162/nnsight-0.7.0-cp314-cp314-win32.whl", hash = "sha256:6b360febd90a1c5330c7e881a11aaa51f76f04e4876fe045c7f0481ded9e0398", size = 267182, upload-time = "2026-05-05T05:40:43.146Z" },
+    { url = "https://files.pythonhosted.org/packages/ef/00/9300ab917f759d2e4cf2c4ceb23b76827fe5e77379c1933f6b67c0bb680c/nnsight-0.7.0-cp314-cp314-win_amd64.whl", hash = "sha256:65a894add8c93b38d781affac4bf10a58050ed3e46497b2038febf9d8633ce21", size = 267606, upload-time = "2026-05-05T05:40:44.303Z" },
+    { url = "https://files.pythonhosted.org/packages/81/36/1f698a62ae01187dd6c4e63a561220ec7999ad489fea796df3733d700d36/nnsight-0.7.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:171a0a1a13f3f8ea8ddc06c98eb3f924c7a8dfcc108be1407c7c071504f0370f", size = 265117, upload-time = "2026-05-05T05:40:45.411Z" },
+    { url = "https://files.pythonhosted.org/packages/ec/2e/0df50b4b98896b02fc76dd08997aab977eef96e0fb3a2646c5c695a3389f/nnsight-0.7.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:eeed67fa7d8ef03a1196533ec341dbe1dfb40bb7615c8dee06a22009eba1be0f", size = 273366, upload-time = "2026-05-05T05:40:46.648Z" },
+    { url = "https://files.pythonhosted.org/packages/7f/da/f47e1102a9216dc487608ce5b7bc058afb0ea9f3f087e85f83ed877d156d/nnsight-0.7.0-cp314-cp314t-win32.whl", hash = "sha256:7140bf2cab8d4bff46e8f6a1a942185fcdf8f18f62ddb173ca7fe000b6e580cc", size = 267290, upload-time = "2026-05-05T05:40:47.956Z" },
+    { url = "https://files.pythonhosted.org/packages/bf/c9/749370592057f790fc60f2542848b09542ca8eb75b4583e1dde01e7e40c9/nnsight-0.7.0-cp314-cp314t-win_amd64.whl", hash = "sha256:544f7e32e277738ca10e85cdd3f07d8965ec3388854fdf14373308580f8fd0f4", size = 267722, upload-time = "2026-05-05T05:40:49.036Z" },
 ]
 
 [[package]]
@@ -1200,16 +1190,16 @@ wheels = [
 
 [[package]]
 name = "parso"
-version = "0.8.6"
+version = "0.8.7"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/81/76/a1e769043c0c0c9fe391b702539d594731a4362334cdf4dc25d0c09761e7/parso-0.8.6.tar.gz", hash = "sha256:2b9a0332696df97d454fa67b81618fd69c35a7b90327cbe6ba5c92d2c68a7bfd", size = 401621, upload-time = "2026-02-09T15:45:24.425Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/30/4b/90c937815137d43ce71ba043cd3566221e9df6b9c805f24b5d138c9d40a7/parso-0.8.7.tar.gz", hash = "sha256:eaaac4c9fdd5e9e8852dc778d2d7405897ec510f2a298071453e5e3a07914bb1", size = 401824, upload-time = "2026-05-01T23:13:02.138Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/b6/61/fae042894f4296ec49e3f193aff5d7c18440da9e48102c3315e1bc4519a7/parso-0.8.6-py2.py3-none-any.whl", hash = "sha256:2c549f800b70a5c4952197248825584cb00f033b29c692671d3bf08bf380baff", size = 106894, upload-time = "2026-02-09T15:45:21.391Z" },
+    { url = "https://files.pythonhosted.org/packages/99/5d/8268b644392ee874ee82a635cd0df1773de230bde356c38de28e298392cc/parso-0.8.7-py2.py3-none-any.whl", hash = "sha256:a8926eb2a1b915486941fdbd31e86a4baf88fe8c210f25f2f35ecec5b574ca1c", size = 107025, upload-time = "2026-05-01T23:12:58.867Z" },
 ]
 
 [[package]]
 name = "persona-data"
-version = "0.2.7"
+version = "0.3.4"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "huggingface-hub" },
@@ -1218,9 +1208,9 @@ dependencies = [
     { name = "python-dotenv" },
     { name = "torch" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/e9/c1/8612da5d216e484f823333fcf6bc1ff97c266c800b597d70f9a19475b78e/persona_data-0.2.7.tar.gz", hash = "sha256:6cd95f82d628a2add68df3935ae47c36a13ab17bafaf495892be5b1192f4693e", size = 9143, upload-time = "2026-05-01T08:54:26.603Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/c5/03/25b6dbcbf4be440cf85e5dc5ee30eff69a5c2db638c5452b50932e758141/persona_data-0.3.4.tar.gz", hash = "sha256:c887b53a6fcddd80bac6485f6639902456cc7fd878af89ba9b8102d02aaa335e", size = 8533, upload-time = "2026-05-04T15:27:17.057Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/b0/ad/65fff568431ae5620f15d5b9c1dff9a941fc77f8c6ebb6f3131c9010d713/persona_data-0.2.7-py3-none-any.whl", hash = "sha256:4395f46381d54e7047bd26a4eaa2e9db1bd29f80e00a813b8be3ede0192aafcb", size = 11832, upload-time = "2026-05-01T08:54:25.778Z" },
+    { url = "https://files.pythonhosted.org/packages/c7/28/b45feb5186615b8a4ff647eaaf35739a35b1afe3fd35b3d09e2734475b2c/persona_data-0.3.4-py3-none-any.whl", hash = "sha256:666c34b8bd7588bfe2d55a62737cd7d87e0d663fbead3a4458d284f3e660bab2", size = 11178, upload-time = "2026-05-04T15:27:17.749Z" },
 ]
 
 [[package]]
@@ -1238,8 +1228,8 @@ dependencies = [
 
 [package.metadata]
 requires-dist = [
-    { name = "persona-data", specifier = ">=0.2.7" },
-    { name = "persona-vectors", specifier = ">=0.4.3" },
+    { name = "persona-data", specifier = ">=0.3.4" },
+    { name = "persona-vectors", editable = "../persona-vectors" },
     { name = "plotly", specifier = ">=6.6.0" },
     { name = "python-dotenv", specifier = ">=1.2.2" },
     { name = "streamlit", specifier = ">=1.44.0" },
@@ -1249,7 +1239,7 @@ requires-dist = [
 [[package]]
 name = "persona-vectors"
 version = "0.4.3"
-source = { registry = "https://pypi.org/simple" }
+source = { editable = "../persona-vectors" }
 dependencies = [
     { name = "kaleido" },
     { name = "nnsight" },
@@ -1265,9 +1255,22 @@ dependencies = [
     { name = "transformers" },
     { name = "umap-learn" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/15/97/5f89f74a98faec09696108a1b2222a5d0bc99c129e035939ba8a70a61314/persona_vectors-0.4.3.tar.gz", hash = "sha256:3b86bed93373868506f1f6006f65e849d560bd20fb771310a33a674f02a90b0b", size = 21614, upload-time = "2026-05-01T16:49:11.757Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/92/de/a666ee51901db501556e9f7014e7f61741df8f2b219ff676179b80e94cd4/persona_vectors-0.4.3-py3-none-any.whl", hash = "sha256:50bfb2de68613ec89a87f5773f4d2f5fe9981e2949f82e4823cd4f4ca093bee3", size = 25563, upload-time = "2026-05-01T16:49:12.694Z" },
+
+[package.metadata]
+requires-dist = [
+    { name = "kaleido", specifier = ">=1.0.0" },
+    { name = "nnsight", specifier = ">=0.6.1" },
+    { name = "nnterp", specifier = ">=1.3.0" },
+    { name = "persona-data", specifier = ">=0.3.4" },
+    { name = "plotly", specifier = ">=6.6.0" },
+    { name = "python-dotenv", specifier = ">=1.2.2" },
+    { name = "safetensors", specifier = ">=0.7.0" },
+    { name = "scikit-learn", specifier = ">=1.6.0" },
+    { name = "torch", specifier = ">=2.10.0" },
+    { name = "torchvision", specifier = ">=0.26.0" },
+    { name = "tqdm", specifier = ">=4.67.3" },
+    { name = "transformers", specifier = ">=5.2.0" },
+    { name = "umap-learn", specifier = ">=0.5.7" },
 ]
 
 [[package]]
@@ -1373,15 +1376,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/90/ad/cba91b3bcf04073e4d1655a5c1710ef3f457f56f7d1b79dcc3d72f4dd912/plotly-6.7.0-py3-none-any.whl", hash = "sha256:ac8aca1c25c663a59b5b9140a549264a5badde2e057d79b8c772ae2920e32ff0", size = 9898444, upload-time = "2026-04-09T20:36:39.812Z" },
 ]
 
-[[package]]
-name = "pluggy"
-version = "1.6.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
-]
-
 [[package]]
 name = "prompt-toolkit"
 version = "3.0.52"
@@ -1626,34 +1620,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b2/e6/94145d714402fd5ade00b5661f2d0ab981219e07f7db9bfa16786cdb9c04/pynndescent-0.6.0-py3-none-any.whl", hash = "sha256:dc8c74844e4c7f5cbd1e0cd6909da86fdc789e6ff4997336e344779c3d5538ef", size = 73511, upload-time = "2026-01-08T21:29:57.306Z" },
 ]
 
-[[package]]
-name = "pytest"
-version = "9.0.3"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32'" },
-    { name = "iniconfig" },
-    { name = "packaging" },
-    { name = "pluggy" },
-    { name = "pygments" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/549bd94f1a0a402dc8cf64563a117c0f3765662e2e668477624baeec44d5/pytest-9.0.3.tar.gz", hash = "sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c", size = 1572165, upload-time = "2026-04-07T17:16:18.027Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d4/24/a372aaf5c9b7208e7112038812994107bc65a84cd00e0354a88c2c77a617/pytest-9.0.3-py3-none-any.whl", hash = "sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9", size = 375249, upload-time = "2026-04-07T17:16:16.13Z" },
-]
-
-[[package]]
-name = "pytest-timeout"
-version = "2.4.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "pytest" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/ac/82/4c9ecabab13363e72d880f2fb504c5f750433b2b6f16e99f4ec21ada284c/pytest_timeout-2.4.0.tar.gz", hash = "sha256:7e68e90b01f9eff71332b25001f85c75495fc4e3a836701876183c4bcfd0540a", size = 17973, upload-time = "2025-05-05T19:44:34.99Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/fa/b6/3127540ecdf1464a00e5a01ee60a1b09175f6913f0644ac748494d9c4b21/pytest_timeout-2.4.0-py3-none-any.whl", hash = "sha256:c42667e5cdadb151aeb5b26d114aff6bdf5a907f176a007a30b940d3d865b5c2", size = 14382, upload-time = "2025-05-05T19:44:33.502Z" },
-]
-
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"
@@ -2564,11 +2530,11 @@ wheels = [
 
 [[package]]
 name = "wcwidth"
-version = "0.6.0"
+version = "0.7.0"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/35/a2/8e3becb46433538a38726c948d3399905a4c7cabd0df578ede5dc51f0ec2/wcwidth-0.6.0.tar.gz", hash = "sha256:cdc4e4262d6ef9a1a57e018384cbeb1208d8abbc64176027e2c2455c81313159", size = 159684, upload-time = "2026-02-06T19:19:40.919Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/2c/ee/afaf0f85a9a18fe47a67f1e4422ed6cf1fe642f0ae0a2f81166231303c52/wcwidth-0.7.0.tar.gz", hash = "sha256:90e3a7ea092341c44b99562e75d09e4d5160fe7a3974c6fb842a101a95e7eed0", size = 182132, upload-time = "2026-05-02T16:04:12.653Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/68/5a/199c59e0a824a3db2b89c5d2dade7ab5f9624dbf6448dc291b46d5ec94d3/wcwidth-0.6.0-py3-none-any.whl", hash = "sha256:1a3a1e510b553315f8e146c54764f4fb6264ffad731b3d78088cdb1478ffbdad", size = 94189, upload-time = "2026-02-06T19:19:39.646Z" },
+    { url = "https://files.pythonhosted.org/packages/41/52/e465037f5375f43533d1a80b6923955201596a99142ed524d77b571a1418/wcwidth-0.7.0-py3-none-any.whl", hash = "sha256:5d69154c429a82910e241c738cd0e2976fac8a2dd47a1a805f4afed1c0f136f2", size = 110825, upload-time = "2026-05-02T16:04:11.033Z" },
 ]
 
 [[package]]