Upload train.py with huggingface_hub

train.py

@@ -8,27 +8,27 @@ from trl import SFTTrainer, SFTConfig
 import trackio
 import os
 
-print("🚀 Starting FunctionGemma 2B Fine-tuning")
+print("🚀 Starting FunctionGemma 270M Fine-tuning")
 
 # Load dataset
 dataset = load_dataset("epinfomax/vn-function-calling-dataset", split="train")
 
 # Training configuration
 config = SFTConfig(
-    output_dir="vn-function-gemma-finetuned",
+    output_dir="vn-function-gemma-270m-finetuned",
     push_to_hub=True,
-    hub_model_id="epinfomax/vn-function-gemma-finetuned",
+    hub_model_id="epinfomax/vn-function-gemma-270m-finetuned",
     hub_strategy="every_save",
-    num_train_epochs=3,
-    per_device_train_batch_size=4,
-    gradient_accumulation_steps=4,
-    learning_rate=2e-5,
-    logging_steps=10,
+    num_train_epochs=5,  # Increased epochs for the smaller model
+    per_device_train_batch_size=8, # Increased batch size for the smaller model
+    gradient_accumulation_steps=2,
+    learning_rate=5e-5, # Slightly higher LR for smaller model
+    logging_steps=5,
     save_strategy="steps",
     save_steps=50,
     report_to="trackio",
     project="vn-function-calling",
-    run_name="function-gemma-2b-baseline"
+    run_name="function-gemma-270m-v1"
 )
 
 # LoRA configuration
@@ -41,7 +41,7 @@ peft_config = LoraConfig(
 
 # Initialize and train
 trainer = SFTTrainer(
-    model="google/function-gemma-2b",
+    model="google/functiongemma-270m-it",
     train_dataset=dataset,
     peft_config=peft_config,
     args=config,