大模型应用：中小显存适配方案：大模型微调底座选型指标与应用实现.52

from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
import torch
import time
import matplotlib.pyplot as plt
import numpy as np

# 2. 全局配置（小白仅需修改TEST_MODE切换分支）
TEST_MODE = "simulate"  # 切换测试分支
TEST_PROMPTS = [  # 统一测试用例（3条电商文案，确保对比公平）
    "生成无线耳机文案：突出续航7天、主动降噪、平价，80字左右",
    "生成护肤品文案：突出敏感肌适用、补水保湿、无香精，80字左右",
    "生成充电宝文案：突出20000mAh、双向快充、轻薄便携，80字左右"
]
MODEL_CONFIGS = [  # 对比底座配置
    {
        "name": "ChatGLM-6B（INT4量化）",
        "model_id": "THUDM/chatglm-6b",
        "color": "#2E86AB"  # 蓝色
    },
    {
        "name": "Qwen-7B（INT4量化）",
        "model_id": "Qwen/Qwen-7B-Chat",
        "color": "#A23B72"  # 紫红色
    }
]
# INT4量化配置（8G显存必加，降低50%+显存占用）
QUANT_CONFIG = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

def load_model_and_tokenizer(model_id):
    """加载模型和分词器（统一逻辑，适配两个底座）"""
    try:
        tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            trust_remote_code=True,
            cache_dir="./model_cache",  # 缓存目录，避免重复下载
            # 国内用户添加镜像地址，解决下载超时（可选）
            # proxy="https://mirror.sjtu.edu.cn/huggingface"
        )
        model = AutoModel.from_pretrained(
            model_id,
            trust_remote_code=True,
            cache_dir="./model_cache",
            quantization_config=QUANT_CONFIG,
            device_map="auto"  # 自动分配GPU/CPU
        ).eval()  # 评估模式，关闭训练层，节省显存
        return model, tokenizer
    except Exception as e:
        print(f"模型加载失败：{e}")
        print("提示：国内用户可替换model_id为镜像地址，或检查网络连接")
        exit()

def calculate_score(response, prompt):
    """文案效果评分（1-5分，量化对比效果，规则可自定义）"""
    score = 0
    # 规则1：包含所有核心卖点（3分，确保满足任务需求）
    keywords = {
        TEST_PROMPTS[0]: ["续航7天", "主动降噪", "平价"],
        TEST_PROMPTS[1]: ["敏感肌", "补水保湿", "无香精"],
        TEST_PROMPTS[2]: ["20000mAh", "双向快充", "轻薄便携"]
    }[prompt]
    for kw in keywords:
        if kw in response:
            score += 1
    # 规则2：语句流畅（1分，避免语序混乱）
    if len(response.split("。")) >= 2 and "，" in response:
        score += 1
    # 规则3：符合电商风格（1分，适配营销场景）
    if any(word in response for word in ["必备", "闭眼入", "性价比", "推荐"]):
        score += 1
    return score

def get_gpu_memory():
    """获取GPU显存占用（GB），无GPU返回0"""
    if torch.cuda.is_available():
        used = torch.cuda.memory_allocated() / (1024**3)
        return round(used, 2)
    else:
        return 0.0

def test_real_model(model_config):
    """实际测试：加载模型生成结果，获取真实指标"""
    model_name = model_config["name"]
    model_id = model_config["model_id"]
    print(f"\n=== 开始测试 {model_name} ===")
    
    # 加载模型并记录显存
    print(f"正在加载模型（首次运行需下载，约13GB）...")
    model, tokenizer = load_model_and_tokenizer(model_id)
    gpu_used = get_gpu_memory()
    print(f"模型加载完成，显存占用：{gpu_used}GB")
    
    # 存储测试指标
    metrics = {
        "显存占用(GB)": gpu_used,
        "生成速度(秒/条)": [],
        "效果评分(1-5分)": [],
        "生成结果": []
    }
    
    # 逐一生成测试用例
    for i, prompt in enumerate(TEST_PROMPTS):
        inputs = tokenizer(prompt, return_tensors="pt").to(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        # 记录生成时间（速度指标）
        start_time = time.time()
        with torch.no_grad():  # 关闭梯度计算，进一步节省显存
            response = model.generate(
                **inputs,
                max_new_tokens=100,  # 生成字数上限
                temperature=0.7,     # 随机性适中
                top_p=0.9,           # 采样阈值，避免杂乱
                repetition_penalty=1.1  # 抑制重复语句
            )
        gen_time = round(time.time() - start_time, 2)
        # 解码结果并评分
        response_text = tokenizer.decode(response[0], skip_special_tokens=True)
        score = calculate_score(response_text, prompt)
        # 保存指标
        metrics["生成速度(秒/条)"].append(gen_time)
        metrics["效果评分(1-5分)"].append(score)
        metrics["生成结果"].append(response_text)
        
        # 打印单条用例结果
        print(f"\n用例{i+1}：")
        print(f"生成时间：{gen_time}秒 | 评分：{score}分")
        print(f"生成文案：{response_text[:100]}...")  # 截取前100字展示
    
    # 计算平均值
    metrics["平均生成速度(秒/条)"] = round(np.mean(metrics["生成速度(秒/条)"]), 2)
    metrics["平均效果评分(1-5分)"] = round(np.mean(metrics["效果评分(1-5分)"]), 1)
    # 打印模型测试总结
    print(f"\n{model_name} 测试总结：")
    print(f"显存占用：{metrics['显存占用(GB)']}GB")
    print(f"平均生成速度：{metrics['平均生成速度(秒/条)']}秒/条")
    print(f"平均效果评分：{metrics['平均效果评分(1-5分)']}分")
    return metrics

def get_simulate_data():
    """模拟测试数据：基于案例逻辑生成，无需加载模型"""
    print("=== 使用模拟数据生成对比结果（无需下载模型）===")
    simulate_results = [
        {
            "显存占用(GB)": 5.8,
            "平均生成速度(秒/条)": 1.2,
            "平均效果评分(1-5分)": 4.2,
            "生成结果": [
                "这款无线耳机续航拉满7天无忧，主动降噪隔绝杂音，平价也能拥有旗舰体验！日常通勤必备，性价比直接拉满，闭眼入不亏～",
                "敏感肌专属补水护肤品，无香精添加温和不刺激，深层补水保湿锁水，换季维稳必备，平价大碗适合日常囤货，干皮姐妹冲！",
                "20000mAh大容量充电宝，双向快充告别电量焦虑，机身轻薄便携不占地，出差旅行必备，性价比拉满，学生党闭眼入～"
            ]
        },
        {
            "显存占用(GB)": 7.3,
            "平均生成速度(秒/条)": 2.1,
            "平均效果评分(1-5分)": 4.0,
            "生成结果": [
                "无线耳机续航长达7天，主动降噪技术过滤环境音，价格亲民性价比高，适合日常使用，办公学习都能搭，值得入手～",
                "专为敏感肌设计的补水保湿护肤品，无香精配方温和修护，深层滋养肌肤，质地清爽不黏腻，日常护理必备单品～",
                "20000mAh充电宝支持双向快充，机身轻薄便携易收纳，满足多设备充电需求，外出游玩必备，实用性拉满～"
            ]
        }
    ]
    return simulate_results

def generate_comparison_chart(results):
    """生成差异对比图表（1行3列柱状图，高清保存）"""
    # 配置中文显示，避免乱码
    plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    # 创建图表（18x6英寸，适合展示）
    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
    model_names = [cfg["name"] for cfg in MODEL_CONFIGS]
    colors = [cfg["color"] for cfg in MODEL_CONFIGS]
    
    # 子图1：显存占用对比（核心适配指标）
    gpu_data = [results[i]["显存占用(GB)"] for i in range(2)]
    bars1 = axes[0].bar(model_names, gpu_data, color=colors, alpha=0.8, edgecolor='black', linewidth=1)
    axes[0].set_title("显存占用对比（8G显存场景）", fontsize=14, fontweight="bold", pad=20)
    axes[0].set_ylabel("显存占用(GB)", fontsize=12)
    axes[0].set_ylim(0, 8)  # 固定Y轴为8G（显存上限）
    axes[0].grid(axis='y', alpha=0.3)
    # 标注数值与占比
    for i, (bar, v) in enumerate(zip(bars1, gpu_data)):
        ratio = (v / 8) * 100
        axes[0].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.2,
                     f"{v}GB\n({ratio:.1f}%)", ha="center", fontsize=11, fontweight="bold")
    
    # 子图2：平均生成速度对比（部署效率指标）
    speed_data = [results[i]["平均生成速度(秒/条)"] for i in range(2)]
    bars2 = axes[1].bar(model_names, speed_data, color=colors, alpha=0.8, edgecolor='black', linewidth=1)
    axes[1].set_title("平均生成速度对比", fontsize=14, fontweight="bold", pad=20)
    axes[1].set_ylabel("平均生成速度(秒/条)", fontsize=12)
    axes[1].set_ylim(0, 3)  # 固定Y轴范围，突出差异
    axes[1].grid(axis='y', alpha=0.3)
    # 标注数值
    for bar, v in zip(bars2, speed_data):
        axes[1].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.05,
                     f"{v}秒", ha="center", fontsize=11, fontweight="bold")
    
    # 子图3：平均效果评分对比（任务适配指标）
    score_data = [results[i]["平均效果评分(1-5分)"] for i in range(2)]
    bars3 = axes[2].bar(model_names, score_data, color=colors, alpha=0.8, edgecolor='black', linewidth=1)
    axes[2].set_title("平均文案效果评分对比", fontsize=14, fontweight="bold", pad=20)
    axes[2].set_ylabel("平均评分(1-5分)", fontsize=12)
    axes[2].set_ylim(0, 5)  # 固定Y轴为满分5分
    axes[2].grid(axis='y', alpha=0.3)
    # 标注数值
    for bar, v in zip(bars3, score_data):
        axes[2].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1,
                     f"{v}分", ha="center", fontsize=11, fontweight="bold")
    
    # 调整布局，避免中文截断
    plt.tight_layout()
    # 保存高清图片（DPI=300，支持插入报告/打印）
    plt.savefig("底座差异对比图表.png", dpi=300, bbox_inches="tight", facecolor='white')
    print("\n✅ 差异对比图表已保存至当前目录：底座差异对比图表.png")

# 4. 主函数（执行入口，小白直接运行即可）
if __name__ == "__main__":
    # 安装依赖提示
    print("=== 依赖安装提示 ===")
    print("请先执行以下命令安装依赖（避免版本冲突）：")
    print("pip install transformers==4.30.2 torch==2.0.1 peft==0.4.0 accelerate==0.20.3")
    print("pip install sentencepiece cpm_kernels matplotlib numpy\n")
    
    # 执行测试
    if TEST_MODE == "real":
        test_results = []
        for cfg in MODEL_CONFIGS:
            result = test_real_model(cfg)
            test_results.append(result)
    else:
        test_results = get_simulate_data()
    
    # 生成对比图表
    generate_comparison_chart(test_results)
    
    # 打印最终结论
    print("\n=== 底座选择结论 ===")
    print(f"ChatGLM-6B（INT4）：显存占用{test_results[0]['显存占用(GB)']}GB，速度{test_results[0]['平均生成速度(秒/条)']}秒/条，评分{test_results[0]['平均效果评分(1-5分)']}分")
    print(f"Qwen-7B（INT4）：显存占用{test_results[1]['显存占用(GB)']}GB，速度{test_results[1]['平均生成速度(秒/条)']}秒/条，评分{test_results[1]['平均效果评分(1-5分)']}分")
    print("结论：8G显存场景下，ChatGLM-6B（INT4量化版）综合适配性更优，兼顾显存、速度与效果。")