Update app.py

app.py

@@ -1,64 +1,365 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import warnings
+import time
+import random
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+warnings.filterwarnings('ignore')
 
+# Global model variables
+model = None
+tokenizer = None
+device = None
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+def load_model():
+    """Load the psychology-tuned model with error handling"""
+    global model, tokenizer, device
+    
+    try:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Loading model on {device}...")
+        
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        # Load base model
+        base_model = AutoModelForCausalLM.from_pretrained(
+            "Qwen/Qwen2.5-0.5B",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None,
+            trust_remote_code=True
+        )
+        
+        # Load PEFT adapter
+        model = PeftModel.from_pretrained(base_model, "phxdev/psychology-qwen-0.5b")
+        model = model.merge_and_unload()
+        
+        if not torch.cuda.is_available():
+            model = model.to(device)
+        
+        print("✅ Model loaded successfully!")
+        return True
+        
+    except Exception as e:
+        print(f"❌ Error loading model: {e}")
+        return False
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+def generate_response(
+    message: str,
+    history: list,
+    temperature: float = 0.8,
+    max_tokens: int = 300,
+    prompt_style: str = "Therapeutic"
+) -> str:
+    """Generate psychology-focused response"""
+    
+    if model is None:
+        return "⚠️ Model is still loading. Please wait a moment and try again."
+    
+    # Define prompt templates
+    prompt_templates = {
+        "Therapeutic": """You're a supportive therapist in session. The client just said: "{message}"
 
-    messages.append({"role": "user", "content": message})
+Respond with empathy and practical guidance. Start with validation, then give 2-3 specific strategies they can try:""",
+        
+        "Supportive Friend": """You're a caring friend who studied psychology. Someone you care about just told you: "{message}"
 
-    response = ""
+Give them warm, understanding support with practical advice:""",
+        
+        "Crisis Support": """You are a crisis counselor. Someone is reaching out for support. They said: "{message}"
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+Provide immediate, caring support and grounding techniques:""",
+        
+        "Anxiety Coach": """You're an anxiety specialist. Help them manage their anxiety with evidence-based techniques.
 
-        response += token
-        yield response
+They're struggling with: "{message}"
 
+Here are specific techniques they can try right now:""",
+        
+        "Mindfulness Guide": """You're a mindfulness teacher. Guide them toward present-moment awareness and self-compassion.
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+They shared: "{message}"
 
+Offer mindful perspective and a gentle practice:"""
+    }
+    
+    # Select and format prompt
+    template = prompt_templates.get(prompt_style, prompt_templates["Therapeutic"])
+    formatted_prompt = template.format(message=message)
+    
+    try:
+        # Tokenize
+        inputs = tokenizer(
+            formatted_prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512
+        ).to(device)
+        
+        # Generate
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=len(inputs.input_ids[0]) + max_tokens,
+                temperature=temperature,
+                top_p=0.9,
+                do_sample=True,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+        
+        # Decode
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        generated_text = response[len(formatted_prompt):].strip()
+        
+        # Clean up response
+        if generated_text.startswith('"') and generated_text.endswith('"'):
+            generated_text = generated_text[1:-1]
+        
+        return generated_text
+        
+    except Exception as e:
+        return f"⚠️ Error generating response: {str(e)}"
 
-if __name__ == "__main__":
-    demo.launch()
+def get_example_prompts():
+    """Return curated example prompts for different scenarios"""
+    examples = {
+        "Work Stress": [
+            "I feel overwhelmed with my workload and don't know how to manage everything.",
+            "My boss is very demanding and I'm afraid of disappointing them.",
+            "I feel like I'm not good enough at my job compared to my colleagues."
+        ],
+        "Anxiety & Worry": [
+            "I can't stop worrying about things that might go wrong.",
+            "I have panic attacks before important meetings or presentations.",
+            "My mind races at night and I can't fall asleep."
+        ],
+        "Relationships": [
+            "I have trouble setting boundaries with people who take advantage of me.",
+            "I feel lonely even when I'm around other people.",
+            "I'm afraid of being rejected if I show my true self."
+        ],
+        "Self-Esteem": [
+            "I constantly criticize myself and focus on my mistakes.",
+            "I feel like everyone else has their life figured out except me.",
+            "I'm afraid to try new things because I might fail."
+        ],
+        "Life Changes": [
+            "I'm going through a major life transition and feel lost.",
+            "I'm grieving the loss of someone important to me.",
+            "I feel stuck in patterns that aren't serving me anymore."
+        ]
+    }
+    return examples
+
+def create_interface():
+    """Create the main Gradio interface"""
+    
+    # Custom CSS for modern, calming design
+    custom_css = """
+    .main-header {
+        text-align: center;
+        padding: 2rem;
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        color: white;
+        border-radius: 10px;
+        margin-bottom: 2rem;
+    }
+    .disclaimer-box {
+        background-color: #f8f9fa;
+        border-left: 4px solid #17a2b8;
+        padding: 1rem;
+        margin: 1rem 0;
+        border-radius: 5px;
+    }
+    .example-category {
+        margin: 0.5rem 0;
+        padding: 0.5rem;
+        background-color: #f1f3f4;
+        border-radius: 5px;
+    }
+    """
+    
+    # Get example prompts
+    examples = get_example_prompts()
+    
+    with gr.Blocks(css=custom_css, title="Psychology AI Assistant", theme=gr.themes.Soft()) as demo:
+        
+        # Header
+        gr.HTML("""
+        <div class="main-header">
+            <h1>🧠 Psychology AI Assistant</h1>
+            <p>A supportive AI trained in psychology and mental health</p>
+            <p><em>Powered by fine-tuned Qwen 2.5-0.5B</em></p>
+        </div>
+        """)
+        
+        # Disclaimer
+        gr.HTML("""
+        <div class="disclaimer-box">
+            <strong>⚠️ Important Disclaimer:</strong> This AI assistant provides supportive guidance based on psychological principles, 
+            but it is not a replacement for professional therapy or medical advice. If you're experiencing a mental health crisis, 
+            please contact a mental health professional or crisis hotline immediately.
+        </div>
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=3):
+                # Main chat interface
+                chatbot = gr.Chatbot(
+                    height=500,
+                    placeholder="👋 Hi! I'm here to provide supportive guidance and practical strategies for mental wellness. What's on your mind today?",
+                    avatar_images=("🧑‍💼", "🧠")
+                )
+                
+                with gr.Row():
+                    message_input = gr.Textbox(
+                        placeholder="Share what's on your mind...",
+                        container=False,
+                        scale=4,
+                        lines=2
+                    )
+                    send_btn = gr.Button("Send", variant="primary", scale=1)
+                
+                # Example prompts
+                gr.HTML("<h3>💡 Try these conversation starters:</h3>")
+                
+                for category, prompts in examples.items():
+                    with gr.Row():
+                        gr.HTML(f"<strong>{category}:</strong>")
+                    with gr.Row():
+                        for prompt in prompts:
+                            example_btn = gr.Button(
+                                prompt[:60] + "..." if len(prompt) > 60 else prompt,
+                                size="sm",
+                                variant="secondary"
+                            )
+                            example_btn.click(
+                                lambda x=prompt: x,
+                                outputs=[message_input]
+                            )
+            
+            with gr.Column(scale=1):
+                # Settings panel
+                gr.HTML("<h3>⚙️ Settings</h3>")
+                
+                prompt_style = gr.Dropdown(
+                    choices=["Therapeutic", "Supportive Friend", "Crisis Support", "Anxiety Coach", "Mindfulness Guide"],
+                    value="Therapeutic",
+                    label="Response Style",
+                    info="Choose the type of support you prefer"
+                )
+                
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.8,
+                    step=0.1,
+                    label="Creativity",
+                    info="Higher = more creative responses"
+                )
+                
+                max_tokens = gr.Slider(
+                    minimum=100,
+                    maximum=500,
+                    value=300,
+                    step=50,
+                    label="Response Length",
+                    info="Maximum response length"
+                )
+                
+                clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
+                
+                # Model info
+                gr.HTML("""
+                <div style="margin-top: 2rem; padding: 1rem; background-color: #f8f9fa; border-radius: 5px;">
+                    <h4>📊 Model Info</h4>
+                    <p><strong>Base Model:</strong> Qwen/Qwen2.5-0.5B</p>
+                    <p><strong>Fine-tuned:</strong> phxdev/psychology-qwen-0.5b</p>
+                    <p><strong>Specialization:</strong> Psychology & Mental Health</p>
+                    <p><strong>Training:</strong> PEFT/LoRA</p>
+                </div>
+                """)
+        
+        # Chat functionality
+        def respond(message, history, temp, max_tok, style):
+            if not message.strip():
+                return history, ""
+            
+            # Add user message
+            history = history + [[message, None]]
+            
+            # Generate response
+            bot_response = generate_response(message, history, temp, max_tok, style)
+            
+            # Add bot response
+            history[-1][1] = bot_response
+            
+            return history, ""
+        
+        def clear_chat():
+            return [], ""
+        
+        # Event handlers
+        send_btn.click(
+            respond,
+            inputs=[message_input, chatbot, temperature, max_tokens, prompt_style],
+            outputs=[chatbot, message_input]
+        )
+        
+        message_input.submit(
+            respond,
+            inputs=[message_input, chatbot, temperature, max_tokens, prompt_style],
+            outputs=[chatbot, message_input]
+        )
+        
+        clear_btn.click(clear_chat, outputs=[chatbot, message_input])
+        
+        # Footer
+        gr.HTML("""
+        <div style="text-align: center; margin-top: 2rem; padding: 1rem; color: #666;">
+            <p>Built with ❤️ using Gradio • Fine-tuned by @phxdev</p>
+            <p>If you're in crisis, please reach out: <a href="https://988lifeline.org/" target="_blank">988 Suicide & Crisis Lifeline</a></p>
+        </div>
+        """)
+    
+    return demo
+
+# Initialize the model
+print("🚀 Starting Psychology AI Assistant...")
+model_loaded = load_model()
+
+if model_loaded:
+    print("✅ Model loaded successfully!")
+    demo = create_interface()
+    
+    if __name__ == "__main__":
+        demo.launch(
+            share=False,
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_error=True
+        )
+else:
+    print("❌ Failed to load model. Creating error interface...")
+    
+    def create_error_interface():
+        with gr.Blocks() as error_demo:
+            gr.HTML("""
+            <div style="text-align: center; padding: 2rem;">
+                <h1>⚠️ Model Loading Error</h1>
+                <p>Sorry, there was an issue loading the psychology model.</p>
+                <p>Please try refreshing the page or contact support.</p>
+            </div>
+            """)
+        return error_demo
+    
+    demo = create_error_interface()
+    
+    if __name__ == "__main__":
+        demo.launch()