ONS2

app.py

@@ -35,14 +35,12 @@ def load_model(hf_token):
 
     try:
         # Try different model versions from smallest to largest
-        # Prioritize instruction-tuned models
         model_options = [
             "google/gemma-2b-it",
             "google/gemma-7b-it",
             "google/gemma-2b",
             "google/gemma-7b",
-            # Add a smaller, potentially public model as a last resort
-            "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+            "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Fallback
         ]
 
         print(f"Attempting to load models with token starting with: {hf_token[:5]}...")
@@ -51,10 +49,8 @@ def load_model(hf_token):
             try:
                 print(f"\n--- Attempting to load model: {model_name} ---")
                 is_gemma = "gemma" in model_name.lower()
-                is_fallback = "tinyllama" in model_name.lower()
-                current_token = hf_token if is_gemma else None # Only use token for Gemma models
+                current_token = hf_token if is_gemma else None
 
-                # Load tokenizer
                 print("Loading tokenizer...")
                 global_tokenizer = AutoTokenizer.from_pretrained(
                     model_name,
@@ -62,13 +58,11 @@ def load_model(hf_token):
                 )
                 print("Tokenizer loaded successfully.")
 
-                # Load model
                 print(f"Loading model {model_name}...")
                 global_model = AutoModelForCausalLM.from_pretrained(
                     model_name,
-                    # torch_dtype=torch.bfloat16, # Use bfloat16 for better performance/compatibility if available - fallback to float16 if needed
                     torch_dtype=torch.float16, # Using float16 for broader compatibility
-                    device_map="auto", # Let HF decide device placement
+                    device_map="auto",
                     token=current_token
                 )
                 print(f"Model {model_name} loaded successfully!")
@@ -76,34 +70,26 @@ def load_model(hf_token):
                 model_loaded = True
                 loaded_model_name = model_name
                 loaded_successfully = True
-                tabs_update = gr.Tabs.update(visible=True) # Show tabs on success
-                if is_fallback:
-                     return f"✅ Fallback model '{model_name}' loaded successfully! Limited capabilities compared to Gemma.", tabs_update
-                else:
-                     return f"✅ Model '{model_name}' loaded successfully!", tabs_update
+                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:
-                 # Handle potential missing dependencies like bitsandbytes if bfloat16 fails
-                 print(f"Import Error loading {model_name}: {import_err}. Check dependencies.")
-                 continue # Try next model
+                 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}")
-                # traceback.print_exc() # Keep for debugging if needed, but can be verbose
-                if "401 Client Error" in str(specific_e) and is_gemma:
+                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.")
-                    # Don't immediately fail, try next model
-                elif "requires you to be logged in" in str(specific_e) and is_gemma:
-                     print("Authentication error likely. Check token and license agreement.")
-                     # Don't immediately fail, try next model
-                # Continue to the next model option
                 continue
 
-        # If loop finishes without loading
         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 (ensure it has read permissions and you've accepted Gemma's license on Hugging Face) and network connection.", initial_tabs_update
+            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
@@ -111,16 +97,14 @@ def load_model(hf_token):
         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. Please check your Hugging Face token and ensure you have accepted the Gemma license agreement on the Hugging Face model page.", initial_tabs_update
+            return "❌ Authentication failed. Check token/license.", initial_tabs_update
         else:
-            return f"❌ An unexpected error occurred during model loading: {error_msg}", initial_tabs_update
+            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"""
-    # Using a dictionary-based approach for cleaner prompt generation
     prompts = {
         "creative": "Write a {style} about {topic}. Be creative and engaging.",
         "informational": "Write an {format_type} about {topic}. Be clear, factual, and informative.",
@@ -138,24 +122,17 @@ def generate_prompt(task_type, **kwargs):
         "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:
-            # Prepare kwargs safely for formatting
-            # Find placeholders like {key}
             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} # Use default if key missing in input kwargs
-
-            # Add any extra kwargs provided that weren't in the template (e.g., for 'custom' type)
-            final_kwargs.update(kwargs)
-
+            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}") # Fallback
+            return kwargs.get("prompt", f"Generate text based on: {kwargs}")
     else:
-        # Fallback for custom or undefined task types
         return kwargs.get("prompt", "Generate text based on the input.")
 
 
@@ -169,689 +146,435 @@ def generate_text(prompt, max_new_tokens=1024, temperature=0.7, top_p=0.9):
     print(f"Prompt (start): {prompt[:150]}...")
 
     if not model_loaded or global_model is None or global_tokenizer is None:
-        print("Model not loaded error.")
         return "⚠️ Model not loaded. Please authenticate first."
-
     if not prompt:
         return "⚠️ Please enter a prompt or configure a task."
 
     try:
-        # Add role/turn indicators if using an instruction-tuned model
-        # Simple check based on model name conventions
+        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()):
-             # Simple chat structure assumed by many instruction models
-             # Using Gemma's specific format if it's a Gemma IT model
              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:"
-        else:
-             # Base models might not need specific turn indicators
-             chat_prompt = prompt
 
         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}")
 
-        # Ensure max_new_tokens isn't excessively large for the model context
-        # Cap generation length for stability
-        effective_max_new_tokens = min(int(max_new_tokens), 2048) # Cap generation length and ensure int
+        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, # Include attention mask
+            "attention_mask": inputs.attention_mask,
             "max_new_tokens": effective_max_new_tokens,
             "do_sample": True,
-            "temperature": float(temperature), # Ensure float
-            "top_p": float(top_p),             # Ensure float
-            "pad_token_id": global_tokenizer.eos_token_id if global_tokenizer.eos_token_id is not None else 50256 # Use EOS token for padding, provide fallback
+            "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}")
 
-        # Generate text
-        with torch.no_grad(): # Disable gradient calculation for inference
+        with torch.no_grad():
             outputs = global_model.generate(**generation_args)
 
-        # Decode only the newly generated tokens
         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() # Remove leading/trailing whitespace
+        return generated_text.strip()
 
     except Exception as e:
         error_msg = str(e)
         print(f"Generation error: {error_msg}")
-        print(f"Error type: {type(e)}")
         traceback.print_exc()
-        # Check for common CUDA errors
         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 variant if possible."
-        elif "probability tensor contains nan" in error_msg:
-             return f"❌ Error: Generation failed (NaN probability). Try adjusting Temperature/Top-P or modifying the prompt."
+             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}\n\nPlease check the logs or try adjusting parameters."
+             return f"❌ Error during text generation: {error_msg}"
+
+# --- UI Components & Layout ---
 
-# Create parameters UI component (reusable function)
 def create_parameter_ui():
     with gr.Accordion("✨ Generation Parameters", open=False):
         with gr.Row():
-             # Renamed max_length to max_new_tokens for clarity with HF generate API
-            max_new_tokens = gr.Slider(
-                minimum=64,
-                maximum=2048, # Set a reasonable max limit
-                value=512, # Default to a moderate length
-                step=64,
-                label="Max New Tokens",
-                info="Max number of tokens to generate.",
-                elem_id="max_new_tokens_slider"
-            )
-            temperature = gr.Slider(
-                minimum=0.1, # Avoid 0 which disables sampling
-                maximum=1.5,
-                value=0.7,
-                step=0.1,
-                label="Temperature",
-                info="Controls randomness. Lower is focused, higher is diverse.",
-                elem_id="temperature_slider"
-            )
-            top_p = gr.Slider(
-                minimum=0.1,
-                maximum=1.0, # Can be 1.0
-                value=0.9,
-                step=0.05,
-                label="Top-P (Nucleus Sampling)",
-                info="Considers tokens with cumulative probability >= top_p.",
-                elem_id="top_p_slider"
-            )
+            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 ---
-# Use the soft theme for a clean look, allow light/dark switching
 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 style="font-size: 2.5em; font-weight: bold;">
-                <span style="font-size: 1.5em;">🤖</span> Gemma Capabilities Demo
-            </h1>
-            <p style="font-size: 1.1em; color: #555;">
-                Explore the text generation capabilities of Google's Gemma models (or a fallback).
-            </p>
-            <p style="font-size: 0.9em; color: #777;">
-                Requires a Hugging Face token with access to Gemma models.
-                <a href="https://huggingface.co/google/gemma-7b-it" target="_blank">[Accept Gemma License Here]</a>
-            </p>
-        </div>
-        """
+        <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 Section ---
-    # REMOVED variant="panel" from gr.Group for compatibility
-    with gr.Group(): # Use default Group appearance
-        gr.Markdown("### 🔑 Authentication") # Added heading inside group
+    # --- 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 your HF token here (hf_...)",
-                    type="password",
-                    value=DEFAULT_HF_TOKEN,
-                    info="Get your token from https://huggingface.co/settings/tokens",
-                    elem_id="hf_token_input"
-                )
+                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", elem_id="auth_button")
-
-        auth_status = gr.Markdown("ℹ️ Enter your Hugging Face token and click 'Load Model'. This might take a minute.", elem_id="auth_status")
-        # Add instructions on getting token inside the auth group
+                auth_button = gr.Button("Load Model", variant="primary")
+        auth_status = gr.Markdown("ℹ️ Enter token & click 'Load Model'. May take time.")
         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. Ensure you've accepted the [Gemma model license](https://huggingface.co/google/gemma-7b-it) on the model page.
-             """
+             "**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 (Initially Hidden) ---
-    # Define the tabs variable here
+    # --- Main Content Tabs ---
     with gr.Tabs(elem_id="main_tabs", visible=False) as tabs:
 
         # --- Text Generation Tab ---
-        with gr.TabItem("📝 Creative & Informational", id="tab_text_gen"):
-            with gr.Row(equal_height=False): # Allow columns to have different heights if needed
-                # Input Column
+        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",
-                        elem_id="text_gen_type"
-                    )
-
-                    # --- Dynamic Options ---
-                    with gr.Group(visible=True, elem_id="creative_options") as creative_options:
-                        style = gr.Dropdown(["short story", "poem", "script", "song lyrics", "joke", "dialogue"], label="Style", value="short story", elem_id="creative_style")
-                        creative_topic = gr.Textbox(label="Topic", placeholder="e.g., a lonely astronaut on Mars", value="a robot discovering music", elem_id="creative_topic", lines=2)
-
-                    with gr.Group(visible=False, elem_id="info_options") as info_options:
-                        format_type = gr.Dropdown(["article", "summary", "explanation", "report", "comparison"], label="Format", value="article", elem_id="info_format")
-                        info_topic = gr.Textbox(label="Topic", placeholder="e.g., the basics of quantum physics", value="the impact of AI on healthcare", elem_id="info_topic", lines=2)
-
-                    with gr.Group(visible=False, elem_id="custom_prompt_group") as custom_prompt_group:
-                        custom_prompt = gr.Textbox(label="Custom Prompt", placeholder="Enter your full prompt here...", lines=5, elem_id="custom_prompt")
-
-                    # Show/hide logic (using gr.update for better practice)
-                    def update_text_gen_visibility(choice):
-                        is_creative = choice == "Creative Writing"
-                        is_info = choice == "Informational Writing"
-                        is_custom = choice == "Custom Prompt"
-                        return {
-                            creative_options: gr.update(visible=is_creative),
-                            info_options: gr.update(visible=is_info),
-                            custom_prompt_group: gr.update(visible=is_custom)
-                        }
-                    text_gen_type.change(update_text_gen_visibility, inputs=text_gen_type, outputs=[creative_options, info_options, custom_prompt_group], queue=False)
-
-                    # Parameters
+                    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()
-                    gr.Spacer(height=15)
-                    generate_text_btn = gr.Button("Generate Text", variant="primary", elem_id="generate_text_btn")
-
-                # Output Column
+                    # 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, elem_id="text_output", show_copy_button=True) # Added copy button
-
-            # Handler
-            def text_generation_handler(gen_type, style, creative_topic, format_type, info_topic, custom_prompt_text, max_tokens, temp, top_p_val):
-                task_map = {
-                    "Creative Writing": ("creative", {"style": style, "topic": creative_topic}),
-                    "Informational Writing": ("informational", {"format_type": format_type, "topic": info_topic}),
-                    "Custom Prompt": ("custom", {"prompt": custom_prompt_text}) # Use 'custom' as type, 'prompt' as key
-                }
-                # Default to custom if type not found (shouldn't happen with Radio)
-                task_type, kwargs = task_map.get(gen_type, ("custom", {"prompt": custom_prompt_text}))
-
-                # Ensure safe values for specific task types
-                if task_type == "creative":
-                    kwargs["style"] = safe_value(style, "story")
-                    kwargs["topic"] = safe_value(creative_topic, "a default topic")
-                elif task_type == "informational":
-                    kwargs["format_type"] = safe_value(format_type, "article")
-                    kwargs["topic"] = safe_value(info_topic, "a default topic")
-                elif task_type == "custom":
-                     kwargs["prompt"] = safe_value(custom_prompt_text, "Write something interesting.")
-
-
+                    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, max_tokens, temp, top_p_val)
-
-            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
-            )
+                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
-            # Simplified examples list for clarity
-            gr.Examples(
-                examples=[
-                    ["Creative Writing", "poem", "the sound of rain on a tin roof", "", "", "", 512, 0.7, 0.9],
-                    ["Informational Writing", "", "", "explanation", "how photosynthesis works", "", 768, 0.6, 0.9],
-                    ["Custom Prompt", "", "", "", "", "Write a short dialogue between a cat and a dog discussing their humans.", 512, 0.8, 0.95],
-                ],
-                # Ensure the order matches the handler's inputs
-                inputs=[text_gen_type, style, creative_topic, format_type, info_topic, custom_prompt, *text_gen_params[:3]], # Pass only the UI elements needed
-                outputs=text_output,
-                label="Try these examples...",
-                #fn=text_generation_handler # fn is deprecated, click event handles execution
-            )
+            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", id="tab_brainstorm"):
-            with gr.Row(equal_height=False):
-                # Input Column
-                with gr.Column(scale=1):
-                    gr.Markdown("### Brainstorming Setup")
-                    brainstorm_category = gr.Dropdown(["project", "business", "creative", "solution", "content", "feature", "product name"], label="Idea Category", value="project", elem_id="brainstorm_category")
-                    brainstorm_topic = gr.Textbox(label="Topic or Problem", placeholder="e.g., reducing plastic waste", value="unique mobile app ideas", elem_id="brainstorm_topic", lines=3)
-                    brainstorm_params = create_parameter_ui()
-                    gr.Spacer(height=15)
-                    brainstorm_btn = gr.Button("Generate Ideas", variant="primary", elem_id="brainstorm_btn")
-
-                # Output Column
-                with gr.Column(scale=1):
-                    gr.Markdown("### Generated Ideas")
-                    brainstorm_output = gr.Textbox(label="Result", lines=25, interactive=False, elem_id="brainstorm_output", show_copy_button=True)
-
-            # Handler
-            def brainstorm_handler(category, topic, max_tokens, temp, top_p_val):
-                category = safe_value(category, "project")
-                topic = safe_value(topic, "innovative concepts")
-                prompt = generate_prompt("brainstorm", category=category, topic=topic)
-                return generate_text(prompt, max_tokens, temp, top_p_val)
-
-            brainstorm_btn.click(brainstorm_handler, inputs=[brainstorm_category, brainstorm_topic, *brainstorm_params], outputs=brainstorm_output)
-
-            gr.Examples(
-                examples=[
-                    ["solution", "making online learning more engaging", 768, 0.8, 0.9],
-                    ["business", "eco-friendly subscription boxes", 768, 0.75, 0.9],
-                    ["creative", "themes for a fantasy novel", 512, 0.85, 0.95],
-                ],
-                inputs=[brainstorm_category, brainstorm_topic, *brainstorm_params[:3]],
-                outputs=brainstorm_output,
-                label="Try these examples...",
-            )
-
-        # --- Code Capabilities Tab ---
-        with gr.TabItem("💻 Code", id="tab_code"):
-             # Language mapping for syntax highlighting (defined once)
-            lang_map = {"Python": "python", "JavaScript": "javascript", "Java": "java", "C++": "cpp", "HTML": "html", "CSS": "css", "SQL": "sql", "Bash": "bash", "Rust": "rust", "Other": "plaintext"}
-
-            with gr.Tabs() as code_tabs:
-                # --- Code Generation ---
-                with gr.TabItem("Generate Code", id="subtab_code_gen"):
-                    with gr.Row(equal_height=False):
-                        # Input Column
+        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("### Code Generation Setup")
-                             code_language_gen = gr.Dropdown(list(lang_map.keys())[:-1], label="Language", value="Python", elem_id="code_language_gen") # Exclude 'Other'
-                             code_task = gr.Textbox(label="Task Description", placeholder="e.g., function to calculate factorial", value="create a Python class for a basic calculator", lines=4, elem_id="code_task")
-                             code_gen_params = create_parameter_ui()
-                             gr.Spacer(height=15)
-                             code_gen_btn = gr.Button("Generate Code", variant="primary", elem_id="code_gen_btn")
-
-                        # Output Column
+                            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, elem_id="code_output") # No copy button needed for gr.Code
-
-                    # Handler
-                    def code_gen_handler(language, task, max_tokens, temp, top_p_val):
-                        language = safe_value(language, "Python")
-                        task = safe_value(task, "hello world program")
-                        prompt = generate_prompt("code_generate", language=language, task=task)
-                        result = generate_text(prompt, max_tokens, temp, top_p_val)
-                        # Try to extract code block if markdown is used
-                        if "```" in result:
+                            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:
-                                 code_block = parts[1]
-                                 # Remove potential language hint (e.g., ```python)
-                                 if '\n' in code_block:
-                                     first_line, rest_of_code = code_block.split('\n', 1)
-                                     if first_line.strip().lower() == language.lower():
-                                         return rest_of_code.strip()
-                                     else:
-                                          # Language hint might be missing or different
-                                          return code_block.strip()
-                                 else:
-                                      # Code block might be single line without language hint after ```
-                                      return code_block.strip()
-                        # Return full result if no markdown block found or extraction failed
-                        return result.strip()
-
-
-                    # Update output language display based on dropdown
-                    def update_code_language_display(lang):
-                         return gr.Code.update(language=lang_map.get(lang, "plaintext")) # Use update method
-
-                    code_language_gen.change(update_code_language_display, inputs=code_language_gen, outputs=code_output, queue=False)
-                    code_gen_btn.click(code_gen_handler, inputs=[code_language_gen, code_task, *code_gen_params], outputs=code_output)
-
-                    gr.Examples(
-                        examples=[
-                             ["JavaScript", "function to validate an email address using regex", 768, 0.6, 0.9],
-                             ["SQL", "query to select users older than 30 from a 'users' table", 512, 0.5, 0.8],
-                             ["HTML", "basic structure for a personal portfolio website", 1024, 0.7, 0.9],
-                        ],
-                        inputs=[code_language_gen, code_task, *code_gen_params[:3]],
-                        outputs=code_output,
-                        label="Try these examples...",
-                    )
-
-                # --- Code Explanation ---
-                with gr.TabItem("Explain Code", id="subtab_code_explain"):
-                     with gr.Row(equal_height=False):
-                         # Input Column
-                         with gr.Column(scale=1):
-                             gr.Markdown("### Code Explanation Setup")
-                             code_language_explain = gr.Dropdown(list(lang_map.keys()), label="Code Language (for context)", value="Python", elem_id="code_language_explain")
-                             code_to_explain = gr.Code(label="Paste Code Here", language="python", lines=15, elem_id="code_to_explain")
-                             explain_code_params = create_parameter_ui()
-                             gr.Spacer(height=15)
-                             explain_code_btn = gr.Button("Explain Code", variant="primary", elem_id="explain_code_btn")
-
-                         # Output Column
-                         with gr.Column(scale=1):
-                             gr.Markdown("### Explanation")
-                             code_explanation = gr.Textbox(label="Result", lines=25, interactive=False, elem_id="code_explanation", show_copy_button=True)
-
-                     # Update code input language display
-                     def update_explain_language_display(lang):
-                          return gr.Code.update(language=lang_map.get(lang, "plaintext"))
-                     code_language_explain.change(update_explain_language_display, inputs=code_language_explain, outputs=code_to_explain, queue=False)
-
-                     # Handler
-                     def explain_code_handler(language, code, max_tokens, temp, top_p_val):
-                         code_content = safe_value(code['code'] if isinstance(code, dict) else code, "# Add code here") # Handle potential dict input from gr.Code
-                         language = safe_value(language, "code") # Use selected language in prompt
-                         prompt = generate_prompt("code_explain", language=language, code=code_content)
-                         return generate_text(prompt, max_tokens, temp, top_p_val)
-
-                     explain_code_btn.click(explain_code_handler, inputs=[code_language_explain, code_to_explain, *explain_code_params], outputs=code_explanation)
-
-                # --- Code Debugging ---
-                with gr.TabItem("Debug Code", id="subtab_code_debug"):
-                     with gr.Row(equal_height=False):
-                         # Input Column
+                                 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("### Code Debugging Setup")
-                             code_language_debug = gr.Dropdown(list(lang_map.keys()), label="Code Language (for context)", value="Python", elem_id="code_language_debug")
-                             code_to_debug = gr.Code(
-                                 label="Paste Potentially Buggy Code Here",
-                                 language="python",
-                                 lines=15,
-                                 value="def calculate_average(numbers):\n  sum = 0\n  for n in numbers:\n    sum += n\n  # Bug: potential division by zero if numbers is empty\n  return sum / len(numbers)", # Example with potential bug
-                                 elem_id="code_to_debug"
-                             )
-                             debug_code_params = create_parameter_ui()
-                             gr.Spacer(height=15)
-                             debug_code_btn = gr.Button("Debug Code", variant="primary", elem_id="debug_code_btn")
-
-                         # Output Column
+                            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("### Debugging Analysis & Fix")
-                             debug_result = gr.Textbox(label="Result", lines=25, interactive=False, elem_id="debug_result", show_copy_button=True)
+                            gr.Markdown("#### Explanation")
+                            code_explanation = gr.Textbox(label="Result", lines=25, interactive=False, show_copy_button=True)
 
-                     # Update code input language display
-                     def update_debug_language_display(lang):
-                          return gr.Code.update(language=lang_map.get(lang, "plaintext"))
-                     code_language_debug.change(update_debug_language_display, inputs=code_language_debug, outputs=code_to_debug, queue=False)
+                     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)
 
-                     # Handler
-                     def debug_code_handler(language, code, max_tokens, temp, top_p_val):
-                         code_content = safe_value(code['code'] if isinstance(code, dict) else code, "# Add potentially buggy code here")
-                         language = safe_value(language, "code")
-                         prompt = generate_prompt("code_debug", language=language, code=code_content)
-                         return generate_text(prompt, max_tokens, temp, top_p_val)
 
-                     debug_code_btn.click(debug_code_handler, inputs=[code_language_debug, code_to_debug, *debug_code_params], outputs=debug_result)
-
-
-        # --- Text Comprehension Tab (Summarize, QA, Translate) ---
-        with gr.TabItem("📚 Comprehension", id="tab_comprehension"):
-            with gr.Tabs() as comprehension_tabs:
-
-                # --- Summarization ---
-                with gr.TabItem("Summarize", id="subtab_summarize"):
-                    with gr.Row(equal_height=False):
-                        # Input Column
+                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("### Summarization Setup")
-                            summarize_text = gr.Textbox(label="Text to Summarize", placeholder="Paste long text here...", lines=15, elem_id="summarize_text")
+                            gr.Markdown("#### Setup")
+                            summarize_text = gr.Textbox(label="Text to Summarize", lines=15, placeholder="Paste long text...")
                             summarize_params = create_parameter_ui()
-                            gr.Spacer(height=15)
-                            summarize_btn = gr.Button("Summarize Text", variant="primary", elem_id="summarize_btn")
-                        # Output Column
+                            # 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, elem_id="summary_output", show_copy_button=True)
-
-                    # Handler
-                    def summarize_handler(text, max_tokens, temp, top_p_val):
-                        text = safe_value(text, "Please provide text to summarize.")
-                        # Use shorter max_tokens default for summary, but ensure it's reasonable
-                        max_tokens = min(max(int(max_tokens), 64), 512) # Ensure int, set min/max bounds
-                        prompt = generate_prompt("summarize", text=text)
-                        return generate_text(prompt, max_tokens, temp, top_p_val)
-
-                    summarize_btn.click(summarize_handler, inputs=[summarize_text, *summarize_params], outputs=summary_output)
-
-                # --- Question Answering ---
-                with gr.TabItem("Q & A", id="subtab_qa"):
-                    with gr.Row(equal_height=False):
-                        # Input Column
+                            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("### Question Answering Setup")
-                            qa_text = gr.Textbox(label="Context Text", placeholder="Paste the text containing the answer...", lines=10, elem_id="qa_text")
-                            qa_question = gr.Textbox(label="Question", placeholder="Ask a question based on the text above...", elem_id="qa_question")
+                            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()
-                            gr.Spacer(height=15)
-                            qa_btn = gr.Button("Get Answer", variant="primary", elem_id="qa_btn")
-                        # Output Column
+                            # 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, elem_id="qa_output", show_copy_button=True)
-
-                    # Handler
-                    def qa_handler(text, question, max_tokens, temp, top_p_val):
-                        text = safe_value(text, "Please provide context text.")
-                        question = safe_value(question, "What is the main point?")
-                        # Use shorter max_tokens default for QA
-                        max_tokens = min(max(int(max_tokens), 32), 256) # Ensure int, set min/max bounds
-                        prompt = generate_prompt("qa", text=text, question=question)
-                        return generate_text(prompt, max_tokens, temp, top_p_val)
-
-                    qa_btn.click(qa_handler, inputs=[qa_text, qa_question, *qa_params], outputs=qa_output)
-
-                # --- Translation ---
-                with gr.TabItem("Translate", id="subtab_translate"):
-                     with gr.Row(equal_height=False):
-                         # Input Column
+                            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("### Translation Setup")
-                             translate_text = gr.Textbox(label="Text to Translate", placeholder="Enter text in any language...", lines=8, elem_id="translate_text")
-                             target_lang = gr.Dropdown(
-                                 ["French", "Spanish", "German", "Japanese", "Chinese", "Russian", "Arabic", "Hindi", "Portuguese", "Italian"],
-                                 label="Translate To", value="French", elem_id="target_lang"
-                             )
+                             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()
-                             gr.Spacer(height=15)
-                             translate_btn = gr.Button("Translate Text", variant="primary", elem_id="translate_btn")
-                         # Output Column
+                             # 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, elem_id="translation_output", show_copy_button=True)
-
-                     # Handler
-                     def translate_handler(text, lang, max_tokens, temp, top_p_val):
-                         text = safe_value(text, "Please enter text to translate.")
-                         lang = safe_value(lang, "French")
-                         prompt = generate_prompt("translate", text=text, target_lang=lang)
-                         # Translation length is often similar to input, allow reasonable max_tokens
-                         max_tokens = max(int(max_tokens), 64) # Ensure int, set min bound
-                         return generate_text(prompt, max_tokens, temp, top_p_val)
-
-                     translate_btn.click(translate_handler, inputs=[translate_text, target_lang, *translate_params], outputs=translation_output)
-
-
-        # --- More Tasks Tab (Consolidating less common ones) ---
-        with gr.TabItem("🛠️ More Tasks", id="tab_more"):
-             with gr.Tabs() as more_tasks_tabs:
-
-                 # --- Content Creation ---
-                 with gr.TabItem("Content Creation", id="tab_content"):
-                     with gr.Row(equal_height=False):
+                             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("### Content 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", elem_id="content_type")
-                            content_topic = gr.Textbox(label="Topic", placeholder="e.g., benefits of remote work", value="sustainable travel tips", elem_id="content_topic", lines=2)
-                            content_audience = gr.Textbox(label="Target Audience", placeholder="e.g., small business owners", value="eco-conscious millennials", elem_id="content_audience")
-                            content_params = create_parameter_ui()
-                            gr.Spacer(height=15)
-                            content_btn = gr.Button("Generate Content", variant="primary", elem_id="content_btn")
+                            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 Content")
-                            content_output = gr.Textbox(label="Result", lines=20, interactive=False, elem_id="content_output", show_copy_button=True)
-
-                     def content_handler(c_type, topic, audience, max_tok, temp, top_p_val):
-                         c_type = safe_value(c_type, "text")
-                         topic = safe_value(topic, "a given subject")
-                         audience = safe_value(audience, "a general audience")
-                         prompt = generate_prompt("content_creation", content_type=c_type, topic=topic, audience=audience)
-                         return generate_text(prompt, max_tok, temp, top_p_val)
-                     content_btn.click(content_handler, inputs=[content_type, content_topic, content_audience, *content_params], outputs=content_output)
-
-
-                 # --- Email Drafting ---
-                 with gr.TabItem("Email Drafting", id="tab_email"):
-                     with gr.Row(equal_height=False):
+                            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("### Email Setup")
-                            email_type = gr.Dropdown(["job inquiry", "meeting request", "follow-up", "thank you note", "customer support response", "sales outreach"], label="Email Type", value="meeting request", elem_id="email_type")
-                            email_context = gr.Textbox(label="Key Points / Context", placeholder="Provide bullet points or context...", lines=5, value="Request a brief meeting next week to discuss project X. Suggest Tuesday or Wednesday afternoon. Mention attached agenda.", elem_id="email_context")
-                            email_params = create_parameter_ui()
-                            gr.Spacer(height=15)
-                            email_btn = gr.Button("Generate Email Draft", variant="primary", elem_id="email_btn")
+                            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("### Generated Email")
-                            email_output = gr.Textbox(label="Result", lines=20, interactive=False, elem_id="email_output", show_copy_button=True)
-
-                     def email_handler(e_type, context, max_tok, temp, top_p_val):
-                         e_type = safe_value(e_type, "professional")
-                         context = safe_value(context, "the following points")
-                         prompt = generate_prompt("email_draft", email_type=e_type, context=context)
-                         return generate_text(prompt, max_tok, temp, top_p_val)
-                     email_btn.click(email_handler, inputs=[email_type, email_context, *email_params], outputs=email_output)
-
-
-                 # --- Document Editing ---
-                 with gr.TabItem("Document Editing", id="tab_edit"):
-                    with gr.Row(equal_height=False):
+                            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("### Editing Setup")
-                            edit_text = gr.Textbox(label="Text to Edit", placeholder="Paste text here...", lines=10, elem_id="edit_text")
-                            edit_type = gr.Dropdown(["improve clarity", "fix grammar & spelling", "make more concise", "make more formal", "make more casual", "simplify language"], label="Improve For", value="improve clarity", elem_id="edit_type")
-                            edit_params = create_parameter_ui()
-                            gr.Spacer(height=15)
-                            edit_btn = gr.Button("Edit Text", variant="primary", elem_id="edit_btn")
+                            gr.Markdown("#### Setup")
+                            extract_text = gr.Textbox(label="Source Text", lines=10, value="Order #123 by Jane (j@ex.com). 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("### Edited Text")
-                            edit_output = gr.Textbox(label="Result", lines=10, interactive=False, elem_id="edit_output", show_copy_button=True)
-
-                    def edit_handler(text, e_type, max_tok, temp, top_p_val):
-                        text = safe_value(text, "Provide text to edit.")
-                        e_type = safe_value(e_type, "clarity and grammar")
-                        prompt = generate_prompt("document_edit", text=text, edit_type=e_type)
-                        # Editing might expand text, give it reasonable token count based on input + max_new
-                        input_tokens_estimate = len(text.split()) # Rough estimate
-                        max_tok = max(int(max_tok), input_tokens_estimate + 64) # Ensure enough room
-                        return generate_text(prompt, max_tok, temp, top_p_val)
-                    edit_btn.click(edit_handler, inputs=[edit_text, edit_type, *edit_params], outputs=edit_output)
-
-
-                 # --- Classification ---
-                 with gr.TabItem("Classification", id="tab_classify"):
-                     with gr.Row(equal_height=False):
-                         with gr.Column(scale=1):
-                             gr.Markdown("### Classification Setup")
-                             classify_text = gr.Textbox(label="Text to Classify", placeholder="Enter text...", lines=8, value="This new sci-fi movie explores themes of AI consciousness and interstellar travel.")
-                             classify_categories = gr.Textbox(label="Categories (comma-separated)", placeholder="e.g., positive, negative, neutral", value="Technology, Entertainment, Science, Politics, Sports, Health")
-                             classify_params = create_parameter_ui()
-                             gr.Spacer(height=15)
-                             classify_btn = gr.Button("Classify Text", variant="primary")
-                         with gr.Column(scale=1):
-                             gr.Markdown("### Classification Result")
-                             classify_output = gr.Textbox(label="Predicted Category", lines=2, interactive=False, show_copy_button=True)
-
-                     def classify_handler(text, cats, max_tok, temp, top_p_val):
-                         text = safe_value(text, "Text to classify needed.")
-                         cats = safe_value(cats, "category1, category2")
-                         # Classification usually needs short output
-                         max_tok = min(max(int(max_tok), 16), 128) # Ensure int, constrain tightly
-                         prompt = generate_prompt("classify", text=text, categories=cats)
-                         # Often the model just outputs the category, so we might not need the prompt structure removal
-                         raw_output = generate_text(prompt, max_tok, temp, top_p_val)
-                         # Post-process to get just the category if possible
-                         lines = raw_output.split('\n')
-                         if lines:
-                              last_line = lines[-1].strip()
-                              # Check if the last line seems like one of the categories
-                              possible_cats = [c.strip().lower() for c in cats.split(',')]
-                              if last_line.lower() in possible_cats:
-                                   return last_line
-                         # Fallback to raw output
-                         return raw_output
-
-                     classify_btn.click(classify_handler, inputs=[classify_text, classify_categories, *classify_params], outputs=classify_output)
-
-
-                 # --- Data Extraction ---
-                 with gr.TabItem("Data Extraction", id="tab_extract"):
-                      with gr.Row(equal_height=False):
-                         with gr.Column(scale=1):
-                             gr.Markdown("### Extraction Setup")
-                             extract_text = gr.Textbox(label="Source Text", placeholder="Paste text containing data...", lines=10, value="Order #12345 placed on 2024-03-15 by Jane Doe (jane.d@email.com). Total amount: $99.95. Shipping to 123 Main St, Anytown, USA.")
-                             extract_data_points = gr.Textbox(label="Data to Extract (comma-separated)", placeholder="e.g., name, email, order number", value="order number, date, customer name, email, total amount, address")
-                             extract_params = create_parameter_ui()
-                             gr.Spacer(height=15)
-                             extract_btn = gr.Button("Extract Data", variant="primary")
-                         with gr.Column(scale=1):
-                             gr.Markdown("### Extracted Data")
-                             extract_output = gr.Textbox(label="Result (e.g., JSON or key-value pairs)", lines=10, interactive=False, show_copy_button=True)
+                            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)
 
-                      def extract_handler(text, points, max_tok, temp, top_p_val):
-                          text = safe_value(text, "Provide text for extraction.")
-                          points = safe_value(points, "key information")
-                          prompt = generate_prompt("data_extract", text=text, data_points=points)
-                          return generate_text(prompt, max_tok, temp, top_p_val)
-                      extract_btn.click(extract_handler, inputs=[extract_text, extract_data_points, *extract_params], outputs=extract_output)
 
+    # --- Authentication Handler & Footer ---
+    footer_status = gr.Markdown(f"...", elem_id="footer-status-md") # Placeholder
 
-    # Define authentication handler AFTER tabs is defined
     def handle_auth(token):
-        # Show loading message immediately
-        yield "⏳ Authenticating and loading model... Please wait.", gr.Tabs.update(visible=False)
-        # Call the actual model loading function
+        yield "⏳ Authenticating & loading model...", gr.Tabs.update(visible=False)
         status_message, tabs_update = load_model(token)
         yield status_message, tabs_update
 
-    # Link button click to the handler
-    auth_button.click(
-        fn=handle_auth,
-        inputs=[hf_token],
-        outputs=[auth_status, tabs], # Use the defined 'tabs' variable here
-        queue=True # Run in queue for potentially long operation
-    )
-
-    # --- Footer ---
-    footer_status = gr.Markdown( # Use a separate Markdown for dynamic updates
-        f"""
-        ---
-        <div style="text-align: center; font-size: 0.9em; color: #777;">
-            <p>Powered by Google's Gemma models via Hugging Face 🤗 Transformers & Gradio.</p>
-            <p>Remember to review generated content. Model outputs may be inaccurate or incomplete.</p>
-            <p>Model Loaded: <strong>{loaded_model_name if model_loaded else 'None'}</strong></p>
-        </div>
-        """
-    )
+    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>""")
 
-    # Update footer when authentication status changes
-    def update_footer_status(status_text):
-         # You could parse status_text, but easier to just use global state here
-         return gr.Markdown.update(value=f"""
-            ---
-            <div style="text-align: center; font-size: 0.9em; color: #777;">
-                <p>Powered by Google's Gemma models via Hugging Face 🤗 Transformers & Gradio.</p>
-                <p>Remember to review generated content. Model outputs may be inaccurate or incomplete.</p>
-                <p>Model Loaded: <strong>{loaded_model_name if model_loaded else 'None'}</strong></p>
-            </div>
-            """)
-    auth_status.change(fn=update_footer_status, inputs=auth_status, outputs=footer_status, queue=False)
+    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 ---
-# Allow built-in theme switching
-# Use queue() to handle multiple requests better
 demo.launch(share=False, allowed_themes=["light", "dark"])