webapp/app.py

"""
Personality Chatbot - Multi-personality LLM with LoRA adapters
Deployed on Hugging Face Spaces
"""

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import os

# Configuration
BASE_MODEL = "Qwen/Qwen2-0.5B-Instruct"
ADAPTERS = {
    "🧠 Brainrot": "qwen-brainrot-lora-stage1-final",
    "🏴‍☠️ Pirate": "pirate-lora-adapter",
    "🧙 Yoda": "yoda-lora-adapter",
    "🤓 Nerd": "nerd-lora-adapter",
}

# Global state
base_model = None
tokenizer = None
current_adapter = None
current_personality = None
device = None

def load_base_model():
    """Load base model and tokenizer once at startup"""
    global base_model, tokenizer, device
    
    print("🔄 Loading base model...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        device_map="auto" if device == "cuda" else None,
        trust_remote_code=True,
    )
    
    if device == "cpu":
        base_model = base_model.to(device)
    
    print(f"✅ Base model loaded on {device}")
    return f"Base model loaded on {device}"

def switch_personality(personality_name):
    """Switch to a different personality adapter"""
    global current_adapter, current_personality
    
    if personality_name == current_personality:
        return f"Already using {personality_name}"
    
    adapter_path = ADAPTERS.get(personality_name)
    if not adapter_path:
        return f"❌ Personality '{personality_name}' not found"
    
    if not os.path.exists(adapter_path):
        return f"❌ Adapter folder '{adapter_path}' not found. Make sure adapters are uploaded."
    
    try:
        print(f"🔄 Loading {personality_name} adapter from {adapter_path}...")
        
        # Load adapter on top of base model
        current_adapter = PeftModel.from_pretrained(
            base_model,
            adapter_path,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        )
        current_adapter.eval()
        
        current_personality = personality_name
        print(f"✅ Switched to {personality_name}")
        return f"✅ Switched to {personality_name}"
    
    except Exception as e:
        return f"❌ Error loading adapter: {str(e)}"

def generate_response(message, history, temperature=0.7, max_tokens=256):
    """Generate response using current personality"""
    if current_adapter is None:
        return "⚠️ Please select a personality first!"
    
    try:
        # Format with chat template
        prompt = f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
        
        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = current_adapter.generate(
                **inputs,
                max_new_tokens=max_tokens,  # Changed from max_length to max_new_tokens
                temperature=temperature,
                do_sample=True,
                top_p=0.9,
                top_k=50,  # Added top-k sampling
                repetition_penalty=1.1,  # Added to reduce repetition
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        
        # Decode and extract only the new tokens (the response)
        response_tokens = outputs[0][inputs['input_ids'].shape[1]:]
        response = tokenizer.decode(response_tokens, skip_special_tokens=True).strip()
        
        # Remove any leftover special tokens or formatting
        response = response.replace("<|im_start|>", "").replace("<|im_end|>", "")
        response = response.replace("assistant\n", "").strip()
        
        return response
    
    except Exception as e:
        return f"❌ Error generating response: {str(e)}"

def handle_personality_change(personality_name):
    """Handle personality dropdown change"""
    status = switch_personality(personality_name)
    return status

# Load base model on startup
print("🚀 Starting application...")
load_base_model()

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft(), title="Personality Chatbot") as demo:
    gr.Markdown(
        """
        # 🎭 Multi-Personality Chatbot
        Chat with AI personalities powered by LoRA adapters on Qwen2-0.5B-Instruct
        
        **Select a personality** and start chatting!
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            personality_dropdown = gr.Dropdown(
                choices=list(ADAPTERS.keys()),
                label="🎭 Select Personality",
                value=list(ADAPTERS.keys())[0],
                interactive=True,
            )
            
            status_box = gr.Textbox(
                label="Status",
                value="Select a personality to begin",
                interactive=False,
                lines=2,
            )
            
            with gr.Accordion("⚙️ Generation Settings", open=False):
                temperature_slider = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=0.8,
                    step=0.1,
                    label="Temperature (creativity)",
                )
                max_tokens_slider = gr.Slider(
                    minimum=30,
                    maximum=256,
                    value=100,
                    step=10,
                    label="Max tokens (response length)",
                )
            
            gr.Markdown(
                """
                ### 📝 Personality Descriptions
                - **🧠 Brainrot**: Internet slang and Gen-Z speak
                - **🏴‍☠️ Pirate**: Arr matey, talks like a pirate!
                - **🧙 Yoda**: Wise Jedi master, speaks in reverse
                - **🤓 Nerd**: Intellectual, loves facts and science
                """
            )
        
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="Chat",
                height=500,
                show_label=True,
            )
            
            msg_box = gr.Textbox(
                label="Your message",
                placeholder="Type your message here...",
                lines=2,
            )
            
            with gr.Row():
                submit_btn = gr.Button("Send 💬", variant="primary")
                clear_btn = gr.Button("Clear 🗑️", variant="secondary")
    
    # Event handlers
    def respond(message, chat_history, temperature, max_tokens):
        if not message.strip():
            return chat_history, ""
        
        bot_response = generate_response(message, chat_history, temperature, max_tokens)
        chat_history.append((message, bot_response))
        return chat_history, ""
    
    # Personality change handler
    personality_dropdown.change(
        fn=handle_personality_change,
        inputs=[personality_dropdown],
        outputs=[status_box],
    )
    
    # Chat handlers
    submit_btn.click(
        fn=respond,
        inputs=[msg_box, chatbot, temperature_slider, max_tokens_slider],
        outputs=[chatbot, msg_box],
    )
    
    msg_box.submit(
        fn=respond,
        inputs=[msg_box, chatbot, temperature_slider, max_tokens_slider],
        outputs=[chatbot, msg_box],
    )
    
    clear_btn.click(
        fn=lambda: ([], ""),
        outputs=[chatbot, msg_box],
    )
    
    # Load first personality on startup
    demo.load(
        fn=handle_personality_change,
        inputs=[personality_dropdown],
        outputs=[status_box],
    )

# Launch
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
    )