app_2.py

import gradio as gr
import torch
import os
from transformers import AutoTokenizer, AutoModelForCausalLM
import random
from gradio.themes.utils import colors, sizes
from gradio.themes import Base, Soft

# Create custom themes for light and dark mode
class GemmaLightTheme(Soft):
    def __init__(self):
        super().__init__(
            primary_hue=colors.blue,
            secondary_hue=colors.indigo,
            neutral_hue=colors.gray,
            spacing_size=sizes.spacing_md,
            radius_size=sizes.radius_md,
            text_size=sizes.text_md,
        )
        
class GemmaDarkTheme(Soft):
    def __init__(self):
        super().__init__(
            primary_hue=colors.blue,
            secondary_hue=colors.indigo,
            neutral_hue=colors.gray,
            spacing_size=sizes.spacing_md,
            radius_size=sizes.radius_md,
            text_size=sizes.text_md,
        )
        self.name = "gemma_dark"
        # Make it dark
        self.background_fill_primary = "#1F1F2E"
        self.background_fill_secondary = "#2A2A3C"
        self.border_color_primary = "#3A3A4C"
        self.border_color_secondary = "#4A4A5C"
        self.color_accent_soft = "#3B5FA3"
        self.color_accent = "#4B82C4"
        self.text_color = "#FFFFFF"
        self.text_color_subdued = "#CCCCCC"
        self.shadow_spread = "8px"
        self.shadow_inset = "0px 1px 2px 0px rgba(0, 0, 0, 0.1) inset"

# Helper function to handle empty values
def safe_value(value, default):
    """Return default if value is empty or None"""
    if value is None or value == "":
        return default
    return value

# Get Hugging Face token from environment variable (as fallback)
DEFAULT_HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN", None)

# Create global variables for model and tokenizer
global_model = None
global_tokenizer = None
model_loaded = False

def load_model(hf_token):
    """Load the model with the provided token"""
    global global_model, global_tokenizer, model_loaded
    
    if not hf_token:
        model_loaded = False
        return "⚠️ Please enter your Hugging Face token to use the model."
    
    try:
        # Try different model versions from smallest to largest
        model_options = [
            "google/gemma-2b-it",    # Try an instruction-tuned 2B model first (smallest)
            "google/gemma-2b",       # Try base 2B model next
            "google/gemma-7b-it",    # Try 7B instruction-tuned model
            "google/gemma-7b",       # Try base 7B model
        ]
        
        print(f"Attempting to load models with token starting with: {hf_token[:5]}...")
        
        # Try to load models in order until one works
        for model_name in model_options:
            try:
                print(f"Attempting to load model: {model_name}")
                
                # Load tokenizer
                print("Loading tokenizer...")
                global_tokenizer = AutoTokenizer.from_pretrained(
                    model_name, 
                    token=hf_token
                )
                print("Tokenizer loaded successfully")
                
                # Load model with minimal configuration
                print(f"Loading model {model_name}...")
                global_model = AutoModelForCausalLM.from_pretrained(
                    model_name,
                    torch_dtype=torch.float16,
                    device_map="auto",
                    token=hf_token
                )
                print(f"Model {model_name} loaded successfully!")
                
                model_loaded = True
                return f"✅ Model {model_name} loaded successfully!"
            except Exception as specific_e:
                print(f"Failed to load {model_name}: {specific_e}")
                import traceback
                traceback.print_exc()
                continue
                
        # If we get here, all model options failed - try one more option with no token
        try:
            print("Trying a public model with no token requirement...")
            model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
            global_tokenizer = AutoTokenizer.from_pretrained(model_name)
            global_model = AutoModelForCausalLM.from_pretrained(
                model_name,
                torch_dtype=torch.float16,
                device_map="auto"
            )
            model_loaded = True
            return f"✅ Fallback model {model_name} loaded successfully! Note: This is not Gemma but a fallback model."
        except Exception as fallback_e:
            print(f"Failed to load fallback model: {fallback_e}")
                
        # If we get here, all model options failed
        model_loaded = False
        return "❌ Could not load any model version. Please check your token and try again."
        
    except Exception as e:
        model_loaded = False
        error_msg = str(e)
        print(f"Error in load_model: {error_msg}")
        import traceback
        traceback.print_exc()
        
        if "401 Client Error" in error_msg:
            return "❌ Authentication failed. Please check your token and make sure you've accepted the model license on Hugging Face."
        else:
            return f"❌ Error loading model: {error_msg}"

def generate_prompt(task_type, **kwargs):
    """Generate appropriate prompts based on task type and parameters"""
    if task_type == "creative":
        style = kwargs.get("style", "")
        topic = kwargs.get("topic", "")
        return f"Write a {style} about {topic}. Be creative and engaging."
    
    elif task_type == "informational":
        format_type = kwargs.get("format_type", "")
        topic = kwargs.get("topic", "")
        return f"Write an {format_type} about {topic}. Be clear, factual, and informative."
    
    elif task_type == "summarize":
        text = kwargs.get("text", "")
        return f"Summarize the following text in a concise way:\n\n{text}"
    
    elif task_type == "translate":
        text = kwargs.get("text", "")
        target_lang = kwargs.get("target_lang", "")
        return f"Translate the following text to {target_lang}:\n\n{text}"
    
    elif task_type == "qa":
        text = kwargs.get("text", "")
        question = kwargs.get("question", "")
        return f"Based on the following text:\n\n{text}\n\nAnswer this question: {question}"
    
    elif task_type == "code_generate":
        language = kwargs.get("language", "")
        task = kwargs.get("task", "")
        return f"Write {language} code to {task}. Include helpful comments."
    
    elif task_type == "code_explain":
        code = kwargs.get("code", "")
        return f"Explain what the following code does in simple terms:\n\n```\n{code}\n```"
    
    elif task_type == "code_debug":
        code = kwargs.get("code", "")
        return f"The following code has an issue. Identify and fix the problem:\n\n```\n{code}\n```"
    
    elif task_type == "brainstorm":
        topic = kwargs.get("topic", "")
        category = kwargs.get("category", "")
        return f"Brainstorm {category} ideas about {topic}. Provide a diverse list of options."
    
    elif task_type == "content_creation":
        content_type = kwargs.get("content_type", "")
        topic = kwargs.get("topic", "")
        audience = kwargs.get("audience", "")
        return f"Create a {content_type} about {topic} for {audience}. Make it engaging and relevant."
    
    elif task_type == "email_draft":
        email_type = kwargs.get("email_type", "")
        context = kwargs.get("context", "")
        return f"Write a professional {email_type} email with the following context:\n\n{context}"
    
    elif task_type == "document_edit":
        text = kwargs.get("text", "")
        edit_type = kwargs.get("edit_type", "")
        return f"Improve the following text for {edit_type}:\n\n{text}"
    
    elif task_type == "explain":
        topic = kwargs.get("topic", "")
        level = kwargs.get("level", "")
        return f"Explain {topic} in a way that's easy to understand for a {level} audience."
    
    elif task_type == "classify":
        text = kwargs.get("text", "")
        categories = kwargs.get("categories", "")
        return f"Classify the following text into one of these categories: {categories}\n\nText: {text}\n\nCategory:"
    
    elif task_type == "data_extract":
        text = kwargs.get("text", "")
        data_points = kwargs.get("data_points", "")
        return f"Extract the following information from the text: {data_points}\n\nText: {text}\n\nExtracted information:"
    
    else:
        return kwargs.get("prompt", "")

def generate_text(prompt, max_length=1024, temperature=0.7, top_p=0.95):
    """Generate text using the Gemma model"""
    global global_model, global_tokenizer, model_loaded
    
    print(f"Generating text with params: max_length={max_length}, temp={temperature}, top_p={top_p}")
    print(f"Prompt: {prompt[:100]}...")
    
    if not model_loaded or global_model is None or global_tokenizer is None:
        print("Model not loaded")
        return "⚠️ Model not loaded. Please authenticate with your Hugging Face token."
    
    if not prompt:
        return "Please enter a prompt to generate text."
    
    try:
        # Keep generation simple to avoid errors
        inputs = global_tokenizer(prompt, return_tensors="pt").to(global_model.device)
        print(f"Input token length: {len(inputs.input_ids[0])}")
        
        # Use even simpler generation parameters
        generation_args = {
            "input_ids": inputs.input_ids,
            "max_length": min(2048, max_length + len(inputs.input_ids[0])),
            "do_sample": True,
        }
        
        # Only add temperature if not too low (can cause issues)
        if temperature >= 0.3:
            generation_args["temperature"] = temperature
            
        print(f"Generation args: {generation_args}")
        
        # Generate text
        outputs = global_model.generate(**generation_args)
        
        # Decode and return the generated text
        generated_text = global_tokenizer.decode(outputs[0], skip_special_tokens=True)
        print(f"Generated text length: {len(generated_text)}")
        return generated_text
    except Exception as e:
        error_msg = str(e)
        print(f"Generation error: {error_msg}")
        print(f"Error type: {type(e)}")
        import traceback
        traceback.print_exc()
        
        if "probability tensor" in error_msg:
            return "Error: There was a problem with the generation parameters. Try using simpler parameters or a different prompt."
        else:
            return f"Error generating text: {error_msg}"

# Create parameters UI component
def create_parameter_ui():
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            max_length = gr.Slider(
                minimum=128, 
                maximum=2048, 
                value=1024, 
                step=64,
                label="Maximum Length",
                info="Control the maximum length of generated text",
                elem_id="max_length_slider"
            )
        with gr.Column(scale=1):
            temperature = gr.Slider(
                minimum=0.3, 
                maximum=1.5, 
                value=0.8, 
                step=0.1,
                label="Temperature",
                info="Higher values create more creative but potentially less coherent outputs",
                elem_id="temperature_slider"
            )
        with gr.Column(scale=1):
            top_p = gr.Slider(
                minimum=0.5, 
                maximum=0.99, 
                value=0.9, 
                step=0.05,
                label="Top-p",
                info="Controls diversity of generated text",
                elem_id="top_p_slider"
            )
    return [max_length, temperature, top_p]

# Create Gradio interface
with gr.Blocks(theme=GemmaLightTheme()) as demo:
    # Header with theme toggle
    with gr.Row(equal_height=True):
        with gr.Column(scale=6):
            gr.Markdown(
                """
                # 🤖 Gemma Capabilities Demo
                
                This interactive demo showcases Google's Gemma model capabilities across different tasks.
                """
            )
        
        with gr.Column(scale=1, min_width=150):
            theme_toggle = gr.Radio(
                ["Light", "Dark"], 
                value="Light",
                label="Theme", 
                info="Switch between light and dark mode",
                elem_id="theme_toggle"
            )
    
    # Add CSS for themes and set up JavaScript for theme toggle
    gr.HTML("""
    <style>
    /* Dark theme styles */
    .dark-theme {
        --background-fill-primary: #1F1F2E !important;
        --background-fill-secondary: #2A2A3C !important;
        --border-color-primary: #3A3A4C !important;
        --border-color-secondary: #4A4A5C !important;
        --color-accent: #4B82C4 !important;
        --color-accent-soft: #3B5FA3 !important;
        --color-text: #FFFFFF !important;
        --color-subtext: #CCCCCC !important;
    }
    </style>
    
    <script>
    // This script will run after the page loads
    (function() {
        // Function to check the theme toggle and apply the appropriate theme
        function applyTheme() {
            const themeToggle = document.getElementById('theme_toggle');
            if (!themeToggle) return;
            
            const inputs = themeToggle.querySelectorAll('input');
            for (let input of inputs) {
                if (input.checked && input.value === 'Dark') {
                    document.body.classList.add('dark-theme');
                    return;
                }
            }
            
            // If we get here, light theme is selected or no selection
            document.body.classList.remove('dark-theme');
        }
        
        // Set up a mutation observer to watch for changes to the radio buttons
        function setupObserver() {
            const themeToggle = document.getElementById('theme_toggle');
            if (!themeToggle) {
                // Try again in a bit if the element isn't ready yet
                setTimeout(setupObserver, 100);
                return;
            }
            
            // Apply theme initially
            applyTheme();
            
            // Watch for changes
            const observer = new MutationObserver(function(mutations) {
                mutations.forEach(function(mutation) {
                    if (mutation.type === 'attributes' || mutation.type === 'childList') {
                        applyTheme();
                    }
                });
            });
            
            observer.observe(themeToggle, { 
                attributes: true,
                childList: true,
                subtree: true
            });
            
            // Also add click listeners to ensure theme changes on click
            const inputs = themeToggle.querySelectorAll('input');
            inputs.forEach(input => {
                input.addEventListener('change', applyTheme);
            });
        }
        
        // Start the setup when the document is ready
        if (document.readyState === 'loading') {
            document.addEventListener('DOMContentLoaded', setupObserver);
        } else {
            setupObserver();
        }
    })();
    </script>
    """)
    
    # Authentication Section
    with gr.Group(elem_id="auth_box"):
        gr.Markdown("## 🔑 Authentication", elem_id="auth_heading")
        
        with gr.Row(equal_height=True):
            with gr.Column(scale=3):
                hf_token = gr.Textbox(
                    label="Hugging Face Token", 
                    placeholder="Enter your token here...",
                    type="password",
                    value=DEFAULT_HF_TOKEN,
                    info="Get your token from https://huggingface.co/settings/tokens",
                    elem_id="hf_token_input"
                )
            
            with gr.Column(scale=1):
                auth_button = gr.Button("Authenticate", variant="primary", elem_id="auth_button")
        
        auth_status = gr.Markdown("Please authenticate to use the model.", elem_id="auth_status")
        
        def authenticate(token):
            return "⏳ Loading model... Please wait, this may take a minute."
            
        def auth_complete(token):
            result = load_model(token)
            return result
            
        # Two-step authentication to show loading message
        auth_button.click(
            fn=authenticate,
            inputs=[hf_token],
            outputs=[auth_status],
            queue=False
        ).then(
            fn=auth_complete,
            inputs=[hf_token],
            outputs=[auth_status]
        )
        
        gr.Markdown(
            """
            ### How to get a token:
            1. Go to [Hugging Face Token Settings](https://huggingface.co/settings/tokens)
            2. Create a new token with read access
            3. Make sure you've accepted the [Gemma model license](https://huggingface.co/google/gemma-3-4b-pt)
            """
        )
    
    # Main content - only show when authenticated
    with gr.Tabs(elem_id="main_tabs") as tabs:
        # Text Generation Tab
        with gr.TabItem("Text Generation", id="tab_text_gen"):
            gr.Markdown(
                """
                ## ✏️ Creative Text Generation
                
                Generate stories, poems, and other creative content. Choose a style and topic or enter your own prompt.
                """
            )
            
            with gr.Group(elem_id="text_gen_box"):
                with gr.Row(equal_height=True):
                    with gr.Column(scale=1):
                        text_gen_type = gr.Radio(
                            ["Creative Writing", "Informational Writing", "Custom Prompt"],
                            label="Writing Type",
                            value="Creative Writing",
                            elem_id="text_gen_type"
                        )
                        
                        # Creative writing options
                        with gr.Group(visible=True, elem_id="creative_options") as creative_options:
                            style = gr.Dropdown(
                                ["short story", "poem", "script", "song lyrics", "joke"],
                                label="Style",
                                value="short story",
                                elem_id="creative_style"
                            )
                            creative_topic = gr.Textbox(
                                label="Topic",
                                placeholder="Enter a topic...",
                                value="a robot discovering emotions",
                                elem_id="creative_topic"
                            )
                        
                        # Informational writing options
                        with gr.Group(visible=False, elem_id="info_options") as info_options:
                            format_type = gr.Dropdown(
                                ["article", "summary", "explanation", "report"],
                                label="Format",
                                value="article",
                                elem_id="info_format"
                            )
                            info_topic = gr.Textbox(
                                label="Topic",
                                placeholder="Enter a topic...",
                                value="artificial intelligence",
                                elem_id="info_topic"
                            )
                        
                        # Custom prompt
                        with gr.Group(visible=False, elem_id="custom_prompt_group") as custom_prompt_group:
                            custom_prompt = gr.Textbox(
                                label="Custom Prompt",
                                placeholder="Enter your custom prompt...",
                                lines=3,
                                elem_id="custom_prompt"
                            )
                        
                        # Show/hide options based on selection
                        def update_text_gen_visibility(choice):
                            return {
                                creative_options: choice == "Creative Writing",
                                info_options: choice == "Informational Writing",
                                custom_prompt_group: choice == "Custom Prompt"
                            }
                        
                        text_gen_type.change(
                            update_text_gen_visibility,
                            inputs=text_gen_type,
                            outputs=[creative_options, info_options, custom_prompt_group]
                        )
                        
                        # Generation parameters
                        with gr.Accordion("Advanced Parameters", open=False, elem_id="text_gen_params"):
                            text_gen_params = create_parameter_ui()
                        
                        generate_text_btn = gr.Button("Generate", variant="primary", size="lg", elem_id="generate_text_btn")
                    
                    with gr.Column(scale=1):
                        text_output = gr.Textbox(
                            label="Generated Text",
                            lines=20,
                            elem_id="text_output"
                        )
                
                # Handle text generation
                def text_generation_handler(
                    gen_type, style, creative_topic, format_type, info_topic, 
                    custom_prompt, max_length, temperature, top_p
                ):
                    if gen_type == "Creative Writing":
                        style = safe_value(style, "short story")
                        creative_topic = safe_value(creative_topic, "a story")
                        prompt = generate_prompt("creative", style=style, topic=creative_topic)
                    elif gen_type == "Informational Writing":
                        format_type = safe_value(format_type, "article")
                        info_topic = safe_value(info_topic, "a topic")
                        prompt = generate_prompt("informational", format_type=format_type, topic=info_topic)
                    else:
                        prompt = safe_value(custom_prompt, "Write something interesting")
                    
                    return generate_text(prompt, max_length, temperature, top_p)
                
                generate_text_btn.click(
                    text_generation_handler,
                    inputs=[
                        text_gen_type, style, creative_topic, format_type, info_topic,
                        custom_prompt, *text_gen_params
                    ],
                    outputs=text_output
                )
                
                # Examples for text generation
                gr.Examples(
                    [
                        ["Creative Writing", "short story", "a robot learning to paint", "article", "artificial intelligence", "", 1024, 0.8, 0.9],
                        ["Creative Writing", "poem", "the beauty of mathematics", "article", "artificial intelligence", "", 768, 0.8, 0.9],
                        ["Informational Writing", "short story", "a robot discovering emotions", "article", "quantum computing", "", 1024, 0.7, 0.9],
                        ["Custom Prompt", "short story", "a robot discovering emotions", "article", "artificial intelligence", "Write a marketing email for a new smartphone with innovative AI features", 1024, 0.8, 0.9]
                    ],
                    fn=text_generation_handler,
                    inputs=[
                        text_gen_type, style, creative_topic, format_type, info_topic,
                        custom_prompt, *text_gen_params
                    ],
                    outputs=text_output,
                    label="Examples"
                )
        
        # Brainstorming Tab
        with gr.TabItem("Brainstorming", id="tab_brainstorm"):
            gr.Markdown(
                """
                ## 🧠 Brainstorming Ideas
                
                Generate creative ideas for projects, solutions, or any topic you're interested in.
                """
            )
            
            with gr.Group(elem_id="brainstorm_box"):
                with gr.Row(equal_height=True):
                    with gr.Column(scale=1):
                        brainstorm_category = gr.Dropdown(
                            ["project", "business", "creative", "solution", "content", "feature", "product"],
                            label="Category",
                            value="project",
                            elem_id="brainstorm_category"
                        )
                        brainstorm_topic = gr.Textbox(
                            label="Topic or Problem",
                            placeholder="What would you like ideas for?",
                            value="sustainable technology",
                            elem_id="brainstorm_topic"
                        )
                        
                        with gr.Accordion("Advanced Parameters", open=False, elem_id="brainstorm_params"):
                            brainstorm_params = create_parameter_ui()
                        
                        brainstorm_btn = gr.Button("Generate Ideas", variant="primary", size="lg", elem_id="brainstorm_btn")
                    
                    with gr.Column(scale=1):
                        brainstorm_output = gr.Textbox(
                            label="Generated Ideas",
                            lines=20,
                            elem_id="brainstorm_output"
                        )
            
            def brainstorm_handler(category, topic, max_length, temperature, top_p):
                category = safe_value(category, "project")
                topic = safe_value(topic, "innovative ideas")
                prompt = generate_prompt("brainstorm", category=category, topic=topic)
                return generate_text(prompt, max_length, temperature, top_p)
            
            brainstorm_btn.click(
                brainstorm_handler,
                inputs=[brainstorm_category, brainstorm_topic, *brainstorm_params],
                outputs=brainstorm_output
            )
            
            # Examples for brainstorming
            gr.Examples(
                [
                    ["project", "educational app for children", 1024, 0.8, 0.9],
                    ["business", "eco-friendly food packaging", 1024, 0.8, 0.9],
                    ["solution", "reducing urban traffic congestion", 1024, 0.8, 0.9],
                ],
                fn=brainstorm_handler,
                inputs=[brainstorm_category, brainstorm_topic, *brainstorm_params],
                outputs=brainstorm_output,
                label="Examples"
            )
        
        # Content Creation Tab
        with gr.TabItem("Content Creation", id="tab_content"):
            gr.Markdown(
                """
                ## 📝 Content Creation
                
                Generate various types of content such as blog posts, social media updates, marketing copy, etc.
                """
            )
            
            with gr.Group(elem_id="content_box"):
                with gr.Row(equal_height=True):
                    with gr.Column(scale=1):
                        content_type = gr.Dropdown(
                            ["blog post", "social media post", "marketing copy", "product description", "press release", "newsletter"],
                            label="Content Type",
                            value="blog post",
                            elem_id="content_type"
                        )
                        content_topic = gr.Textbox(
                            label="Topic",
                            placeholder="What is your content about?",
                            value="the future of artificial intelligence",
                            elem_id="content_topic"
                        )
                        content_audience = gr.Textbox(
                            label="Target Audience",
                            placeholder="Who is your audience?",
                            value="tech enthusiasts",
                            elem_id="content_audience"
                        )
                        
                        with gr.Accordion("Advanced Parameters", open=False, elem_id="content_params"):
                            content_params = create_parameter_ui()
                        
                        content_btn = gr.Button("Generate Content", variant="primary", size="lg", elem_id="content_btn")
                    
                    with gr.Column(scale=1):
                        content_output = gr.Textbox(
                            label="Generated Content",
                            lines=20,
                            elem_id="content_output"
                        )
                
                def content_creation_handler(content_type, topic, audience, max_length, temperature, top_p):
                    content_type = safe_value(content_type, "blog post")
                    topic = safe_value(topic, "interesting topic")
                    audience = safe_value(audience, "general audience")
                    prompt = generate_prompt("content_creation", content_type=content_type, topic=topic, audience=audience)
                    return generate_text(prompt, max_length, temperature, top_p)
                
                content_btn.click(
                    content_creation_handler,
                    inputs=[content_type, content_topic, content_audience, *content_params],
                    outputs=content_output
                )
                
                # Examples for content creation
                gr.Examples(
                    [
                        ["blog post", "sustainable living tips", "environmentally conscious consumers", 1536, 0.8, 0.9],
                        ["social media post", "product launch announcement", "existing customers", 512, 0.8, 0.9],
                        ["marketing copy", "new fitness app", "health-focused individuals", 1024, 0.8, 0.9],
                    ],
                    fn=content_creation_handler,
                    inputs=[content_type, content_topic, content_audience, *content_params],
                    outputs=content_output,
                    label="Examples"
                )
        
        # Email Drafting Tab
        with gr.TabItem("Email Drafting", id="tab_email"):
            gr.Markdown(
                """
                ## ✉️ Email Drafting
                
                Generate professional email drafts for various purposes.
                """
            )
            
            with gr.Group(elem_id="email_box"):
                with gr.Row(equal_height=True):
                    with gr.Column(scale=1):
                        email_type = gr.Dropdown(
                            ["job application", "customer support", "business proposal", "networking", "follow-up", "thank you", "meeting request"],
                            label="Email Type",
                            value="job application",
                            elem_id="email_type"
                        )
                        email_context = gr.Textbox(
                            label="Context and Details",
                            placeholder="Provide necessary context for the email...",
                            lines=5,
                            value="Applying for a software developer position at Tech Solutions Inc. I have 3 years of experience with Python and JavaScript.",
                            elem_id="email_context"
                        )
                        
                        with gr.Accordion("Advanced Parameters", open=False, elem_id="email_params"):
                            email_params = create_parameter_ui()
                        
                        email_btn = gr.Button("Generate Email", variant="primary", size="lg", elem_id="email_btn")
                    
                    with gr.Column(scale=1):
                        email_output = gr.Textbox(
                            label="Generated Email",
                            lines=20,
                            elem_id="email_output"
                        )
            
            def email_draft_handler(email_type, context, max_length, temperature, top_p):
                email_type = safe_value(email_type, "professional")
                context = safe_value(context, "general communication")
                prompt = generate_prompt("email_draft", email_type=email_type, context=context)
                return generate_text(prompt, max_length, temperature, top_p)
            
            email_btn.click(
                email_draft_handler,
                inputs=[email_type, email_context, *email_params],
                outputs=email_output
            )
            
            # Examples for email drafting
            gr.Examples(
                [
                    ["job application", "Applying for a marketing specialist position at ABC Marketing. I have 5 years of experience in digital marketing.", 1024, 0.8, 0.9],
                    ["business proposal", "Proposing a collaboration between our companies for a joint product development effort.", 1024, 0.8, 0.9],
                    ["follow-up", "Following up after our meeting last Thursday about the project timeline and resources.", 1024, 0.8, 0.9],
                ],
                fn=email_draft_handler,
                inputs=[email_type, email_context, *email_params],
                outputs=email_output,
                label="Examples"
            )
        
        # Document Editing Tab
        with gr.TabItem("Document Editing", id="tab_edit"):
            gr.Markdown(
                """
                ## ✂️ Document Editing
                
                Improve the clarity, grammar, and style of your writing.
                """
            )
            
            with gr.Group(elem_id="edit_box"):
                with gr.Row(equal_height=True):
                    with gr.Column(scale=1):
                        edit_text = gr.Textbox(
                            label="Text to Edit",
                            placeholder="Paste your text here...",
                            lines=10,
                            value="The company have been experiencing rapid growth over the past few years and is expecting to continue this trend in the coming years. They believe that it's success is due to the quality of their products and the dedicated team.",
                            elem_id="edit_text"
                        )
                        edit_type = gr.Dropdown(
                            ["grammar and clarity", "conciseness", "formal tone", "casual tone", "simplification", "academic style", "persuasive style"],
                            label="Edit Type",
                            value="grammar and clarity",
                            elem_id="edit_type"
                        )
                        
                        with gr.Accordion("Advanced Parameters", open=False, elem_id="edit_params"):
                            edit_params = create_parameter_ui()
                        
                        edit_btn = gr.Button("Edit Document", variant="primary", size="lg", elem_id="edit_btn")
                    
                    with gr.Column(scale=1):
                        edit_output = gr.Textbox(
                            label="Edited Text",
                            lines=10,
                            elem_id="edit_output"
                        )
            
            def document_edit_handler(text, edit_type, max_length, temperature, top_p):
                text = safe_value(text, "Please provide text to edit.")
                edit_type = safe_value(edit_type, "grammar and clarity")
                prompt = generate_prompt("document_edit", text=text, edit_type=edit_type)
                return generate_text(prompt, max_length, temperature, top_p)
            
            edit_btn.click(
                document_edit_handler,
                inputs=[edit_text, edit_type, *edit_params],
                outputs=edit_output
            )
        
        # Learning & Explanation Tab
        with gr.TabItem("Learning & Explanation", id="tab_explain"):
            gr.Markdown(
                """
                ## 🎓 Learning & Explanation
                
                Get easy-to-understand explanations of complex topics.
                """
            )
            
            with gr.Group(elem_id="explain_box"):
                with gr.Row(equal_height=True):
                    with gr.Column(scale=1):
                        explain_topic = gr.Textbox(
                            label="Topic to Explain",
                            placeholder="What topic would you like explained?",
                            value="quantum computing",
                            elem_id="explain_topic"
                        )
                        explain_level = gr.Dropdown(
                            ["beginner", "child", "teenager", "college student", "professional", "expert"],
                            label="Audience Level",
                            value="beginner",
                            elem_id="explain_level"
                        )
                        
                        with gr.Accordion("Advanced Parameters", open=False, elem_id="explain_params"):
                            explain_params = create_parameter_ui()
                        
                        explain_btn = gr.Button("Generate Explanation", variant="primary", size="lg", elem_id="explain_btn")
                    
                    with gr.Column(scale=1):
                        explain_output = gr.Textbox(
                            label="Explanation",
                            lines=20,
                            elem_id="explain_output"
                        )
            
            def explanation_handler(topic, level, max_length, temperature, top_p):
                topic = safe_value(topic, "an interesting concept")
                level = safe_value(level, "beginner")
                prompt = generate_prompt("explain", topic=topic, level=level)
                return generate_text(prompt, max_length, temperature, top_p)
            
            explain_btn.click(
                explanation_handler,
                inputs=[explain_topic, explain_level, *explain_params],
                outputs=explain_output
            )
            
            # Examples for explanation
            gr.Examples(
                [
                    ["blockchain technology", "beginner", 1024, 0.8, 0.9],
                    ["photosynthesis", "child", 1024, 0.8, 0.9],
                    ["machine learning", "college student", 1024, 0.8, 0.9],
                ],
                fn=explanation_handler,
                inputs=[explain_topic, explain_level, *explain_params],
                outputs=explain_output,
                label="Examples"
            )
        
        # Classification & Categorization Tab
        with gr.TabItem("Classification", id="tab_classify"):
            gr.Markdown(
                """
                ## 🏷️ Classification & Categorization
                
                Classify text into different categories or themes.
                """
            )
            
            with gr.Group(elem_id="classify_box"):
                with gr.Row(equal_height=True):
                    with gr.Column(scale=1):
                        classify_text = gr.Textbox(
                            label="Text to Classify",
                            placeholder="Enter the text you want to classify...",
                            lines=8,
                            value="The latest smartphone features a powerful processor, excellent camera, and impressive battery life, making it a top choice for tech enthusiasts.",
                            elem_id="classify_text"
                        )
                        classify_categories = gr.Textbox(
                            label="Categories (comma-separated)",
                            placeholder="List categories separated by commas...",
                            value="technology, health, finance, entertainment, education, sports",
                            elem_id="classify_categories"
                        )
                        
                        with gr.Accordion("Advanced Parameters", open=False, elem_id="classify_params"):
                            classify_params = create_parameter_ui()
                        
                        classify_btn = gr.Button("Classify Text", variant="primary", size="lg", elem_id="classify_btn")
                    
                    with gr.Column(scale=1):
                        classify_output = gr.Textbox(
                            label="Classification Result",
                            lines=5,
                            elem_id="classify_output"
                        )
            
            def classification_handler(text, categories, max_length, temperature, top_p):
                text = safe_value(text, "Please provide text to classify.")
                categories = safe_value(categories, "general, specific, other")
                prompt = generate_prompt("classify", text=text, categories=categories)
                return generate_text(prompt, max_length, temperature, top_p)
            
            classify_btn.click(
                classification_handler,
                inputs=[classify_text, classify_categories, *classify_params],
                outputs=classify_output
            )
            
            # Examples for classification
            gr.Examples(
                [
                    ["The stock market saw significant gains today as tech companies reported strong quarterly earnings.", "technology, health, finance, entertainment, education, sports", 256, 0.5, 0.9],
                    ["The team scored in the final minutes to secure their victory in the championship game.", "technology, health, finance, entertainment, education, sports", 256, 0.5, 0.9],
                    ["The new educational app helps students master complex math concepts through interactive exercises.", "technology, health, finance, entertainment, education, sports", 256, 0.5, 0.9],
                ],
                fn=classification_handler,
                inputs=[classify_text, classify_categories, *classify_params],
                outputs=classify_output,
                label="Examples"
            )
        
        # Data Extraction Tab
        with gr.TabItem("Data Extraction", id="tab_extract"):
            gr.Markdown(
                """
                ## 📊 Data Extraction
                
                Extract specific data points from text.
                """
            )
            
            with gr.Group(elem_id="extract_box"):
                with gr.Row(equal_height=True):
                    with gr.Column(scale=1):
                        extract_text = gr.Textbox(
                            label="Text to Process",
                            placeholder="Enter the text containing data to extract...",
                            lines=10,
                            value="John Smith, born on May 15, 1985, is a software engineer at Tech Solutions Inc. He can be reached at [email protected] or (555) 123-4567. John graduated from MIT in 2007 with a degree in Computer Science.",
                            elem_id="extract_text"
                        )
                        extract_data_points = gr.Textbox(
                            label="Data Points to Extract (comma-separated)",
                            placeholder="Specify what data to extract...",
                            value="name, email, phone number, birth date, company, education",
                            elem_id="extract_data_points"
                        )
                        
                        with gr.Accordion("Advanced Parameters", open=False, elem_id="extract_params"):
                            extract_params = create_parameter_ui()
                        
                        extract_btn = gr.Button("Extract Data", variant="primary", size="lg", elem_id="extract_btn")
                    
                    with gr.Column(scale=1):
                        extract_output = gr.Textbox(
                            label="Extracted Data",
                            lines=10,
                            elem_id="extract_output"
                        )
            
            def data_extraction_handler(text, data_points, max_length, temperature, top_p):
                text = safe_value(text, "Please provide text with data to extract.")
                data_points = safe_value(data_points, "key information")
                prompt = generate_prompt("data_extract", text=text, data_points=data_points)
                return generate_text(prompt, max_length, temperature, top_p)
            
            extract_btn.click(
                data_extraction_handler,
                inputs=[extract_text, extract_data_points, *extract_params],
                outputs=extract_output
            )
            
            # Examples for data extraction
            gr.Examples(
                [
                    ["Sarah Johnson is the CEO of Green Innovations, founded in 2012. The company reported $8.5 million in revenue for 2023. Contact her at [email protected].", "name, position, company, founding year, revenue, contact", 768, 0.5, 0.9],
                    ["The new iPhone 15 Pro features a 6.1-inch display, A17 Pro chip, 48MP camera, and starts at $999 for the 128GB model.", "product name, screen size, processor, camera, price, storage capacity", 768, 0.5, 0.9],
                ],
                fn=data_extraction_handler,
                inputs=[extract_text, extract_data_points, *extract_params],
                outputs=extract_output,
                label="Examples"
            )
        
        # Text Comprehension Tab
        with gr.TabItem("Text Comprehension", id="tab_comprehension"):
            gr.Markdown(
                """
                ## 📚 Text Comprehension
                
                Test Gemma's ability to understand and process text. Try summarization, Q&A, or translation.
                """
            )
            
            with gr.Tabs() as comprehension_tabs:
                # Summarization
                with gr.TabItem("Summarization", id="subtab_summarize"):
                    with gr.Group(elem_id="summarize_box"):
                        with gr.Row(equal_height=True):
                            with gr.Column(scale=1):
                                summarize_text = gr.Textbox(
                                    label="Text to Summarize",
                                    placeholder="Paste text here...",
                                    lines=10,
                                    elem_id="summarize_text"
                                )
                                
                                with gr.Accordion("Advanced Parameters", open=False, elem_id="summarize_params"):
                                    summarize_params = create_parameter_ui()
                                
                                summarize_btn = gr.Button("Summarize", variant="primary", size="lg", elem_id="summarize_btn")
                            
                            with gr.Column(scale=1):
                                summary_output = gr.Textbox(
                                    label="Summary",
                                    lines=10,
                                    elem_id="summary_output"
                                )
                
                # Question Answering
                with gr.TabItem("Question Answering", id="subtab_qa"):
                    with gr.Group(elem_id="qa_box"):
                        with gr.Row(equal_height=True):
                            with gr.Column(scale=1):
                                qa_text = gr.Textbox(
                                    label="Context Text",
                                    placeholder="Paste text here...",
                                    lines=10,
                                    elem_id="qa_text"
                                )
                                qa_question = gr.Textbox(
                                    label="Question",
                                    placeholder="Ask a question about the text...",
                                    elem_id="qa_question"
                                )
                                
                                with gr.Accordion("Advanced Parameters", open=False, elem_id="qa_params"):
                                    qa_params = create_parameter_ui()
                                
                                qa_btn = gr.Button("Answer", variant="primary", size="lg", elem_id="qa_btn")
                            
                            with gr.Column(scale=1):
                                qa_output = gr.Textbox(
                                    label="Answer",
                                    lines=10,
                                    elem_id="qa_output"
                                )
                
                # Translation
                with gr.TabItem("Translation", id="subtab_translate"):
                    with gr.Group(elem_id="translate_box"):
                        with gr.Row(equal_height=True):
                            with gr.Column(scale=1):
                                translate_text = gr.Textbox(
                                    label="Text to Translate",
                                    placeholder="Enter text to translate...",
                                    lines=5,
                                    elem_id="translate_text"
                                )
                                target_lang = gr.Dropdown(
                                    ["French", "Spanish", "German", "Japanese", "Chinese", "Russian", "Arabic", "Hindi"],
                                    label="Target Language",
                                    value="French",
                                    elem_id="target_lang"
                                )
                                
                                with gr.Accordion("Advanced Parameters", open=False, elem_id="translate_params"):
                                    translate_params = create_parameter_ui()
                                
                                translate_btn = gr.Button("Translate", variant="primary", size="lg", elem_id="translate_btn")
                            
                            with gr.Column(scale=1):
                                translation_output = gr.Textbox(
                                    label="Translation",
                                    lines=5,
                                    elem_id="translation_output"
                                )
        
        # Code Capabilities Tab
        with gr.TabItem("Code Capabilities", id="tab_code"):
            gr.Markdown(
                """
                ## 💻 Code Generation and Understanding
                
                Test Gemma's ability to generate, explain, and debug code in various programming languages.
                """
            )
            
            with gr.Tabs() as code_tabs:
                # Code Generation
                with gr.TabItem("Code Generation", id="subtab_code_gen"):
                    with gr.Group(elem_id="code_gen_box"):
                        with gr.Row(equal_height=True):
                            with gr.Column(scale=1):
                                code_language = gr.Dropdown(
                                    ["Python", "JavaScript", "Java", "C++", "HTML/CSS", "SQL", "Bash"],
                                    label="Programming Language",
                                    value="Python",
                                    elem_id="code_language"
                                )
                                code_task = gr.Textbox(
                                    label="Coding Task",
                                    placeholder="Describe what you want the code to do...",
                                    value="Create a function to find prime numbers up to n",
                                    elem_id="code_task"
                                )
                                
                                with gr.Accordion("Advanced Parameters", open=False, elem_id="code_gen_params"):
                                    code_gen_params = create_parameter_ui()
                                
                                code_gen_btn = gr.Button("Generate Code", variant="primary", size="lg", elem_id="code_gen_btn")
                            
                            with gr.Column(scale=1):
                                code_output = gr.Code(
                                    label="Generated Code",
                                    language="python",
                                    elem_id="code_output"
                                )
                    
                    def code_gen_handler(language, task, max_length, temperature, top_p):
                        language = safe_value(language, "Python")
                        task = safe_value(task, "write a hello world program")
                        prompt = generate_prompt("code_generate", language=language, task=task)
                        result = generate_text(prompt, max_length, temperature, top_p)
                        return result
                    
                    # Update language in code output component
                    def update_code_language(lang):
                        lang_map = {
                            "Python": "python",
                            "JavaScript": "javascript",
                            "Java": "java",
                            "C++": "cpp",
                            "HTML/CSS": "html",
                            "SQL": "sql",
                            "Bash": "bash"
                        }
                        return gr.Code.update(language=lang_map.get(lang, "python"))
                    
                    code_language.change(update_code_language, inputs=code_language, outputs=code_output)
                    
                    code_gen_btn.click(
                        code_gen_handler,
                        inputs=[code_language, code_task, *code_gen_params],
                        outputs=code_output
                    )
                
                # Code Explanation
                with gr.TabItem("Code Explanation", id="subtab_code_explain"):
                    with gr.Group(elem_id="code_explain_box"):
                        with gr.Row(equal_height=True):
                            with gr.Column(scale=1):
                                code_to_explain = gr.Code(
                                    label="Code to Explain",
                                    language="python",
                                    value="def quicksort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr) // 2]\n    left = [x for x in arr if x < pivot]\n    middle = [x for x in arr if x == pivot]\n    right = [x for x in arr if x > pivot]\n    return quicksort(left) + middle + quicksort(right)",
                                    elem_id="code_to_explain"
                                )
                                
                                with gr.Accordion("Advanced Parameters", open=False, elem_id="explain_code_params"):
                                    explain_code_params = create_parameter_ui()
                                
                                explain_code_btn = gr.Button("Explain Code", variant="primary", size="lg", elem_id="explain_code_btn")
                            
                            with gr.Column(scale=1):
                                code_explanation = gr.Textbox(
                                    label="Explanation",
                                    lines=10,
                                    elem_id="code_explanation"
                                )
                    
                    def explain_code_handler(code, max_length, temperature, top_p):
                        code = safe_value(code, "print('Hello, world!')")
                        prompt = generate_prompt("code_explain", code=code)
                        return generate_text(prompt, max_length, temperature, top_p)
                    
                    explain_code_btn.click(
                        explain_code_handler,
                        inputs=[code_to_explain, *explain_code_params],
                        outputs=code_explanation
                    )
                
                # Code Debugging
                with gr.TabItem("Code Debugging", id="subtab_code_debug"):
                    with gr.Group(elem_id="code_debug_box"):
                        with gr.Row(equal_height=True):
                            with gr.Column(scale=1):
                                code_to_debug = gr.Code(
                                    label="Code to Debug",
                                    language="python",
                                    value="def fibonacci(n):\n    if n <= 0:\n        return []\n    elif n == 1:\n        return [0]\n    elif n == 2:\n        return [0, 1]\n    \n    fib = [0, 1]\n    for i in range(2, n):\n        fib.append(fib[i-1] - fib[i-2])  # Bug is here (should be +)\n    \n    return fib\n\nprint(fibonacci(10))",
                                    elem_id="code_to_debug"
                                )
                                
                                with gr.Accordion("Advanced Parameters", open=False, elem_id="debug_code_params"):
                                    debug_code_params = create_parameter_ui()
                                
                                debug_code_btn = gr.Button("Debug Code", variant="primary", size="lg", elem_id="debug_code_btn")
                            
                            with gr.Column(scale=1):
                                debug_result = gr.Textbox(
                                    label="Debugging Result",
                                    lines=10,
                                    elem_id="debug_result"
                                )
                    
                    def debug_code_handler(code, max_length, temperature, top_p):
                        code = safe_value(code, "print('Hello, world!')")
                        prompt = generate_prompt("code_debug", code=code)
                        return generate_text(prompt, max_length, temperature, top_p)
                    
                    debug_code_btn.click(
                        debug_code_handler,
                        inputs=[code_to_debug, *debug_code_params],
                        outputs=debug_result
                    )
    
    # Footer
    with gr.Group(elem_id="footer"):
        gr.Markdown(
            """
            ## About Gemma
            
            Gemma is a family of lightweight, state-of-the-art open models from Google, built from the same research and technology used to create the Gemini models. 
            It's designed to be efficient and accessible for various applications.
            
            [Learn more about Gemma](https://huggingface.co/google/gemma-3-4b-pt)
            
            <div style="text-align: center; margin-top: 20px; padding: 10px;">
                <p>© 2023 Gemma Capabilities Demo | Made with ❤️ using Gradio</p>
            </div>
            """
        )

    # Add CSS for better styling
    demo.load(js="""
    () => {
        // Add custom CSS for better styling
        const style = document.createElement('style');
        style.textContent = `
            .tabs {
                box-shadow: 0 2px 5px rgba(0,0,0,0.1);
                border-radius: 10px;
                overflow: hidden;
            }
            
            /* Make buttons more noticeable */
            button.primary {
                transition: transform 0.2s, box-shadow 0.2s;
            }
            
            button.primary:hover {
                transform: translateY(-2px);
                box-shadow: 0 4px 8px rgba(0,0,0,0.2);
            }
            
            /* Add hover effect to tabs */
            .tab-nav button {
                transition: background-color 0.3s;
            }
            
            /* Make textboxes more readable */
            textarea, .input-box {
                font-size: 16px !important;
            }
            
            /* Improve box styling */
            .container {
                border-radius: 10px;
                overflow: hidden;
            }
        `;
        document.head.appendChild(style);
    }
    """)

    # Load default token if available
    if DEFAULT_HF_TOKEN:
        demo.load(lambda x: authenticate(DEFAULT_HF_TOKEN), outputs=auth_status).then(
            lambda x: auth_complete(DEFAULT_HF_TOKEN), outputs=auth_status
        )

demo.launch(share=False)