from transformers import AutoModelForCausalLM, AutoTokenizer  
from llmcompressor.modifiers.quantization import GPTQModifier  
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier    
from llmcompressor.transformers import oneshot  
from datasets import load_dataset  

# 加载模型  
model_stub = "Qwen/Qwen2.5-7B-Instruct"  
model_name = model_stub.split("/")[-1]  

num_samples = 512  
max_seq_len = 8192  

tokenizer = AutoTokenizer.from_pretrained(model_stub)  

model = AutoModelForCausalLM.from_pretrained(  
    model_stub,  
    device_map="auto",  
    torch_dtype="auto",  
)  

def preprocess_fn(example):  
    return {"text": tokenizer.apply_chat_template(example["messages"], add_generation_prompt=False, tokenize=False)}  

ds = load_dataset("neuralmagic/LLM_compression_calibration", split="train")  
ds = ds.map(preprocess_fn)  

# 配置量化算法和方案  
recipe = [  
    SmoothQuantModifier(  
      smoothing_strength=0.8,  
      mappings=[  
          [["re:.*q_proj", "re:.*k_proj", "re:.*v_proj"], "re:.*input_layernorm"],  
          [["re:.*gate_proj", "re:.*up_proj"], "re:.*post_attention_layernorm"],  
          [["re:.*down_proj"], "re:.*up_proj"],  
      ],  
    ),  
    GPTQModifier(  
        ignore=["lm_head"],  
        sequential_targets=["Qwen2DecoderLayer"],  
        dampening_frac=0.01,  
        targets="Linear",  
        scheme="W8A8",  
    ),  
]  

# 应用量化  
oneshot(  
    model=model,  
    dataset=ds,   
    recipe=recipe,  
    max_seq_length=max_seq_len,  
    num_calibration_samples=num_samples,  
)  

# 以压缩张量格式保存到磁盘  
save_path = model_name + "-quantized.w8a8"  
model.save_pretrained(save_path)  
tokenizer.save_pretrained(save_path)  
print(f"模型和分词器保存至：{save_path}")