#!/usr/bin/env python3
"""
DeepMath — merge LoRA adapters into a single, directly-loadable full model.

DeepMath is produced by a TWO-STAGE LoRA pipeline (LLaMA-Factory, template `deepseekr1`):

    DeepSeek-R1-Distill-Qwen-7B  (base)
      + SFT LoRA adapter   --(trained on cleaned NuminaMath-CoT)-->
        DeepMath-SFT  (merged intermediate)
      + DPO LoRA adapter --(trained on math preference pairs)---->
        DeepMath (final)

This repository ships the two adapters (small). Run this script ONCE on a
machine with the base model available to reconstruct the final full model:

    python merge.py \
        --base  deepseek-ai/DeepSeek-R1-Distill-Qwen-7B \
        --sft   ./sft_adapter \
        --dpo   ./adapter \
        --out   ./DeepMath-merged

Notes
-----
* Merging is pure PEFT `merge_and_unload` (no LLaMA-Factory needed at merge time).
* CPU is fine but slow; a GPU with ~16GB+ is comfortable. Weights are bf16.
* The original training merged with `llamafactory-cli export --template deepseekr1`;
  both approaches yield the same weights for these adapters.
"""

import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel


def merge(base_path: str, sft_adapter: str, dpo_adapter: str, out: str, device: str, dtype: str):
    torch_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[dtype]
    print(f"[1/4] Loading base model: {base_path}  ({dtype}, device={device})")
    model = AutoModelForCausalLM.from_pretrained(
        base_path, torch_dtype=torch_dtype, device_map=device
    )
    tokenizer = AutoTokenizer.from_pretrained(base_path)

    print(f"[2/4] Merging SFT adapter: {sft_adapter}")
    model = PeftModel.from_pretrained(model, sft_adapter)
    model = model.merge_and_unload()

    print(f"[3/4] Merging DPO adapter: {dpo_adapter}")
    model = PeftModel.from_pretrained(model, dpo_adapter)
    model = model.merge_and_unload()

    print(f"[4/4] Saving merged model -> {out}")
    model.save_pretrained(out, safe_serialization=True)
    tokenizer.save_pretrained(out)
    print("Done. Load with: AutoModelForCausalLM.from_pretrained('%s')" % out)


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Merge DeepMath SFT+DPO LoRA adapters into the base model.")
    ap.add_argument("--base", default="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
                    help="Base model path or HF id.")
    ap.add_argument("--sft", default="./sft_adapter", help="Path to the SFT LoRA adapter.")
    ap.add_argument("--dpo", default="./adapter", help="Path to the final DPO LoRA adapter.")
    ap.add_argument("--out", default="./DeepMath-merged", help="Output directory for the merged model.")
    ap.add_argument("--device", default="auto", help="device_map: 'auto', 'cpu', 'cuda:0', ...")
    ap.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"], help="Weight dtype.")
    args = ap.parse_args()
    merge(args.base, args.sft, args.dpo, args.out, args.device, args.dtype)