Upload train.py with huggingface_hub

train.py

@@ -8,11 +8,9 @@ from transformers import (
     TrainingArguments,
 )
 from peft import LoraConfig
-from trl import SFTTrainer, SFTConfig
+from trl import SFTTrainer
 
 # --- CONFIGURATION ---
-# Base model: Using a quantized Llama 3 or Mistral is recommended for consumer GPUs.
-# Ensure you have access to the model on Hugging Face (might need login).
 MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct" 
 DATASET_NAME = "ceperaltab/elixir-golden-dataset"
 OUTPUT_DIR = "elixir-model-qwen"
@@ -21,7 +19,6 @@ def main():
     print(f"Loading dataset from {DATASET_NAME}...")
     # 1. Load Dataset
     try:
-        # Load directly from HF Hub
         dataset = load_dataset(DATASET_NAME, split="train")
     except Exception as e:
         print(f"Error loading dataset: {e}")
@@ -46,9 +43,9 @@ def main():
     # 4. Load Tokenizer
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
     tokenizer.pad_token = tokenizer.eos_token
-    tokenizer.padding_side = "right" # Critical for fp16 training
+    tokenizer.padding_side = "right"
 
-    # 5. LoRA Config (Parameter Efficient Fine-Tuning)
+    # 5. LoRA Config
     peft_config = LoraConfig(
         lora_alpha=16,
         lora_dropout=0.1,
@@ -59,41 +56,41 @@ def main():
     )
 
     # 6. Formatting Function for Chat Dataset
-    # Converts {"messages": [...]} into the model's expected prompt format
     def formatting_prompts_func(examples):
         output_texts = []
         for messages in examples['messages']:
-            # Apply chat template (e.g., <|begin_of_text|><|start_header_id|>user...)
-            # We don't tokenize yet, SFTTrainer handles it
             text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
             output_texts.append(text)
         return output_texts
 
     print("Starting SFTTrainer setup...")
-    # 7. Trainer
+    
+    # 7. Training Arguments (using transformers TrainingArguments for stability)
+    training_args = TrainingArguments(
+        output_dir=OUTPUT_DIR,
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=4,
+        learning_rate=2e-4,
+        logging_steps=10,
+        num_train_epochs=1,
+        optim="paged_adamw_32bit",
+        fp16=True,
+        group_by_length=True,
+        save_strategy="epoch",
+        report_to="none",
+        push_to_hub=True,
+        hub_model_id=f"ceperaltab/{OUTPUT_DIR}",
+    )
+
+    # 8. Trainer - use older stable API
     trainer = SFTTrainer(
         model=model,
         train_dataset=dataset,
         peft_config=peft_config,
         formatting_func=formatting_prompts_func,
         tokenizer=tokenizer,
-        args=SFTConfig(
-            output_dir=OUTPUT_DIR,
-            max_seq_length=2048, # Moved here
-            per_device_train_batch_size=2,
-            gradient_accumulation_steps=4,
-            learning_rate=2e-4,
-            logging_steps=10,
-            num_train_epochs=1,
-            optim="paged_adamw_32bit",
-            fp16=True,
-            group_by_length=True,
-            save_strategy="epoch",
-            report_to="none",
-            push_to_hub=True,
-            hub_model_id=f"ceperaltab/{OUTPUT_DIR}",
-            dataset_text_field="text", # SFTConfig requires this or packing, though we use formatting_func
-        ),
+        args=training_args,
+        max_seq_length=2048,  # Passed directly to SFTTrainer (old API)
     )
 
     print("Starting training...")