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

base_model_path = '../Finllm/Finllm_computing'
lora_model_path = '../saves/Finllm_computing_train/lora/sft/checkpoint-4000'
new_model_path = '../Finllm/Finllm_computing-merged'


# Loading the base model and tokenzier
model = AutoModel.from_pretrained(
       base_model_path, torch_dtype=torch.float16, trust_remote_code=True
    )

tokenzier = AutoTokenizer.from_pretrained(
       base_model_path, trust_remote_code=True
    )

# Loading the LoRA adapter
lora_model = PeftModel.from_pretrained(model, lora_model_path, torch_dtype=torch.float16,)

# Applying the LoRA
new_model = lora_model.merge_and_unload()

# Saving the target model to new model paths
new_model.save_pretrained(new_model_path)
tokenzier.save_pretrained(new_model_path)