Upload scripts/train_orpo_n8n_thinking.py with huggingface_hub

scripts/train_orpo_n8n_thinking.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+# /// script
+# dependencies = [
+#     "trl>=0.12.0",
+#     "transformers>=4.46.0",
+#     "accelerate>=0.24.0",
+#     "peft>=0.7.0",
+#     "trackio",
+#     "bitsandbytes",
+# ]
+# ///
+
+"""
+ORPO training for n8n workflows with chain-of-thought reasoning.
+
+Fine-tunes stmasson/mistral-7b-n8n-workflows on the n8n-workflows-thinking dataset
+to generate structured reasoning (<thinking>) before producing n8n workflow JSON.
+
+ORPO (Odds Ratio Preference Optimization) combines SFT and preference learning
+in a single training objective, making it more efficient than DPO for this use case.
+"""
+
+import trackio
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import ORPOTrainer, ORPOConfig
+
+
+# Load ORPO dataset
+print("Loading n8n-workflows-thinking dataset (ORPO split)...")
+train_dataset = load_dataset(
+    "stmasson/n8n-workflows-thinking",
+    data_files="data/orpo/train.jsonl",
+    split="train"
+)
+eval_dataset = load_dataset(
+    "stmasson/n8n-workflows-thinking",
+    data_files="data/orpo/validation.jsonl",
+    split="train"
+)
+
+print(f"Train: {len(train_dataset)} examples")
+print(f"Eval: {len(eval_dataset)} examples")
+
+# Remove metadata column (not needed for training)
+train_dataset = train_dataset.remove_columns(["metadata"])
+eval_dataset = eval_dataset.remove_columns(["metadata"])
+
+# LoRA configuration for efficient training on 7B model
+lora_config = LoraConfig(
+    r=32,
+    lora_alpha=64,
+    lora_dropout=0.05,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+    task_type="CAUSAL_LM",
+)
+
+# ORPO training configuration
+config = ORPOConfig(
+    # Hub settings - CRITICAL for saving
+    output_dir="mistral-7b-n8n-thinking-orpo",
+    push_to_hub=True,
+    hub_model_id="stmasson/mistral-7b-n8n-thinking-orpo",
+    hub_strategy="every_save",
+    hub_private_repo=False,
+
+    # ORPO-specific parameter
+    beta=0.1,  # Weight for the odds ratio loss
+
+    # Training parameters
+    num_train_epochs=2,
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=16,  # Effective batch size = 16
+    learning_rate=5e-5,
+    max_length=4096,  # Long context for workflows + thinking
+    max_prompt_length=512,
+
+    # Memory optimization
+    gradient_checkpointing=True,
+    bf16=True,
+
+    # Logging & checkpointing
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=200,
+    save_total_limit=3,
+
+    # Evaluation
+    eval_strategy="steps",
+    eval_steps=200,
+
+    # Optimization
+    warmup_ratio=0.1,
+    lr_scheduler_type="cosine",
+    optim="adamw_8bit",  # Memory-efficient optimizer
+
+    # Monitoring with Trackio
+    report_to="trackio",
+    project="n8n-thinking-training",
+    run_name="mistral-7b-orpo-reasoning",
+)
+
+# Initialize trainer
+print("Initializing ORPO trainer...")
+trainer = ORPOTrainer(
+    model="stmasson/mistral-7b-n8n-workflows",
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    peft_config=lora_config,
+    args=config,
+)
+
+print("Starting ORPO training...")
+print(f"  Model: stmasson/mistral-7b-n8n-workflows")
+print(f"  Dataset: stmasson/n8n-workflows-thinking (ORPO)")
+print(f"  Output: stmasson/mistral-7b-n8n-thinking-orpo")
+
+trainer.train()
+
+print("Pushing final model to Hub...")
+trainer.push_to_hub()
+
+# Finish Trackio tracking
+trackio.finish()
+
+print("Training complete!")
+print("Model: https://huggingface.co/stmasson/mistral-7b-n8n-thinking-orpo")
+print("Metrics: https://huggingface.co/spaces/stmasson/trackio")