from typing import Union

import torch
from peft import PeftModel
from peft.tuners.lora import LoraModel
from transformers import LlamaForCausalLM as ModelCls
from transformers import LlamaTokenizerFast as TkCls

PeftCls = Union[PeftModel, LoraModel]

orig_model = "TheBloke/Llama-2-7B-Chat-fp16"
lora_model = "models/Llama-7B-TwAddr-LoRA"
output_dir = "models/Llama-7B-TwAddr-Merged"

model = ModelCls.from_pretrained(
    orig_model,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
)

# Due to generation config validation.
model.generation_config.temperature = 1.0
model.generation_config.top_p = 1.0

model: PeftCls = PeftModel.from_pretrained(
    model,
    lora_model,
    torch_dtype=torch.float16,
)

model = model.merge_and_unload()
model.save_pretrained(
    output_dir,
    safe_serialization=True,
)

# Tokenizer 也要跟著另外存一份
tk: TkCls = TkCls.from_pretrained(orig_model)
tk.save_pretrained(output_dir)