Update app.py

actual last attempt

app.py

@@ -65,35 +65,6 @@ model.eval()
 # ----------------------
 
 
-def build_messages_for_userlm(
-    system_prompt: str, history: List[Tuple[str, str]]
-) -> List[Dict[str, str]]:
-    """Build messages for UserLM generation.
-
-    In history tuples: (user_msg, assistant_msg) where:
-    - user_msg: what UserLM previously generated
-    - assistant_msg: what the human (playing assistant) said
-
-    For UserLM training, these roles were flipped, so we need to reconstruct
-    the conversation as UserLM saw it during training.
-    """
-    messages: List[Dict[str, str]] = []
-
-    # System prompt defines the user's intent
-    if system_prompt.strip():
-        messages.append({"role": "system", "content": system_prompt.strip()})
-
-    # Add conversation history in the format UserLM expects
-    # UserLM was trained to generate "user" role messages given prior context
-    for user_msg, assistant_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
-
-    return messages
-
-
 def is_valid_length(text: str, min_words: int = 3, max_words: int = 50) -> bool:
     """Check if generated text meets length requirements (Guardrail 3)."""
     word_count = len(text.split())
@@ -101,21 +72,19 @@ def is_valid_length(text: str, min_words: int = 3, max_words: int = 50) -> bool:
 
 
 def is_verbatim_repetition(
-    new_text: str, history: List[Tuple[str, str]], system_prompt: str
+    new_text: str, history: List[Dict], system_prompt: str
 ) -> bool:
-    """Check if text is exact repetition of prior user turn or system prompt (Guardrail 4).
-
-    History format: (assistant_msg, user_msg) - so user messages are in position 1
-    """
+    """Check if text is exact repetition. History is now list of message dicts."""
     new_text_normalized = new_text.strip().lower()
 
     if new_text_normalized == system_prompt.strip().lower():
         return True
 
-    # User messages are now in position 1 of the tuple
-    for _, user_msg in history:
-        if user_msg and new_text_normalized == user_msg.strip().lower():
-            return True
+    # Check against previous user messages
+    for msg in history:
+        if msg.get("role") == "user" and msg.get("content"):
+            if new_text_normalized == msg["content"].strip().lower():
+                return True
 
     return False
 
@@ -123,7 +92,7 @@ def is_verbatim_repetition(
 @spaces.GPU
 def generate_user_message(
     messages: List[Dict[str, str]],
-    history: List[Tuple[str, str]],
+    history: List[Dict],
     system_prompt: str,
     max_new_tokens: int = 256,
     temperature: float = 1.0,
@@ -174,38 +143,35 @@ def generate_user_message(
 
 def generate_next_turn(
     assistant_response: str,
-    chat_history: List[Tuple[str, str]],
+    chat_history: List[Dict],
     system_prompt: str,
     max_new_tokens: int,
     temperature: float,
     top_p: float,
 ):
     """
-    Tuple: (human_assistant, userlm_user)
-    Testing if Gradio shows first element on RIGHT not LEFT
+    History format: List of {"role": "user"/"assistant", "content": "..."}
+    - "user" role = UserLM (displays LEFT)
+    - "assistant" role = Human (displays RIGHT)
     """
 
-    # If we have an assistant response, fill in position 0 of last turn
-    if assistant_response.strip() and len(chat_history) > 0:
-        _, last_user_msg = chat_history[-1]
-        chat_history = chat_history[:-1] + [(assistant_response.strip(), last_user_msg)]
+    # If we have an assistant response, add it to history
+    if assistant_response.strip():
+        chat_history.append(
+            {"role": "assistant", "content": assistant_response.strip()}
+        )
 
-    # Build messages for UserLM - user msgs are now in position 1
+    # Build messages for UserLM from history
     messages = []
     if system_prompt.strip():
         messages.append({"role": "system", "content": system_prompt.strip()})
-
-    for asst_msg, user_msg in chat_history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if asst_msg:
-            messages.append({"role": "assistant", "content": asst_msg})
+    messages.extend(chat_history)
 
     # Generate next user message
     try:
         user_msg = generate_user_message(
             messages,
-            [(u, a) for a, u in chat_history],  # Swap for repetition check
+            chat_history,
             system_prompt,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
@@ -214,14 +180,14 @@ def generate_next_turn(
     except Exception as e:
         user_msg = f"(Generation error: {e})"
 
-    # Add: (empty_assistant, new_user_msg)
-    new_history = chat_history + [("", user_msg)]
+    # Add new user message to history
+    new_history = chat_history + [{"role": "user", "content": user_msg}]
 
     return "", new_history, "Generate Next User Message"
 
 
 def clear_conversation():
-    return [], DEFAULT_SYSTEM_PROMPT, "Generate First User Message"
+    return [], DEFAULT_SYSTEM_PROMPT, [], "Generate First User Message", []
 
 
 # ----------------------
@@ -270,7 +236,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         top_p = gr.Slider(0.0, 1.0, value=0.8, step=0.01, label="top_p")
 
     with gr.Row():
-        submit_btn = gr.Button("Generate First User Message", variant="primary")
+        submit_btn = gr.Button("Generate User Message", variant="primary")
         clear_btn = gr.Button("Clear")
 
     state = gr.State([])
@@ -291,23 +257,26 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         )
 
     def _submit(asst_text, history, system_prompt, mnt, temp, tp):
-        return generate_next_turn(asst_text, history, system_prompt, mnt, temp, tp)
+        new_msg, new_history = generate_next_turn(
+            asst_text, history, system_prompt, mnt, temp, tp
+        )
+        return new_msg, new_history, new_history
 
     submit_btn.click(
         fn=_submit,
         inputs=[msg, state, system_box, max_new_tokens, temperature, top_p],
-        outputs=[msg, state, submit_btn],
+        outputs=[msg, state, chatbot],
     )
     msg.submit(
         fn=_submit,
         inputs=[msg, state, system_box, max_new_tokens, temperature, top_p],
-        outputs=[msg, state, submit_btn],
+        outputs=[msg, state, chatbot],
     )
 
-    # Keep chatbot display in sync with state
-    state.change(lambda x: x, inputs=[state], outputs=[chatbot])
-
-    clear_btn.click(fn=clear_conversation, outputs=[state, system_box, submit_btn])
+    clear_btn.click(
+        fn=clear_conversation,
+        outputs=[state, system_box, chatbot],
+    )
 
 if __name__ == "__main__":
     demo.queue().launch()