app.py

import gradio as gr
import torch
import os
from transformers import AutoTokenizer, AutoModelForCausalLM
import random
import traceback # Keep traceback for detailed error logging

# 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
loaded_model_name = "None" # Keep track of which model was loaded

def load_model(hf_token):
    """Load the model with the provided token"""
    global global_model, global_tokenizer, model_loaded, loaded_model_name

    # Initially assume tabs should be hidden until successful load
    initial_tabs_update = gr.Tabs.update(visible=False)

    if not hf_token:
        model_loaded = False
        loaded_model_name = "None"
        return "⚠️ Please enter your Hugging Face token to use the model.", initial_tabs_update

    try:
        # Try different model versions from smallest to largest
        model_options = [
            "google/gemma-2b-it",
            "google/gemma-7b-it",
            "google/gemma-2b",
            "google/gemma-7b",
            "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Fallback
        ]

        print(f"Attempting to load models with token starting with: {hf_token[:5]}...")
        loaded_successfully = False
        for model_name in model_options:
            try:
                print(f"\n--- Attempting to load model: {model_name} ---")
                is_gemma = "gemma" in model_name.lower()
                current_token = hf_token if is_gemma else None

                print("Loading tokenizer...")
                global_tokenizer = AutoTokenizer.from_pretrained(
                    model_name,
                    token=current_token
                )
                print("Tokenizer loaded successfully.")

                print(f"Loading model {model_name}...")
                global_model = AutoModelForCausalLM.from_pretrained(
                    model_name,
                    torch_dtype=torch.float16, # Using float16 for broader compatibility
                    device_map="auto",
                    token=current_token
                )
                print(f"Model {model_name} loaded successfully!")

                model_loaded = True
                loaded_model_name = model_name
                loaded_successfully = True
                tabs_update = gr.Tabs.update(visible=True)
                status_msg = f"✅ Model '{model_name}' loaded successfully!"
                if "tinyllama" in model_name.lower():
                    status_msg = f"✅ Fallback model '{model_name}' loaded successfully! Limited capabilities compared to Gemma."
                return status_msg, tabs_update

            except ImportError as import_err:
                 print(f"Import Error loading {model_name}: {import_err}. Check dependencies (e.g., bitsandbytes, accelerate).")
                 continue
            except Exception as specific_e:
                print(f"Failed to load {model_name}: {specific_e}")
                if "401 Client Error" in str(specific_e) or "requires you to be logged in" in str(specific_e) and is_gemma:
                    print("Authentication error likely. Check token and license agreement.")
                continue

        if not loaded_successfully:
            model_loaded = False
            loaded_model_name = "None"
            print("Could not load any model version.")
            return "❌ Could not load any model. Please check your token, license acceptance, dependencies, and network connection.", initial_tabs_update

    except Exception as e:
        model_loaded = False
        loaded_model_name = "None"
        error_msg = str(e)
        print(f"Error in load_model: {error_msg}")
        traceback.print_exc()
        if "401 Client Error" in error_msg or "requires you to be logged in" in error_msg :
            return "❌ Authentication failed. Check token/license.", initial_tabs_update
        else:
            return f"❌ Unexpected error during model loading: {error_msg}", initial_tabs_update


def generate_prompt(task_type, **kwargs):
    """Generate appropriate prompts based on task type and parameters"""
    prompts = {
        "creative": "Write a {style} about {topic}. Be creative and engaging.",
        "informational": "Write an {format_type} about {topic}. Be clear, factual, and informative.",
        "summarize": "Summarize the following text concisely:\n\n{text}",
        "translate": "Translate the following text to {target_lang}:\n\n{text}",
        "qa": "Based on the following text:\n\n{text}\n\nAnswer this question: {question}",
        "code_generate": "Write {language} code to {task}. Include comments explaining the code.",
        "code_explain": "Explain the following {language} code in simple terms:\n\n```\n{code}\n```",
        "code_debug": "Identify and fix the potential bug(s) in the following {language} code. Explain the fix:\n\n```\n{code}\n```",
        "brainstorm": "Brainstorm {category} ideas about {topic}. Provide a diverse list.",
        "content_creation": "Create a {content_type} about {topic} targeting {audience}. Make it engaging.",
        "email_draft": "Draft a professional {email_type} email regarding the following:\n\n{context}",
        "document_edit": "Improve the following text for {edit_type}:\n\n{text}",
        "explain": "Explain {topic} clearly for a {level} audience.",
        "classify": "Classify the following text into one of these categories: {categories}\n\nText: {text}\n\nCategory:",
        "data_extract": "Extract the following data points ({data_points}) from the text below:\n\nText: {text}\n\nExtracted Data:",
    }
    prompt_template = prompts.get(task_type)
    if prompt_template:
        try:
            keys_in_template = [k[1:-1] for k in prompt_template.split('{') if '}' in k for k in [k.split('}')[0]]]
            final_kwargs = {key: kwargs.get(key, f"[{key}]") for key in keys_in_template}
            final_kwargs.update(kwargs) # Add extras
            return prompt_template.format(**final_kwargs)
        except KeyError as e:
            print(f"Warning: Missing key for prompt template '{task_type}': {e}")
            return kwargs.get("prompt", f"Generate text based on: {kwargs}")
    else:
        return kwargs.get("prompt", "Generate text based on the input.")


def generate_text(prompt, max_new_tokens=1024, temperature=0.7, top_p=0.9):
    """Generate text using the loaded model"""
    global global_model, global_tokenizer, model_loaded, loaded_model_name

    print(f"\n--- Generating Text ---")
    print(f"Model: {loaded_model_name}")
    print(f"Params: max_new_tokens={max_new_tokens}, temp={temperature}, top_p={top_p}")
    print(f"Prompt (start): {prompt[:150]}...")

    if not model_loaded or global_model is None or global_tokenizer is None:
        return "⚠️ Model not loaded. Please authenticate first."
    if not prompt:
        return "⚠️ Please enter a prompt or configure a task."

    try:
        chat_prompt = prompt # Default to raw prompt
        if loaded_model_name and ("it" in loaded_model_name.lower() or "instruct" in loaded_model_name.lower() or "chat" in loaded_model_name.lower()):
             if "gemma" in loaded_model_name.lower():
                 # Use Gemma's specific format
                 chat_prompt = f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
             elif "tinyllama" in loaded_model_name.lower():
                 # Use TinyLlama's chat format
                 chat_prompt = f"<|system|>\nYou are a friendly chatbot.</s>\n<|user|>\n{prompt}</s>\n<|assistant|>\n"
             else: # Generic instruction format
                 chat_prompt = f"User: {prompt}\nAssistant:"

        inputs = global_tokenizer(chat_prompt, return_tensors="pt", add_special_tokens=True).to(global_model.device)
        input_length = inputs.input_ids.shape[1]
        print(f"Input token length: {input_length}")

        effective_max_new_tokens = min(int(max_new_tokens), 2048)

        # Handle potential None for eos_token_id
        eos_token_id = global_tokenizer.eos_token_id
        if eos_token_id is None:
            print("Warning: eos_token_id is None, using default 50256.")
            eos_token_id = 50256 # A common default EOS token ID

        generation_args = {
            "input_ids": inputs.input_ids,
            "attention_mask": inputs.attention_mask,
            "max_new_tokens": effective_max_new_tokens,
            "do_sample": True,
            "temperature": float(temperature),
            "top_p": float(top_p),
            "pad_token_id": eos_token_id # Use determined EOS or default
        }

        print(f"Generation args: {generation_args}")

        with torch.no_grad():
            outputs = global_model.generate(**generation_args)

        generated_ids = outputs[0, input_length:]
        generated_text = global_tokenizer.decode(generated_ids, skip_special_tokens=True)

        print(f"Generated text length: {len(generated_text)}")
        print(f"Generated text (start): {generated_text[:150]}...")
        return generated_text.strip()

    except Exception as e:
        error_msg = str(e)
        print(f"Generation error: {error_msg}")
        traceback.print_exc()
        if "CUDA out of memory" in error_msg:
             return f"❌ Error: CUDA out of memory. Try reducing 'Max New Tokens' or using a smaller model."
        elif "probability tensor contains nan" in error_msg or "invalid value encountered" in error_msg:
             return f"❌ Error: Generation failed (invalid probability). Try adjusting Temperature/Top-P or modifying the prompt."
        else:
             return f"❌ Error during text generation: {error_msg}"

# --- UI Components & Layout ---

def create_parameter_ui():
    with gr.Accordion("✨ Generation Parameters", open=False):
        with gr.Row():
            max_new_tokens = gr.Slider(minimum=64, maximum=2048, value=512, step=64, label="Max New Tokens", info="Max tokens to generate.")
            temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Temperature", info="Controls randomness.")
            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P", info="Nucleus sampling probability.")
    return [max_new_tokens, temperature, top_p]

# Language map (defined once)
lang_map = {"Python": "python", "JavaScript": "javascript", "Java": "java", "C++": "cpp", "HTML": "html", "CSS": "css", "SQL": "sql", "Bash": "bash", "Rust": "rust", "Other": "plaintext"}

# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft(), fill_height=True, title="Gemma Capabilities Demo") as demo:

    # Header
    gr.Markdown(
        """
        <div style="text-align: center; margin-bottom: 20px;"><h1><span style="font-size: 1.5em;">🤖</span> Gemma Capabilities Demo</h1>
        <p>Explore text generation with Google's Gemma models (or a fallback).</p>
        <p style="font-size: 0.9em;"><a href="https://huggingface.co/google/gemma-7b-it" target="_blank">[Accept Gemma License Here]</a></p></div>"""
    )

    # --- Authentication ---
    with gr.Group(): # Removed variant="panel"
        gr.Markdown("### 🔑 Authentication")
        with gr.Row():
            with gr.Column(scale=4):
                hf_token = gr.Textbox(label="Hugging Face Token", placeholder="Paste token (hf_...)", type="password", value=DEFAULT_HF_TOKEN, info="Needed for Gemma models.")
            with gr.Column(scale=1, min_width=150):
                auth_button = gr.Button("Load Model", variant="primary")
        auth_status = gr.Markdown("ℹ️ Enter token & click 'Load Model'. May take time.")
        gr.Markdown(
             "**Token Info:** Get from [HF Settings](https://huggingface.co/settings/tokens) (read access). Ensure Gemma license is accepted.",
             elem_id="token-info" # Optional ID for styling if needed later
         )

    # --- Main Content Tabs ---
    with gr.Tabs(elem_id="main_tabs", visible=False) as tabs:

        # --- Text Generation Tab ---
        with gr.TabItem("📝 Creative & Informational"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("#### Configure Task")
                    text_gen_type = gr.Radio(["Creative Writing", "Informational Writing", "Custom Prompt"], label="Writing Type", value="Creative Writing")
                    with gr.Group(visible=True) as creative_options:
                        style = gr.Dropdown(["short story", "poem", "script", "song lyrics", "joke", "dialogue"], label="Style", value="short story")
                        creative_topic = gr.Textbox(label="Topic", placeholder="e.g., a lonely astronaut", value="a robot discovering music", lines=2)
                    with gr.Group(visible=False) as info_options:
                        format_type = gr.Dropdown(["article", "summary", "explanation", "report", "comparison"], label="Format", value="article")
                        info_topic = gr.Textbox(label="Topic", placeholder="e.g., quantum physics basics", value="AI impact on healthcare", lines=2)
                    with gr.Group(visible=False) as custom_prompt_group:
                        custom_prompt = gr.Textbox(label="Custom Prompt", placeholder="Enter full prompt...", lines=5)
                    text_gen_params = create_parameter_ui()
                    # Removed gr.Spacer
                    generate_text_btn = gr.Button("Generate Text", variant="primary")
                with gr.Column(scale=1):
                    gr.Markdown("#### Generated Output")
                    text_output = gr.Textbox(label="Result", lines=25, interactive=False, show_copy_button=True)

            # Visibility logic
            def update_text_gen_visibility(choice):
                return { creative_options: gr.update(visible=choice == "Creative Writing"),
                         info_options: gr.update(visible=choice == "Informational Writing"),
                         custom_prompt_group: gr.update(visible=choice == "Custom Prompt") }
            text_gen_type.change(update_text_gen_visibility, text_gen_type, [creative_options, info_options, custom_prompt_group], queue=False)

            # Click handler
            def text_gen_click(gen_type, style, c_topic, fmt_type, i_topic, custom_pr, *params):
                task_map = {"Creative Writing": ("creative", {"style": style, "topic": c_topic}),
                            "Informational Writing": ("informational", {"format_type": fmt_type, "topic": i_topic}),
                            "Custom Prompt": ("custom", {"prompt": custom_pr})}
                task_type, kwargs = task_map.get(gen_type, ("custom", {"prompt": custom_pr}))
                # Apply safe_value inside handler where needed
                if task_type == "creative": kwargs = {"style": safe_value(style, "story"), "topic": safe_value(c_topic, "[topic]")}
                elif task_type == "informational": kwargs = {"format_type": safe_value(fmt_type, "article"), "topic": safe_value(i_topic, "[topic]")}
                else: kwargs = {"prompt": safe_value(custom_pr, "Write something.")}
                final_prompt = generate_prompt(task_type, **kwargs)
                return generate_text(final_prompt, *params)
            generate_text_btn.click(text_gen_click, [text_gen_type, style, creative_topic, format_type, info_topic, custom_prompt, *text_gen_params], text_output)

            # Examples
            gr.Examples( examples=[ ["Creative Writing", "poem", "sound of rain", "", "", "", 512, 0.7, 0.9],
                                     ["Informational Writing", "", "", "explanation", "photosynthesis", "", 768, 0.6, 0.9],
                                     ["Custom Prompt", "", "", "", "", "Dialogue: cat and dog discuss humans.", 512, 0.8, 0.95] ],
                         inputs=[text_gen_type, style, creative_topic, format_type, info_topic, custom_prompt, *text_gen_params[:3]], # Pass UI elements
                         outputs=text_output, label="Try examples...")


        # --- Brainstorming Tab ---
        with gr.TabItem("🧠 Brainstorming"):
             with gr.Row():
                 with gr.Column(scale=1):
                     gr.Markdown("#### Setup")
                     brainstorm_category = gr.Dropdown(["project", "business", "creative", "solution", "content", "feature", "product name"], label="Category", value="project")
                     brainstorm_topic = gr.Textbox(label="Topic/Problem", placeholder="e.g., reducing plastic waste", value="unique mobile app ideas", lines=3)
                     brainstorm_params = create_parameter_ui()
                     # Removed gr.Spacer
                     brainstorm_btn = gr.Button("Generate Ideas", variant="primary")
                 with gr.Column(scale=1):
                     gr.Markdown("#### Generated Ideas")
                     brainstorm_output = gr.Textbox(label="Result", lines=25, interactive=False, show_copy_button=True)

             def brainstorm_click(category, topic, *params):
                 prompt = generate_prompt("brainstorm", category=safe_value(category, "project"), topic=safe_value(topic, "ideas"))
                 return generate_text(prompt, *params)
             brainstorm_btn.click(brainstorm_click, [brainstorm_category, brainstorm_topic, *brainstorm_params], brainstorm_output)
             gr.Examples([ ["solution", "engaging online learning", 768, 0.8, 0.9],
                           ["business", "eco-friendly subscription boxes", 768, 0.75, 0.9],
                           ["creative", "fantasy novel themes", 512, 0.85, 0.95] ],
                          inputs=[brainstorm_category, brainstorm_topic, *brainstorm_params[:3]], outputs=brainstorm_output, label="Try examples...")


        # --- Code Tab ---
        with gr.TabItem("💻 Code"):
            with gr.Tabs():
                with gr.TabItem("Generate"):
                     with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown("#### Setup")
                            code_lang_gen = gr.Dropdown(list(lang_map.keys())[:-1], label="Language", value="Python")
                            code_task = gr.Textbox(label="Task", placeholder="e.g., function for factorial", value="Python class for calculator", lines=4)
                            code_gen_params = create_parameter_ui()
                            # Removed gr.Spacer
                            code_gen_btn = gr.Button("Generate Code", variant="primary")
                        with gr.Column(scale=1):
                            gr.Markdown("#### Generated Code")
                            code_output = gr.Code(label="Result", language="python", lines=25, interactive=False)

                     def gen_code_click(lang, task, *params):
                         prompt = generate_prompt("code_generate", language=safe_value(lang, "Python"), task=safe_value(task, "hello world"))
                         result = generate_text(prompt, *params)
                         # Basic code block extraction
                         if "```" in result:
                             parts = result.split("```")
                             if len(parts) >= 2:
                                 block = parts[1]
                                 if '\n' in block: first_line, rest = block.split('\n', 1); return rest.strip() if first_line.strip().lower() == lang.lower() else block.strip()
                                 else: return block.strip()
                         return result.strip()
                     def update_gen_lang_display(lang): return gr.Code.update(language=lang_map.get(lang, "plaintext"))
                     code_lang_gen.change(update_gen_lang_display, code_lang_gen, code_output, queue=False)
                     code_gen_btn.click(gen_code_click, [code_lang_gen, code_task, *code_gen_params], code_output)
                     gr.Examples([ ["JavaScript", "email validation regex function", 768, 0.6, 0.9],
                                   ["SQL", "select users > 30 yrs old", 512, 0.5, 0.8],
                                   ["HTML", "basic portfolio structure", 1024, 0.7, 0.9] ],
                                 inputs=[code_lang_gen, code_task, *code_gen_params[:3]], outputs=code_output, label="Try examples...")

                with gr.TabItem("Explain"):
                     with gr.Row():
                         with gr.Column(scale=1):
                            gr.Markdown("#### Setup")
                            code_lang_explain = gr.Dropdown(list(lang_map.keys()), label="Language", value="Python")
                            code_to_explain = gr.Code(label="Code to Explain", language="python", lines=15)
                            explain_code_params = create_parameter_ui()
                            # Removed gr.Spacer
                            explain_code_btn = gr.Button("Explain Code", variant="primary")
                         with gr.Column(scale=1):
                            gr.Markdown("#### Explanation")
                            code_explanation = gr.Textbox(label="Result", lines=25, interactive=False, show_copy_button=True)

                     def explain_code_click(lang, code, *params):
                         code_content = safe_value(code['code'] if isinstance(code, dict) else code, "# Empty code")
                         prompt = generate_prompt("code_explain", language=safe_value(lang, "code"), code=code_content)
                         return generate_text(prompt, *params)
                     def update_explain_lang_display(lang): return gr.Code.update(language=lang_map.get(lang, "plaintext"))
                     code_lang_explain.change(update_explain_lang_display, code_lang_explain, code_to_explain, queue=False)
                     explain_code_btn.click(explain_code_click, [code_lang_explain, code_to_explain, *explain_code_params], code_explanation)


                with gr.TabItem("Debug"):
                     with gr.Row():
                         with gr.Column(scale=1):
                            gr.Markdown("#### Setup")
                            code_lang_debug = gr.Dropdown(list(lang_map.keys()), label="Language", value="Python")
                            code_to_debug = gr.Code(label="Buggy Code", language="python", lines=15, value="def avg(nums):\n  # Potential div by zero\n  return sum(nums)/len(nums)")
                            debug_code_params = create_parameter_ui()
                            # Removed gr.Spacer
                            debug_code_btn = gr.Button("Debug Code", variant="primary")
                         with gr.Column(scale=1):
                            gr.Markdown("#### Debugging Analysis")
                            debug_result = gr.Textbox(label="Result", lines=25, interactive=False, show_copy_button=True)

                     def debug_code_click(lang, code, *params):
                         code_content = safe_value(code['code'] if isinstance(code, dict) else code, "# Empty code")
                         prompt = generate_prompt("code_debug", language=safe_value(lang, "code"), code=code_content)
                         return generate_text(prompt, *params)
                     def update_debug_lang_display(lang): return gr.Code.update(language=lang_map.get(lang, "plaintext"))
                     code_lang_debug.change(update_debug_lang_display, code_lang_debug, code_to_debug, queue=False)
                     debug_code_btn.click(debug_code_click, [code_lang_debug, code_to_debug, *debug_code_params], debug_result)


        # --- Comprehension Tab ---
        with gr.TabItem("📚 Comprehension"):
            with gr.Tabs():
                with gr.TabItem("Summarize"):
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown("#### Setup")
                            summarize_text = gr.Textbox(label="Text to Summarize", lines=15, placeholder="Paste long text...")
                            summarize_params = create_parameter_ui()
                            # Removed gr.Spacer
                            summarize_btn = gr.Button("Summarize Text", variant="primary")
                        with gr.Column(scale=1):
                            gr.Markdown("#### Summary")
                            summary_output = gr.Textbox(label="Result", lines=15, interactive=False, show_copy_button=True)
                    def summarize_click(text, *params):
                        prompt = generate_prompt("summarize", text=safe_value(text, "[empty text]"))
                        # Adjust max tokens for summary specifically if needed
                        p_list = list(params); p_list[0] = min(max(int(p_list[0]), 64), 512)
                        return generate_text(prompt, *p_list)
                    summarize_btn.click(summarize_click, [summarize_text, *summarize_params], summary_output)


                with gr.TabItem("Q & A"):
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown("#### Setup")
                            qa_text = gr.Textbox(label="Context Text", lines=10, placeholder="Paste text containing answer...")
                            qa_question = gr.Textbox(label="Question", placeholder="Ask question about text...")
                            qa_params = create_parameter_ui()
                            # Removed gr.Spacer
                            qa_btn = gr.Button("Get Answer", variant="primary")
                        with gr.Column(scale=1):
                            gr.Markdown("#### Answer")
                            qa_output = gr.Textbox(label="Result", lines=10, interactive=False, show_copy_button=True)
                    def qa_click(text, question, *params):
                        prompt = generate_prompt("qa", text=safe_value(text, "[context]"), question=safe_value(question,"[question]"))
                        p_list = list(params); p_list[0] = min(max(int(p_list[0]), 32), 256)
                        return generate_text(prompt, *p_list)
                    qa_btn.click(qa_click, [qa_text, qa_question, *qa_params], qa_output)


                with gr.TabItem("Translate"):
                     with gr.Row():
                         with gr.Column(scale=1):
                             gr.Markdown("#### Setup")
                             translate_text = gr.Textbox(label="Text to Translate", lines=8, placeholder="Enter text...")
                             target_lang = gr.Dropdown(["French", "Spanish", "German", "Japanese", "Chinese", "Russian", "Arabic", "Hindi", "Portuguese", "Italian"], label="Translate To", value="French")
                             translate_params = create_parameter_ui()
                             # Removed gr.Spacer
                             translate_btn = gr.Button("Translate Text", variant="primary")
                         with gr.Column(scale=1):
                             gr.Markdown("#### Translation")
                             translation_output = gr.Textbox(label="Result", lines=8, interactive=False, show_copy_button=True)
                     def translate_click(text, lang, *params):
                         prompt = generate_prompt("translate", text=safe_value(text,"[text]"), target_lang=safe_value(lang,"French"))
                         p_list = list(params); p_list[0] = max(int(p_list[0]), 64)
                         return generate_text(prompt, *p_list)
                     translate_btn.click(translate_click, [translate_text, target_lang, *translate_params], translation_output)


        # --- More Tasks Tab ---
        with gr.TabItem("🛠️ More Tasks"):
             with gr.Tabs():
                 with gr.TabItem("Content Creation"):
                     with gr.Row():
                         with gr.Column(scale=1):
                             gr.Markdown("#### Setup")
                             content_type = gr.Dropdown(["blog post outline", "social media post (Twitter)", "social media post (LinkedIn)", "marketing email subject line", "product description", "press release intro"], label="Content Type", value="blog post outline")
                             content_topic = gr.Textbox(label="Topic", value="sustainable travel tips", lines=2)
                             content_audience = gr.Textbox(label="Audience", value="eco-conscious millennials")
                             content_params = create_parameter_ui()
                             # Removed gr.Spacer
                             content_btn = gr.Button("Generate Content", variant="primary")
                         with gr.Column(scale=1):
                             gr.Markdown("#### Generated Content")
                             content_output = gr.Textbox(label="Result", lines=20, interactive=False, show_copy_button=True)
                     def content_click(c_type, topic, audience, *params):
                         prompt = generate_prompt("content_creation", content_type=safe_value(c_type,"text"), topic=safe_value(topic,"[topic]"), audience=safe_value(audience,"[audience]"))
                         return generate_text(prompt, *params)
                     content_btn.click(content_click, [content_type, content_topic, content_audience, *content_params], content_output)

                 with gr.TabItem("Email Drafting"):
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown("#### Setup")
                            email_type = gr.Dropdown(["job inquiry", "meeting request", "follow-up", "thank you", "support response", "sales outreach"], label="Email Type", value="meeting request")
                            email_context = gr.Textbox(label="Context/Points", lines=5, value="Request meeting next week re: project X. Suggest Tue/Wed afternoon.")
                            email_params = create_parameter_ui()
                            # Removed gr.Spacer
                            email_btn = gr.Button("Generate Draft", variant="primary")
                        with gr.Column(scale=1):
                            gr.Markdown("#### Generated Draft")
                            email_output = gr.Textbox(label="Result", lines=20, interactive=False, show_copy_button=True)
                    def email_click(e_type, context, *params):
                        prompt = generate_prompt("email_draft", email_type=safe_value(e_type,"email"), context=safe_value(context,"[context]"))
                        return generate_text(prompt, *params)
                    email_btn.click(email_click, [email_type, email_context, *email_params], email_output)

                 with gr.TabItem("Doc Editing"):
                     with gr.Row():
                         with gr.Column(scale=1):
                             gr.Markdown("#### Setup")
                             edit_text = gr.Textbox(label="Text to Edit", lines=10, placeholder="Paste text...")
                             edit_type = gr.Dropdown(["improve clarity", "fix grammar/spelling", "make concise", "make formal", "make casual", "simplify"], label="Improve For", value="improve clarity")
                             edit_params = create_parameter_ui()
                             # Removed gr.Spacer
                             edit_btn = gr.Button("Edit Text", variant="primary")
                         with gr.Column(scale=1):
                             gr.Markdown("#### Edited Text")
                             edit_output = gr.Textbox(label="Result", lines=10, interactive=False, show_copy_button=True)
                     def edit_click(text, e_type, *params):
                         prompt = generate_prompt("document_edit", text=safe_value(text,"[text]"), edit_type=safe_value(e_type,"clarity"))
                         p_list = list(params); input_tokens = len(safe_value(text,"").split()); p_list[0] = max(int(p_list[0]), input_tokens + 64)
                         return generate_text(prompt, *p_list)
                     edit_btn.click(edit_click, [edit_text, edit_type, *edit_params], edit_output)

                 with gr.TabItem("Classification"):
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown("#### Setup")
                            classify_text = gr.Textbox(label="Text to Classify", lines=8, value="Sci-fi movie explores AI consciousness.")
                            classify_categories = gr.Textbox(label="Categories (comma-sep)", value="Tech, Entertainment, Science, Politics")
                            classify_params = create_parameter_ui()
                            # Removed gr.Spacer
                            classify_btn = gr.Button("Classify Text", variant="primary")
                        with gr.Column(scale=1):
                            gr.Markdown("#### Classification")
                            classify_output = gr.Textbox(label="Predicted Category", lines=2, interactive=False, show_copy_button=True)
                    def classify_click(text, cats, *params):
                        prompt = generate_prompt("classify", text=safe_value(text,"[text]"), categories=safe_value(cats,"cat1, cat2"))
                        p_list = list(params); p_list[0] = min(max(int(p_list[0]), 16), 128)
                        raw = generate_text(prompt, *p_list)
                        # Basic post-processing attempt
                        lines = raw.split('\n'); last = lines[-1].strip(); possible = [c.strip().lower() for c in cats.split(',')]; return last if last.lower() in possible else raw
                    classify_btn.click(classify_click, [classify_text, classify_categories, *classify_params], classify_output)

                 with gr.TabItem("Data Extraction"):
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown("#### Setup")
                            extract_text = gr.Textbox(label="Source Text", lines=10, value="Order #123 by Jane ([email protected]). Total: $99. Shipped: 123 Main St.")
                            extract_data_points = gr.Textbox(label="Data Points (comma-sep)", value="order num, name, email, total, address")
                            extract_params = create_parameter_ui()
                            # Removed gr.Spacer
                            extract_btn = gr.Button("Extract Data", variant="primary")
                        with gr.Column(scale=1):
                            gr.Markdown("#### Extracted Data")
                            extract_output = gr.Textbox(label="Result (JSON or Key-Value)", lines=10, interactive=False, show_copy_button=True)
                    def extract_click(text, points, *params):
                        prompt = generate_prompt("data_extract", text=safe_value(text,"[text]"), data_points=safe_value(points,"info"))
                        return generate_text(prompt, *params)
                    extract_btn.click(extract_click, [extract_text, extract_data_points, *extract_params], extract_output)


    # --- Authentication Handler & Footer ---
    footer_status = gr.Markdown(f"...", elem_id="footer-status-md") # Placeholder

    def handle_auth(token):
        yield "⏳ Authenticating & loading model...", gr.Tabs.update(visible=False)
        status_message, tabs_update = load_model(token)
        yield status_message, tabs_update

    def update_footer_status(status_text): # Updates footer based on global state
        return gr.Markdown.update(value=f"""
           <hr><div style="text-align: center; font-size: 0.9em; color: #777;">
           <p>Powered by Hugging Face 🤗 Transformers & Gradio. Model: <strong>{loaded_model_name if model_loaded else 'None'}</strong>.</p>
           <p>Review outputs carefully. Models may generate inaccurate information.</p></div>""")

    auth_button.click(handle_auth, hf_token, [auth_status, tabs], queue=True)
    # Update footer whenever auth status text changes
    auth_status.change(update_footer_status, auth_status, footer_status, queue=False)
    # Initial footer update on load
    demo.load(update_footer_status, auth_status, footer_status, queue=False)


# --- Launch App ---
demo.launch(share=False)