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

def merge_and_export_model(model_path: str, output_path: str):
    """合并LoRA权重并导出模型"""
    print("🔄 合并LoRA权重...")
    
    # 加载基座模型
    base_model = AutoModelForCausalLM.from_pretrained(
        "Qwen/Qwen2.5-1.5B",
        torch_dtype=torch.float16,
        device_map="auto"
    )
    
    # 加载LoRA适配器
    model = PeftModel.from_pretrained(base_model, model_path)
    
    # 合并权重
    merged_model = model.merge_and_unload()
    
    # 保存合并后的模型
    merged_model.save_pretrained(
        output_path,
        torch_dtype=torch.float16,
        safe_serialization=True
    )
    
    # 保存分词器
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tokenizer.save_pretrained(output_path)
    
    print(f"✅ 模型已导出到: {output_path}")
    
    return merged_model

# 执行导出
merged_model = merge_and_export_model(
    "./sft_sentiment_qwen",
    "./sft_sentiment_qwen_merged"
)