Spaces:
Sleeping
Sleeping
File size: 3,760 Bytes
c1c9e88 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
#!/usr/bin/env python3
"""
模型合并脚本 - 将LoRA权重合并到基础模型中
用于推理和部署
"""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import argparse
def merge_lora_model(base_model_path, lora_model_path, output_path):
"""
合并LoRA权重到基础模型
Args:
base_model_path: 基础模型路径
lora_model_path: LoRA模型路径（训练输出）
output_path: 合并后模型保存路径
"""
print("📥 Loading base model...")
# 加载基础模型（不使用量化）
base_model = AutoModelForCausalLM.from_pretrained(
base_model_path,
torch_dtype=torch.float16,
device_map="auto",
trust_remote_code=True,
)
print("📥 Loading LoRA model...")
# 加载LoRA模型
model = PeftModel.from_pretrained(base_model, lora_model_path)
print("🔄 Merging LoRA weights...")
# 合并权重
model = model.merge_and_unload()
print("💾 Saving merged model...")
# 保存合并后的模型
model.save_pretrained(output_path, safe_serialization=True)
# 复制tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model_path)
tokenizer.save_pretrained(output_path)
print(f"✅ Model merged and saved to {output_path}")
def test_merged_model(model_path):
"""测试合并后的模型"""
print("🧪 Testing merged model...")
# 加载模型和tokenizer
model = AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype=torch.float16,
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(model_path)
# 测试提示
test_prompt = "### Human: Create an advertisement for a revolutionary AI-powered smartwatch\n### Assistant:"
inputs = tokenizer(test_prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=200,
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tokenizer.eos_token_id,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
generated_text = response[len(test_prompt):].strip()
print(f"\n📝 Test Prompt: Create an advertisement for a revolutionary AI-powered smartwatch")
print(f"📄 Generated Response:\n{generated_text}")
def main():
parser = argparse.ArgumentParser(description="Merge LoRA weights with base model")
parser.add_argument("--base_model", required=True, help="Path to base model")
parser.add_argument("--lora_model", required=True, help="Path to LoRA model (training output)")
parser.add_argument("--output", required=True, help="Output path for merged model")
parser.add_argument("--test", action="store_true", help="Test the merged model")
args = parser.parse_args()
# 合并模型
merge_lora_model(args.base_model, args.lora_model, args.output)
# 测试模型（可选）
if args.test:
test_merged_model(args.output)
if __name__ == "__main__":
# 示例用法
print("📋 Merge LoRA Model Script")
print("\n使用方法:")
print("python merge_model.py --base_model microsoft/DialoGPT-medium --lora_model ./results --output ./merged_model --test")
print("\n或者直接运行默认配置:")
# 默认配置
merge_lora_model(
base_model_path="microsoft/DialoGPT-medium", # 替换为实际的OpenAI OSS 120B模型
lora_model_path="./results",
output_path="./merged_model"
)
# 测试合并后的模型
test_merged_model("./merged_model") |