ONS

app.py

@@ -25,10 +25,13 @@ def load_model(hf_token):
     """Load the model with the provided token"""
     global global_model, global_tokenizer, model_loaded, loaded_model_name
 
+    # Initially assume tabs should be hidden until successful load
+    initial_tabs_update = gr.Tabs.update(visible=False)
+
     if not hf_token:
         model_loaded = False
         loaded_model_name = "None"
-        return "⚠️ Please enter your Hugging Face token to use the model.", gr.Tabs.update(visible=False)
+        return "⚠️ Please enter your Hugging Face token to use the model.", initial_tabs_update
 
     try:
         # Try different model versions from smallest to largest
@@ -63,7 +66,8 @@ def load_model(hf_token):
                 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
+                    # 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
                     token=current_token
                 )
@@ -72,11 +76,16 @@ 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.", gr.Tabs.update(visible=True)
+                     return f"✅ Fallback model '{model_name}' loaded successfully! Limited capabilities compared to Gemma.", tabs_update
                 else:
-                     return f"✅ Model '{model_name}' loaded successfully!", gr.Tabs.update(visible=True)
+                     return f"✅ Model '{model_name}' loaded successfully!", 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
             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
@@ -94,7 +103,7 @@ def load_model(hf_token):
             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.", gr.Tabs.update(visible=False)
+            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
 
     except Exception as e:
         model_loaded = False
@@ -104,9 +113,9 @@ def load_model(hf_token):
         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.", gr.Tabs.update(visible=False)
+            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
         else:
-            return f"❌ An unexpected error occurred during model loading: {error_msg}", gr.Tabs.update(visible=False)
+            return f"❌ An unexpected error occurred during model loading: {error_msg}", initial_tabs_update
 
 
 def generate_prompt(task_type, **kwargs):
@@ -133,12 +142,14 @@ def generate_prompt(task_type, **kwargs):
     prompt_template = prompts.get(task_type)
     if prompt_template:
         try:
-            # Ensure all required keys are present with defaults if necessary
-            # This prevents KeyError if a kwarg is missing
-            required_keys = [k[1:-1] for k in prompt_template.replace('{',' ').replace('}',' ').split() if '{' in k and '}' in k]
-            final_kwargs = {key: kwargs.get(key, f"[{key}]") for key in required_keys}
-            # Add remaining kwargs that might not be in the template explicitly
+            # 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)
+
             return prompt_template.format(**final_kwargs)
         except KeyError as e:
             print(f"Warning: Missing key for prompt template '{task_type}': {e}")
@@ -166,9 +177,14 @@ def generate_text(prompt, max_new_tokens=1024, temperature=0.7, top_p=0.9):
 
     try:
         # Add role/turn indicators if using an instruction-tuned model
-        if "it" in loaded_model_name.lower() or "chat" in loaded_model_name.lower():
+        # Simple check based on model name conventions
+        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
-             chat_prompt = f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
+             # Using Gemma's specific format if it's a Gemma IT model
+             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"
+             else: # Generic instruction format
+                 chat_prompt = f"User: {prompt}\nAssistant:"
         else:
              # Base models might not need specific turn indicators
              chat_prompt = prompt
@@ -177,13 +193,9 @@ def generate_text(prompt, max_new_tokens=1024, temperature=0.7, top_p=0.9):
         input_length = inputs.input_ids.shape[1]
         print(f"Input token length: {input_length}")
 
-        # Adjust max_length based on input, prevent it from being too small
-        # max_length = max(input_length + 64, input_length + max_new_tokens) # Ensure at least some generation
-        # Use max_new_tokens directly as it's clearer for users
         # Ensure max_new_tokens isn't excessively large for the model context
-        # Gemma 2B/7B context is often 8192, TinyLlama 2048
-        # Let's cap generation length for stability
-        effective_max_new_tokens = min(max_new_tokens, 2048) # Cap generation length
+        # Cap generation length for stability
+        effective_max_new_tokens = min(int(max_new_tokens), 2048) # Cap generation length and ensure int
 
         generation_args = {
             "input_ids": inputs.input_ids,
@@ -192,7 +204,7 @@ def generate_text(prompt, max_new_tokens=1024, temperature=0.7, top_p=0.9):
             "do_sample": True,
             "temperature": float(temperature), # Ensure float
             "top_p": float(top_p),             # Ensure float
-            "pad_token_id": global_tokenizer.eos_token_id # Use EOS token for padding
+            "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
         }
 
         print(f"Generation args: {generation_args}")
@@ -201,7 +213,6 @@ def generate_text(prompt, max_new_tokens=1024, temperature=0.7, top_p=0.9):
         with torch.no_grad(): # Disable gradient calculation for inference
             outputs = global_model.generate(**generation_args)
 
-        # Decode response, skipping special tokens and the prompt
         # Decode only the newly generated tokens
         generated_ids = outputs[0, input_length:]
         generated_text = global_tokenizer.decode(generated_ids, skip_special_tokens=True)
@@ -215,7 +226,13 @@ def generate_text(prompt, max_new_tokens=1024, temperature=0.7, top_p=0.9):
         print(f"Generation error: {error_msg}")
         print(f"Error type: {type(e)}")
         traceback.print_exc()
-        return f"❌ Error during text generation: {error_msg}\n\nPlease check the logs or try adjusting parameters (e.g., reduce Max Tokens)."
+        # 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."
+        else:
+             return f"❌ Error during text generation: {error_msg}\n\nPlease check the logs or try adjusting parameters."
 
 # Create parameters UI component (reusable function)
 def create_parameter_ui():
@@ -237,7 +254,7 @@ def create_parameter_ui():
                 value=0.7,
                 step=0.1,
                 label="Temperature",
-                info="Controls randomness. Lower is more focused, higher is more diverse.",
+                info="Controls randomness. Lower is focused, higher is diverse.",
                 elem_id="temperature_slider"
             )
             top_p = gr.Slider(
@@ -253,7 +270,7 @@ def create_parameter_ui():
 
 # --- Gradio Interface ---
 # Use the soft theme for a clean look, allow light/dark switching
-with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
+with gr.Blocks(theme=gr.themes.Soft(), fill_height=True, title="Gemma Capabilities Demo") as demo:
 
     # Header
     gr.Markdown(
@@ -274,7 +291,9 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
     )
 
     # --- Authentication Section ---
-    with gr.Group(variant="panel"): # Use panel variant for visual grouping
+    # REMOVED variant="panel" from gr.Group for compatibility
+    with gr.Group(): # Use default Group appearance
+        gr.Markdown("### 🔑 Authentication") # Added heading inside group
         with gr.Row():
             with gr.Column(scale=4):
                 hf_token = gr.Textbox(
@@ -282,38 +301,31 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                     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"
                 )
             with gr.Column(scale=1, min_width=150):
-                # Add spacer for alignment if needed, or adjust scale
-                # gr.Spacer(height=10) # Add space above button if needed
                 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
+        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.
+             """
+         )
 
-        # Define authentication flow (simplified)
-        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
-            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, gr.get_component("main_tabs")], # Update status and hide/show main_tabs by element id
-            queue=True # Run in queue for potentially long operation
-        )
 
     # --- Main Content Tabs (Initially Hidden) ---
-    # Use gr.Tabs with visible=False initially
+    # Define the tabs variable here
     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():
+            with gr.Row(equal_height=False): # Allow columns to have different heights if needed
                 # Input Column
                 with gr.Column(scale=1):
                     gr.Markdown("### Configure Task")
@@ -336,12 +348,15 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                     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
+                    # 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=choice == "Creative Writing"),
-                            info_options: gr.update(visible=choice == "Informational Writing"),
-                            custom_prompt_group: gr.update(visible=choice == "Custom Prompt")
+                            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)
 
@@ -353,19 +368,28 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                 # Output Column
                 with gr.Column(scale=1):
                     gr.Markdown("### Generated Output")
-                    text_output = gr.Textbox(label="Result", lines=25, interactive=False, elem_id="text_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})
+                    "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 k, v in kwargs.items():
-                    kwargs[k] = safe_value(v, f"[{k}]")
+
+                # 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.")
+
 
                 final_prompt = generate_prompt(task_type, **kwargs)
                 return generate_text(final_prompt, max_tokens, temp, top_p_val)
@@ -377,21 +401,24 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
             )
 
             # 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],
                 ],
-                inputs=[text_gen_type, style, creative_topic, format_type, info_topic, custom_prompt, *text_gen_params],
+                # 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 # Need to provide function for examples to run
+                #fn=text_generation_handler # fn is deprecated, click event handles execution
             )
 
+
         # --- Brainstorming Tab ---
         with gr.TabItem("🧠 Brainstorming", id="tab_brainstorm"):
-            with gr.Row():
+            with gr.Row(equal_height=False):
                 # Input Column
                 with gr.Column(scale=1):
                     gr.Markdown("### Brainstorming Setup")
@@ -404,7 +431,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                 # 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")
+                    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):
@@ -421,22 +448,24 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                     ["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],
+                inputs=[brainstorm_category, brainstorm_topic, *brainstorm_params[:3]],
                 outputs=brainstorm_output,
                 label="Try these examples...",
-                fn=brainstorm_handler
             )
 
         # --- 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():
+                    with gr.Row(equal_height=False):
                         # Input Column
                         with gr.Column(scale=1):
                              gr.Markdown("### Code Generation Setup")
-                             code_language_gen = gr.Dropdown(["Python", "JavaScript", "Java", "C++", "HTML", "CSS", "SQL", "Bash", "Rust"], label="Language", value="Python", elem_id="code_language_gen")
+                             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)
@@ -445,9 +474,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                         # Output Column
                         with gr.Column(scale=1):
                              gr.Markdown("### Generated Code")
-                             # Use language mapping for syntax highlighting
-                             lang_map = {"Python": "python", "JavaScript": "javascript", "Java": "java", "C++": "cpp", "HTML": "html", "CSS": "css", "SQL": "sql", "Bash": "bash", "Rust": "rust"}
-                             code_output = gr.Code(label="Result", language="python", lines=25, interactive=False, elem_id="code_output")
+                             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):
@@ -457,19 +484,27 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                         result = generate_text(prompt, max_tokens, temp, top_p_val)
                         # Try to extract code block if markdown is used
                         if "```" in result:
-                             code_block = result.split("```")
-                             if len(code_block) > 1:
-                                 # Return content of the first code block, stripping language hint if present
-                                 content = code_block[1]
-                                 if content.lower().startswith(language.lower()):
-                                     content = content[len(language):].lstrip()
-                                 return content
-                        return result # Return full result if no block found
+                             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(language=lang_map.get(lang, "plaintext")) # Update component property
+                         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)
@@ -480,20 +515,18 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                              ["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],
+                        inputs=[code_language_gen, code_task, *code_gen_params[:3]],
                         outputs=code_output,
                         label="Try these examples...",
-                        fn=code_gen_handler
                     )
 
                 # --- Code Explanation ---
                 with gr.TabItem("Explain Code", id="subtab_code_explain"):
-                     with gr.Row():
+                     with gr.Row(equal_height=False):
                          # Input Column
                          with gr.Column(scale=1):
                              gr.Markdown("### Code Explanation Setup")
-                             # Allow user to select language for context, even if not strictly necessary for model
-                             code_language_explain = gr.Dropdown(["Python", "JavaScript", "Java", "C++", "HTML", "CSS", "SQL", "Bash", "Rust", "Other"], label="Code Language (for context)", value="Python", elem_id="code_language_explain")
+                             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)
@@ -502,29 +535,29 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                          # Output Column
                          with gr.Column(scale=1):
                              gr.Markdown("### Explanation")
-                             code_explanation = gr.Textbox(label="Result", lines=25, interactive=False, elem_id="code_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(language=lang_map.get(lang, "plaintext"))
+                          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 = safe_value(code, "# Add code here")
+                         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)
+                         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():
+                     with gr.Row(equal_height=False):
                          # Input Column
                          with gr.Column(scale=1):
                              gr.Markdown("### Code Debugging Setup")
-                             code_language_debug = gr.Dropdown(["Python", "JavaScript", "Java", "C++", "SQL", "Bash", "Other"], label="Code Language (for context)", value="Python", elem_id="code_language_debug")
+                             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",
@@ -539,18 +572,18 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                          # Output Column
                          with gr.Column(scale=1):
                              gr.Markdown("### Debugging Analysis & Fix")
-                             debug_result = gr.Textbox(label="Result", lines=25, interactive=False, elem_id="debug_result")
+                             debug_result = gr.Textbox(label="Result", lines=25, interactive=False, elem_id="debug_result", show_copy_button=True)
 
                      # Update code input language display
                      def update_debug_language_display(lang):
-                          return gr.Code(language=lang_map.get(lang, "plaintext"))
+                          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)
 
                      # Handler
                      def debug_code_handler(language, code, max_tokens, temp, top_p_val):
-                         code = safe_value(code, "# Add potentially buggy code here")
+                         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)
+                         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)
@@ -562,7 +595,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
 
                 # --- Summarization ---
                 with gr.TabItem("Summarize", id="subtab_summarize"):
-                    with gr.Row():
+                    with gr.Row(equal_height=False):
                         # Input Column
                         with gr.Column(scale=1):
                             gr.Markdown("### Summarization Setup")
@@ -573,13 +606,13 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                         # Output Column
                         with gr.Column(scale=1):
                             gr.Markdown("### Summary")
-                            summary_output = gr.Textbox(label="Result", lines=15, interactive=False, elem_id="summary_output")
+                            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
-                        max_tokens = min(max_tokens, 512)
+                        # 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)
 
@@ -587,7 +620,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
 
                 # --- Question Answering ---
                 with gr.TabItem("Q & A", id="subtab_qa"):
-                    with gr.Row():
+                    with gr.Row(equal_height=False):
                         # Input Column
                         with gr.Column(scale=1):
                             gr.Markdown("### Question Answering Setup")
@@ -599,14 +632,14 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                         # Output Column
                         with gr.Column(scale=1):
                             gr.Markdown("### Answer")
-                            qa_output = gr.Textbox(label="Result", lines=10, interactive=False, elem_id="qa_output")
+                            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_tokens, 256)
+                        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)
 
@@ -614,7 +647,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
 
                 # --- Translation ---
                 with gr.TabItem("Translate", id="subtab_translate"):
-                     with gr.Row():
+                     with gr.Row(equal_height=False):
                          # Input Column
                          with gr.Column(scale=1):
                              gr.Markdown("### Translation Setup")
@@ -629,13 +662,15 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                          # Output Column
                          with gr.Column(scale=1):
                              gr.Markdown("### Translation")
-                             translation_output = gr.Textbox(label="Result", lines=8, interactive=False, elem_id="translation_output")
+                             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)
@@ -647,7 +682,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
 
                  # --- Content Creation ---
                  with gr.TabItem("Content Creation", id="tab_content"):
-                     with gr.Row():
+                     with gr.Row(equal_height=False):
                         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")
@@ -658,7 +693,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                             content_btn = gr.Button("Generate Content", variant="primary", elem_id="content_btn")
                         with gr.Column(scale=1):
                             gr.Markdown("### Generated Content")
-                            content_output = gr.Textbox(label="Result", lines=20, interactive=False, elem_id="content_output")
+                            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")
@@ -671,7 +706,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
 
                  # --- Email Drafting ---
                  with gr.TabItem("Email Drafting", id="tab_email"):
-                     with gr.Row():
+                     with gr.Row(equal_height=False):
                         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")
@@ -681,7 +716,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                             email_btn = gr.Button("Generate Email Draft", variant="primary", elem_id="email_btn")
                         with gr.Column(scale=1):
                             gr.Markdown("### Generated Email")
-                            email_output = gr.Textbox(label="Result", lines=20, interactive=False, elem_id="email_output")
+                            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")
@@ -693,7 +728,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
 
                  # --- Document Editing ---
                  with gr.TabItem("Document Editing", id="tab_edit"):
-                    with gr.Row():
+                    with gr.Row(equal_height=False):
                         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")
@@ -703,20 +738,22 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                             edit_btn = gr.Button("Edit Text", variant="primary", elem_id="edit_btn")
                         with gr.Column(scale=1):
                             gr.Markdown("### Edited Text")
-                            edit_output = gr.Textbox(label="Result", lines=10, interactive=False, elem_id="edit_output")
+                            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)
-                        # Generate potentially longer text as editing might expand it
-                        return generate_text(prompt, max(max_tok, len(text.split()) + 128), temp, top_p_val)
+                        # 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():
+                     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.")
@@ -726,21 +763,33 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                              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)
+                             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_tok, 64)
+                         max_tok = min(max(int(max_tok), 16), 128) # Ensure int, constrain tightly
                          prompt = generate_prompt("classify", text=text, categories=cats)
-                         return generate_text(prompt, max_tok, temp, top_p_val)
+                         # 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():
+                      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.")
@@ -750,7 +799,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                              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)
+                             extract_output = gr.Textbox(label="Result (e.g., JSON or key-value pairs)", lines=10, interactive=False, show_copy_button=True)
 
                       def extract_handler(text, points, max_tok, temp, top_p_val):
                           text = safe_value(text, "Provide text for extraction.")
@@ -760,27 +809,49 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
                       extract_btn.click(extract_handler, inputs=[extract_text, extract_data_points, *extract_params], outputs=extract_output)
 
 
+    # 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
+        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 ---
-    gr.Markdown(
-        """
+    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: <span id="footer-model-name">None</span></p> <!-- Placeholder for dynamic update -->
+            <p>Model Loaded: <strong>{loaded_model_name if model_loaded else 'None'}</strong></p>
         </div>
         """
     )
-    # Add JS to update the footer model name (optional, but nice)
-    auth_status.change(lambda status: gr.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>
-        """), inputs=auth_status, outputs=gr.Markdown(elem_id="footer-model-display")) # Need a dummy output or separate MD for this
+
+    # 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)
+
 
 # --- Launch App ---
 # Allow built-in theme switching
+# Use queue() to handle multiple requests better
 demo.launch(share=False, allowed_themes=["light", "dark"])