Upload train.py with huggingface_hub

train.py

@@ -1,34 +1,54 @@
 # /// script
-# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "transformers", "accelerate"]
+# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "transformers", "accelerate", "jinja2"]
 # ///
 
 from datasets import load_dataset
 from peft import LoraConfig
 from trl import SFTTrainer, SFTConfig
+from transformers import AutoTokenizer
 import trackio
 import os
 
-print("🚀 Starting FunctionGemma 270M Fine-tuning")
+print("🚀 Starting FunctionGemma 270M Fine-tuning (V2 with Template Fix)")
+
+model_id = "google/functiongemma-270m-it"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 # Load dataset
 dataset = load_dataset("epinfomax/vn-function-calling-dataset", split="train")
 
+def format_conversation(example):
+    # Modern transformers template supports 'tools' argument
+    # We render the template to a string so SFTTrainer doesn't have to guess
+    text = tokenizer.apply_chat_template(
+        example["messages"],
+        tools=example["tools"],
+        tokenize=False,
+        add_generation_prompt=False
+    )
+    return {"text": text}
+
+print("🔄 Pre-processing dataset with chat template...")
+dataset = dataset.map(format_conversation, remove_columns=dataset.column_names)
+
 # Training configuration
 config = SFTConfig(
+    dataset_text_field="text", # Use the pre-rendered text
+    max_seq_length=1024,
     output_dir="vn-function-gemma-270m-finetuned",
     push_to_hub=True,
     hub_model_id="epinfomax/vn-function-gemma-270m-finetuned",
     hub_strategy="every_save",
-    num_train_epochs=5,  # Increased epochs for the smaller model
-    per_device_train_batch_size=8, # Increased batch size for the smaller model
+    num_train_epochs=5,
+    per_device_train_batch_size=8,
     gradient_accumulation_steps=2,
-    learning_rate=5e-5, # Slightly higher LR for smaller model
+    learning_rate=5e-5,
     logging_steps=5,
     save_strategy="steps",
     save_steps=50,
     report_to="trackio",
     project="vn-function-calling",
-    run_name="function-gemma-270m-v1"
+    run_name="function-gemma-270m-v2-fixed"
 )
 
 # LoRA configuration
@@ -41,7 +61,7 @@ peft_config = LoraConfig(
 
 # Initialize and train
 trainer = SFTTrainer(
-    model="google/functiongemma-270m-it",
+    model=model_id,
     train_dataset=dataset,
     peft_config=peft_config,
     args=config,