app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import re
import random

print("🤖 Loading Simple AI Assistant...")

# === MODEL CONFIGURATION (DEPLOYMENT-READY) ===
# Using only Mistral-7B-AWQ model as requested
MODEL_CONFIG = {
    "id": "TheBloke/Mistral-7B-Instruct-v0.2-AWQ", 
    "name": "Mistral-AWQ",
    "description": "High-quality instruction model",
    "use_safetensors": True
}

model = None
tokenizer = None
model_name = None
MODEL_ID = None

# Load the single Mistral-7B-AWQ model
try:
    print(f"🔄 Loading {MODEL_CONFIG['description']}...")
    MODEL_ID = MODEL_CONFIG["id"]
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    
    # Try advanced model with fallback parameters
    try:
        # First try with autoawq support
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            device_map="auto" if torch.cuda.is_available() else "cpu",
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
            use_safetensors=True  # Prefer safetensors for security
        )
        print("✅ AWQ model loaded successfully with device mapping!")
    except Exception as awq_error:
        print(f"⚠️ AWQ device mapping failed: {awq_error}")
        print("🔄 Trying standard loading without device mapping...")
        # Fallback without device mapping
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
            use_safetensors=True
        )
        print("✅ AWQ model loaded successfully with standard loading!")
    
    model_name = MODEL_CONFIG["name"]
    print(f"✅ {MODEL_CONFIG['description']} ready for use!")
    
except Exception as e:
    print(f"❌ Failed to load {MODEL_CONFIG['description']}: {e}")
    model = None

if model is None:
    raise RuntimeError("❌ Could not load any model!")

# Add pad token if needed
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Load sentiment analysis for emotion detection
try:
    print("🔄 Loading emotion detection...")
    emotion_detector = pipeline(
        "sentiment-analysis",
        model="distilbert-base-uncased-finetuned-sst-2-english",
        return_all_scores=True
    )
    print("✅ Emotion detection loaded!")
except Exception as e:
    print(f"⚠️ Emotion detection failed: {e}")
    emotion_detector = None

print("✅ Simple AI Assistant ready!")

# Simple AI Assistant System Prompt
SIMPLE_SYSTEM_PROMPT = """You are a helpful AI assistant. Answer questions directly and clearly. Be friendly and concise. If someone seems upset, be understanding. If they seem happy, match their energy. Keep responses to 1-2 sentences unless more detail is needed."""

def check_crisis_keywords(message):
    """Check for crisis-related keywords that require immediate intervention"""
    crisis_keywords = [
        'suicide', 'kill myself', 'end my life', 'hurt myself', 'self harm', 'self-harm',
        'want to die', 'better off dead', 'no point living', 'end it all'
    ]
    message_lower = message.lower()
    return any(keyword in message_lower for keyword in crisis_keywords)

def get_crisis_response():
    """Return crisis intervention response"""
    return """I'm very concerned about what you're sharing with me. Please reach out for immediate help:
    
🆘 **Crisis Hotlines:**
• National Suicide Prevention Lifeline: 988 or 1-800-273-8255
• Crisis Text Line: Text HOME to 741741
• International Association for Suicide Prevention: https://www.iasp.info/resources/Crisis_Centres/

**Please contact emergency services (911) or go to your nearest emergency room if you're in immediate danger.**

You matter, and there are people who want to help you through this. Please reach out to a mental health professional - they have the training and resources to support you in ways I cannot."""

def is_inappropriate_response(response, user_message):
    """Check if the generated response is inappropriate and should be blocked - ENHANCED VERSION"""
    response_lower = response.lower()
    user_lower = user_message.lower()
    
    # Comprehensive list of inappropriate phrases for different contexts
    general_inappropriate = [
        "did you die", "are you dead", "are you okay", "lol", "haha", "funny", "hilarious",
        "cheer up", "look on the bright side", "at least", "could be worse", "think positive",
        "stop complaining", "get over it", "move on", "that's life", "suck it up",
        "just smile", "be happy", "don't worry", ":d", "xd", "lmao", "rofl"
    ]
    
    # Platitudes that minimize feelings
    empty_platitudes = [
        "don't get discouraged", "it gets easier", "stay strong", "everything happens for a reason",
        "things will get better", "this too shall pass", "everything will be fine", 
        "you'll be fine", "it's all good", "no worries", "don't stress", "relax",
        "think positive", "stay positive", "look on the bright side", "count your blessings"
    ]
    
    # Casual dismissive phrases
    dismissive_phrases = [
        "i know many people", "lots of people", "everyone goes through", "we all", "you and me both",
        "been there", "i've been there", "happens to everyone", "normal thing", "no big deal"
    ]
    
    # Combine all inappropriate patterns
    all_inappropriate = general_inappropriate + empty_platitudes + dismissive_phrases
    
    # Check for any inappropriate phrases
    if any(phrase in response_lower for phrase in all_inappropriate):
        return True
    
    # Specific checks for injury/medical situations
    injury_keywords = ["broken", "injured", "hurt", "pain", "fell", "accident", "bleeding", "fracture"]
    if any(word in user_lower for word in injury_keywords):
        # Extra strict for injury responses
        if any(phrase in response_lower for phrase in ["don't worry", "you'll be fine", "no big deal", "happens"]):
            return True
        
        # Block inappropriate medical advice or casual suggestions
        inappropriate_medical = [
            "get a new", "just wear", "try to get", "always try", "just use", 
            "wear a glove", "use the other", "it's not that bad", "could be worse"
        ]
        if any(phrase in response_lower for phrase in inappropriate_medical):
            return True
            
        # Must acknowledge the injury specifically for broken hand cases
        if "broken" in user_lower and "hand" in user_lower:
            if not any(word in response_lower for word in ["broken", "hand", "pain", "injury", "hurt", "doctor", "medical"]):
                return True
    
    # Specific checks for mental health situations
    mental_health_keywords = ["depressed", "sad", "crying", "devastated", "hopeless", "overwhelmed", "anxious"]
    if any(word in user_lower for word in mental_health_keywords):
        # Block casual responses to serious mental health concerns
        if any(phrase in response_lower for phrase in ["just think positive", "snap out of it", "get over it"]):
            return True
    
    # Block responses that are too short or nonsensical
    if len(response.strip()) < 10 or response.count(" ") < 3:
        return True
    
    # Block responses with excessive repetition
    words = response_lower.split()
    if len(words) > 5 and len(set(words)) < len(words) * 0.4:  # More than 60% repetition
        return True
        
    # Block responses that seem to make light of serious situations
    if any(word in user_lower for word in ["help me", "emergency", "urgent", "crisis"]):
        if any(phrase in response_lower for phrase in ["lol", "haha", "funny", "joke"]):
            return True
    
    return False

def format_aura_response(raw_response):
    """Format the response to align with Aura's personality"""
    # Add gentle, empathetic tone if the response seems too direct
    empathetic_starters = [
        "I hear you, and ",
        "That sounds really ",
        "I can imagine that feels ",
        "Thank you for sharing that with me. ",
        "It makes complete sense that you'd feel "
    ]
    
    # If response doesn't start with empathetic language, add some
    if not any(starter.lower() in raw_response.lower()[:50] for starter in empathetic_starters):
        if len(raw_response) > 0:
            return f"I hear you. {raw_response}"
    
    return raw_response

# === EMOTION DETECTION ===
def detect_emotion(message):
    """Detect user emotion for appropriate response tone"""
    if not emotion_detector:
        return "neutral", 0.5
    
    try:
        results = emotion_detector(message)[0]
        for result in results:
            if result['label'] == 'POSITIVE':
                return "positive", result['score']
            elif result['label'] == 'NEGATIVE':
                return "negative", result['score']
        return "neutral", 0.5
    except:
        return "neutral", 0.5

# === EMOJI SELECTION ===
def get_emoji(emotion, confidence):
    """Get appropriate emoji based on emotion"""
    if confidence < 0.6:
        return "😊"
    
    if emotion == "positive":
        return random.choice(["😊", "😄", "🎉", "👍", "✨"])
    elif emotion == "negative":
        return random.choice(["😔", "💙", "🫂", "😞", "💗"])
    else:
        return random.choice(["😊", "👋", "🤔", "💭"])

# === SIMPLE RESPONSE FUNCTION ===
def respond(message, history, max_length=80, temperature=0.7, top_p=0.9, top_k=50, repetition_penalty=1.1):
    """Generate simple, direct responses with appropriate emotion"""
    try:
        # 1. Crisis detection
        if check_crisis_keywords(message):
            return get_crisis_response()
        
        # 2. Detect emotion
        emotion, confidence = detect_emotion(message)
        print(f"Detected emotion: {emotion} (confidence: {confidence:.2f})")
        
        # 3. Build conversation for model
        messages = [
            {"role": "system", "content": SIMPLE_SYSTEM_PROMPT},
            {"role": "user", "content": message}
        ]
        
        # Add recent history (max 2 exchanges)
        if history:
            recent_history = history[-2:]
            full_messages = [{"role": "system", "content": SIMPLE_SYSTEM_PROMPT}]
            for user_msg, bot_msg in recent_history:
                full_messages.append({"role": "user", "content": user_msg})
                if bot_msg:
                    full_messages.append({"role": "assistant", "content": bot_msg})
            full_messages.append({"role": "user", "content": message})
            messages = full_messages
        
        # 4. Handle different model types
        if "mistral" in MODEL_ID.lower() or model_name == "Mistral-AWQ":
            # Use Mistral chat template
            conversation = tokenizer.apply_chat_template(
                messages, 
                tokenize=False, 
                add_generation_prompt=True
            )
        else:
            # Simple format for DialoGPT
            conversation = f"{message}{tokenizer.eos_token}"
        
        # 5. Tokenize
        inputs = tokenizer(
            conversation,
            return_tensors="pt",
            truncation=True,
            max_length=1024,
            padding=True
        )
        
        # 6. Generate response
        with torch.no_grad():
            outputs = model.generate(
                inputs['input_ids'].to(model.device),
                attention_mask=inputs.get('attention_mask', None),
                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=tokenizer.pad_token_id or tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id
            )
        
        # 7. Decode response
        raw_response = tokenizer.decode(
            outputs[:, inputs['input_ids'].shape[-1]:][0],
            skip_special_tokens=True
        ).strip()
        
        # 8. Clean up response
        response = raw_response.replace("Human:", "").replace("Assistant:", "").strip()
        response = re.sub(r'^(User|Bot|AI|Assistant):\s*', '', response)
        
        # 9. Add emotional tone if needed
        if emotion == "negative" and confidence > 0.7:
            if not any(word in response.lower() for word in ["sorry", "understand", "difficult"]):
                response = f"I understand that's tough. {response}"
        elif emotion == "positive" and confidence > 0.7:
            if not any(word in response.lower() for word in ["great", "wonderful", "amazing"]):
                response = f"That's great! {response}"
        
        # 10. Add emoji
        emoji = get_emoji(emotion, confidence)
        
        # 11. Ensure proper formatting
        if response and not response.endswith(('!', '?', '.')):
            response += '.'
        
        final_response = f"{response} {emoji}"
        
        return final_response
        
    except Exception as e:
        print(f"Error: {e}")
        emotion, _ = detect_emotion(message)
        emoji = get_emoji(emotion, 0.5)
        return f"I'm here to help! What can I assist you with? {emoji}"

def add_empathy_to_response(response, user_message):
    """Add Aura's empathetic touch to the raw response with high variety"""
    import random
    
    # Clean up the response first
    cleaned_response = response.replace("Human:", "").replace("Aura:", "").strip()
    
    # Detect context and emotions to provide varied, appropriate responses
    user_lower = user_message.lower()
    
    # Context-specific empathetic starters with high variety
    if "not good enough" in user_lower or "failure" in user_lower or "worthless" in user_lower:
        starters = [
            "Those feelings are so valid, and I want you to know that ",
            "It takes real courage to share those feelings with me. ",
            "I can hear how much pain you're carrying right now. ",
            "That inner voice can be so harsh sometimes. ",
            "What you're feeling makes complete sense, and "
        ]
    elif "tough day" in user_lower or "bad day" in user_lower or "difficult" in user_lower:
        starters = [
            "Some days really are harder than others, aren't they? ",
            "It sounds like today has been especially challenging. ",
            "Days like this can feel so heavy. ",
            "I'm sorry you're going through such a rough time. ",
            "Thank you for trusting me with how you're feeling today. "
        ]
    elif "overwhelmed" in user_lower or "stressed" in user_lower:
        starters = [
            "That feeling of overwhelm can be so intense. ",
            "It sounds like there's a lot weighing on you right now. ",
            "When everything feels like too much, it's really hard. ",
            "I can imagine how exhausting that must feel. ",
            "Sometimes life piles up in ways that feel impossible to manage. "
        ]
    elif "listen" in user_lower or "someone" in user_lower:
        starters = [
            "I'm right here with you, and I'm listening. ",
            "You have my full attention, and I want you to know ",
            "Thank you for reaching out. I'm here, and ",
            "Sometimes we all need someone to just be present with us. ",
            "I'm grateful you felt safe sharing this with me. "
        ]
    elif any(word in user_lower for word in ['hurt', 'pain', 'sad', 'crying']):
        starters = [
            "I can hear how much you're hurting right now. ",
            "Pain like this is so real and so valid. ",
            "It takes strength to sit with feelings like these. ",
            "Your pain matters, and I want you to know ",
            "I'm holding space for everything you're feeling. "
        ]
    elif "fractured" in user_lower or "injury" in user_lower or "hurt" in user_lower:
        starters = [
            "Oh no, that sounds incredibly painful! ",
            "I'm so sorry that happened to you. ",
            "That must be really scary and painful. ",
            "Injuries like that are no small thing. ",
            "I can only imagine how much that must hurt. "
        ]
    else:
        # General varied empathetic responses
        starters = [
            "Thank you for sharing that with me. ",
            "I can hear that this is important to you. ",
            "It sounds like you're going through something difficult. ",
            "I'm here with you in this moment. ",
            "What you're experiencing sounds really challenging. ",
            "I want you to know that your feelings are completely valid. ",
            "It takes courage to open up like this. ",
            "I'm grateful you felt safe sharing this with me. "
        ]
    
    # Choose a random starter for variety
    starter = random.choice(starters)
    
    # Add thoughtful questions sometimes to show deeper engagement
    questions = [
        " What's been the hardest part for you?",
        " How long have you been carrying this?",
        " Is there anything that's helped, even a little?",
        " What do you need most right now?",
        " How are you taking care of yourself through this?",
        ""
    ]
    
    # Add question 30% of the time for engagement variety
    question = random.choice(questions) if random.random() < 0.3 else ""
    
    # Combine everything
    if cleaned_response:
        return f"{starter}{cleaned_response}{question}"
    else:
        return f"{starter}I'm here to listen to whatever you need to share{question}"

def get_fallback_aura_response(user_message):
    """Provide appropriate fallback responses based on user input with variety"""
    import random
    user_lower = user_message.lower()
    
    # Special handling for injury situations (especially broken hand)
    if any(injury_word in user_lower for injury_word in ["broken", "fractured", "injured", "hurt", "fell", "accident"]):
        if "hand" in user_lower or "arm" in user_lower or "wrist" in user_lower:
            responses = [
                "Oh no, that sounds incredibly painful and frightening! 😟 Falling and breaking your hand must be so overwhelming to deal with. Have you been able to see a doctor? How are you managing the pain right now?",
                "I'm so sorry that happened to you - that must be really scary and painful. A broken hand is no small injury. Have you been able to get medical attention? How are you feeling right now?",
                "That sounds absolutely awful and so painful. I can only imagine how much you're hurting and how scary it must have been when you fell. What kind of medical care have you been able to get?",
                "I'm so sorry you're going through this. Breaking your hand from a fall sounds incredibly painful and traumatic. Have you been able to see a doctor about it? How are you coping with everything right now?"
            ]
        else:
            responses = [
                "Oh no, that sounds really painful and scary. I'm so sorry that happened to you. Have you been able to get the medical attention you need?",
                "I'm sorry you got hurt - that must be really frightening and painful to go through. How are you doing right now?",
                "That sounds like it must have been really scary and painful. I'm here with you. Have you been able to get help with your injury?"
            ]
        return random.choice(responses)
    
    # Special handling for depression with physical injury
    if "depressed" in user_lower and any(injury_word in user_lower for injury_word in ["fell", "broken", "hand", "hurt"]):
        responses = [
            "I'm here with you in this difficult moment. Dealing with both emotional pain and a physical injury like a broken hand must feel overwhelming. Your feelings make complete sense - this is a lot to handle.",
            "Thank you for reaching out when you're going through so much. Having depression and dealing with a physical injury at the same time must be incredibly hard. I'm here to listen to whatever you're feeling.",
            "I can hear that you're struggling with both emotional and physical pain right now. That combination must feel so heavy. You don't have to carry this alone - I'm here with you."
        ]
        return random.choice(responses)
    
    if "not good enough" in user_lower:
        responses = [
            "Those feelings of not being good enough are so painful, and I want you to know they're completely valid. It takes real courage to share something so vulnerable.",
            "I can hear how much you're struggling with those feelings about yourself. That inner critic can be so harsh, can't it?",
            "What you're feeling makes complete sense, and I'm grateful you trusted me with something so personal. You're not alone in feeling this way.",
            "That's such a heavy feeling to carry. I want you to know that sharing this with me shows real strength, even when you don't feel strong."
        ]
        return random.choice(responses)
    elif "what" in user_lower and "do" in user_lower:
        responses = [
            "I can hear that you're looking for some direction. Rather than telling you what to do, I'm wondering - what feels most important to you right now?",
            "It sounds like you're feeling a bit lost about next steps. Sometimes when we're overwhelmed, it helps to start with what you're feeling rather than what you should do. What's going on for you?",
            "That feeling of not knowing what to do can be so frustrating. Instead of focusing on actions right now, what would feel most supportive for you in this moment?"
        ]
        return random.choice(responses)
    elif "tough day" in user_lower or "bad day" in user_lower:
        responses = [
            "Some days are just harder than others, aren't they? I'm here with you through this difficult time.",
            "I'm sorry today has been so rough for you. Would it help to talk about what's been making it especially difficult?",
            "Days like this can feel so heavy. Thank you for reaching out - you don't have to go through this alone."
        ]
        return random.choice(responses)
    elif len(user_message.strip()) <= 2:  # Short messages like "??"
        responses = [
            "I'm here with you. Sometimes it's hard to find words, and that's completely okay. Take your time.",
            "You don't need to have the perfect words. I'm here to listen to whatever you're feeling, however you want to express it.",
            "Sometimes feelings are too big for words, and that's alright. I'm here with you in this moment."
        ]
        return random.choice(responses)
    else:
        # General varied empathetic fallback responses
        responses = [
            "Thank you for sharing that with me. I can hear that something important is going on for you, and I want you to know I'm here to listen.",
            "I'm grateful you felt safe reaching out. Whatever you're going through, you don't have to face it alone.",
            "It sounds like you're going through something difficult right now. I'm here with you, and your feelings matter to me.",
            "I can sense that this is weighing on you. Thank you for trusting me with what you're experiencing.",
            "Whatever brought you here today, I want you to know that I'm here to listen without judgment. You're not alone."
        ]
        return random.choice(responses)

# Create Gradio interface  
with gr.Blocks(title="Mistral AI Assistant") as demo:
    gr.Markdown("# 🤖 Mistral AI Assistant")
    gr.Markdown("""
    **Powered by Mistral-7B-Instruct-v0.2-AWQ:**
    - Answers your questions directly and clearly
    - Detects your emotions and responds appropriately
    - Uses emojis to match the conversation tone
    - Keeps responses concise and useful
    - High-quality responses from advanced language model
    """)
    
    chatbot = gr.Chatbot(
        height=500
        # Use default tuples format for compatibility
    )
    
    msg = gr.Textbox(
        placeholder="Ask me anything! I'll help you out 😊",
        container=False,
        scale=7
    )
    
    with gr.Row():
        clear = gr.Button("Clear Chat", variant="secondary")
    
    # Simplified settings
    with gr.Accordion("⚙️ Settings", open=False):
        gr.Markdown("*The assistant is optimized for speed and quality by default.*")
        with gr.Row():
            max_length = gr.Slider(
                minimum=50, maximum=150, value=80, step=10,
                label="Response Length",
                info="Shorter = faster responses"
            )
            temperature = gr.Slider(
                minimum=0.1, maximum=1.0, value=0.7, step=0.1,
                label="Creativity",
                info="Higher = more creative"
            )
    
    def user(user_message, history):
        return "", history + [[user_message, None]]
    
    def bot(history, max_len, temp):
        if history and history[-1][1] is None:
            user_message = history[-1][0]
            bot_response = respond(user_message, history[:-1], max_len, temp)
            history[-1][1] = bot_response
        return history
    
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, max_length, temperature], chatbot
    )
    
    clear.click(lambda: [], None, chatbot, queue=False)
    
    # Example conversations
    gr.Examples(
        examples=[
            "What's the weather like today?",
            "I'm feeling stressed about work",
            "Can you help me with Python code?",
            "I just got a promotion!",
            "How do I make pasta?"
        ],
        inputs=msg,
        label="Try these examples:"
    )

if __name__ == "__main__":
    demo.queue()
    # Check if running on Hugging Face Spaces
    import os
    if "SPACE_ID" in os.environ:
        # Running on HF Spaces - public by default
        demo.launch()
        print(f"🌐 Your chatbot is publicly available at: https://huggingface.co/spaces/{os.environ.get('SPACE_AUTHOR_NAME', 'your-username')}/{os.environ.get('SPACE_REPO_NAME', 'chatbot')}")
    else:
        # Running locally - create public link
        demo.launch(share=True)
        print("🌐 Public link generated above ⬆️")