Update app.py

app.py

@@ -132,6 +132,8 @@ def parse_llm_output(full_output, input_prompt_list):
             if potential_response:
                 potential_response = re.sub(r'^<\/?s?>', '', potential_response).strip()
                 potential_response = re.sub(r'^(assistant|ASSISTANT|System|SYSTEM)[:\s]*', '', potential_response).strip()
+                # Remove special tags
+                potential_response = re.sub(r'<end_of_turn>|<start_of_turn>model', '', potential_response).strip()
                 if potential_response:
                     return potential_response
     cleaned_text = generated_text
@@ -141,9 +143,12 @@ def parse_llm_output(full_output, input_prompt_list):
             pass
     cleaned_text = re.sub(r'^<\/?s?>', '', cleaned_text).strip()
     cleaned_text = re.sub(r'^(assistant|ASSISTANT|System|SYSTEM)[:\s]*', '', cleaned_text).strip()
+    # Remove special tags from full output
+    cleaned_text = re.sub(r'<end_of_turn>|<start_of_turn>model', '', cleaned_text).strip()
     if not cleaned_text and generated_text:
         print("Warning: Parsing resulted in empty string, returning original generation.")
-        return generated_text
+        # Still clean special tags from original
+        return re.sub(r'<end_of_turn>|<start_of_turn>model', '', generated_text).strip()
     if last_input_content and last_occurrence_index == -1:
         print("Warning: Could not find last input prompt in LLM output. Returning cleaned full output.")
     return cleaned_text
@@ -239,6 +244,8 @@ def generate_response(messages):
                 pad_token_id=pipe.tokenizer.eos_token_id if pipe.tokenizer.eos_token_id else None
             )
         parsed_output = parse_llm_output(outputs, messages)
+        # Extra cleanup for specific model tags
+        parsed_output = re.sub(r'<end_of_turn>|<start_of_turn>model', '', parsed_output).strip()
         print("Response generated.")
         return parsed_output if parsed_output else "..."
     except Exception as e:
@@ -326,7 +333,7 @@ def create_interface():
     .gradio-container { font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; }
     .main-container { max-width: 1200px; margin: auto; padding: 0; }
     .header { background: linear-gradient(90deg, #2c3e50, #4ca1af); color: white; padding: 20px; border-radius: 10px 10px 0 0; margin-bottom: 20px; text-align: center; }
-    .setup-section { tôackground-color: #f9f9f9; border-radius: 10px; padding: 20px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); margin-bottom: 20px; }
+    .setup-section { background-color: #f9f9f9; border-radius: 10px; padding: 20px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); margin-bottom: 20px; }
     .chat-section { background-color: white; border-radius: 10px; padding: 20px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); }
     .status-bar { background: #e9ecef; padding: 10px 15px; border-radius: 5px; margin: 15px 0; font-weight: 500; border: 1px solid #ced4da; }
     .chat-container { border: 1px solid #eaeaea; border-radius: 10px; height: 500px !important; overflow-y: auto; background-color: #ffffff; padding: 10px; }
@@ -337,6 +344,11 @@ def create_interface():
     .footer { text-align: center; margin-top: 20px; font-size: 0.9em; color: #666; }
     .typing-indicator { color: #aaa; font-style: italic; }
     """
+    
+    # Define avatar images with full URLs to ensure they work
+    user_avatar = "https://api.dicebear.com/6.x/bottts/svg?seed=user"
+    bot_avatar = "https://api.dicebear.com/6.x/bottts/svg?seed=bot"
+    
     with gr.Blocks(css=css, title="AI Persona Simulator") as interface:
         with gr.Row(elem_classes="main-container"):
             with gr.Column():
@@ -361,14 +373,15 @@ def create_interface():
                         label="Conversation",
                         height=450,
                         elem_classes="chat-container",
-                        avatar_images=(None, "🤖"),
-                        type="messages"
+                        avatar_images=(user_avatar, bot_avatar),
+                        show_label=False
                     )
                     with gr.Row():
                         msg_input = gr.Textbox(label="Your message", placeholder="Type your message here and press Enter...", elem_classes="message-input", scale=4)
                         send_button = gr.Button("Send", variant="primary", elem_classes="send-button", scale=1)
                 with gr.Column(elem_classes="footer"):
                     gr.Markdown(f"Powered by {MODEL_ID}")
+                    
         def set_persona_flow(name, context):
             if not name:
                 yield "Status: Please enter a character name.", "", "", "*No persona created yet*", []
@@ -385,50 +398,53 @@ def create_interface():
             final_status, final_prompt, final_profile, final_character_display, final_history = "Error", "", "", f"### Error creating {name}", []
             try:
                 for status_update, prompt_update, profile_update, history_update in persona_chat.set_persona(name, context):
-                    # Strict validation of history
-                    filtered_history = [
-                        msg for msg in (history_update or [])
-                        if isinstance(msg, dict) and "content" in msg and isinstance(msg["content"], str) and msg["content"]
-                    ]
+                    # For Gradio's Chatbot, convert to tuple list format
+                    gradio_history = []
+                    for i in range(0, len(history_update), 2):
+                        if i+1 < len(history_update):
+                            user_msg = history_update[i].get("content", "")
+                            bot_msg = history_update[i+1].get("content", "")
+                            gradio_history.append([user_msg, bot_msg])
+                    
                     # Log chatbot state
                     current_chatbot_state = chatbot.value if hasattr(chatbot, 'value') else []
                     print(f"set_persona_flow: Current chatbot state: {current_chatbot_state}")
-                    # Clean and validate chatbot state
-                    cleaned_chatbot_state = [
-                        msg for msg in current_chatbot_state
-                        if isinstance(msg, dict) and "content" in msg and isinstance(msg["content"], str) and msg["content"]
-                    ]
-                    # Combine histories
-                    combined_history = cleaned_chatbot_state + filtered_history
-                    print(f"set_persona_flow: Processing yield - status: {status_update}, filtered_history: {filtered_history}, combined_history: {combined_history}")
-                    final_status = status_update
-                    final_prompt = prompt_update
-                    final_profile = profile_update
-                    final_history = combined_history
+                    
                     character_display = f"### Preparing chat with {name}..."
                     if "Ready to chat" in status_update:
                         character_display = f"### Chatting with {name}"
                     elif "Error" in status_update:
                         character_display = f"### Error creating {name}"
-                    yield status_update, final_prompt, final_profile, character_display, final_history
+                    
+                    yield status_update, prompt_update, profile_update, character_display, gradio_history
                     time.sleep(0.1)
             except Exception as e:
                 error_msg = f"Failed to set persona (interface error): {str(e)}"
                 print(f"set_persona_flow: Exception: {error_msg}")
-                yield error_msg, final_prompt, final_profile, f"### Error creating {name}", final_history
+                yield error_msg, final_prompt, final_profile, f"### Error creating {name}", []
+                
         def send_message_flow(message, history):
-            if history is None:
-                history = []
             if not message.strip():
                 return "", history
+            
             if not persona_chat.messages or persona_chat.messages[0]['role'] != 'system':
-                history.append({"role": "user", "content": message})
-                history.append({"role": "assistant", "content": "Error: Please create a valid persona first."})
+                history.append([message, "Error: Please create a valid persona first."])
                 return "", history
-            history.append({"role": "user", "content": message})
+            
+            # Add user message to history
+            history.append([message, None])  # Add placeholder for bot response
+            
+            # Get response from AI
             response_text = persona_chat.chat(message)
-            history.append({"role": "assistant", "content": response_text})
+            
+            # Clean any special tags that might still be in the response
+            response_text = re.sub(r'<end_of_turn>|<start_of_turn>model', '', response_text).strip()
+            
+            # Update the last message with the actual response
+            history[-1][1] = response_text
+            
             return "", history
+            
         set_persona_button.click(
             set_persona_flow,
             inputs=[name_input, context_input],