Update app.py

app.py

@@ -1,7 +1,6 @@
-# AI Persona Simulator - Stateless GPU Optimized Version
+# AI Persona Simulator - Final Optimized Version
 import gradio as gr
 import torch
-from transformers import pipeline
 from duckduckgo_search import DDGS
 import re
 import time
@@ -24,7 +23,6 @@ MAX_GPU_MEMORY = "60GiB"
 @GPU(memory=60)
 def load_model():
     """Load the Gemma 3 1B model without quantization for full precision."""
-    print(f"Attempting to load model: {MODEL_ID} without quantization")
     try:
         pipe = pipeline(
             "text-generation",
@@ -33,7 +31,7 @@ def load_model():
             device_map="auto",
             model_kwargs={"use_cache": True}
         )
-        print(f"Model {MODEL_ID} loaded successfully on device: {pipe.device} (full precision)")
+        print(f"Model {MODEL_ID} loaded successfully on device: {pipe.device}")
         return pipe
     except Exception as e:
         print(f"FATAL Error loading model '{MODEL_ID}': {e}")
@@ -53,8 +51,7 @@ CRITERIA:
 3. NO manipulation/exploitation attempts
 4. NO illegal/harmful scenarios
 5. NO inappropriate relationships
-Respond ONLY with "TRUE" if acceptable, "FALSE" if not."""
-            },
+Respond ONLY with "TRUE" if acceptable, "FALSE" if not."""},
             {"role": "user", "content": f"Character Name: {name}\nContext: {context}"}
         ]
         
@@ -64,6 +61,7 @@ Respond ONLY with "TRUE" if acceptable, "FALSE" if not."""
             add_generation_prompt=True,
             tokenize=False
         )
+        
         with torch.amp.autocast('cuda', dtype=torch.bfloat16):
             outputs = pipe(
                 text,
@@ -72,6 +70,7 @@ Respond ONLY with "TRUE" if acceptable, "FALSE" if not."""
                 temperature=0.1,
                 pad_token_id=pipe.tokenizer.eos_token_id
             )
+            
         result = parse_llm_output(outputs, validation_prompt).strip().upper()
         return result == "TRUE"
     except Exception as e:
@@ -118,54 +117,7 @@ def search_person(name, context=""):
     print(f"Found {len(results)} potential search results.")
     return results
 
-def create_synthetic_profile(name, context):
-    """Create a synthetic profile when search returns no results."""
-    profile = {
-        "title": f"Synthetic Profile for {name}",
-        "href": "",
-        "body": f"{name} is a person described with the context: '{context}'. "
-    }
-    if "grade" in context.lower():
-        grade_match = re.search(r'(\d+)(?:st|nd|rd|th)?\s+grade', context.lower())
-        if grade_match:
-            try:
-                grade = int(grade_match.group(1))
-                age = 5 + grade
-                profile["body"] += f"Based on being in {grade}th grade, {name} is likely around {age} years old. "
-                profile["body"] += f"Typical interests for this age might include friends, hobbies, school subjects, and developing independence. "
-            except ValueError:
-                profile["body"] += f"The grade mentioned ('{grade_match.group(1)}') could not be parsed to estimate age. "
-    profile["body"] += "Since no public information was found, this profile is based solely on the provided context."
-    return [profile]
-
-def extract_text_from_search_results(search_results):
-    """Extract relevant text from search results."""
-    if isinstance(search_results, str):
-        return f"Could not extract text due to search error: {search_results}"
-    
-    combined_text = ""
-    seen_bodies = set()
-    count = 0
-    max_results_to_process = 5
-    
-    for result in search_results:
-        if count >= max_results_to_process:
-            break
-        if isinstance(result, dict) and 'body' in result and result['body']:
-            body = result['body'].strip()
-            if body not in seen_bodies:
-                combined_text += body + "\n"
-                seen_bodies.add(body)
-                count += 1
-    
-    if not combined_text:
-        return "No relevant text found in search results."
-    
-    combined_text = re.sub(r'\s+', ' ', combined_text).strip()
-    max_length = 2000
-    return combined_text[:max_length] + "..." if len(combined_text) > max_length else combined_text
-
-# --- LLM Generation Functions ---
+# --- Text Processing Functions ---
 def parse_llm_output(full_output, input_prompt_list):
     """Attempts to parse only the newly generated text from the LLM output."""
     if isinstance(full_output, list) and len(full_output) > 0:
@@ -207,11 +159,9 @@ def parse_llm_output(full_output, input_prompt_list):
         print("Warning: Parsing resulted in empty string, returning original generation.")
         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
 
+# --- LLM Generation Functions ---
 @GPU(memory=60)
 def generate_enhanced_persona(name, bio_text, context=""):
     """Use the LLM to enhance the persona profile."""
@@ -263,7 +213,12 @@ Additional context for the simulation: {context}
 Maintain this persona consistently. Respond naturally based on the profile. Do not mention that you are an AI or a simulation. If asked about details not in the profile, you can be evasive or state you don't know/remember, consistent with the persona."""
     
     prompt = [
-        {"role": "system", "content": """You are an expert AI prompt engineer specializing in character simulation. Create a concise system prompt that instructs the LLM to embody the character based on the profile. The prompt must: 1. Define core personality and speaking style. 2. Specify how to handle unknown topics. 3. Prohibit breaking character or mentioning AI nature. Output ONLY the system prompt itself."""},
+        {"role": "system", "content": """You are an expert AI prompt engineer specializing in character simulation. Create a concise system prompt that instructs the LLM to embody the character based on the profile. The prompt must:
+1. Define core personality and speaking style
+2. Specify how to handle unknown topics
+3. Prohibit breaking character or mentioning AI nature
+Output ONLY the system prompt itself."""
+        },
         {"role": "user", "content": f"""Create a system prompt for an AI to simulate the character '{name}'. Context for simulation: {context} Character Profile:
 {enhanced_profile}
 Generate the system prompt based *only* on the profile and context provided."""}
@@ -300,6 +255,7 @@ def generate_response(messages):
     print("Generating response...")
     if not messages:
         return "Error: No message history provided."
+    
     try:
         tokenizer = pipe.tokenizer
         text = tokenizer.apply_chat_template(
@@ -323,7 +279,7 @@ def generate_response(messages):
     except Exception as e:
         error_msg = f"Error during response generation: {str(e)}"
         print(error_msg)
-        return f"Sorry, I encountered an error trying to respond."
+        return f"Sorry, I encountered an error: {str(e)}"
 
 # --- Persona Chat Class with Safety ---
 class PersonaChat:
@@ -453,15 +409,6 @@ def create_interface():
         padding: 20px; 
         box-shadow: 0 4px 12px rgba(0, 0, 0, 0.05);
     }
-    .status-bar { 
-        background: #f1f3f5; 
-        padding: 12px 15px; 
-        border-radius: 5px; 
-        margin: 15px 0; 
-        font-weight: 500; 
-        border: 1px solid #e2e6ea;
-        color: #212529;
-    }
     .chat-container { 
         border: 1px solid #eaeaea; 
         border-radius: 10px; 
@@ -473,30 +420,13 @@ def create_interface():
     .message-input { 
         margin-top: 10px; 
     }
-    .send-button { 
+    .send-button, .persona-button { 
         background-color: #1e3c72 !important; 
         color: white !important;
         border-radius: 8px;
         padding: 10px 20px;
         font-weight: bold;
     }
-    .persona-button { 
-        background-color: #2a5298 !important; 
-        color: white !important;
-        border-radius: 8px;
-        padding: 10px 20px;
-        font-weight: bold;
-    }
-    .system-prompt-display { 
-        background-color: #f5f5f5; 
-        border-radius: 8px; 
-        padding: 15px; 
-        margin-top: 15px; 
-        border: 1px solid #e0e0e0; 
-        font-family: monospace; 
-        white-space: pre-wrap; 
-        word-wrap: break-word;
-    }
     .footer { 
         text-align: center; 
         margin-top: 30px; 
@@ -505,12 +435,6 @@ def create_interface():
         padding: 15px;
         border-top: 1px solid #eee;
     }
-    .typing-indicator { 
-        color: #aaa; 
-        font-style: italic; 
-    }
-    
-    /* Mobile styles */
     @media (max-width: 768px) {
         .chat-container { height: 300px !important; }
         .main-container { padding: 10px; }
@@ -539,10 +463,16 @@ def create_interface():
                 with gr.Column():
                     gr.Markdown("### Chat with Character")
                     character_name_display = gr.Markdown("*No persona created yet*")
-                    chatbot = gr.Chatbot(height=450, show_label=False)
+                    chatbot = gr.Chatbot(
+                        height=450, 
+                        show_label=False,
+                        bubble_full_width=False,
+                        type="messages",
+                        avatar_images=("https://api.dicebear.com/6.x/bottts/svg?seed=user ", 
+                                      "https://api.dicebear.com/6.x/bottts/svg?seed=bot ")
+                    )
                     msg_input = gr.Textbox(label="Your message", placeholder="Type your message here and press Enter...")
                     send_button = gr.Button("Send Message")
-                    
                     gr.Markdown("Powered by Gemma 3 1B • Ethically Designed • Safe & Secure")
 
         def set_persona_flow(name, context):
@@ -552,11 +482,9 @@ def create_interface():
             
             initial_status = f"Creating persona for '{name}'..."
             initial_character_display = f"### Preparing to chat with {name}..."
-            initial_prompt = "System prompt will appear here..."
-            initial_profile = "Enhanced profile will appear here..."
             initial_history = []
             
-            yield initial_status, initial_prompt, initial_profile, initial_character_display, initial_history
+            yield initial_status, "", "", initial_character_display, initial_history
             
             try:
                 for status_update, prompt_update, profile_update, history_update in persona_chat.set_persona(name, context):
@@ -565,30 +493,33 @@ def create_interface():
                         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])
+                            gradio_history.append({"role": "user", "content": user_msg})
+                            gradio_history.append({"role": "assistant", "content": bot_msg})
+                    
                     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, 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, initial_prompt, initial_profile, f"### Error creating {name}", []
+                print(error_msg)
+                yield error_msg, "", "", f"### Error creating {name}", []
         
         def send_message_flow(message, history):
             if not message.strip():
                 return "", history
+            
             if not persona_chat.messages or persona_chat.messages[0]['role'] != 'system':
-                history.append([message, "Error: Please create a valid persona first."])
+                history.append({"role": "assistant", "content": "Error: Please create a valid persona first."})
                 return "", history
             
-            history.append([message, None])
-            response_text = persona_chat.chat(message)
-            response_text = re.sub(r'<end_of_turn>|<start_of_turn>model', '', response_text).strip()
-            history[-1][1] = response_text
+            history.append({"role": "user", "content": message})
+            response = persona_chat.chat(message)
+            history.append({"role": "assistant", "content": response})
             return "", history
         
         set_persona_button.click(