Upload train.py with huggingface_hub

train.py

@@ -5,9 +5,10 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
+    TrainingArguments,
 )
 from peft import LoraConfig
-from trl import SFTTrainer, SFTConfig
+from trl import SFTTrainer
 
 # --- CONFIGURATION ---
 MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct" 
@@ -54,12 +55,19 @@ def main():
         target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
     )
 
+    # 6. Formatting Function for Chat Dataset (TRL v0.8.6 API)
+    def formatting_prompts_func(examples):
+        output_texts = []
+        for messages in examples['messages']:
+            text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+            output_texts.append(text)
+        return output_texts
+
     print("Starting SFTTrainer setup...")
     
-    # 6. SFTConfig (new TRL API - all config goes here)
-    sft_config = SFTConfig(
+    # 7. Training Arguments (TRL v0.8.6 uses TrainingArguments from transformers)
+    training_args = TrainingArguments(
         output_dir=OUTPUT_DIR,
-        max_seq_length=2048,
         per_device_train_batch_size=2,
         gradient_accumulation_steps=4,
         learning_rate=2e-4,
@@ -74,14 +82,15 @@ def main():
         hub_model_id=f"ceperaltab/{OUTPUT_DIR}",
     )
 
-    # 7. Trainer - new API: use processing_class instead of tokenizer
-    # The trainer automatically handles conversational datasets with "messages" field
+    # 8. SFTTrainer (TRL v0.8.6 API)
     trainer = SFTTrainer(
         model=model,
-        args=sft_config,
         train_dataset=dataset,
-        processing_class=tokenizer,  # New API: processing_class replaces tokenizer
         peft_config=peft_config,
+        formatting_func=formatting_prompts_func,
+        max_seq_length=2048,
+        tokenizer=tokenizer,
+        args=training_args,
     )
 
     print("Starting training...")