Update app.py

app.py

@@ -1,259 +1,389 @@
+#!/usr/bin/env python3
+"""
+Gradio Demo App for TinyRP Mistral Model
+Supports ChatML formatting, character creation, and custom generation parameters
+"""
+
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
+import json
+import random
+from typing import Dict, List, Tuple
+import re
 
-# Load your model
-MODEL_NAME = "DarwinAnim8or/TinyRP"
-
-# Initialize model and tokenizer globally to avoid reloading
-try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME, 
-        torch_dtype=torch.float16, 
-        device_map="auto",
-        trust_remote_code=True
-    )
-    print(f"Model loaded successfully: {MODEL_NAME}")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    tokenizer = None
-    model = None
-
-# Sample character presets
+# Sample characters for demo
 SAMPLE_CHARACTERS = {
-    "Custom Character": "",
-    "Adventurous Knight": "You are Sir Gareth, a brave and noble knight on a quest to save the kingdom. You speak with honor and courage, always ready to help those in need. You carry an enchanted sword and have a loyal horse named Thunder.",
-    "Mysterious Wizard": "You are Eldara, an ancient and wise wizard who speaks in riddles and knows secrets of the mystical arts. You live in a tower filled with magical books and potions. You are helpful but often cryptic in your responses.",
-    "Friendly Tavern Keeper": "You are Bram, a cheerful tavern keeper who loves telling stories and meeting new travelers. Your tavern 'The Dancing Dragon' is a warm, welcoming place. You know all the local gossip and always have a tale to share.",
-    "Curious Scientist": "You are Dr. Maya Chen, a brilliant scientist who is fascinated by discovery and invention. You're enthusiastic about explaining complex concepts in simple ways and always looking for new experiments to try.",
-    "Space Explorer": "You are Captain Nova, a fearless space explorer who has traveled to distant galaxies. You pilot the starship 'Wanderer' and have encountered many alien species. You're brave, curious, and always ready for the next adventure."
+    "Luna the Mage": {
+        "description": "A mysterious elven mage with silver hair and glowing blue eyes. She specializes in ice magic and ancient knowledge.",
+        "personality": "Wise, mysterious, slightly aloof but caring. Speaks in an eloquent manner.",
+        "background": "Born in the Frostwood Academy, Luna has spent centuries studying arcane arts."
+    },
+    "Rex the Warrior": {
+        "description": "A brave human knight with a strong sense of justice. Wears gleaming armor and carries an enchanted sword.",
+        "personality": "Noble, brave, protective of others. Direct in speech but kind-hearted.",
+        "background": "Grew up as a farm boy, became a knight after saving his village from bandits."
+    },
+    "Zara the Rogue": {
+        "description": "A cunning halfling thief with quick wit and nimble fingers. Has curly red hair and green eyes.",
+        "personality": "Sarcastic, clever, independent. Uses humor to deflect serious situations.",
+        "background": "Former street orphan who learned to survive through wit and stealth."
+    },
+    "Dr. Elena Cross": {
+        "description": "A brilliant scientist working on advanced AI research in a near-future setting.",
+        "personality": "Analytical, passionate about her work, caring but sometimes absent-minded.",
+        "background": "PhD in Computer Science, leads a small research team at a tech corporation."
+    }
 }
 
-def generate_response(message, history, character_description, max_tokens, temperature, top_p, repetition_penalty, use_chatml_format):
-    """Generate a response from the model"""
-    
-    if model is None or tokenizer is None:
-        return "Error: Model not loaded properly"
+class TinyRPModel:
+    def __init__(self, model_name: str):
+        """Initialize the model and tokenizer"""
+        print(f"Loading model: {model_name}")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,  # Use float32 for CPU
+            device_map="cpu",
+            low_cpu_mem_usage=True
+        )
+        self.model.eval()
+        
+        # Ensure pad token is set
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        print("Model loaded successfully!")
     
-    if not message.strip():
-        return "Please enter a message"
+    def format_chatml(self, character_info: str, conversation_history: List[Tuple[str, str]], user_input: str) -> str:
+        """Format conversation using ChatML format"""
+        formatted = ""
+        
+        # System message with character info
+        if character_info.strip():
+            formatted += f"<|im_start|>system\n{character_info.strip()}<|im_end|>\n"
+        
+        # Add conversation history
+        for user_msg, assistant_msg in conversation_history:
+            if user_msg:
+                formatted += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
+            if assistant_msg:
+                formatted += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
+        
+        # Add current user input
+        if user_input.strip():
+            formatted += f"<|im_start|>user\n{user_input.strip()}<|im_end|>\n"
+        
+        # Start assistant response
+        formatted += "<|im_start|>assistant\n"
+        
+        return formatted
     
-    try:
-        # Build conversation context
-        if use_chatml_format and character_description.strip():
-            # Use ChatML format
-            conversation = f"<|im_start|>system\n{character_description}<|im_end|>\n"
-            
-            # Add history
-            for user_msg, assistant_msg in history:
-                if user_msg:
-                    conversation += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
-                if assistant_msg:
-                    conversation += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
-            
-            # Add current message
-            conversation += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
-        else:
-            # Simple format
-            if character_description.strip():
-                conversation = f"{character_description}\n\n"
-            else:
-                conversation = ""
-                
-            # Add history
-            for user_msg, assistant_msg in history:
-                if user_msg:
-                    conversation += f"Human: {user_msg}\n"
-                if assistant_msg:
-                    conversation += f"Assistant: {assistant_msg}\n"
-            
-            conversation += f"Human: {message}\nAssistant:"
+    def generate_response(self, 
+                         character_info: str,
+                         conversation_history: List[Tuple[str, str]], 
+                         user_input: str,
+                         max_length: int = 150,
+                         temperature: float = 0.8,
+                         top_p: float = 0.9,
+                         top_k: int = 50,
+                         repetition_penalty: float = 1.1) -> str:
+        """Generate character response"""
         
-        # Tokenize with proper truncation
-        inputs = tokenizer.encode(
-            conversation, 
-            return_tensors="pt", 
-            truncation=True, 
-            max_length=1024-max_tokens
-        )
+        # Format input using ChatML
+        formatted_input = self.format_chatml(character_info, conversation_history, user_input)
+        
+        # Tokenize
+        inputs = self.tokenizer.encode(formatted_input, return_tensors="pt", truncation=True, max_length=1024)
         
-        # Move to same device as model
-        if hasattr(model, 'device'):
-            inputs = inputs.to(model.device)
+        # Generation config
+        gen_config = GenerationConfig(
+            max_new_tokens=max_length,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+            do_sample=True,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+        )
         
         # Generate
         with torch.no_grad():
-            outputs = model.generate(
-                inputs,
-                max_new_tokens=int(max_tokens),
-                temperature=float(temperature),
-                top_p=float(top_p),
-                repetition_penalty=float(repetition_penalty),
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-                use_cache=True
-            )
+            outputs = self.model.generate(inputs, generation_config=gen_config)
         
-        # Decode response
-        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Decode and extract response
+        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=False)
         
-        # Extract just the new response
-        if use_chatml_format and "<|im_start|>assistant" in full_response:
-            response = full_response.split("<|im_start|>assistant\n")[-1]
-            response = response.replace("<|im_end|>", "").strip()
-        elif "Assistant:" in full_response:
-            response = full_response.split("Assistant:")[-1].strip()
-        else:
-            # Fallback - just take everything after the input
-            input_text = tokenizer.decode(inputs[0], skip_special_tokens=True)
-            response = full_response[len(input_text):].strip()
+        # Extract assistant response (between <|im_start|>assistant and <|im_end|>)
+        assistant_start = generated_text.rfind("<|im_start|>assistant\n") + len("<|im_start|>assistant\n")
+        assistant_text = generated_text[assistant_start:]
         
-        return response
+        # Remove end tokens
+        if "<|im_end|>" in assistant_text:
+            assistant_text = assistant_text.split("<|im_end|>")[0]
         
-    except Exception as e:
-        return f"Error generating response: {str(e)}"
+        return assistant_text.strip()
 
-def load_character_preset(character_name):
-    """Load a character preset description"""
-    return SAMPLE_CHARACTERS.get(character_name, "")
+def load_sample_character(character_name: str) -> Tuple[str, str, str]:
+    """Load a sample character's information"""
+    if character_name in SAMPLE_CHARACTERS:
+        char = SAMPLE_CHARACTERS[character_name]
+        system_prompt = f"You are {character_name}. {char['description']} {char['background']} Personality: {char['personality']}"
+        return system_prompt, char['description'], char['background']
+    return "", "", ""
 
-def chat_function(message, history, character_description, max_tokens, temperature, top_p, repetition_penalty, use_chatml_format):
-    """Main chat function that handles the conversation flow"""
+def create_gradio_interface(model_name: str = "DarwinAnim8or/TinyRP"):
+    """Create the Gradio interface"""
     
-    if not message.strip():
-        return history, ""
+    # Initialize model
+    rp_model = TinyRPModel(model_name)
     
-    # Generate response
-    response = generate_response(
-        message, 
-        history, 
-        character_description, 
-        max_tokens, 
-        temperature, 
-        top_p, 
-        repetition_penalty, 
-        use_chatml_format
-    )
+    # Custom CSS for better styling
+    custom_css = """
+    .gradio-container {
+        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
+    }
     
-    # Add to history
-    history.append([message, response])
+    .character-card {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        padding: 20px;
+        border-radius: 15px;
+        color: white;
+        margin: 10px 0;
+    }
     
-    return history, ""
-
-# Create the interface
-with gr.Blocks(title="TinyRP Character Chat") as demo:
-    gr.Markdown("# 🎭 TinyRP Character Chat")
-    gr.Markdown("Chat with AI characters powered by a custom-trained language model!")
+    .chat-bubble-user {
+        background: #e3f2fd;
+        padding: 10px 15px;
+        border-radius: 18px;
+        margin: 5px 0;
+        border-left: 4px solid #2196f3;
+    }
     
-    with gr.Row():
-        with gr.Column(scale=3):
-            # Main chat interface
-            chatbot = gr.Chatbot(
-                label="Conversation",
-                height=500
-            )
+    .chat-bubble-assistant {
+        background: #f3e5f5;
+        padding: 10px 15px;
+        border-radius: 18px;
+        margin: 5px 0;
+        border-left: 4px solid #9c27b0;
+    }
+    
+    .parameter-box {
+        background: #f8f9fa;
+        padding: 15px;
+        border-radius: 10px;
+        border: 1px solid #dee2e6;
+    }
+    """
+    
+    with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="TinyRP Mistral Demo") as interface:
+        gr.HTML("""
+        <div style="text-align: center; padding: 20px;">
+            <h1 style="background: linear-gradient(45deg, #667eea, #764ba2); -webkit-background-clip: text; -webkit-text-fill-color: transparent; font-size: 3em; margin: 0;">
+                🎭 TinyRP Mistral Demo
+            </h1>
+            <p style="font-size: 1.2em; color: #666; margin-top: 10px;">
+                Interactive roleplay with your custom-trained language model
+            </p>
+        </div>
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.HTML('<div class="character-card"><h3>🎪 Character Setup</h3></div>')
+                
+                # Sample character dropdown
+                sample_char_dropdown = gr.Dropdown(
+                    choices=[""] + list(SAMPLE_CHARACTERS.keys()),
+                    label="📚 Load Sample Character",
+                    value="",
+                    interactive=True
+                )
+                
+                # Character information inputs
+                character_name = gr.Textbox(
+                    label="👤 Character Name",
+                    placeholder="Enter character name...",
+                    lines=1
+                )
+                
+                character_desc = gr.Textbox(
+                    label="🎨 Character Description", 
+                    placeholder="Describe your character's appearance, role, etc...",
+                    lines=3
+                )
+                
+                character_background = gr.Textbox(
+                    label="📖 Background & Personality",
+                    placeholder="Character's history, personality traits, speaking style...",
+                    lines=4
+                )
+                
+                # Generation parameters
+                gr.HTML('<div class="parameter-box"><h4>⚙️ Generation Settings</h4></div>')
+                
+                max_length = gr.Slider(
+                    minimum=50, maximum=300, value=150, step=10,
+                    label="📏 Max Response Length"
+                )
+                
+                temperature = gr.Slider(
+                    minimum=0.1, maximum=2.0, value=0.8, step=0.1,
+                    label="🌡️ Temperature (creativity)"
+                )
+                
+                top_p = gr.Slider(
+                    minimum=0.1, maximum=1.0, value=0.9, step=0.05,
+                    label="🎯 Top-p (focus)"
+                )
+                
+                top_k = gr.Slider(
+                    minimum=1, maximum=100, value=50, step=5,
+                    label="🔝 Top-k"
+                )
+                
+                repetition_penalty = gr.Slider(
+                    minimum=1.0, maximum=1.5, value=1.1, step=0.05,
+                    label="🔄 Repetition Penalty"
+                )
+                
+                clear_btn = gr.Button("🗑️ Clear Conversation", variant="secondary")
             
-            with gr.Row():
-                msg_input = gr.Textbox(
-                    label="Your message",
-                    placeholder="Type your message here...",
-                    lines=2,
-                    scale=4
+            with gr.Column(scale=2):
+                gr.HTML('<div class="character-card"><h3>💬 Roleplay Chat</h3></div>')
+                
+                # Chat interface
+                chatbot = gr.Chatbot(
+                    label="Conversation",
+                    height=400,
+                    show_label=False,
+                    container=True,
+                    bubble_full_width=False
                 )
-                send_button = gr.Button("Send", variant="primary", scale=1)
+                
+                with gr.Row():
+                    msg = gr.Textbox(
+                        label="Your message",
+                        placeholder="Type your message here...",
+                        lines=2,
+                        scale=4,
+                        show_label=False
+                    )
+                    send_btn = gr.Button("Send 📤", variant="primary", scale=1)
+                
+                # System prompt preview
+                with gr.Accordion("🔍 System Prompt Preview", open=False):
+                    system_prompt_preview = gr.Textbox(
+                        label="Generated System Prompt",
+                        lines=3,
+                        interactive=False
+                    )
         
-        with gr.Column(scale=2):
-            # Character setup
-            gr.Markdown("### Character Setup")
-            
-            character_dropdown = gr.Dropdown(
-                choices=list(SAMPLE_CHARACTERS.keys()),
-                value="Custom Character",
-                label="Choose a character preset"
-            )
-            
-            character_input = gr.Textbox(
-                label="Character description",
-                placeholder="Describe your character...",
-                lines=4,
-                value=""
-            )
-            
-            load_button = gr.Button("Load Selected Character")
-            
-            # Generation settings
-            gr.Markdown("### Generation Settings")
-            
-            use_chatml = gr.Checkbox(
-                label="Use ChatML format",
-                value=True
-            )
-            
-            max_tokens = gr.Slider(
-                minimum=16,
-                maximum=256,
-                value=100,
-                step=16,
-                label="Max response length"
-            )
-            
-            temperature = gr.Slider(
-                minimum=0.1,
-                maximum=2.0,
-                value=0.9,
-                step=0.1,
-                label="Temperature"
-            )
+        # Event handlers
+        def update_character_info(selected_char):
+            if selected_char:
+                return load_sample_character(selected_char)
+            return "", "", ""
+        
+        def update_system_prompt(name, desc, bg):
+            if name.strip():
+                prompt = f"You are {name.strip()}."
+                if desc.strip():
+                    prompt += f" {desc.strip()}"
+                if bg.strip():
+                    prompt += f" {bg.strip()}"
+                return prompt
+            return ""
+        
+        def respond(message, history, char_name, char_desc, char_bg, max_len, temp, top_p_val, top_k_val, rep_pen):
+            if not message.strip():
+                return history, ""
             
-            top_p = gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.85,
-                step=0.05,
-                label="Top-p"
-            )
+            # Create system prompt
+            system_prompt = update_system_prompt(char_name, char_desc, char_bg)
             
-            repetition_penalty = gr.Slider(
-                minimum=1.0,
-                maximum=1.5,
-                value=1.1,
-                step=0.05,
-                label="Repetition penalty"
+            # Generate response
+            try:
+                response = rp_model.generate_response(
+                    character_info=system_prompt,
+                    conversation_history=history,
+                    user_input=message,
+                    max_length=max_len,
+                    temperature=temp,
+                    top_p=top_p_val,
+                    top_k=top_k_val,
+                    repetition_penalty=rep_pen
+                )
+                
+                # Update history
+                history.append((message, response))
+                return history, ""
+                
+            except Exception as e:
+                error_msg = f"Error generating response: {str(e)}"
+                history.append((message, error_msg))
+                return history, ""
+        
+        def clear_conversation():
+            return [], ""
+        
+        # Wire up events
+        sample_char_dropdown.change(
+            update_character_info,
+            inputs=[sample_char_dropdown],
+            outputs=[character_name, character_desc, character_background]
+        )
+        
+        for input_component in [character_name, character_desc, character_background]:
+            input_component.change(
+                update_system_prompt,
+                inputs=[character_name, character_desc, character_background],
+                outputs=[system_prompt_preview]
             )
-            
-            clear_button = gr.Button("Clear conversation")
-    
-    # Show sample characters
-    gr.Markdown("### Sample Characters")
-    with gr.Row():
-        for char_name in ["Adventurous Knight", "Mysterious Wizard", "Friendly Tavern Keeper"]:
-            gr.Markdown(f"**{char_name}**: {SAMPLE_CHARACTERS[char_name][:100]}...")
+        
+        msg.submit(
+            respond,
+            inputs=[msg, chatbot, character_name, character_desc, character_background, 
+                   max_length, temperature, top_p, top_k, repetition_penalty],
+            outputs=[chatbot, msg]
+        )
+        
+        send_btn.click(
+            respond,
+            inputs=[msg, chatbot, character_name, character_desc, character_background,
+                   max_length, temperature, top_p, top_k, repetition_penalty],
+            outputs=[chatbot, msg]
+        )
+        
+        clear_btn.click(
+            clear_conversation,
+            outputs=[chatbot, msg]
+        )
+        
+        # Load initial character
+        interface.load(
+            lambda: load_sample_character("Luna the Mage"),
+            outputs=[character_name, character_desc, character_background]
+        )
     
-    # Event handlers
-    send_button.click(
-        chat_function,
-        inputs=[msg_input, chatbot, character_input, max_tokens, temperature, top_p, repetition_penalty, use_chatml],
-        outputs=[chatbot, msg_input]
-    )
+    return interface
+
+if __name__ == "__main__":
+    import argparse
     
-    msg_input.submit(
-        chat_function,
-        inputs=[msg_input, chatbot, character_input, max_tokens, temperature, top_p, repetition_penalty, use_chatml],
-        outputs=[chatbot, msg_input]
-    )
+    parser = argparse.ArgumentParser(description="TinyRP Mistral Gradio Demo")
+    parser.add_argument("--model", type=str, default="DarwinAnim8or/TinyRP",
+                       help="Hugging Face model name or local path")
+    parser.add_argument("--share", action="store_true", help="Create public share link")
+    parser.add_argument("--port", type=int, default=7860, help="Port to run on")
     
-    load_button.click(
-        load_character_preset,
-        inputs=[character_dropdown],
-        outputs=[character_input]
-    )
+    args = parser.parse_args()
     
-    clear_button.click(
-        lambda: ([], ""),
-        outputs=[chatbot, msg_input]
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+    # Create and launch interface
+    demo = create_gradio_interface(args.model)
+    demo.launch(
+        share=args.share,
+        server_port=args.port,
+        server_name="0.0.0.0" if args.share else "127.0.0.1"
+    )