Upload folder using huggingface_hub

app.py

@@ -1,15 +1,9 @@
-from config import AppConfig
-from engine import FunctionGemmaEngine
 from ui import build_interface
 
 if __name__ == "__main__":
-    # Initialize Config
-    config = AppConfig()
-    
-    # Initialize Logic Engine
-    app_engine = FunctionGemmaEngine(config)
-    
     # Build and Launch UI
-    demo = build_interface(app_engine)
-    print("Starting Gradio App...")
+    # Note: Engine creation is now handled per-session inside build_interface
+    demo = build_interface()
+    print("Starting Gradio App with Multi-User Support...")
+    demo.queue() # Enable queueing for concurrent request handling
     demo.launch()

engine.py

@@ -3,12 +3,14 @@ import torch
 import time
 import json
 import queue
+import uuid
 import matplotlib.pyplot as plt
 from functools import partial
 from typing import Generator, Optional, List, Dict, Any, Tuple
 from datasets import Dataset, load_dataset
 from trl import SFTConfig, SFTTrainer
 from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl
+from huggingface_hub import HfApi # Added for Hub Upload
 
 from config import AppConfig
 from tools import DEFAULT_TOOLS
@@ -29,10 +31,6 @@ class AbortCallback(TrainerCallback):
             control.should_training_stop = True
 
 class LogStreamingCallback(TrainerCallback):
-    """
-    Intercepts training logs and pushes them to a queue.
-    Sends tuple: (formatted_string, raw_dict)
-    """
     def __init__(self, log_queue: queue.Queue):
         self.log_queue = log_queue
         
@@ -63,7 +61,6 @@ class LogStreamingCallback(TrainerCallback):
             
                 log_parts.append(f"{label}: {val_str}")
         
-        # Structure for plotting
         log_payload = logs.copy()
         log_payload['step'] = state.global_step
         
@@ -72,6 +69,11 @@ class LogStreamingCallback(TrainerCallback):
 class FunctionGemmaEngine:
     def __init__(self, config: AppConfig):
         self.config = config
+        
+        self.session_id = str(uuid.uuid4())[:8]
+        self.output_dir = self.config.ARTIFACTS_DIR.joinpath(f"session_{self.session_id}")
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
         self.model = None
         self.tokenizer = None
         self.loaded_model_name = None 
@@ -104,17 +106,15 @@ class FunctionGemmaEngine:
     # --- Model & Data Management ---
     
     def _load_model_weights(self):
-        """Internal helper to load model based on current config."""
-        print(f"Loading model: {self.config.MODEL_NAME}...")
+        print(f"[{self.session_id}] Loading model: {self.config.MODEL_NAME}...")
         self.model, self.tokenizer = load_model_and_tokenizer(self.config.MODEL_NAME)
         self.loaded_model_name = self.config.MODEL_NAME
 
     def refresh_data_and_model(self) -> str:
-        """Full reset: Reloads model and clears dataset."""
         self.imported_dataset = []
         try:
             self._load_model_weights()
-            return f"Model loaded: {self.loaded_model_name}\nData cleared.\nReady."
+            return f"Model loaded: {self.loaded_model_name}\nData cleared.\nReady (Session {self.session_id})."
         except Exception as e:
             self.model = None
             self.tokenizer = None
@@ -141,7 +141,6 @@ class FunctionGemmaEngine:
         output_buffer = ""
         last_plot = None
 
-        # 1. Check if model name changed since last load
         if self.config.MODEL_NAME != self.loaded_model_name:
             output_buffer += f"🔄 Model changed. Switching from '{self.loaded_model_name}' to '{self.config.MODEL_NAME}'...\n"
             yield output_buffer, None
@@ -231,7 +230,6 @@ class FunctionGemmaEngine:
         
         train_thread.join()
         
-        # Flush logs
         while not log_queue.empty():
             payload = log_queue.get()
             if isinstance(payload, tuple):
@@ -287,7 +285,7 @@ class FunctionGemmaEngine:
     def _execute_trainer(self, dataset, log_queue: queue.Queue, epochs: int, learning_rate: float) -> List[Dict]:
         torch_dtype = self.model.dtype
         args = SFTConfig(
-            output_dir=str(self.config.OUTPUT_DIR),
+            output_dir=str(self.output_dir),
             max_length=512,
             packing=False,
             num_train_epochs=epochs,
@@ -319,8 +317,6 @@ class FunctionGemmaEngine:
         
     def _generate_loss_plot(self, history: list):
         if not history: return None
-        
-        # CHANGED: Close previous figures to prevent memory warning
         plt.close('all')
         
         train_steps = [x['step'] for x in history if 'loss' in x]
@@ -372,7 +368,32 @@ class FunctionGemmaEngine:
                 yield f"Error during inference: {e}"
 
     def get_zip_path(self) -> Optional[str]:
-        if not self.config.OUTPUT_DIR.exists(): return None
-        timestamp = int(time.time())
-        base_name = str(self.config.ARTIFACTS_DIR.joinpath(f"functiongemma_finetuned_{timestamp}"))
-        return zip_directory(str(self.config.OUTPUT_DIR), base_name)
+        if not self.output_dir.exists(): return None
+        base_name = str(self.config.ARTIFACTS_DIR.joinpath(f"functiongemma_finetuned_{self.session_id}"))
+        return zip_directory(str(self.output_dir), base_name)
+
+    def upload_model_to_hub(self, repo_name: str, oauth_token: str) -> str:
+        """Uploads the trained model to Hugging Face Hub."""
+        if not self.output_dir.exists() or not any(self.output_dir.iterdir()):
+            return "❌ No trained model found in current session. Run training first."
+        
+        try:
+            api = HfApi(token=oauth_token)
+            
+            # Create Repo (if needed)
+            print(f"Creating/Checking repo {repo_name}...")
+            repo_url = api.create_repo(
+                repo_id=repo_name, 
+                exist_ok=True
+            )
+            
+            # Upload
+            print(f"Uploading to {repo_url.repo_id}...")
+            api.upload_folder(
+                folder_path=str(self.output_dir),
+                repo_id=repo_url.repo_id,
+                repo_type="model"
+            )
+            return f"✅ Success! Model uploaded to: {repo_url}"
+        except Exception as e:
+            return f"❌ Upload failed: {str(e)}"

requirements.txt

@@ -2,5 +2,6 @@ accelerate
 datasets
 gradio
 matplotlib
+oauth
 transformers
 trl

ui.py

@@ -1,21 +1,64 @@
 import gradio as gr
+from config import AppConfig
 from engine import FunctionGemmaEngine
 
-def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
+def build_interface() -> gr.Blocks:
     
-    # Wrapper: Update config with selected model, then Run Training
-    def run_training_wrapper(epochs, lr, test_size, shuffle, model_name):
+    # --- State Management Wrappers ---
+    
+    def init_session():
+        config = AppConfig()
+        new_engine = FunctionGemmaEngine(config)
+        return (
+            new_engine,
+            new_engine.get_tools_json(),
+            new_engine.config.MODEL_NAME,
+            f"Ready. (Session {new_engine.session_id})"
+        )
+
+    def run_training_wrapper(engine, epochs, lr, test_size, shuffle, model_name):
         engine.config.MODEL_NAME = model_name.strip()
         yield from engine.run_training_pipeline(epochs, lr, test_size, shuffle)
 
-    # Wrapper: Update config with selected model, then Reset/Reload
-    def handle_reset(model_name):
+    def handle_reset(engine, model_name):
         engine.config.MODEL_NAME = model_name.strip()
         return engine.refresh_data_and_model()
 
+    def update_tools_wrapper(engine, json_val):
+        return engine.update_tools(json_val)
+
+    def import_file_wrapper(engine, file_obj):
+        return engine.load_csv(file_obj)
+
+    def stop_wrapper(engine):
+        engine.trigger_stop()
+        return "Stopping..."
+
+    def zip_wrapper(engine):
+        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):
+        if oauth_token is None:
+            return "❌ Error: You must log in (top right) to upload models."
+        if not repo_name:
+            return "❌ Error: Please enter a repository name."
+        
+        return engine.upload_model_to_hub(
+            repo_name=repo_name, 
+            oauth_token=oauth_token.token, 
+        )
+
+    # --- UI Layout ---
     with gr.Blocks(title="FunctionGemma Modkit") as demo:
-        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.")
+        engine_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.LoginButton(value="(Optional) Sign in to Hugging Face, if you want to push fine-tuned model to your repo.")
         
         with gr.Tabs():
             
@@ -27,7 +70,6 @@ def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
                     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(
-                            value=engine.get_tools_json(), 
                             language="json", 
                             label="Tool Definitions (JSON Schema)",
                             lines=15
@@ -52,10 +94,9 @@ def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
                 with gr.Group():
                     gr.Markdown("**Hyperparameters**")
                     with gr.Row():
-                        # Dropdown that allows custom typing
+                        default_models = AppConfig().AVAILABLE_MODELS
                         param_model = gr.Dropdown(
-                            choices=engine.config.AVAILABLE_MODELS,
-                            value=engine.config.MODEL_NAME,
+                            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/gemma-3-1b-it')",
@@ -82,91 +123,107 @@ def build_interface(engine: FunctionGemmaEngine) -> gr.Blocks:
                         )
 
                 with gr.Row():
-                    run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary", scale=2)
+                    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():
-                    # Left column: Text Logs
                     output_display = gr.Textbox(
                         lines=20, 
                         label="Logs & Results", 
-                        value="Ready.",
+                        value="Initializing...", 
                         interactive=False,
                         autoscroll=True
                     )
-                    # Right column: Plot
                     loss_plot = gr.Plot(label="Training Metrics")
 
             # --- TAB 3: EXPORT ---
             with gr.TabItem("3. Export"):
                 gr.Markdown("### 📦 Export Trained Model")
-                gr.Markdown("Download the fine-tuned LoRA adapters or full model weights (depending on configuration) as a ZIP file.")
                 
                 with gr.Row():
-                    zip_btn = gr.Button("⬇️ Prepare Model ZIP", variant="primary", scale=1)
-                    download_file = gr.File(label="Download Archive", interactive=False, scale=2)
+                    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")
+                        download_file = gr.File(label="Download Archive", interactive=False)
+                    
+                    with gr.Column():
+                        gr.Markdown("#### Option B: Upload to Hugging Face Hub")
+                        gr.Markdown("Publish to your HF profile. **Requires Login**.")
+                        
+                        with gr.Group():
+                            repo_id_input = gr.Textbox(
+                                label="Repository Name", 
+                                placeholder="my-function-gemma-v1",
+                                info="Will be created under your username (e.g. user/repo)"
+                            )
+                            upload_hub_btn = gr.Button("☁️ Upload to Hub", variant="primary")
+                        
+                        upload_status = gr.Markdown("")
 
         # --- EVENT WIRING ---
 
-        # Tab 1: Tools
+        demo.load(
+            fn=init_session,
+            inputs=None,
+            outputs=[engine_state, tools_editor, param_model, output_display]
+        )
+
         update_tools_btn.click(
-            fn=engine.update_tools,
-            inputs=[tools_editor],
+            fn=update_tools_wrapper,
+            inputs=[engine_state, tools_editor],
             outputs=[tools_status]
         )
 
-        # Tab 1: File Import
         import_file.upload(
-            fn=engine.load_csv, 
-            inputs=[import_file], 
+            fn=import_file_wrapper, 
+            inputs=[engine_state, import_file], 
             outputs=[import_status]
         )
 
-        # Tab 2: Training (Uses Wrapper)
         run_training_btn.click(
             fn=lambda: (
-                gr.update(visible=False), # Hide Run
-                gr.update(interactive=False), # Disable Reset
-                gr.update(visible=True)   # Show Stop
+                gr.update(visible=False), 
+                gr.update(interactive=False), 
+                gr.update(visible=True)
             ),
             outputs=[run_training_btn, clear_reload_btn, stop_training_btn]
         ).then(
             fn=run_training_wrapper,
-            inputs=[param_epochs, param_lr, param_test_size, param_shuffle, param_model],
+            inputs=[engine_state, param_epochs, param_lr, param_test_size, param_shuffle, param_model],
             outputs=[output_display, loss_plot],
         ).then(
             fn=lambda: (
-                gr.update(visible=True),  # Show Run
-                gr.update(interactive=True), # Enable Reset
-                gr.update(visible=False)  # Hide Stop
+                gr.update(visible=True), 
+                gr.update(interactive=True), 
+                gr.update(visible=False)
             ),
             outputs=[run_training_btn, clear_reload_btn, stop_training_btn]
         )
         
-        # Tab 2: Stop
         stop_training_btn.click(
-            fn=lambda: (engine.trigger_stop(), "Stopping...")[1],
+            fn=stop_wrapper,
+            inputs=[engine_state],
             outputs=None
         )
 
-        # Tab 2: Reset (Uses Wrapper to capture model name)
         clear_reload_btn.click(
             fn=handle_reset, 
-            inputs=[param_model], 
+            inputs=[engine_state, param_model], 
             outputs=[output_display]
         )
         
-        # Tab 3: Download
-        def handle_zip():
-            path = engine.get_zip_path()
-            if path:
-                return gr.update(value=path, visible=True)
-            return gr.update(value=None, visible=False)
-
         zip_btn.click(
-            fn=handle_zip,
+            fn=zip_wrapper,
+            inputs=[engine_state],
             outputs=[download_file]
         )
+        
+        upload_hub_btn.click(
+            fn=upload_wrapper,
+            inputs=[engine_state, repo_id_input],
+            outputs=[upload_status]
+        )
 
     return demo