Create app.py

app.py

@@ -0,0 +1,557 @@
+import gradio as gr
+import transformers
+import torch
+from transformers import pipeline
+from duckduckgo_search import DDGS
+import re
+import time
+
+def search_person(name, context=""):
+    """Search for information about a person using DuckDuckGo."""
+    results = []
+    search_terms = []
+    
+    # Create search terms based on provided context
+    if "grade" in context.lower():
+        grade_match = re.search(r'(\d+)(st|nd|rd|th)?\s+grade', context.lower())
+        if grade_match:
+            grade = grade_match.group(1)
+            search_terms.append(f"{name} student {grade} grade")
+            
+    # Add basic search terms
+    search_terms.extend([
+        f"{name} {context}" if context else name,
+        f"{name} interests",
+        f"{name} personality"
+    ])
+    
+    try:
+        with DDGS() as ddgs:
+            for term in search_terms:
+                search_results = list(ddgs.text(term, max_results=3))
+                results.extend(search_results)
+    except Exception as e:
+        return f"Error during search: {str(e)}"
+    
+    # If no results found but we have context, create synthetic information
+    if not results and context:
+        return create_synthetic_profile(name, context)
+    
+    return results
+
+def create_synthetic_profile(name, context):
+    """Create a synthetic profile when search returns no results."""
+    profile = {
+        "body": f"{name} is a person described as: {context}."
+    }
+    
+    # Extract age/grade information
+    if "grade" in context.lower():
+        grade_match = re.search(r'(\d+)(st|nd|rd|th)?\s+grade', context.lower())
+        if grade_match:
+            grade = grade_match.group(1)
+            age = 5 + int(grade)  # Approximate age based on grade
+            profile["body"] += f" {name} is approximately {age} years old and in {grade}th grade."
+            profile["body"] += f" Like most {grade}th graders, {name} is likely interested in friends, learning new things, and developing their own identity."
+    
+    return [profile]
+
+def extract_text_from_search_results(search_results):
+    """Extract relevant text from search results."""
+    combined_text = ""
+    for result in search_results:
+        if isinstance(result, dict) and 'body' in result:
+            combined_text += result['body'] + "\n\n"
+    
+    # Clean up the text
+    combined_text = re.sub(r'\s+', ' ', combined_text)
+    return combined_text
+
+def load_model():
+    """Load the LLM model."""
+    model_id = "nvidia/Llama-3.1-Nemotron-8B-UltraLong-4M-Instruct"
+    pipe = pipeline(
+        "text-generation",
+        model=model_id,
+        model_kwargs={"torch_dtype": torch.bfloat16},
+        device_map="auto",
+    )
+    return pipe
+
+def generate_enhanced_persona(model, name, bio_text, context=""):
+    """Use the LLM to enhance the persona profile."""
+    
+    enhancement_prompt = [
+        {"role": "system", "content": """You are an expert AI character developer. 
+Your task is to create a detailed character profile based on limited information.
+Output ONLY the enhanced profile with no additional explanations or formatting."""},
+        {"role": "user", "content": f"""Here's some information I found about {name}:
+
+{bio_text}
+
+Additional context: {context}
+
+Based on this information, create a detailed, rich character profile for {name}. 
+Include personality traits, speaking style, interests, background details, quirks, and mannerisms.
+If this is a child in school, include age-appropriate details about school life, friends, family dynamics, and interests.
+Be creative but make the profile coherent with the known facts.
+Focus on what makes {name} unique and authentic.
+Structure your response as a character profile that could be used by an actor playing this role.
+
+DO NOT prefix your response with anything. Start directly with the character profile.
+DO NOT include any disclaimers, explanations, or notes in your response.
+DO NOT use bullet points or section titles."""}
+    ]
+    
+    try:
+        outputs = model(enhancement_prompt, max_new_tokens=1024)
+        if isinstance(outputs, list) and len(outputs) > 0:
+            if isinstance(outputs[0], dict) and "generated_text" in outputs[0]:
+                if isinstance(outputs[0]["generated_text"], list) and len(outputs[0]["generated_text"]) > 0:
+                    last_message = outputs[0]["generated_text"][-1]
+                    if isinstance(last_message, dict) and "content" in last_message:
+                        return last_message["content"]
+                elif isinstance(outputs[0]["generated_text"], str):
+                    return outputs[0]["generated_text"]
+        
+        # Fallback parsing
+        if isinstance(outputs, str):
+            return outputs
+        return bio_text  # Fall back to original if parsing fails
+    except Exception as e:
+        print(f"Error generating enhanced persona: {str(e)}")
+        return bio_text
+
+def generate_system_prompt_with_llm(model, name, enhanced_profile, context=""):
+    """Use the LLM to generate an optimized system prompt."""
+    
+    prompt_generation_message = [
+        {"role": "system", "content": """You are an expert AI prompt engineer. 
+Your task is to create an optimal system prompt that will make an LLM simulate a specific person accurately.
+Output ONLY the system prompt with no additional explanations."""},
+        {"role": "user", "content": f"""Here is a detailed profile for {name}:
+
+{enhanced_profile}
+
+Additional context: {context}
+
+Create a comprehensive and effective system prompt that would make an LLM perfectly simulate {name}.
+The prompt should:
+1. Include key personality traits and communication style
+2. Specify how to handle questions about unknown topics
+3. Provide guidance on maintaining consistent character behavior
+4. Include instructions for age-appropriate responses if this is a child
+5. Give specific examples of phrases or expressions this person might use
+
+Format the system prompt for direct use - don't include any explanations outside the prompt itself."""}
+    ]
+    
+    try:
+        outputs = model(prompt_generation_message, max_new_tokens=1024)
+        if isinstance(outputs, list) and len(outputs) > 0:
+            if isinstance(outputs[0], dict) and "generated_text" in outputs[0]:
+                if isinstance(outputs[0]["generated_text"], list) and len(outputs[0]["generated_text"]) > 0:
+                    last_message = outputs[0]["generated_text"][-1]
+                    if isinstance(last_message, dict) and "content" in last_message:
+                        return last_message["content"]
+                elif isinstance(outputs[0]["generated_text"], str):
+                    return outputs[0]["generated_text"]
+        
+        # Fallback parsing
+        if isinstance(outputs, str):
+            return outputs
+        
+        # If all parsing fails, generate a basic system prompt
+        return f"""You are now simulating {name}. Use the following information to respond as if you were {name}:
+
+{enhanced_profile}
+
+{context}
+
+Always stay in character as {name} and respond directly as {name} would respond."""
+    except Exception as e:
+        print(f"Error generating system prompt: {str(e)}")
+        # Fallback to basic prompt
+        return f"""You are now simulating {name}. Use the following information to respond as if you were {name}:
+
+{enhanced_profile}
+
+{context}
+
+Always stay in character as {name} and respond directly as {name} would respond."""
+
+def generate_response(model, messages):
+    """Generate a response using the LLM."""
+    outputs = model(messages, max_new_tokens=512)
+    # Extract the content from the generated response
+    if isinstance(outputs, list) and len(outputs) > 0:
+        if isinstance(outputs[0], dict) and "generated_text" in outputs[0]:
+            if isinstance(outputs[0]["generated_text"], list) and len(outputs[0]["generated_text"]) > 0:
+                last_message = outputs[0]["generated_text"][-1]
+                if isinstance(last_message, dict) and "content" in last_message:
+                    return last_message["content"]
+            elif isinstance(outputs[0]["generated_text"], str):
+                return outputs[0]["generated_text"]
+    
+    # Fallback parsing for different output formats
+    if isinstance(outputs, str):
+        return outputs
+    return "I couldn't generate a proper response. Please try again."
+
+class PersonaChat:
+    def __init__(self):
+        self.model = None
+        self.system_prompt = "You are a helpful assistant."
+        self.persona_name = "Assistant" 
+        self.persona_context = ""
+        self.messages = []
+        self.enhanced_profile = ""
+        
+    def load_model_if_needed(self):
+        if self.model is None:
+            self.model = load_model()
+    
+    def set_persona(self, name, context=""):
+        self.load_model_if_needed()
+        self.persona_name = name
+        self.persona_context = context
+        
+        # Show loading indicator
+        status = f"Searching for information about {name}..."
+        yield status, "", [{"role": "system", "content": "Starting persona creation..."}]
+        
+        search_results = search_person(name, context)
+        if isinstance(search_results, str) and search_results.startswith("Error"):
+            yield f"Error: {search_results}", "", [{"role": "system", "content": f"Error: {search_results}"}]
+            return
+        
+        # Extract text from search results
+        bio_text = extract_text_from_search_results(search_results)
+        
+        # Use LLM to enhance the persona
+        status = f"Creating enhanced profile for {name}..."
+        yield status, "", [{"role": "system", "content": status}]
+        
+        self.enhanced_profile = generate_enhanced_persona(self.model, name, bio_text, context)
+        
+        # Use LLM to generate the optimized system prompt
+        status = f"Generating optimal system prompt for {name}..."
+        yield status, "", [{"role": "system", "content": status}]
+        
+        self.system_prompt = generate_system_prompt_with_llm(self.model, name, self.enhanced_profile, context)
+        self.messages = [{"role": "system", "content": self.system_prompt}]
+        
+        yield f"Persona set to {name}. Ready to chat!", self.system_prompt, self.messages
+    
+    def chat(self, user_message):
+        """Process a chat message and return the response."""
+        self.load_model_if_needed()
+        
+        try:
+            # Create message format for the model
+            if isinstance(user_message, str):
+                formatted_message = {"role": "user", "content": user_message}
+            else:
+                formatted_message = user_message
+            
+            # Add the message to history
+            self.messages.append(formatted_message)
+            
+            # Generate response using the model
+            response = generate_response(self.model, self.messages)
+            
+            # Format the response
+            assistant_message = {"role": "assistant", "content": response}
+            self.messages.append(assistant_message)
+            
+            return response
+            
+        except Exception as e:
+            error_msg = f"Error generating response: {str(e)}"
+            print(error_msg)
+            return error_msg
+
+def create_interface():
+    persona_chat = PersonaChat()
+    
+    # Custom CSS for better UI
+    css = """
+    .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 {
+        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: #f0f0f0;
+        padding: 10px 15px;
+        border-radius: 5px;
+        margin: 15px 0;
+        font-weight: 500;
+    }
+    
+    .chat-container {
+        border: 1px solid #eaeaea;
+        border-radius: 10px;
+        height: 500px !important;
+        overflow-y: auto;
+        background-color: #f9f9f9;
+    }
+    
+    .message-input {
+        margin-top: 10px;
+    }
+    
+    .send-button {
+        background-color: #2c3e50 !important;
+    }
+    
+    .persona-button {
+        background-color: #4ca1af !important;
+    }
+    
+    .system-prompt {
+        background-color: #f5f5f5;
+        border-radius: 8px;
+        padding: 10px;
+        margin-top: 15px;
+        border: 1px solid #e0e0e0;
+    }
+    
+    .footer {
+        text-align: center;
+        margin-top: 20px;
+        font-size: 0.9rem;
+        color: #666;
+    }
+    
+    /* Avatar styling */
+    .user-message {
+        background-color: #e1f5fe;
+        border-radius: 15px 15px 0 15px;
+        padding: 10px 15px;
+        margin: 8px 0;
+        max-width: 80%;
+        float: right;
+        clear: both;
+    }
+    
+    .bot-message {
+        background-color: #f0f0f0;
+        border-radius: 15px 15px 15px 0;
+        padding: 10px 15px;
+        margin: 8px 0;
+        max-width: 80%;
+        float: left;
+        clear: both;
+    }
+    
+    /* Loading animation */
+    @keyframes pulse {
+        0% { opacity: 0.6; }
+        50% { opacity: 1; }
+        100% { opacity: 0.6; }
+    }
+    
+    .loading {
+        animation: pulse 1.5s infinite;
+        padding: 10px;
+        background-color: #eee;
+        border-radius: 5px;
+        display: inline-block;
+    }
+    """
+    
+    with gr.Blocks(css=css, title="AI Persona Simulator") as interface:
+        with gr.Row(elem_classes="main-container"):
+            with gr.Column():
+                # Header
+                with gr.Column(elem_classes="header"):
+                    gr.Markdown("# AI Persona Simulator")
+                    gr.Markdown("Create lifelike character simulations with advanced AI")
+                
+                # Setup Section
+                with gr.Column(elem_classes="setup-section"):
+                    gr.Markdown("### Create Your Persona")
+                    gr.Markdown("Enter details about the character you want to simulate")
+                    
+                    with gr.Row():
+                        name_input = gr.Textbox(
+                            label="Character Name",
+                            placeholder="e.g. Erenalp",
+                            elem_classes="input-field"
+                        )
+                        
+                    with gr.Row():    
+                        context_input = gr.Textbox(
+                            label="Character Context", 
+                            placeholder="e.g. in 7th grade, loves math and video games, has a pet cat",
+                            lines=2,
+                            elem_classes="input-field"
+                        )
+                    
+                    with gr.Row():
+                        set_persona_button = gr.Button(
+                            "Create Persona", 
+                            variant="primary",
+                            elem_classes="persona-button"
+                        )
+                
+                    status_output = gr.Textbox(
+                        label="Status", 
+                        interactive=False,
+                        elem_classes="status-bar"
+                    )
+                    
+                    with gr.Accordion("Character System Prompt", open=False, elem_classes="system-prompt-section"):
+                        system_prompt_display = gr.TextArea(
+                            label="", 
+                            interactive=False,
+                            lines=10,
+                            elem_classes="system-prompt"
+                        )
+                
+                # Chat Section
+                with gr.Column(elem_classes="chat-section"):
+                    gr.Markdown("### Chat with Your Character")
+                    
+                    # Display character name dynamically
+                    character_name_display = gr.Markdown(
+                        elem_id="character-name",
+                        value="Start by creating a persona above"
+                    )
+                    
+                    chatbot = gr.Chatbot(
+                        label="",
+                        height=450,
+                        elem_classes="chat-container",
+                        avatar_images=("👤", "🤖"),
+                        type="messages"
+                    )
+                    
+                    with gr.Row():
+                        msg_input = gr.Textbox(
+                            label="Your message",
+                            placeholder="Type your message here...",
+                            elem_classes="message-input"
+                        )
+                        send_button = gr.Button(
+                            "Send", 
+                            variant="primary",
+                            elem_classes="send-button"
+                        )
+                
+                # Footer
+                with gr.Column(elem_classes="footer"):
+                    gr.Markdown("Powered by Llama-3.1-Nemotron-8B-UltraLong-4M-Instruct")
+        
+        # Functions
+        def update_character_name(name):
+            if name:
+                return f"### Chatting with {name}"
+            return "### Start by creating a persona above"
+        
+        def set_persona_generator(name, context):
+            initial_status = f"Creating persona for {name}..."
+            initial_character_display = f"### Creating persona for {name}..."
+            initial_prompt = ""
+            initial_history = [{"role": "system", "content": "Initializing..."}]
+            
+            # Initial yield
+            yield initial_status, initial_prompt, initial_history, initial_character_display
+            
+            # Process persona creation
+            for status, prompt, history in persona_chat.set_persona(name, context):
+                character_display = f"### Creating persona for {name}..."
+                if "Ready to chat" in status:
+                    character_display = f"### Chatting with {name}"
+                yield status, prompt, history, character_display
+        
+        def send_message(message, history):
+            if not message.strip():
+                return "", history
+            
+            if not persona_chat.messages:
+                new_history = list(history) if history else []
+                new_history.append({"role": "user", "content": message})
+                new_history.append({"role": "assistant", "content": "Please create a persona first using the form above."})
+                return "", new_history
+            
+            try:
+                # Show typing indicator
+                new_history = list(history) if history else []
+                new_history.append({"role": "user", "content": message})
+                new_history.append({"role": "assistant", "content": "..."})
+                yield "", new_history
+                
+                # Generate actual response
+                response = persona_chat.chat(message)
+                new_history[-1]["content"] = response
+                yield "", new_history
+                
+            except Exception as e:
+                print(f"Error in send_message: {str(e)}")
+                new_history[-1]["content"] = "Sorry, there was an error processing your message."
+                yield "", new_history
+        
+        # Event handlers
+        set_persona_button.click(
+            set_persona_generator,
+            inputs=[name_input, context_input],
+            outputs=[status_output, system_prompt_display, chatbot, character_name_display]
+        )
+        
+        name_input.change(
+            update_character_name,
+            inputs=[name_input],
+            outputs=[character_name_display]
+        )
+        
+        send_button.click(
+            send_message,
+            inputs=[msg_input, chatbot],
+            outputs=[msg_input, chatbot]
+        )
+        
+        msg_input.submit(
+            send_message,
+            inputs=[msg_input, chatbot],
+            outputs=[msg_input, chatbot]
+        )
+    
+    return interface
+
+# Install required packages if not already installed
+# !pip install gradio transformers torch duckduckgo_search
+
+# Create and launch the interface
+demo = create_interface()
+demo.launch(share=True)