Refactor ui.py

README.md

@@ -8,6 +8,8 @@ sdk_version: 6.0.1
 app_file: app.py
 pinned: false
 hf_oauth: true
+hf_oauth_scopes:
+ - write-repos
 license: apache-2.0
 ---
 

ui.py

@@ -1,52 +1,74 @@
 import gradio as gr
+from typing import Optional, Tuple, Generator, List, Any
 from config import AppConfig
 from engine import FunctionGemmaEngine
-from typing import Optional
 
-def build_interface() -> gr.Blocks:
-    
-    # --- State Management Wrappers ---
-    
-    def init_session(profile: Optional[gr.OAuthProfile] = None):
+# --- Controller / Logic Layer ---
+
+class UIController:
+    """
+    Handles the business logic and interaction with the Engine.
+    Stateless methods that operate on the passed Engine state.
+    """
+
+    @staticmethod
+    def init_session(profile: Optional[gr.OAuthProfile] = None) -> Tuple[Any, ...]:
         config = AppConfig()
         new_engine = FunctionGemmaEngine(config)
-
         username = profile.username if profile else None
-        repo_update, push_update = update_hub_interactive(new_engine, username)
+        username="bebechien"
+        
+        # Calculate initial interactivity state
+        repo_update, push_update, zip_update = UIController.update_hub_interactive(new_engine, username)
 
         return (
             new_engine,
             new_engine.get_tools_json(),
             new_engine.config.MODEL_NAME,
             f"Ready. (Session {new_engine.session_id})",
-            repo_update, push_update, username
+            repo_update, 
+            push_update,
+            zip_update, 
+            username
         )
 
-    def run_training_wrapper(engine, epochs, lr, test_size, shuffle, model_name):
+    @staticmethod
+    def run_training(engine: FunctionGemmaEngine, epochs: int, lr: float, 
+                     test_size: float, shuffle: bool, model_name: str) -> Generator:
+        if not engine:
+            yield "⚠️ Engine not initialized.", None
+            return
+            
         engine.config.MODEL_NAME = model_name.strip()
         yield from engine.run_training_pipeline(epochs, lr, test_size, shuffle)
 
-    def handle_reset(engine, model_name):
+    @staticmethod
+    def handle_reset(engine: FunctionGemmaEngine, model_name: str) -> str:
         engine.config.MODEL_NAME = model_name.strip()
         return engine.refresh_model()
 
-    def update_tools_wrapper(engine, json_val):
+    @staticmethod
+    def update_tools(engine: FunctionGemmaEngine, json_val: str) -> str:
         return engine.update_tools(json_val)
 
-    def import_file_wrapper(engine, file_obj):
+    @staticmethod
+    def import_file(engine: FunctionGemmaEngine, file_obj: Any) -> str:
         return engine.load_csv(file_obj)
 
-    def stop_wrapper(engine):
+    @staticmethod
+    def stop_process(engine: FunctionGemmaEngine) -> str:
         engine.trigger_stop()
         return "Stopping..."
 
-    def zip_wrapper(engine):
+    @staticmethod
+    def zip_model(engine: FunctionGemmaEngine) -> Any:
         path = engine.get_zip_path()
         if path:
             return gr.update(value=path, visible=True)
         return gr.update(value=None, visible=False)
 
-    def upload_wrapper(engine, repo_name, oauth_token: gr.OAuthToken | None):
+    @staticmethod
+    def upload_model(engine: FunctionGemmaEngine, repo_name: str, oauth_token: Optional[gr.OAuthToken]) -> str:
         if oauth_token is None:
             return "❌ Error: You must log in (top right) to upload models."
         if not repo_name:
@@ -57,258 +79,252 @@ def build_interface() -> gr.Blocks:
             oauth_token=oauth_token.token, 
         )
 
-    def update_repo_preview(username, repo_name):
-        """Updates the markdown preview to show 'username/repo_name'."""
+    @staticmethod
+    def update_repo_preview(username: Optional[str], repo_name: str) -> str:
         if not username:
             return "⚠️ Sign in to see the target repository path."
-    
         clean_repo = repo_name.strip() if repo_name else "..."
         return f"Target Repository: **`{username}/{clean_repo}`**"
         
-    def update_hub_interactive(engine, username: Optional[str] = None):
+    @staticmethod
+    def update_hub_interactive(engine: Optional[FunctionGemmaEngine], username: Optional[str] = None):
         is_logged_in = username is not None
-        has_model_tuned = engine is not None and engine.has_model_tuned
+        has_model_tuned = engine is not None and getattr(engine, 'has_model_tuned', False)
         
-        return gr.update(interactive=is_logged_in), gr.update(interactive=is_logged_in and has_model_tuned)
+        return (
+            gr.update(interactive=is_logged_in), 
+            gr.update(interactive=is_logged_in and has_model_tuned),
+            gr.update(interactive=has_model_tuned)
+        )
+
+# --- View / Layout Layer ---
 
-    # --- UI Layout ---
+def _render_header():
+    with gr.Column():
+        gr.Markdown("# 🤖 FunctionGemma Modkit: Fine-Tuning")
+        gr.Markdown("Fine-tune FunctionGemma to understand your custom functions.<br>"
+                    "See [README](https://huggingface.co/spaces/google/functiongemma-modkit/blob/main/README.md) for more details.")
+        gr.Markdown("(Optional) Sign in to Hugging Face if you plan to push your fine-tuned model to the Hub later (3. Export).")
+        with gr.Row():
+            gr.LoginButton(value="Sign in with Hugging Face")
+            with gr.Column(scale=3):
+                gr.Markdown("")
+
+def _render_dataset_tab(engine_state):
+    with gr.TabItem("1. Preparing Dataset"):
+        gr.Markdown("### 🛠️ Tool Schema & Data Import")
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("**Step 1: Define Functions**<br>Edit the JSON schema below to define the tools the model should learn.")
+                tools_editor = gr.Code(language="json", label="Tool Definitions (JSON Schema)", lines=15)
+                update_tools_btn = gr.Button("💾 Update Tool Schema")
+                tools_status = gr.Markdown("")
+
+            with gr.Column(scale=1):
+                gr.Markdown("**Step 2: Upload Data (Optional)**<br>To train on your own data, upload a CSV file to replace the [default dataset](https://huggingface.co/datasets/bebechien/SimpleToolCalling).")
+                gr.Markdown("**Example CSV Row:** No header required.<br>Format: `[User Prompt, Tool Name, Tool Args JSON]`\n```csv\n\"What is the weather in London?\", \"get_weather\", \"{\"\"location\"\": \"\"London, UK\"\"}\"\n```")
+                import_file = gr.File(label="Upload Dataset (.csv)", file_types=[".csv"], height=100)
+                import_status = gr.Markdown("")
+    
+    # Return controls needed for wiring
+    return {
+        "tools_editor": tools_editor,
+        "update_tools_btn": update_tools_btn,
+        "tools_status": tools_status,
+        "import_file": import_file,
+        "import_status": import_status
+    }
+
+def _render_training_tab(engine_state):
+    with gr.TabItem("2. Training"):
+        gr.Markdown("### 🚀 Fine-Tuning Configuration")
+        with gr.Group():
+            gr.Markdown("**Hyperparameters**")
+            with gr.Row():
+                default_models = AppConfig().AVAILABLE_MODELS
+                param_model = gr.Dropdown(
+                    choices=default_models, allow_custom_value=True, label="Base Model", info="Select a preset OR type a custom Hugging Face model ID (e.g. 'google/functiongemma-270m-it')", interactive=True
+                )
+                param_epochs = gr.Slider(1, 20, value=5, step=1, label="Epochs", info="Total training passes")
+            with gr.Row():
+                param_lr = gr.Number(value=5e-5, label="Learning Rate", info="e.g. 5e-5")
+                param_test_size = gr.Slider(0.1, 0.9, value=0.2, step=0.05, label="Test Split", info="Validation ratio (0.2 = 20%)")
+                param_shuffle = gr.Checkbox(value=True, label="Shuffle Data", info="Randomize before split")
+
+        with gr.Row():
+            run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary", scale=1)
+            stop_training_btn = gr.Button("🛑 Stop", variant="stop", visible=False, scale=1)
+            clear_reload_btn = gr.Button("🔄 Reload Model & Reset Data", variant="secondary", scale=1)
+
+        with gr.Row():
+            output_display = gr.Textbox(lines=20, label="Logs", value="Initializing...", interactive=False, autoscroll=True)
+            loss_plot = gr.Plot(label="Training Metrics")
+
+    return {
+        "params": [param_epochs, param_lr, param_test_size, param_shuffle, param_model],
+        "buttons": [run_training_btn, stop_training_btn, clear_reload_btn],
+        "outputs": [output_display, loss_plot],
+        "model_input": param_model # specifically needed for initialization
+    }
+
+def _render_export_tab(engine_state, username_state):
+    with gr.TabItem("3. Export"):
+        gr.Markdown("### 📦 Export Trained Model")
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("#### Option A: Download ZIP")
+                gr.Markdown("Download the model weights locally.")
+                zip_btn = gr.Button("⬇️ Prepare Model ZIP", variant="secondary", interactive=False)
+                download_file = gr.File(label="Download Archive", interactive=False)
+            
+            with gr.Column():
+                gr.Markdown("#### Option B: Save to Hugging Face Hub")
+                gr.Markdown("Publish your fine-tuned model to your personal Hugging Face account.")
+                repo_name_input = gr.Textbox(
+                    label="Target Repository Name", value="my-functiongemma-v1", placeholder="e.g., my-functiongemma-v1", interactive=False
+                )
+                push_to_hub_btn = gr.Button("Save to Hugging Face Hub", variant="secondary", interactive=False)
+                repo_id_preview = gr.Markdown("Target Repository: (Waiting for input...)")
+                upload_status = gr.Markdown("")
+
+    return {
+        "zip_controls": [zip_btn, download_file],
+        "hub_controls": [repo_name_input, push_to_hub_btn, repo_id_preview, upload_status]
+    }
+
+# --- Main Build Function ---
+
+def build_interface() -> gr.Blocks:
     with gr.Blocks(title="FunctionGemma Modkit") as demo:
         engine_state = gr.State()
         username_state = gr.State()
 
-        with gr.Column():
-            gr.Markdown("# 🤖 FunctionGemma Modkit: Fine-Tuning")
-            gr.Markdown("Fine-tune FunctionGemma to understand your custom functions.<br>See [README](https://huggingface.co/spaces/google/functiongemma-modkit/blob/main/README.md) for more details.")
-            gr.Markdown("(Optional) Sign in to Hugging Face if you plan to push your fine-tuned model to the Hub later (3. Export).")
-            with gr.Row():
-                gr.LoginButton(value="Sign in with Hugging Face")
-                with gr.Column(scale=3):
-                    gr.Markdown("")
+        _render_header()
         
         with gr.Tabs():
-            
-            # --- TAB 1: PREPARING DATASET ---
-            with gr.TabItem("1. Preparing Dataset"):
-                gr.Markdown("### 🛠️ Tool Schema & Data Import")
-                
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        gr.Markdown("**Step 1: Define Functions**<br>Edit the JSON schema below to define the tools the model should learn.")
-                        tools_editor = gr.Code(
-                            language="json", 
-                            label="Tool Definitions (JSON Schema)",
-                            lines=15
-                        )
-                        update_tools_btn = gr.Button("💾 Update Tool Schema")
-                        tools_status = gr.Markdown("")
-
-                    with gr.Column(scale=1):
-                        gr.Markdown("**Step 2: Upload Data (Optional)**<br>To train on your own data, upload a CSV file to replace the [default dataset](https://huggingface.co/datasets/bebechien/SimpleToolCalling).")
-                        gr.Markdown("**Example CSV Row:** No header required.<br>Format: `[User Prompt, Tool Name, Tool Args JSON]`\n```csv\n\"What is the weather in London?\", \"get_weather\", \"{\"\"location\"\": \"\"London, UK\"\"}\"\n```")
-                        import_file = gr.File(
-                            label="Upload Dataset (.csv)", 
-                            file_types=[".csv"], 
-                            height=100
-                        )
-                        import_status = gr.Markdown("")
-
-            # --- TAB 2: TRAINING ---
-            with gr.TabItem("2. Training"):
-                gr.Markdown("### 🚀 Fine-Tuning Configuration")
-                
-                with gr.Group():
-                    gr.Markdown("**Hyperparameters**")
-                    with gr.Row():
-                        default_models = AppConfig().AVAILABLE_MODELS
-                        param_model = gr.Dropdown(
-                            choices=default_models,
-                            allow_custom_value=True, 
-                            label="Base Model",
-                            info="Select a preset OR type a custom Hugging Face model ID (e.g. 'google/functiongemma-270m-it')",
-                            interactive=True
-                        )
-                        param_epochs = gr.Slider(
-                            minimum=1, maximum=20, value=5, step=1, 
-                            label="Epochs", info="Total training passes"
-                        )
-                    with gr.Row():
-                        param_lr = gr.Number(
-                            value=5e-5, 
-                            label="Learning Rate", 
-                            info="e.g. 5e-5"
-                        )
-                        param_test_size = gr.Slider(
-                            minimum=0.1, maximum=0.9, value=0.2, step=0.05, 
-                            label="Test Split", info="Validation ratio (0.2 = 20%)"
-                        )
-                        param_shuffle = gr.Checkbox(
-                            value=True,
-                            label="Shuffle Data",
-                            info="Randomize before split"
-                        )
-
-                with gr.Row():
-                    run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary", scale=1)
-                    stop_training_btn = gr.Button("🛑 Stop", variant="stop", visible=False, scale=1)
-                    clear_reload_btn = gr.Button("🔄 Reload Model & Reset Data", variant="secondary", scale=1)
-
-                with gr.Row():
-                    output_display = gr.Textbox(
-                        lines=20, 
-                        label="Logs & Results", 
-                        value="Initializing...", 
-                        interactive=False,
-                        autoscroll=True
-                    )
-                    loss_plot = gr.Plot(label="Training Metrics")
-
-            # --- TAB 3: EXPORT ---
-            with gr.TabItem("3. Export"):
-                gr.Markdown("### 📦 Export Trained Model")
-                
-                with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("#### Option A: Download ZIP")
-                        gr.Markdown("Download the model weights locally.")
-                        zip_btn = gr.Button("⬇️ Prepare Model ZIP", variant="secondary", interactive=False)
-                        download_file = gr.File(label="Download Archive", interactive=False)
-                    
-                    with gr.Column():
-                        gr.Markdown("#### Option B: Save to Hugging Face Hub")
-                        gr.Markdown("Publish your fine-tuned model to your personal Hugging Face account.")
-                        
-                        repo_name_input = gr.Textbox(
-                            label="TargetRepository Name", 
-                            value="my-functiongemma-v1",
-                            placeholder="e.g., my-functiongemma-v1",
-                            interactive=False
-                        )
-                        push_to_hub_btn = gr.Button("Save to Hugging Face Hub", variant="secondary", interactive=False)
-                        repo_id_preview = gr.Markdown("Target Repository: (Waiting for input...)")
-                        
-                        upload_status = gr.Markdown("")
-
-        # --- EVENT WIRING ---
+            data_ui = _render_dataset_tab(engine_state)
+            train_ui = _render_training_tab(engine_state)
+            export_ui = _render_export_tab(engine_state, username_state)
+
+        # Helpers for UI State
+        # 'action_buttons' now ONLY contains buttons that should always be enabled after a process
+        # Zip and Push buttons are excluded here because their state depends on has_model_tuned
+        action_buttons = [train_ui["buttons"][2], train_ui["buttons"][0]] # [Reload, Run]
+        
+        repo_input = export_ui["hub_controls"][0]
+        push_btn = export_ui["hub_controls"][1]
+        zip_btn = export_ui["zip_controls"][0]
+
+        def lock_ui():
+            """Locks all buttons (including Zip/Push) during processing"""
+            return [gr.update(interactive=False) for _ in action_buttons] + \
+                   [gr.update(interactive=False), gr.update(interactive=False)]
         
-        action_buttons = [
-            clear_reload_btn,
-            run_training_btn,
-            zip_btn
-        ]
+        def unlock_ui():
+            """Unlocks general action buttons only. Zip/Push are handled by update_hub_interactive"""
+            return [gr.update(interactive=True) for _ in action_buttons]
 
-        def set_interactivity(interactive: bool):
-            return [gr.update(interactive=interactive) for _ in action_buttons]
+        # --- Event Wiring ---
 
-        demo.load(
-            fn=lambda: set_interactivity(False), outputs=action_buttons
-        ).then(
-            fn=init_session,
+        # 1. Initialization
+        demo.load(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
+            fn=UIController.init_session,
             inputs=None,
-            outputs=[engine_state, tools_editor, param_model, output_display, repo_name_input, push_to_hub_btn, username_state]
+            outputs=[
+                engine_state, 
+                data_ui["tools_editor"], 
+                train_ui["model_input"], 
+                train_ui["outputs"][0], # log output
+                repo_input, 
+                push_btn, 
+                zip_btn,                # Update Zip state based on initial engine state
+                username_state
+            ]
         ).then(
-            fn=update_repo_preview,
-            inputs=[username_state, repo_name_input],
-            outputs=[repo_id_preview]
-        ).then(
-            fn=lambda: [gr.update(interactive=True)]*2, outputs=[clear_reload_btn, run_training_btn]
-        )
+            fn=UIController.update_repo_preview,
+            inputs=[username_state, repo_input],
+            outputs=[export_ui["hub_controls"][2]] 
+        ).then(unlock_ui, outputs=action_buttons)
 
-        update_tools_btn.click(
-            fn=update_tools_wrapper,
-            inputs=[engine_state, tools_editor],
-            outputs=[tools_status]
+        # 2. Data Tab
+        data_ui["update_tools_btn"].click(
+            fn=UIController.update_tools,
+            inputs=[engine_state, data_ui["tools_editor"]],
+            outputs=[data_ui["tools_status"]]
         )
 
-        import_file.upload(
-            fn=import_file_wrapper, 
-            inputs=[engine_state, import_file], 
-            outputs=[import_status]
+        data_ui["import_file"].upload(
+            fn=UIController.import_file,
+            inputs=[engine_state, data_ui["import_file"]],
+            outputs=[data_ui["import_status"]]
         )
+
+        # 3. Training Tab
+        run_btn, stop_btn, reload_btn = train_ui["buttons"]
         
-        run_training_btn.click(
+        run_btn.click(
             fn=lambda: (
                 gr.update(visible=False), 
-                gr.update(interactive=False), 
-                gr.update(interactive=False), 
-                gr.update(visible=True)
+                gr.update(interactive=False), # Lock Reload
+                gr.update(interactive=False), # Lock Zip 
+                gr.update(visible=True)       # Show Stop
             ),
-            outputs=[run_training_btn, clear_reload_btn, zip_btn, stop_training_btn]
+            outputs=[run_btn, reload_btn, zip_btn, stop_btn]
         ).then(
-            fn=run_training_wrapper,
-            inputs=[engine_state, param_epochs, param_lr, param_test_size, param_shuffle, param_model],
-            outputs=[output_display, loss_plot],
+            fn=UIController.run_training,
+            inputs=[engine_state, *train_ui["params"]],
+            outputs=train_ui["outputs"],
         ).then(
             fn=lambda: (
                 gr.update(visible=True), 
                 gr.update(interactive=True), 
-                gr.update(interactive=True), 
                 gr.update(visible=False)
             ),
-            outputs=[run_training_btn, clear_reload_btn, zip_btn, stop_training_btn]
+            outputs=[run_btn, reload_btn, stop_btn]
         ).then(
-            fn=update_hub_interactive,
+            # Final check determines if Zip/Push should unlock
+            fn=UIController.update_hub_interactive,
             inputs=[engine_state, username_state],
-            outputs=[repo_name_input, push_to_hub_btn]
-        )
-        
-        stop_training_btn.click(
-            fn=stop_wrapper,
-            inputs=[engine_state],
-            outputs=None
+            outputs=[repo_input, push_btn, zip_btn] 
         )
 
-        clear_reload_btn.click(
-            fn=lambda: set_interactivity(False), outputs=action_buttons
-        ).then(
-            fn=lambda: gr.update(interactive=False), outputs=push_to_hub_btn
-        ).then(
-            fn=handle_reset, 
-            inputs=[engine_state, param_model], 
-            outputs=[output_display]
-        ).then(
-            fn=lambda: [gr.update(interactive=True)]*2, outputs=[clear_reload_btn, run_training_btn]
-        ).then(
-            fn=update_hub_interactive,
+        stop_btn.click(fn=UIController.stop_process, inputs=[engine_state])
+
+        reload_btn.click(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
+            fn=UIController.handle_reset,
+            inputs=[engine_state, train_ui["model_input"]],
+            outputs=[train_ui["outputs"][0]]
+        ).then(unlock_ui, outputs=action_buttons).then(
+            fn=UIController.update_hub_interactive,
             inputs=[engine_state, username_state],
-            outputs=[repo_name_input, push_to_hub_btn]
+            outputs=[repo_input, push_btn, zip_btn]
         )
-        
-        zip_btn.click(
-            fn=lambda: set_interactivity(False), outputs=action_buttons
-        ).then(
-            fn=lambda: gr.update(interactive=False), outputs=push_to_hub_btn
-        ).then(
-            fn=zip_wrapper,
+
+        # 4. Export Tab
+        zip_btn.click(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
+            fn=UIController.zip_model,
             inputs=[engine_state],
-            outputs=[download_file]
-        ).then(
-            fn=lambda: set_interactivity(True), outputs=action_buttons
-        ).then(
-            fn=update_hub_interactive,
+            outputs=[export_ui["zip_controls"][1]]
+        ).then(unlock_ui, outputs=action_buttons).then(
+            fn=UIController.update_hub_interactive,
             inputs=[engine_state, username_state],
-            outputs=[repo_name_input, push_to_hub_btn]
+            outputs=[repo_input, push_btn, zip_btn]
         )
-        
-        repo_name_input.change(
-            fn=update_repo_preview,
-            inputs=[username_state, repo_name_input],
-            outputs=[repo_id_preview]
+
+        repo_input.change(
+            fn=UIController.update_repo_preview,
+            inputs=[username_state, repo_input],
+            outputs=[export_ui["hub_controls"][2]]
         )
 
-        push_to_hub_btn.click(
-            fn=lambda: set_interactivity(False), outputs=action_buttons
-        ).then(
-            fn=lambda: gr.update(interactive=False), outputs=push_to_hub_btn
-        ).then(
-            fn=upload_wrapper,
-            inputs=[engine_state, repo_name_input],
-            outputs=[upload_status]
-        ).then(
-            fn=lambda: set_interactivity(True), outputs=action_buttons
-        ).then(
-            fn=update_hub_interactive,
+        push_btn.click(lock_ui, outputs=action_buttons + [push_btn, zip_btn]).then(
+            fn=UIController.upload_model,
+            inputs=[engine_state, repo_input],
+            outputs=[export_ui["hub_controls"][3]]
+        ).then(unlock_ui, outputs=action_buttons).then(
+            fn=UIController.update_hub_interactive,
             inputs=[engine_state, username_state],
-            outputs=[repo_name_input, push_to_hub_btn]
+            outputs=[repo_input, push_btn, zip_btn]
         )
 
     return demo