adikuma's picture
initial upload: cleanup code and 688-pair seed dataset
fd0b01f verified
Raw
History Blame Contribute Delete
1.07 kB
# 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