Upload DPO-Training/merge_lora.py with huggingface_hub

DPO-Training/merge_lora.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""
+Merge LoRA adapters with base model after DPO training.
+
+Usage:
+    python merge_lora.py --lora_path ./qwen3-0.6b-dpo --output_path ./qwen3-0.6b-dpo-merged
+"""
+
+import argparse
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+import os
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Merge LoRA adapters with base model")
+    parser.add_argument(
+        "--base_model", default="Qwen/Qwen3-0.6B", help="Base model name or path"
+    )
+    parser.add_argument(
+        "--lora_path", default="./qwen3-0.6b-dpo", help="Path to LoRA adapters"
+    )
+    parser.add_argument(
+        "--output_path",
+        default="./qwen3-0.6b-dpo-merged",
+        help="Output path for merged model",
+    )
+    parser.add_argument(
+        "--device", default="auto", help="Device to use (auto/cuda/cpu)"
+    )
+    parser.add_argument(
+        "--push_to_hub", default=None, help="Push to HuggingFace Hub (repo name)"
+    )
+
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("Merging LoRA Adapters")
+    print("=" * 60)
+    print(f"Base model: {args.base_model}")
+    print(f"LoRA path: {args.lora_path}")
+    print(f"Output: {args.output_path}")
+    print("=" * 60)
+
+    # Load tokenizer
+    print("\n📥 Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(args.base_model, trust_remote_code=True)
+
+    # Load base model
+    print("📥 Loading base model...")
+    base_model = AutoModelForCausalLM.from_pretrained(
+        args.base_model,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map=args.device,
+        trust_remote_code=True,
+    )
+
+    # Load LoRA adapters
+    print("📥 Loading LoRA adapters...")
+    model = PeftModel.from_pretrained(base_model, args.lora_path)
+
+    # Merge adapters
+    print("🔧 Merging adapters...")
+    model = model.merge_and_unload()
+
+    # Save merged model
+    print(f"💾 Saving merged model to {args.output_path}...")
+    os.makedirs(args.output_path, exist_ok=True)
+    model.save_pretrained(args.output_path)
+    tokenizer.save_pretrained(args.output_path)
+
+    print("\n" + "=" * 60)
+    print("✅ Merge Complete!")
+    print("=" * 60)
+    print(f"Merged model saved to: {args.output_path}")
+
+    # Push to hub if requested
+    if args.push_to_hub:
+        print(f"\n📤 Pushing to HuggingFace Hub: {args.push_to_hub}")
+        model.push_to_hub(args.push_to_hub)
+        tokenizer.push_to_hub(args.push_to_hub)
+        print(f"✅ Pushed to: https://huggingface.co/{args.push_to_hub}")
+
+    print("\n🎉 Done!")
+
+
+if __name__ == "__main__":
+    main()