Update code

README.md

@@ -1,5 +1,7 @@
 # Persona UI
 
+[![Deploy to Hugging Face Spaces](https://huggingface.co/landing/badge.svg)](https://huggingface.co/spaces/implicit-personalization/persona-ui)
+
 Streamlit interface for persona vector extraction, analysis, and chat.
 
 > [!WARNING]
@@ -74,6 +76,21 @@ parent/
 streamlit run app.py
 ```
 
+## Hugging Face Spaces Deployment
+
+This app can be deployed to Hugging Face Spaces using Docker.
+
+### Prerequisites
+
+No secrets needed! The dependencies are published on PyPI.
+
+### Build Locally (Optional)
+
+```bash
+docker build -t persona-ui .
+docker run -p 8501:8501 persona-ui
+```
+
 ## Configuration
 
 Copy `.env.example` to `.env` and fill in:

pyproject.toml

@@ -10,14 +10,15 @@ dependencies = [
     "streamlit>=1.44.0",
     "plotly>=6.6.0",
     "python-dotenv>=1.2.2",
+    "transformers>=5.5.0",
 ]
 
+
 [tool.uv.sources]
 # Local development:
-# persona-vectors = { path = "../persona-vectors", editable = true }
 # persona-data = { path = "../persona-data", editable = true }
-persona-vectors = { git = "ssh://git@github.com/implicit-personalization/persona-vectors.git" }
-persona-data = { git = "ssh://git@github.com/implicit-personalization/persona-data.git" }
+persona-vectors = { path = "../persona-vectors", editable = true }
+persona-data = { git = "https://github.com/implicit-personalization/persona-data.git" }
 
 # [build-system]
 # requires = ["uv_build>=0.11.3,<0.12"]

tabs/chat.py

@@ -38,7 +38,8 @@ def _render_collapsible_markdown(content: str) -> None:
 def _render_chat_message(message: dict[str, str]) -> None:
     if not message.get("content"):
         return
-    with st.chat_message(message["role"]):
+    with st.container(border=True):
+        st.caption(message["role"])
         _render_collapsible_markdown(message["content"])
 
 
@@ -46,47 +47,24 @@ def _render_inline_system_prompt(
     prompt_key: str,
     prompt_mode: str,
     active_system_prompt: str | None,
-    edit_key: str,
     height: int = 200,
 ) -> str | None:
-    """Render the system prompt as an inline editable item at the top of the chat."""
+    """Render the system prompt as an always-editable text area at the top of the chat."""
     if prompt_mode == "empty":
         return active_system_prompt
 
     if prompt_key not in st.session_state:
         st.session_state[prompt_key] = active_system_prompt or ""
 
-    current_prompt = st.session_state[prompt_key] or None
-    is_editing = st.session_state.get(edit_key) == -1
-
     with st.container(border=True):
         st.caption("System prompt")
-        if is_editing:
-            new_val = st.text_area(
-                "system_prompt_edit",
-                value=current_prompt or "",
-                height=height,
-                label_visibility="collapsed",
-                key=f"{prompt_key}_inline_edit",
-            )
-            c1, c2 = st.columns(2)
-            with c1:
-                if st.button("Save", key=f"{edit_key}_sys_save", type="primary"):
-                    st.session_state[prompt_key] = new_val
-                    st.session_state[edit_key] = None
-                    st.rerun()
-            with c2:
-                if st.button("Cancel", key=f"{edit_key}_sys_cancel"):
-                    st.session_state[edit_key] = None
-                    st.rerun()
-        else:
-            if current_prompt:
-                _render_collapsible_markdown(current_prompt)
-            else:
-                st.markdown("*(empty)*")
-            if st.button("Edit", key=f"{edit_key}_sys_edit"):
-                st.session_state[edit_key] = -1
-                st.rerun()
+        st.text_area(
+            "system_prompt_edit",
+            value=st.session_state[prompt_key],
+            height=height,
+            label_visibility="collapsed",
+            key=prompt_key,
+        )
 
     return st.session_state.get(prompt_key) or None
 
@@ -105,7 +83,8 @@ def _render_editable_message(
 
     is_editing = st.session_state.get(edit_key) == msg_index
 
-    with st.chat_message(message["role"]):
+    with st.container(border=True):
+        st.caption(message["role"])
         if is_editing:
             new_content = st.text_area(
                 "Edit",
@@ -305,14 +284,9 @@ def _render_compare_mode(
     """Render the full side-by-side comparison UI."""
     left_col, right_col = st.columns(2)
 
-    def render_panel(
-        side: str, column
-    ) -> tuple[dict[str, object], Any, str | None, str]:
+    def render_panel(side: str) -> tuple[dict[str, object], Any, str | None, str]:
         panel_key = widget_key(context_key, f"cmp_{side}")
-        state = st.session_state.get(panel_key)
-        if state is None:
-            state = _default_chat_state()
-            st.session_state[panel_key] = state
+        state = _panel_state(panel_key)
         prompt_key = widget_key(panel_key, "custom_prompt")
         show_all_key = widget_key(panel_key, "show_all")
         edit_key = widget_key(panel_key, "edit_idx")
@@ -374,7 +348,6 @@ def _render_compare_mode(
                 prompt_key,
                 prompt_mode,
                 active_system_prompt,
-                edit_key,
                 height=150,
             )
         _render_chat_window(
@@ -390,11 +363,9 @@ def _render_compare_mode(
         return state, chat_log, active_system_prompt, pending_regen_key
 
     with left_col:
-        left_state, left_log, left_prompt, left_pending = render_panel("left", left_col)
+        left_state, left_log, left_prompt, left_pending = render_panel("left")
     with right_col:
-        right_state, right_log, right_prompt, right_pending = render_panel(
-            "right", right_col
-        )
+        right_state, right_log, right_prompt, right_pending = render_panel("right")
 
     panels = [
         (left_state, left_log, left_prompt, left_pending),
@@ -454,12 +425,9 @@ def _render_compare_mode(
                     executor.submit(
                         generate_chat_reply,
                         model=model,
-                        messages=(
-                            [{"role": "system", "content": panel_prompt}]
-                            if panel_prompt
-                            else []
-                        )
-                        + panel_state["messages"],
+                        messages=_build_chat_messages(
+                            panel_prompt, panel_state["messages"]
+                        ),
                         remote=remote,
                         past_key_values=panel_state["past_key_values"],
                         **gen_kwargs,
@@ -479,12 +447,9 @@ def _render_compare_mode(
                     results.append(
                         generate_chat_reply(
                             model=model,
-                            messages=(
-                                [{"role": "system", "content": panel_prompt}]
-                                if panel_prompt
-                                else []
-                            )
-                            + panel_state["messages"],
+                            messages=_build_chat_messages(
+                                panel_prompt, panel_state["messages"]
+                            ),
                             remote=remote,
                             past_key_values=panel_state["past_key_values"],
                             **gen_kwargs,
@@ -507,36 +472,22 @@ def _render_compare_mode(
         with panel_log:
             _render_chat_message({"role": "assistant", "content": result.text})
 
+    # Rerun so the newly appended turns are redrawn through the editable history
+    # renderer instead of only appearing in the one-off generation pass.
+    st.rerun()
 
-# ── Main tab entry point ───────────────────────────────────────────────────────
-
-
-def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
-    """Render the chat tab."""
 
-    st.title("Chat")
+# ── Main tab entry point ───────────────────────────────────────────────────────
 
-    context_key = chat_session_key(model_name, dataset_source)
-    chat_state = get_chat_state(model_name, remote, dataset_source)
-    try:
-        dataset, dataset_status = load_dataset(
-            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}")
-        st.info("Check the selected dataset source or upload both JSONL files.")
-        return
 
-    personas = list(dataset)
-    if not personas:
-        st.warning("No personas found in the selected dataset.")
-        st.info("Try a different dataset source or upload a non-empty personas file.")
-        return
+def _render_generation_settings(
+    context_key: str, remote: bool
+) -> tuple[dict, bool]:
+    """Render the Advanced generation settings expander.
 
-    # ── Generation settings ───────────────────────────────────────────────────
+    Returns ``(gen_kwargs, advanced_generation)`` where ``advanced_generation``
+    is True when any setting differs from its default.
+    """
     with st.expander("Advanced", expanded=False):
         config_col1, config_col2 = st.columns([2, 1])
         with config_col1:
@@ -643,6 +594,35 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
         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:
+    """Render the chat tab."""
+
+    st.title("Chat")
+
+    context_key = chat_session_key(model_name, dataset_source)
+    chat_state = get_chat_state(model_name, remote, dataset_source)
+    try:
+        dataset, dataset_status = load_dataset(
+            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}")
+        st.info("Check the selected dataset source or upload both JSONL files.")
+        return
+
+    personas = list(dataset)
+    if not personas:
+        st.warning("No personas found in the selected dataset.")
+        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)
 
     # ── Mode toggle ───────────────────────────────────────────────────────────
     compare_mode = st.toggle(
@@ -731,7 +711,6 @@ def render_chat_tab(remote: bool, model_name: str, dataset_source: str) -> None:
             prompt_key,
             prompt_mode,
             active_system_prompt,
-            edit_key,
             height=200,
         )
 

tabs/extract.py

@@ -111,7 +111,6 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         st.info("Select at least one persona.")
         return
 
-    runs = None
     max_questions = 0
 
     with st.expander("Advanced", expanded=False):
@@ -190,9 +189,6 @@ def render_extract_tab(remote: bool, model_name: str, dataset_source: str) -> No
         )
         st.session_state[_LAST_MAX_QUESTIONS_KEY] = max_questions
 
-    if runs is None:
-        return
-
     run_clicked = st.button("Run extraction", type="primary")
     if not run_clicked:
         return

utils/chat.py

@@ -82,11 +82,9 @@ def _format_generation_prompt(
     Tries the tokenizer's chat template first, falls back to normalized messages,
     then to a plain-text format if both template attempts fail.
     """
-    normalized_messages = messages
-
     try:
         prompt = tokenizer.apply_chat_template(
-            normalized_messages,
+            messages,
             tokenize=False,
             add_generation_prompt=True,
         )
@@ -94,11 +92,11 @@ def _format_generation_prompt(
         logger.debug(
             "Chat template failed on raw messages, trying normalized", exc_info=True
         )
-        normalized_messages = normalize_messages(messages)
+        messages = normalize_messages(messages)
 
         try:
             prompt = tokenizer.apply_chat_template(
-                normalized_messages,
+                messages,
                 tokenize=False,
                 add_generation_prompt=True,
             )
@@ -108,7 +106,7 @@ def _format_generation_prompt(
                 exc_info=True,
             )
             prompt = _format_plain_messages(
-                normalized_messages,
+                messages,
                 add_generation_prompt=True,
             )
 

uv.lock

@@ -1561,7 +1561,7 @@ wheels = [
 [[package]]
 name = "persona-data"
 version = "0.1.0"
-source = { git = "ssh://git@github.com/implicit-personalization/persona-data.git#3763bd6e42472b589b4e32acd3e47b711a0af1f5" }
+source = { git = "https://github.com/implicit-personalization/persona-data#4316f47026a40ad1c5337c3830141267527be2fc" }
 dependencies = [
     { name = "huggingface-hub" },
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
@@ -1580,21 +1580,23 @@ dependencies = [
     { name = "plotly" },
     { name = "python-dotenv" },
     { name = "streamlit" },
+    { name = "transformers" },
 ]
 
 [package.metadata]
 requires-dist = [
-    { name = "persona-data", git = "ssh://git@github.com/implicit-personalization/persona-data.git" },
-    { name = "persona-vectors", git = "ssh://git@github.com/implicit-personalization/persona-vectors.git" },
+    { name = "persona-data", git = "https://github.com/implicit-personalization/persona-data.git" },
+    { 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" },
+    { name = "transformers", specifier = ">=5.5.0" },
 ]
 
 [[package]]
 name = "persona-vectors"
 version = "0.1.0"
-source = { git = "ssh://git@github.com/implicit-personalization/persona-vectors.git#fa6b4b61eaaba9ce64ee8614766bf75879148bbb" }
+source = { editable = "../persona-vectors" }
 dependencies = [
     { name = "kaleido" },
     { name = "nnsight" },
@@ -1612,6 +1614,23 @@ dependencies = [
     { name = "umap-learn" },
 ]
 
+[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", git = "https://github.com/implicit-personalization/persona-data" },
+    { 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]]
 name = "pexpect"
 version = "4.9.0"