55 lines
2.5 KiB
Python
55 lines
2.5 KiB
Python
# 导入所需的库
|
|
from transformers import AutoModelForCausalLM # 用于加载预训练的语言模型
|
|
from peft import LoraConfig, get_peft_model, PeftModel # 用于处理LoRA权重
|
|
import argparse # 用于解析命令行参数
|
|
import shutil # 用于文件操作，如复制
|
|
import os # 用于文件路径操作
|
|
import torch # 用于深度学习操作
|
|
|
|
def main():
|
|
# 创建参数解析器
|
|
parser = argparse.ArgumentParser()
|
|
# 添加命令行参数
|
|
parser.add_argument("--base_model_path", type=str, required=True,
|
|
help="Path to pretrained model or model identifier from huggingface.co/models")
|
|
parser.add_argument("--adapter_model_path", type=str, required=True, help="Path to adapter model")
|
|
parser.add_argument("--output_path", type=str, required=True, help="Path to save the output model")
|
|
parser.add_argument("--save_dtype", type=str, choices=['bf16', 'fp32', 'fp16'],
|
|
default='fp32', help="In which dtype to save, fp32, bf16 or fp16.")
|
|
# 解析命令行参数
|
|
args = parser.parse_args()
|
|
|
|
name2dtype = {'bf16': torch.bfloat16, 'fp32': torch.float32, 'fp16': torch.float16}
|
|
# 加载基座模型
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
args.base_model_path, device_map='cpu',
|
|
trust_remote_code=True, torch_dtype=name2dtype[args.save_dtype]
|
|
)
|
|
# 在基座模型的基础上加载 adapter 权重
|
|
model = PeftModel.from_pretrained(model, args.adapter_model_path, trust_remote_code=True)
|
|
# 融合模型和 adapter
|
|
model = model.merge_and_unload()
|
|
# 保存融合后的模型权重
|
|
model.save_pretrained(args.output_path, safe_serialization=False)
|
|
|
|
# Copy tokenizer, config and other non-weight files from base model
|
|
# Skip model weight files (.safetensors, .bin, .pt) and index files
|
|
_SKIP_SUFFIXES = ('.safetensors', '.bin', '.pt', '.pth')
|
|
_SKIP_NAMES = {'model.safetensors.index.json', 'pytorch_model.bin.index.json'}
|
|
|
|
for fname in os.listdir(args.base_model_path):
|
|
src = os.path.join(args.base_model_path, fname)
|
|
if not os.path.isfile(src):
|
|
continue
|
|
if fname in _SKIP_NAMES or fname.endswith(_SKIP_SUFFIXES):
|
|
continue
|
|
dst = os.path.join(args.output_path, fname)
|
|
if not os.path.exists(dst):
|
|
shutil.copy(src, dst)
|
|
print(f'Copied {fname}')
|
|
|
|
print(f'Merged model weight is saved to {args.output_path}')
|
|
|
|
if __name__ == "__main__":
|
|
main()
|