Upload train.py with huggingface_hub

train.py

@@ -5,11 +5,16 @@
 from datasets import load_dataset
 from peft import LoraConfig
 from trl import SFTTrainer, SFTConfig
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, TrainingArguments
+import trl
+import transformers
 import trackio
 import os
+import inspect
 
-print("🚀 Starting FunctionGemma 270M Fine-tuning (V3 - Config Fix)")
+print(f"🚀 Starting FunctionGemma 270M Fine-tuning (V4 - Diagnostic)")
+print(f"📦 TRL Version: {trl.__version__}")
+print(f"📦 Transformers Version: {transformers.__version__}")
 
 model_id = "google/functiongemma-270m-it"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -18,7 +23,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 dataset = load_dataset("epinfomax/vn-function-calling-dataset", split="train")
 
 def format_conversation(example):
-    # Pre-render the conversation using the model's chat template
     text = tokenizer.apply_chat_template(
         example["messages"],
         tools=example["tools"],
@@ -31,40 +35,63 @@ 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",
-    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,
-    per_device_train_batch_size=4, # Reduced for stability
-    gradient_accumulation_steps=4,
-    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-v3-fixed"
-)
+# Trying max_seq_length again but checking if it exists in SFTConfig first
+sft_config_args = {
+    "dataset_text_field": "text",
+    "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,
+    "per_device_train_batch_size": 4,
+    "gradient_accumulation_steps": 4,
+    "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-v4-diag"
+}
+
+# Check which parameter to use
+sft_fields = SFTConfig.__dataclass_fields__
+if "max_seq_length" in sft_fields:
+    print("✅ Using max_seq_length in SFTConfig")
+    sft_config_args["max_seq_length"] = 1024
+elif "max_length" in sft_fields:
+    print("✅ Using max_length in SFTConfig")
+    sft_config_args["max_length"] = 1024
+else:
+    print("⚠️ Neither max_seq_length nor max_length found in SFTConfig fields!")
+    print("Fields:", list(sft_fields.keys()))
+
+config = SFTConfig(**sft_config_args)
+
+# Initialize and train
+print("🎯 Initializing SFTTrainer...")
+trainer_kwargs = {
+    "model": model_id,
+    "train_dataset": dataset,
+    "peft_config": peft_config,
+    "args": config,
+}
+
+# Check SFTTrainer init signature
+trainer_params = inspect.signature(SFTTrainer.__init__).parameters
+if "max_seq_length" in trainer_params and "max_seq_length" not in sft_config_args:
+    print("✅ Adding max_seq_length to SFTTrainer")
+    trainer_kwargs["max_seq_length"] = 1024
 
-# LoRA configuration
 peft_config = LoraConfig(
     r=16,
     lora_alpha=32,
     target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
     task_type="CAUSAL_LM",
 )
+trainer_kwargs["peft_config"] = peft_config
 
-# Initialize and train
-trainer = SFTTrainer(
-    model=model_id,
-    train_dataset=dataset,
-    peft_config=peft_config,
-    args=config,
-    max_seq_length=1024, # Moved here from SFTConfig
-)
+trainer = SFTTrainer(**trainer_kwargs)
 
 trainer.train()
 trainer.push_to_hub()