Upload scripts/train_sft_n8n_multitask.py with huggingface_hub

scripts/train_sft_n8n_multitask.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+# /// script
+# dependencies = [
+#     "trl>=0.12.0",
+#     "transformers>=4.46.0",
+#     "accelerate>=0.24.0",
+#     "peft>=0.7.0",
+#     "trackio",
+#     "bitsandbytes",
+#     "sentencepiece",
+#     "protobuf",
+# ]
+# ///
+
+"""
+SFT training for n8n agentic multi-task workflows.
+
+Continues fine-tuning from stmasson/mistral-7b-n8n-thinking-orpo (ORPO-trained model)
+on the n8n-agentic-multitask dataset for complex multi-step tasks:
+- generate: Create n8n workflows from descriptions
+- edit: Modify existing workflows
+- fix: Repair broken workflows
+- improve: Optimize and enhance workflows
+- explain: Describe what workflows do
+- debug: Diagnose workflow issues
+
+The model learns to use <thinking> tags for chain-of-thought reasoning
+before producing structured JSON outputs.
+"""
+
+import trackio
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, PeftModel
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from trl import SFTTrainer, SFTConfig
+
+
+# Load multitask dataset
+print("Loading n8n-agentic-multitask dataset...")
+train_dataset = load_dataset(
+    "stmasson/n8n-agentic-multitask",
+    data_files="data/multitask_large/train.jsonl",
+    split="train"
+)
+eval_dataset = load_dataset(
+    "stmasson/n8n-agentic-multitask",
+    data_files="data/multitask_large/val.jsonl",
+    split="train"
+)
+
+print(f"Train: {len(train_dataset)} examples")
+print(f"Eval: {len(eval_dataset)} examples")
+
+# Load tokenizer from ORPO-trained model
+MODEL_NAME = "stmasson/mistral-7b-n8n-thinking-orpo"
+BASE_MODEL = "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
+
+# Step 1: Load base model WITHOUT quantization to merge ORPO adapter
+print(f"Loading base model {BASE_MODEL} (full precision for merge)...")
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    attn_implementation="sdpa",
+)
+
+print(f"Loading ORPO adapter from {MODEL_NAME}...")
+model = PeftModel.from_pretrained(base_model, MODEL_NAME)
+
+print("Merging ORPO adapter into base model...")
+model = model.merge_and_unload()
+print("ORPO adapter merged successfully!")
+
+# Step 2: Prepare for LoRA training with gradient checkpointing
+model.gradient_checkpointing_enable()
+model.enable_input_require_grads()
+
+# New LoRA configuration for SFT training
+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",
+)
+
+# SFT training configuration
+config = SFTConfig(
+    # Hub settings
+    output_dir="mistral-7b-n8n-agentic-multitask",
+    push_to_hub=True,
+    hub_model_id="stmasson/mistral-7b-n8n-agentic-multitask",
+    hub_strategy="every_save",
+    hub_private_repo=False,
+
+    # Training parameters
+    num_train_epochs=1,  # Large dataset, 1 epoch is enough
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=32,  # Effective batch size = 32
+    learning_rate=2e-5,  # Lower LR for continued fine-tuning
+    max_length=4096,  # Longer context for complex workflows
+
+    # Memory optimization
+    gradient_checkpointing=True,
+    bf16=True,
+
+    # Logging & checkpointing
+    logging_steps=25,
+    save_strategy="steps",
+    save_steps=500,
+    save_total_limit=3,
+
+    # Evaluation
+    eval_strategy="steps",
+    eval_steps=500,
+
+    # Optimization
+    warmup_ratio=0.03,
+    lr_scheduler_type="cosine",
+    optim="adamw_8bit",
+
+    # Monitoring
+    report_to="trackio",
+    project="n8n-agentic-training",
+    run_name="mistral-7b-multitask-sft",
+)
+
+# Initialize trainer
+print("Initializing SFT trainer...")
+trainer = SFTTrainer(
+    model=model,
+    processing_class=tokenizer,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    peft_config=lora_config,
+    args=config,
+)
+
+print("Starting SFT training...")
+print(f"  Base: stmasson/mistral-7b-n8n-thinking-orpo (merged)")
+print(f"  Dataset: stmasson/n8n-agentic-multitask")
+print(f"  Output: stmasson/mistral-7b-n8n-agentic-multitask")
+print(f"  Tasks: generate, edit, fix, improve, explain, debug")
+
+trainer.train()
+
+print("Pushing final model to Hub...")
+trainer.push_to_hub()
+
+# Finish Trackio
+trackio.finish()
+
+print("Training complete!")
+print("Model: https://huggingface.co/stmasson/mistral-7b-n8n-agentic-multitask")
+print("Metrics: https://huggingface.co/spaces/stmasson/trackio")