Upload scripts/train_orpo_n8n_thinking.py with huggingface_hub

scripts/train_orpo_n8n_thinking.py

@@ -23,6 +23,7 @@ in a single training objective, making it more efficient than DPO for this use c
 import trackio
 from datasets import load_dataset
 from peft import LoraConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from trl import ORPOTrainer, ORPOConfig
 
 
@@ -46,6 +47,20 @@ print(f"Eval: {len(eval_dataset)} examples")
 train_dataset = train_dataset.remove_columns(["metadata"])
 eval_dataset = eval_dataset.remove_columns(["metadata"])
 
+# Load model and tokenizer
+MODEL_NAME = "stmasson/mistral-7b-n8n-workflows"
+print(f"Loading tokenizer from {MODEL_NAME}...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+print(f"Loading model from {MODEL_NAME}...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype="auto",
+    device_map="auto",
+)
+
 # LoRA configuration for efficient training on 7B model
 lora_config = LoraConfig(
     r=32,
@@ -103,7 +118,8 @@ config = ORPOConfig(
 # Initialize trainer
 print("Initializing ORPO trainer...")
 trainer = ORPOTrainer(
-    model="stmasson/mistral-7b-n8n-workflows",
+    model=model,
+    processing_class=tokenizer,
     train_dataset=train_dataset,
     eval_dataset=eval_dataset,
     peft_config=lora_config,