Upload train.py with huggingface_hub

train.py

@@ -5,10 +5,9 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
-    TrainingArguments,
 )
 from peft import LoraConfig
-from trl import SFTTrainer
+from trl import SFTTrainer, SFTConfig
 
 # --- CONFIGURATION ---
 MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct" 
@@ -55,19 +54,12 @@ def main():
         target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
     )
 
-    # 6. Formatting Function for Chat Dataset
-    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...")
     
-    # 7. Training Arguments (using transformers TrainingArguments for stability)
-    training_args = TrainingArguments(
+    # 6. SFTConfig (new TRL API - all config goes here)
+    sft_config = SFTConfig(
         output_dir=OUTPUT_DIR,
+        max_seq_length=2048,
         per_device_train_batch_size=2,
         gradient_accumulation_steps=4,
         learning_rate=2e-4,
@@ -82,15 +74,14 @@ def main():
         hub_model_id=f"ceperaltab/{OUTPUT_DIR}",
     )
 
-    # 8. Trainer - use older stable API
+    # 7. Trainer - new API: use processing_class instead of tokenizer
+    # The trainer automatically handles conversational datasets with "messages" field
     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,
-        tokenizer=tokenizer,
-        args=training_args,
-        max_seq_length=2048,  # Passed directly to SFTTrainer (old API)
     )
 
     print("Starting training...")