Upload folder using huggingface_hub

README.md

@@ -1,5 +1,5 @@
 ---
-title: Functiongemma Modkit
+title: FunctionGemma Modkit
 emoji: 📊
 colorFrom: gray
 colorTo: indigo
@@ -8,7 +8,9 @@ sdk_version: 6.0.1
 app_file: app.py
 pinned: false
 license: apache-2.0
-short_description: FunctionGemma Modkit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FunctionGemma Modkit
+
+This project provides a set of tools to fine-tune FunctionGemma to understand your personal needs.
+

app.py

@@ -0,0 +1,311 @@
+import gradio as gr
+import os
+import json
+import torch
+
+from typing import Final, Optional, List
+from pathlib import Path
+from huggingface_hub import login
+from trl import SFTConfig, SFTTrainer
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from datasets import Dataset, load_dataset
+from transformers.utils import get_json_schema
+
+
+ARTIFACTS_DIR: Final[Path] = Path("artifacts")
+
+def authenticate_hf(token: Optional[str]) -> None:
+    """Logs into the Hugging Face Hub."""
+    if token:
+        print("Logging into Hugging Face Hub...")
+        login(token=token)
+    else:
+        print("Skipping Hugging Face login: HF_TOKEN not set.")
+
+def load_model(model_name: str):
+    print(f"Loading Transformer model: {model_name}")
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        print("Model loaded successfully.")
+        return (model, tokenizer)
+    except Exception as e:
+        print(f"Error loading Transformer model {model_name}: {e}")
+        raise
+
+# --- Tool Definitions ---
+def search_knowledge_base(query: str) -> str:
+    """
+    Search internal company documents, policies and project data.
+
+    Args:
+        query: query string
+    """
+    return "Interal Result"
+
+def search_google(query: str) -> str:
+    """
+    Search public information.
+
+    Args:
+        query: query string
+    """
+    return "Public Result"
+
+
+TOOLS = [get_json_schema(search_knowledge_base), get_json_schema(search_google)]
+
+DEFAULT_SYSTEM_MSG = "You are a model that can do function calling with the following functions"
+
+
+def create_conversation(sample):
+  return {
+      "messages": [
+          {"role": "developer", "content": DEFAULT_SYSTEM_MSG},
+          {"role": "user", "content": sample["user_content"]},
+          {"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": sample["tool_name"], "arguments": json.loads(sample["tool_arguments"])}}]},
+      ],
+      "tools": TOOLS
+  }
+
+
+def train_with_dataset(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    dataset: Dataset,
+    output_dir: Path,
+    learning_rate: float = 5e-5
+) -> None:
+
+    torch_dtype = model.dtype
+
+    args = SFTConfig(
+        output_dir=output_dir,                  # directory to save and repository id
+        max_length=512,                         # max sequence length for model and packing of the dataset
+        packing=False,                          # Groups multiple samples in the dataset into a single sequence
+        num_train_epochs=5,                     # number of training epochs
+        per_device_train_batch_size=4,          # batch size per device during training
+        gradient_checkpointing=False,           # Caching is incompatible with gradient checkpointing
+        optim="adamw_torch_fused",              # use fused adamw optimizer
+        logging_steps=1,                        # log every step
+        #save_strategy="epoch",                  # save checkpoint every epoch
+        eval_strategy="epoch",                  # evaluate checkpoint every epoch
+        learning_rate=learning_rate,            # learning rate
+        fp16=True if torch_dtype == torch.float16 else False,   # use float16 precision
+        bf16=True if torch_dtype == torch.bfloat16 else False,  # use bfloat16 precision
+        lr_scheduler_type="constant",           # use constant learning rate scheduler
+        push_to_hub=False,                       # push model to hub
+        report_to="none",                # report metrics to tensorboard
+        dataset_kwargs={
+            "add_special_tokens": False, # Template with special tokens
+            "append_concat_token": True, # Add EOS token as separator token between examples
+        }
+    )
+
+    # Create Trainer object
+    trainer = SFTTrainer(
+        model=model,
+        args=args,
+        train_dataset=dataset['train'],
+        eval_dataset=dataset['test'],
+        processing_class=tokenizer,
+    )
+
+    trainer.train()
+
+    print("Training finished. Model weights are updated in memory.")
+
+    # Save the final fine-tuned model
+    trainer.save_model()
+
+    print(f"Model saved locally to: {output_dir}")
+
+class AppConfig:
+    """
+    Central configuration class for the Fine-Tuner application.
+    """
+    ARTIFACTS_DIR: Final[Path] = ARTIFACTS_DIR
+    HF_TOKEN: Final[str | None] = os.getenv('HF_TOKEN')
+    MODEL_NAME: Final[str] = '../hf/270m'
+    DEFAULT_DATASET: Final[str] = 'bebechien/SimpleToolCalling'
+    OUTPUT_DIR: Final[Path] = ARTIFACTS_DIR.joinpath("functiongemma-270m-it-modkit-demo")
+
+
+class FunctionGemmaTuner:
+    def __init__(self, config: AppConfig = AppConfig):
+        self.config = config
+
+        os.makedirs(self.config.ARTIFACTS_DIR, exist_ok=True)
+        print(f"Created artifact directory: {self.config.ARTIFACTS_DIR}")
+
+        authenticate_hf(self.config.HF_TOKEN)
+
+        self._initial_load()
+
+    def _initial_load(self):
+        """Helper to run the refresh function once at startup."""
+        print("--- Running Initial Data Load ---")
+        self.refresh_data_and_model()
+        print("--- Initial Load Complete ---")
+
+    def refresh_data_and_model(self):
+        print("\n" + "=" * 50)
+        print("RELOADING MODEL and RE-FETCHING DATA")
+
+        # Reset dataset state
+        self.imported_dataset = []
+
+        # 1. Reload the base model
+        try:
+            self.model, self.tokenizer = load_model(self.config.MODEL_NAME)
+        except Exception as e:
+            gr.Error(f"Model load failed: {e}")
+            self.model = None
+            self.tokenizer = None
+            return gr.update(value=f"CRITICAL ERROR: Model failed to load. {e}")
+
+        status_value: str = f"Model and data reloaded. Click 'Run Fine-Tuning' to begin."
+
+        # Return Gradio updates for CheckboxGroup and Textbox
+        return gr.update(value=status_value)
+
+    # --- Import Dataset/Export ---
+    def import_additional_dataset(self, file_path: str) -> str:
+        if not file_path:
+            return "Please upload a CSV file."
+        new_dataset, num_imported = [], 0
+        try:
+            with open(file_path, 'r', newline='', encoding='utf-8') as f:
+                reader = csv.reader(f)
+                try:
+                    header = next(reader)
+                    if not (header and header[0].lower().strip() == 'anchor'):
+                        f.seek(0)
+                except StopIteration:
+                    return "Error: Uploaded file is empty."
+
+                for row in reader:
+                    if len(row) == 3:
+                        new_dataset.append([s.strip() for s in row])
+                        num_imported += 1
+            if num_imported == 0:
+                raise ValueError("No valid [Anchor, Positive, Negative] rows found in the CSV.")
+            self.imported_dataset = new_dataset
+            return f"Successfully imported {num_imported} additional training triplets."
+        except Exception as e:
+            gr.Error(f"Import failed. Ensure the CSV format is: [Anchor, Positive, Negative]. Error: {e}")
+            return "Import failed. Check console for details."
+
+    def download_model(self) -> Optional[str]:
+        if not os.path.exists(self.config.OUTPUT_DIR):
+            gr.Warning(f"The model directory '{self.config.OUTPUT_DIR}' does not exist. Please run training first.")
+            return None
+        timestamp = int(time.time())
+        try:
+            base_name = os.path.join(self.config.ARTIFACTS_DIR, f"embedding_gemma_finetuned_{timestamp}")
+            archive_path = shutil.make_archive(
+                base_name=base_name,
+                format='zip',
+                root_dir=self.config.OUTPUT_DIR,
+            )
+            gr.Info(f"Model files successfully zipped to: {archive_path}")
+            return archive_path
+        except Exception as e:
+            gr.Error(f"Failed to create the model ZIP file. Error: {e}")
+            return None
+
+    def training(self, test_size: float = 0.5) -> str:
+        """
+        Generates a training dataset from user selection and runs the fine-tuning process.
+        """
+        if self.model is None:
+             raise gr.Error("Training failed: Model is not loaded.")
+
+        if not self.imported_dataset:
+            print("No imported dataset, use the default")
+            dataset = load_dataset(self.config.DEFAULT_DATASET, split="train")
+        else:
+            dataset_as_dicts = [{
+                "user_content": row[0], "tool_name": row[1], "tool_arguments": row[2]}
+                for row in self.imported_dataset
+            ]
+            dataset = Dataset.from_list(dataset_as_dicts)
+
+        dataset = dataset.map(create_conversation, batched=False)
+        dataset = dataset.train_test_split(test_size=test_size, shuffle=False)
+        print(dataset)
+        print("--- dataset input ---")
+        print(json.dumps(dataset["train"][0], indent=2))
+        debug_msg = self.tokenizer.apply_chat_template(dataset["train"][0]["messages"], tools=dataset["train"][0]["tools"], add_generation_prompt=False, tokenize=False)
+        print("--- Formatted prompt ---")
+        print(debug_msg)
+
+        result = "### Success Rate (Before Training):\n" + f"{self.check_success_rate(dataset["test"])}\n\n"
+        print("-" * 50 + "\nStarting Fine-tuning...")
+        train_with_dataset(model=self.model, tokenizer=self.tokenizer, dataset=dataset, output_dir=self.config.OUTPUT_DIR)
+        print("Fine-tuning Complete.\n" + "-" * 50)
+
+        result += "### Success Rate (After Training):\n" + f"{self.check_success_rate(dataset["test"])}\n\n"
+        return result
+
+    def check_success_rate(self, test_dataset):
+        result = []
+        success_count = 0
+        for idx, item in enumerate(test_dataset):
+            messages = [
+                item["messages"][0],
+                item["messages"][1],
+            ]
+
+            inputs = self.tokenizer.apply_chat_template(messages, tools=TOOLS, add_generation_prompt=True, return_dict=True, return_tensors="pt")
+
+            out = self.model.generate(**inputs.to(self.model.device), pad_token_id=self.tokenizer.eos_token_id, max_new_tokens=128)
+            output = self.tokenizer.decode(out[0][len(inputs["input_ids"][0]) :], skip_special_tokens=False)
+
+            result.append(f"{idx+1} Prompt: {item['messages'][1]['content']}")
+            result.append(f"  Output: {output}")
+            if item['messages'][2]['tool_calls'][0]['function']['name'] in output:
+                result.append("  `-> ✅ correct!")
+                success_count += 1
+            else:
+                result.append(" `-> ❌ wrong tool")
+
+        result.append(f"Success : {success_count} / {len(test_dataset)}")
+
+        return result
+
+    def build_interface(self) -> gr.Blocks:
+        with gr.Blocks(title="FunctionGemma Modkit") as demo:
+            gr.Markdown("# 🤖 FunctionGemma Modkit: Fine-Tuning")
+            gr.Markdown("This project provides a set of tools to fine-tune FunctionGemma to understand your personal needs.<br>See [README](https://huggingface.co/spaces/google/functiongemma-modkit/blob/main/README.md) for more details.")
+            self._build_training_interface()
+        return demo
+
+    def _build_training_interface(self):
+        with gr.Column():
+            gr.Markdown("## Fine-Tuning")
+            with gr.Row():
+                output = gr.Textbox(lines=14, label="Training and Search Results", value="Click 'Run Fine-Tuning' to begin.")
+            with gr.Row():
+                clear_reload_btn = gr.Button("Clear & Reload Model/Data")
+                run_training_btn = gr.Button("🚀 Run Fine-Tuning", variant="primary")
+            gr.Markdown("--- \n ## Dataset & Model Management")
+            import_file = gr.File(label="Upload Additional Dataset (.csv)", file_types=[".csv"], height=50)
+            with gr.Row():
+                download_model_btn = gr.Button("⬇️ Download Fine-Tuned Model")
+            download_status = gr.Markdown("Ready.")
+            with gr.Row():
+                model_output = gr.File(label="Download Model ZIP", height=50, visible=False, interactive=False)
+
+            run_training_btn.click(fn=self.training, outputs=output)
+            clear_reload_btn.click(fn=self.refresh_data_and_model, inputs=None, outputs=[output], queue=False)
+            import_file.change(fn=self.import_additional_dataset, inputs=[import_file], outputs=download_status)
+            download_model_btn.click(lambda: [gr.update(value=None, visible=False), "Zipping..."], None, [model_output, download_status], queue=False).then(self.download_model, None, model_output).then(lambda p: [gr.update(visible=p is not None, value=p), "ZIP ready." if p else "Zipping failed."], [model_output], [model_output, download_status])
+
+
+if __name__ == "__main__":
+    app = FunctionGemmaTuner(AppConfig)
+    demo = app.build_interface()
+    print("Starting Gradio App...")
+    demo.launch()
+

requirements.txt

@@ -0,0 +1,5 @@
+accelerate
+datasets
+gradio
+transformers
+trl