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

    # --- FIX: Use gr.update() for visibility ---
    initial_tabs_update = gr.update(visible=False) # Generic update targeted by outputs list

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

    try:
        model_options = [
            "google/gemma-2b-it", "google/gemma-7b-it",
            "google/gemma-2b", "google/gemma-7b",
            "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
        ]
        print(f"Attempting loading with token: {hf_token[:5]}...")
        loaded_successfully = False
        for model_name in model_options:
            try:
                print(f"\n--- Attempting: {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("Loading model...")
                global_model = AutoModelForCausalLM.from_pretrained(
                    model_name, torch_dtype=torch.float16,
                    device_map="auto", token=current_token
                )
                print(f"Success: {model_name}")
                model_loaded = True
                loaded_model_name = model_name
                loaded_successfully = True
                # --- FIX: Use gr.update() for visibility ---
                tabs_update = gr.update(visible=True) # Generic update targeted by outputs list
                status_msg = f"✅ Model '{model_name}' loaded!"
                if "tinyllama" in model_name.lower():
                    status_msg = f"✅ Fallback '{model_name}' loaded!"
                return status_msg, tabs_update
            except ImportError as ie:
                 print(f"Import Error ({model_name}): {ie}. Check dependencies.")
                 continue
            except Exception as e:
                print(f"Failed ({model_name}): {e}")
                if "401" in str(e) or "logged in" in str(e) and is_gemma: print("Auth error likely.")
                continue
        if not loaded_successfully:
            model_loaded = False; loaded_model_name = "None"
            return "❌ Failed to load any model. Check token/license/deps/network.", initial_tabs_update
    except Exception as e:
        model_loaded = False; loaded_model_name = "None"
        print(f"Outer load error: {e}"); traceback.print_exc()
        if "401" in str(e) or "logged in" in str(e): return "❌ Auth failed.", initial_tabs_update
        else: return f"❌ Unexpected load error: {e}", 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 ---")
    # ... (rest of the function remains the same as the previous valid version) ...
    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():
                 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():
                 chat_prompt = f"<|system|>\nYou are a helpful assistant.</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)

        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

        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
        }
        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 use 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). Adjust Temp/Top-P or prompt."
        else:
             return f"❌ Error during text generation: {error_msg}"


# --- UI Components & Layout ---

def create_parameter_ui():
    # ... (function remains the same) ...
    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
    # ... (remains the same) ...
    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 ---
    # ... (remains the same) ...
    with gr.Group():
        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"
         )

    # --- Main Content Tabs ---
    # Define tabs instance first
    with gr.Tabs(elem_id="main_tabs", visible=False) as tabs:
        # ... (All TabItem definitions remain the same as the previous working version) ...
        # --- 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()
                    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)

            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)

            def text_gen_click(gen_type, style, c_topic, fmt_type, i_topic, custom_pr, *params):
                task_map = {"Creative Writing": ("creative", {}), "Informational Writing": ("informational", {}), "Custom Prompt": ("custom", {})}
                task_type, kwargs = task_map.get(gen_type, ("custom", {}))
                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)

            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]],
                         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()
                     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()
                            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 extraction...
                         if "```" in result: parts = result.split("```"); block = parts[1] if len(parts)>1 else ''; return block.split('\n',1)[1].strip() if '\n' in block and block.split('\n',1)[0].strip().lower() == lang.lower() else 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([["JS", "email validation", 768, 0.6, 0.9], ["SQL", "users > 30", 512, 0.5, 0.8], ["HTML", "portfolio", 1024, 0.7, 0.9]], [code_lang_gen, code_task, *code_gen_params[:3]], code_output, label="Try...") # Abbreviated 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(); 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, "#"); 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(); 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, "#"); 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", lines=15, placeholder="Paste..."); summarize_params = create_parameter_ui(); summarize_btn = gr.Button("Summarize", 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,"[text]")); p = list(params); p[0]=min(max(int(p[0]),64),512); return generate_text(prompt, *p)
                    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", lines=10, placeholder="Paste context..."); qa_question = gr.Textbox(label="Question", placeholder="Ask..."); qa_params = create_parameter_ui(); qa_btn = gr.Button("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, q, *params): prompt = generate_prompt("qa", text=safe_value(text,"[ctx]"), question=safe_value(q,"[q]")); p = list(params); p[0]=min(max(int(p[0]),32),256); return generate_text(prompt, *p)
                    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", lines=8, placeholder="Enter text..."); target_lang = gr.Dropdown(["French", "Spanish", "German", "Japanese", "Chinese", "Russian", "Arabic", "Hindi", "Portuguese", "Italian"], label="To", value="French"); translate_params = create_parameter_ui(); translate_btn = gr.Button("Translate", 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(params); p[0]=max(int(p[0]),64); return generate_text(prompt, *p)
                     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"): # Abbreviated names for brevity
                     with gr.Row():
                         with gr.Column(scale=1): gr.Markdown("#### Setup"); content_type = gr.Dropdown(["blog outline", "tweet", "linkedin post", "email subject", "product desc", "press release intro"], label="Type", value="blog outline"); content_topic = gr.Textbox(label="Topic", value="sustainable travel", lines=2); content_audience = gr.Textbox(label="Audience", value="millennials"); content_params = create_parameter_ui(); content_btn = gr.Button("Generate", variant="primary")
                         with gr.Column(scale=1): gr.Markdown("#### Result"); content_output = gr.Textbox(lines=20, interactive=False, show_copy_button=True)
                     def content_click(t, top, aud, *p): prompt = generate_prompt("content_creation", content_type=safe_value(t,"text"), topic=safe_value(top,"[topic]"), audience=safe_value(aud,"[audience]")); return generate_text(prompt, *p)
                     content_btn.click(content_click, [content_type, content_topic, content_audience, *content_params], content_output)
                 with gr.TabItem("Email"):
                     with gr.Row():
                         with gr.Column(scale=1): gr.Markdown("#### Setup"); email_type = gr.Dropdown(["job inquiry", "meeting request", "follow-up", "thank you", "support reply", "sales outreach"], label="Type", value="meeting request"); email_context = gr.Textbox(label="Context", lines=5, value="Meet next week re: project X. Tue/Wed PM?"); email_params = create_parameter_ui(); email_btn = gr.Button("Generate", variant="primary")
                         with gr.Column(scale=1): gr.Markdown("#### Draft"); email_output = gr.Textbox(lines=20, interactive=False, show_copy_button=True)
                     def email_click(t, ctx, *p): prompt = generate_prompt("email_draft", email_type=safe_value(t,"email"), context=safe_value(ctx,"[context]")); return generate_text(prompt, *p)
                     email_btn.click(email_click, [email_type, email_context, *email_params], email_output)
                 with gr.TabItem("Edit"):
                     with gr.Row():
                         with gr.Column(scale=1): gr.Markdown("#### Setup"); edit_text = gr.Textbox(label="Text", lines=10, placeholder="Paste..."); edit_type = gr.Dropdown(["clarity", "grammar/spelling", "concise", "formal", "casual", "simplify"], label="Improve For", value="clarity"); edit_params = create_parameter_ui(); edit_btn = gr.Button("Edit", variant="primary")
                         with gr.Column(scale=1): gr.Markdown("#### Edited"); edit_output = gr.Textbox(lines=10, interactive=False, show_copy_button=True)
                     def edit_click(txt, et, *p): prompt = generate_prompt("document_edit", text=safe_value(txt,"[text]"), edit_type=safe_value(et,"clarity")); p_list = list(p); p_list[0] = max(int(p_list[0]), len(safe_value(txt,"").split()) + 64); return generate_text(prompt, *p_list)
                     edit_btn.click(edit_click, [edit_text, edit_type, *edit_params], edit_output)
                 with gr.TabItem("Classify"):
                     with gr.Row():
                         with gr.Column(scale=1): gr.Markdown("#### Setup"); classify_text = gr.Textbox(label="Text", lines=8, value="Sci-fi movie explores AI."); classify_categories = gr.Textbox(label="Categories", value="Tech, Entertainment, Science"); classify_params = create_parameter_ui(); classify_btn = gr.Button("Classify", variant="primary")
                         with gr.Column(scale=1): gr.Markdown("#### Category"); classify_output = gr.Textbox(lines=2, interactive=False, show_copy_button=True)
                     def classify_click(txt, cats, *p): prompt = generate_prompt("classify", text=safe_value(txt,"[text]"), categories=safe_value(cats,"c1,c2")); p_list = list(p); p_list[0] = min(max(int(p_list[0]),16),128); raw=generate_text(prompt,*p_list); lines=raw.split('\n');last=lines[-1].strip();poss=[c.strip().lower() for c in cats.split(',')]; return last if last.lower() in poss else raw
                     classify_btn.click(classify_click, [classify_text, classify_categories, *classify_params], classify_output)
                 with gr.TabItem("Extract"):
                     with gr.Row():
                         with gr.Column(scale=1): gr.Markdown("#### Setup"); extract_text = gr.Textbox(label="Source", lines=10, value="Order #123 by Jane ([email protected]). Total: $99."); extract_data_points = gr.Textbox(label="Extract", value="order num, name, email, total"); extract_params = create_parameter_ui(); extract_btn = gr.Button("Extract", variant="primary")
                         with gr.Column(scale=1): gr.Markdown("#### Data"); extract_output = gr.Textbox(lines=10, interactive=False, show_copy_button=True)
                     def extract_click(txt, pts, *p): prompt = generate_prompt("data_extract", text=safe_value(txt,"[text]"), data_points=safe_value(pts,"info")); return generate_text(prompt, *p)
                     extract_btn.click(extract_click, [extract_text, extract_data_points, *extract_params], extract_output)


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

    # Define authentication handler AFTER tabs is defined
    def handle_auth(token):
        # --- FIX: Use gr.update() for visibility ---
        yield "⏳ Authenticating & loading model...", gr.update(visible=False)
        # Call the actual model loading function
        status_message, tabs_update_obj = load_model(token) # Get the update object
        yield status_message, tabs_update_obj # Yield the object

    # Define footer update handler
    def update_footer_status(status_text): # Updates footer based on global state
        # --- FIX: Use gr.update() for Markdown ---
        return gr.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>""")

    # Link button click to the handler
    auth_button.click(
        fn=handle_auth,
        inputs=hf_token,
        outputs=[auth_status, tabs], # Target auth_status and the tabs instance
        queue=True
    )
    # Update footer whenever auth status text changes
    auth_status.change(
        fn=update_footer_status,
        inputs=auth_status, # Trigger based on auth_status text
        outputs=footer_status, # Update the footer_status Markdown
        queue=False
    )
    # Initial footer update on load
    demo.load(
        fn=update_footer_status,
        inputs=auth_status, # Use initial auth_status text
        outputs=footer_status,
        queue=False
    )


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