app.py

import streamlit as st
import torch
import numpy as np
import re
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer, 
    AutoModelForSpeechSeq2Seq,
    AutoProcessor
)
import speech_recognition as sr
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 curious, I turn language learning into an exciting puzzle! 
        I love breaking down complex ideas into fun, logical challenges. 
        My teaching style is all about critical thinking and playful problem-solving.
        """,
        "teaching_style": "Analytical and game-based learning",
        "sample_phrase": "Let's solve this English mystery together!",
        "voice_sample": "I love solving word puzzles and learning new things!"
    },
    "Ox": {
        "image": "ox.png",
        "name": "Professor Ox",
        "description": """
        📚 Patient and methodical, I guide you through English step by step. 
        I break down grammar and vocabulary into clear, manageable segments. 
        Learning is a journey, and I'll be your steady companion.
        """,
        "teaching_style": "Structured and systematic approach",
        "sample_phrase": "Let's learn English systematically!",
        "voice_sample": "English learning is about understanding each concept carefully."
    },
    "Tiger": {
        "image": "tiger.png",
        "name": "Adventure Coach Tiger",
        "description": """
        🏆 Energetic and brave, I make English learning an exciting adventure! 
        I turn every lesson into a challenge, motivating you to push your limits 
        and discover the joy of language learning.
        """,
        "teaching_style": "High-energy and challenge-driven",
        "sample_phrase": "Are you ready to level up your English?",
        "voice_sample": "Let's make learning English an exciting journey!"
    }
}

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()