# merge the lora adapter into the base model and save the result as a normal
# transformers checkpoint. ONNX export is happier with a single dense model.

from pathlib import Path

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

from cleanup.config import TrainConfig


def merge_adapter(cfg: TrainConfig, adapter_dir: Path, out_dir: Path) -> Path:
    adapter_dir = Path(adapter_dir)
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    print(f"[merge] loading base {cfg.base_model}")
    model = AutoModelForCausalLM.from_pretrained(
        cfg.base_model,
        torch_dtype=torch.float32,
    )
    print(f"[merge] loading adapter {adapter_dir}")
    merged = PeftModel.from_pretrained(model, adapter_dir)
    merged = merged.merge_and_unload()
    merged.save_pretrained(out_dir)
    tokenizer = AutoTokenizer.from_pretrained(adapter_dir, use_fast=True)
    tokenizer.save_pretrained(out_dir)
    print(f"[merge] wrote merged transformers checkpoint to {out_dir}")
    return out_dir