import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

def quantize_model(model_path: str, output_path: str):
    """量化模型以减少内存占用"""
    print("⚡ 量化模型...")
    
    # 量化配置
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4"
    )
    
    # 加载量化模型
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        quantization_config=quantization_config,
        device_map="auto"
    )
    
    # 保存量化模型
    model.save_pretrained(output_path)
    
    print(f"✅ 量化模型已保存到: {output_path}")
    print(f"   内存占用减少约75%")
    
    return model