Upload scripts/train_n8n_dpo.py with huggingface_hub

scripts/train_n8n_dpo.py

@@ -0,0 +1,244 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "transformers>=4.45.0",
+#     "trl>=0.12.0",
+#     "peft>=0.13.0",
+#     "datasets>=3.0.0",
+#     "accelerate>=1.0.0",
+#     "bitsandbytes>=0.44.0",
+#     "wandb>=0.18.0",
+#     "huggingface_hub>=0.26.0",
+#     "torch>=2.4.0",
+#     "einops>=0.8.0",
+#     "sentencepiece>=0.2.0",
+# ]
+# [tool.uv]
+# extra-index-url = ["https://download.pytorch.org/whl/cu124"]
+# ///
+"""
+Script d'entraînement DPO pour le modèle n8n Expert.
+À exécuter APRÈS l'entraînement SFT.
+
+Usage sur HuggingFace Jobs:
+    hf jobs uv run \
+        --script train_n8n_dpo.py \
+        --flavor h100x1 \
+        --name n8n-expert-dpo \
+        --timeout 12h \
+        --env BASE_MODEL=stmasson/n8n-expert-14b-sft
+
+Variables d'environnement:
+    - HF_TOKEN: Token HuggingFace
+    - BASE_MODEL: Modèle SFT à utiliser comme base
+    - WANDB_API_KEY: (optionnel) Pour le tracking
+"""
+
+import os
+import torch
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import LoraConfig, PeftModel
+from trl import DPOTrainer, DPOConfig
+from huggingface_hub import login
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+# Modèle SFT fine-tuné
+BASE_MODEL = os.environ.get("BASE_MODEL", "stmasson/n8n-expert-14b-sft")
+ORIGINAL_MODEL = os.environ.get("ORIGINAL_MODEL", "Qwen/Qwen2.5-14B-Instruct")
+
+# Dataset DPO
+DATASET_REPO = "stmasson/n8n-workflows-thinking"
+DPO_FILE = "n8n_dpo_train.jsonl"
+
+# Output
+OUTPUT_DIR = "./n8n-expert-dpo"
+HF_REPO = os.environ.get("HF_REPO", "stmasson/n8n-expert-14b-dpo")
+
+# Hyperparamètres DPO
+NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", "2"))
+BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "1"))
+GRAD_ACCUM = int(os.environ.get("GRAD_ACCUM", "16"))
+LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "5e-6"))
+BETA = float(os.environ.get("DPO_BETA", "0.1"))
+MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "8192"))
+MAX_PROMPT_LENGTH = int(os.environ.get("MAX_PROMPT_LENGTH", "2048"))
+
+# LoRA (plus léger pour DPO)
+LORA_R = int(os.environ.get("LORA_R", "32"))
+LORA_ALPHA = int(os.environ.get("LORA_ALPHA", "64"))
+
+# ============================================================================
+# AUTHENTIFICATION
+# ============================================================================
+
+print("=" * 60)
+print("ENTRAÎNEMENT DPO - N8N EXPERT")
+print("=" * 60)
+
+hf_token = os.environ.get("HF_TOKEN")
+if hf_token:
+    login(token=hf_token)
+    print("Authentifié sur HuggingFace")
+
+wandb_key = os.environ.get("WANDB_API_KEY")
+if wandb_key:
+    import wandb
+    wandb.login(key=wandb_key)
+    report_to = "wandb"
+else:
+    report_to = "none"
+
+# ============================================================================
+# CHARGEMENT DU MODÈLE
+# ============================================================================
+
+print(f"\nChargement du modèle SFT: {BASE_MODEL}")
+
+# Charger le modèle de base
+model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+    device_map="auto",
+    trust_remote_code=True,
+)
+
+# Charger le modèle de référence (pour DPO)
+ref_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+    device_map="auto",
+    trust_remote_code=True,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+print("Modèle chargé")
+
+# ============================================================================
+# CONFIGURATION LORA
+# ============================================================================
+
+print(f"\nConfiguration LoRA: r={LORA_R}, alpha={LORA_ALPHA}")
+
+lora_config = LoraConfig(
+    r=LORA_R,
+    lora_alpha=LORA_ALPHA,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM"
+)
+
+# ============================================================================
+# CHARGEMENT DU DATASET DPO
+# ============================================================================
+
+print(f"\nChargement du dataset DPO: {DATASET_REPO}")
+
+dataset = load_dataset(
+    DATASET_REPO,
+    data_files={"train": DPO_FILE},
+    split="train"
+)
+
+print(f"Exemples DPO: {len(dataset)}")
+
+# Fonction de formatage pour DPO
+def format_dpo_example(example):
+    """
+    Format attendu par DPOTrainer:
+    - prompt: le prompt de l'utilisateur
+    - chosen: la bonne réponse
+    - rejected: la mauvaise réponse
+    """
+    return {
+        "prompt": example["prompt"],
+        "chosen": example["chosen"],
+        "rejected": example["rejected"],
+    }
+
+# Le dataset devrait déjà être au bon format
+print("\nExemple de données DPO:")
+print(f"Prompt: {dataset[0]['prompt'][:200]}...")
+print(f"Chosen: {dataset[0]['chosen'][:200]}...")
+print(f"Rejected: {dataset[0]['rejected'][:200]}...")
+
+# ============================================================================
+# CONFIGURATION D'ENTRAÎNEMENT DPO
+# ============================================================================
+
+print(f"\nConfiguration DPO:")
+print(f"  - Beta: {BETA}")
+print(f"  - Epochs: {NUM_EPOCHS}")
+print(f"  - Batch size: {BATCH_SIZE}")
+print(f"  - Gradient accumulation: {GRAD_ACCUM}")
+print(f"  - Learning rate: {LEARNING_RATE}")
+
+dpo_config = DPOConfig(
+    output_dir=OUTPUT_DIR,
+    num_train_epochs=NUM_EPOCHS,
+    per_device_train_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    learning_rate=LEARNING_RATE,
+    beta=BETA,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.1,
+    bf16=True,
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=200,
+    save_total_limit=3,
+    max_length=MAX_LENGTH,
+    max_prompt_length=MAX_PROMPT_LENGTH,
+    gradient_checkpointing=True,
+    gradient_checkpointing_kwargs={"use_reentrant": False},
+    report_to=report_to,
+    run_name="n8n-expert-dpo",
+    hub_model_id=HF_REPO if hf_token else None,
+    push_to_hub=bool(hf_token),
+)
+
+# ============================================================================
+# ENTRAÎNEMENT DPO
+# ============================================================================
+
+print("\nInitialisation du DPO trainer...")
+
+trainer = DPOTrainer(
+    model=model,
+    ref_model=ref_model,
+    args=dpo_config,
+    train_dataset=dataset,
+    peft_config=lora_config,
+    tokenizer=tokenizer,
+)
+
+print("\n" + "=" * 60)
+print("DÉMARRAGE DE L'ENTRAÎNEMENT DPO")
+print("=" * 60)
+
+trainer.train()
+
+# ============================================================================
+# SAUVEGARDE
+# ============================================================================
+
+print("\nSauvegarde du modèle...")
+trainer.save_model(f"{OUTPUT_DIR}/final")
+
+if hf_token:
+    print(f"Push vers {HF_REPO}...")
+    trainer.push_to_hub()
+    print(f"Modèle disponible sur: https://huggingface.co/{HF_REPO}")
+
+print("\n" + "=" * 60)
+print("ENTRAÎNEMENT DPO TERMINÉ")
+print("=" * 60)