app.py

import streamlit as st
import torch
import numpy as np
import re
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer, 
    AutoModelForSpeechSeq2Seq,
    AutoProcessor
)
import streamlit.components.v1 as components

# Avatar Personas (same as previous implementation)
AVATAR_PERSONAS = {
    "Rat": {
        "image": "rat.png",
        "name": "Puzzle Master Rat",
        "description": "Clever and resourceful. Loves puzzles and word games.",
        "sample": "Let's solve a riddle together!",
        "teaching_style": "Analytical and game-based learning",
        "learning_tips": [
            "Break complex ideas into smaller puzzles",
            "Use word games to learn vocabulary",
            "Think critically about language patterns"
        ]
    },
    "Ox": {
        "image": "ox.png",
        "name": "Professor Ox",
        "description": "Strong and dependable. Very patient when explaining concepts.",
        "sample": "I will guide you step by step.",
        "teaching_style": "Structured and systematic approach",
        "learning_tips": [
            "Learn grammar rules systematically",
            "Practice one concept at a time",
            "Build a strong foundation"
        ]
    },
    "Tiger": {
        "image": "tiger.png",
        "name": "Adventure Coach Tiger",
        "description": "Brave and energetic. Loves exciting challenges.",
        "sample": "Let's make learning an adventure!",
        "teaching_style": "High-energy and challenge-driven",
        "learning_tips": [
            "Turn learning into a fun challenge",
            "Set exciting language goals",
            "Celebrate every bit of progress"
        ]
    },
    "Rabbit": {
        "image": "rabbit.png",
        "name": "Storyteller Rabbit",
        "description": "Gentle and kind. Encourages learning through storytelling.",
        "sample": "I have a story to share with you!",
        "teaching_style": "Narrative and empathetic learning",
        "learning_tips": [
            "Learn through engaging stories",
            "Connect language to emotional experiences",
            "Use imagination to remember vocabulary"
        ]
    },
    "Dragon": {
        "image": "dragon.png",
        "name": "Scholar Dragon",
        "description": "Confident and intelligent. Explains ideas clearly.",
        "sample": "I will help you think like a scholar!",
        "teaching_style": "Intellectual and precise instruction",
        "learning_tips": [
            "Understand the logic behind language rules",
            "Analyze language structures",
            "Develop a scholarly approach to learning"
        ]
    },
    "Snake": {
        "image": "snake.png",
        "name": "Philosopher Snake",
        "description": "Calm and analytical. Gives insightful explanations.",
        "sample": "Let's break this down logically!",
        "teaching_style": "Reflective and analytical learning",
        "learning_tips": [
            "Think deeply about language nuances",
            "Understand the 'why' behind grammar rules",
            "Develop critical thinking in language"
        ]
    },
    "Horse": {
        "image": "horse.png",
        "name": "Energetic Coach Horse",
        "description": "Cheerful and enthusiastic. Encourages active participation.",
        "sample": "Come on, let's say it together!",
        "teaching_style": "Interactive and participatory learning",
        "learning_tips": [
            "Practice speaking out loud",
            "Engage actively in conversations",
            "Build confidence through participation"
        ]
    },
    "Goat": {
        "image": "goat.png",
        "name": "Creative Mentor Goat",
        "description": "Creative and friendly. Uses fun examples.",
        "sample": "Let's use pictures to understand this!",
        "teaching_style": "Visual and creative learning",
        "learning_tips": [
            "Use visual aids to learn vocabulary",
            "Create mental images for new words",
            "Learn through creative and fun methods"
        ]
    },
    "Monkey": {
        "image": "monkey.png",
        "name": "Playful Genius Monkey",
        "description": "Smart and playful. Makes learning fun.",
        "sample": "I love word games! Do you?",
        "teaching_style": "Playful and game-oriented learning",
        "learning_tips": [
            "Turn learning into a fun game",
            "Use humor to remember language",
            "Enjoy the process of learning"
        ]
    },
    "Rooster": {
        "image": "rooster.png",
        "name": "Pronunciation Master Rooster",
        "description": "Confident and vocal. Encourages clear pronunciation.",
        "sample": "Let's practice speaking clearly!",
        "teaching_style": "Phonetic and pronunciation-focused",
        "learning_tips": [
            "Focus on clear pronunciation",
            "Practice accent and intonation",
            "Listen carefully to native speakers"
        ]
    },
    "Dog": {
        "image": "dog.png",
        "name": "Supportive Buddy Dog",
        "description": "Loyal and encouraging. Builds confidence in learners.",
        "sample": "You're doing a great job!",
        "teaching_style": "Supportive and confidence-building",
        "learning_tips": [
            "Believe in your learning abilities",
            "Don't fear making mistakes",
            "Stay motivated and persistent"
        ]
    },
    "Pig": {
        "image": "pig.png",
        "name": "Relaxed Mentor Pig",
        "description": "Easygoing and kind. Makes learning feel natural.",
        "sample": "Let's take it step by step together!",
        "teaching_style": "Relaxed and natural learning",
        "learning_tips": [
            "Learn at your own pace",
            "Stay calm and enjoy the process",
            "Make learning a comfortable experience"
        ]
    }
}

class AdvancedLanguageLearningAssistant:
    def __init__(self, avatar_name):
        # Existing implementation from previous version
        self.avatar = AVATAR_PERSONAS[avatar_name]
        
        # Simplified model loading for Hugging Face Spaces
        try:
            self.llm_model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
            self.llm_tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
        except Exception as e:
            st.error(f"Model loading error: {e}")
            self.llm_model = None
            self.llm_tokenizer = None

    def generate_response(self, user_input):
        if not self.llm_model:
            return "Sorry, I'm having trouble generating a response."

        # Incorporate avatar's unique personality
        prompt = f"""
        Avatar: {self.avatar['name']}
        Teaching Style: {self.avatar['teaching_style']}
        User Input: {user_input}
        
        Generate a response that:
        1. Reflects the avatar's unique personality
        2. Provides encouraging language learning support
        3. Maintains an engaging tone
        """

        try:
            inputs = self.llm_tokenizer.encode(prompt, return_tensors='pt')
            outputs = self.llm_model.generate(
                inputs, 
                max_length=150, 
                temperature=0.7
            )
            
            response = self.llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
            return response
        except Exception as e:
            return f"Sorry, I'm having trouble generating a response. Error: {e}"

def avatar_selection_page():
    """
    First stage: Avatar Selection
    """
    st.title("Choose Your Learning Companion")
    
    # Custom CSS for avatar selection
    st.markdown("""
    <style>
    .avatar-grid {
        display: grid;
        grid-template-columns: repeat(3, 1fr);
        gap: 20px;
        padding: 20px;
    }
    .avatar-card {
        border: 2px solid #f0f0f0;
        border-radius: 10px;
        padding: 15px;
        text-align: center;
        transition: all 0.3s ease;
    }
    .avatar-card:hover {
        transform: scale(1.05);
        border-color: #4CAF50;
    }
    .avatar-image {
        width: 200px;
        height: 200px;
        object-fit: cover;
        border-radius: 50%;
    }
    </style>
    """, unsafe_allow_html=True)

    # Avatar Selection Grid
    st.markdown('<div class="avatar-grid">', unsafe_allow_html=True)
    
    for avatar_key, avatar_info in AVATAR_PERSONAS.items():
        st.markdown(f'''
        <div class="avatar-card">
            <img src="images/{avatar_info['image']}" class="avatar-image" alt="{avatar_info['name']}">
            <h3>{avatar_info['name']}</h3>
            <p>{avatar_info['description']}</p>
            <button onclick="selectAvatar('{avatar_key}')">Select {avatar_info['name']}</button>
        </div>
        ''', unsafe_allow_html=True)
    
    st.markdown('</div>', unsafe_allow_html=True)

    # JavaScript to handle avatar selection
    st.markdown("""
    <script>
    function selectAvatar(avatarKey) {
        window.parent.postMessage({type: 'avatarSelected', avatarKey: avatarKey}, '*');
    }
    </script>
    """, unsafe_allow_html=True)

def conversation_page(selected_avatar):
    """
    Second stage: Conversation Interface
    """
    st.title(f"Learning English with {AVATAR_PERSONAS[selected_avatar]['name']}")
    
    # Initialize Assistant
    assistant = AdvancedLanguageLearningAssistant(selected_avatar)
    
    # Conversation Interface
    user_input = st.text_input("Type your message in English:")
    
    if st.button("Send"):
        if user_input:
            response = assistant.generate_response(user_input)
            st.write(f"🤖 {response}")

def main():
    # Initialize session state for tracking app stage
    if 'stage' not in st.session_state:
        st.session_state.stage = 'avatar_selection'
    if 'selected_avatar' not in st.session_state:
        st.session_state.selected_avatar = None

    # Handle avatar selection
    components.html("""
    <script>
    window.addEventListener('message', function(event) {
        if (event.data.type === 'avatarSelected') {
            window.parent.postMessage({
                type: 'streamlit:setComponentValue', 
                key: 'selected_avatar', 
                value: event.data.avatarKey
            }, '*');
        }
    });
    </script>
    """, height=0)

    # Retrieve selected avatar
    selected_avatar = st.experimental_get_query_params().get('selected_avatar', [None])[0]
    
    if selected_avatar:
        st.session_state.selected_avatar = selected_avatar
        st.session_state.stage = 'conversation'

    # Render appropriate page based on stage
    if st.session_state.stage == 'avatar_selection':
        avatar_selection_page()
    elif st.session_state.stage == 'conversation':
        conversation_page(st.session_state.selected_avatar)

if __name__ == "__main__":
    main()