import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from llmcompressor import oneshot
from llmcompressor.modifiers.gptq import GPTQModifier
model_id = "microsoft/Phi-4-reasoning-plus"
output_dir = "./Phi-4-reasoning-plus-w4a16-llmcompressor-v0.12.0"
NUM_CALIBRATION_SAMPLES = 128
MAX_SEQUENCE_LENGTH = 2048
model = AutoModelForCausalLM.from_pretrained(
model_id,
dtype=torch.bfloat16,
device_map="cpu",
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
ds = load_dataset(
"HuggingFaceH4/ultrachat_200k",
split=f"train_sft[:{NUM_CALIBRATION_SAMPLES}]",
)
ds = ds.map(
lambda example: {"text": "\n".join(m["content"] for m in example["messages"])},
remove_columns=ds.column_names,
)
recipe = GPTQModifier(scheme="W4A16", targets="Linear", ignore=["lm_head"])
oneshot(
model=model,
dataset=ds,
recipe=recipe,
max_seq_length=MAX_SEQUENCE_LENGTH,
tokenizer=tokenizer,
output_dir=output_dir,
trust_remote_code_model=True,
)
inputs = tokenizer("What are we having for dinner?", return_tensors="pt")
output = model.generate(**inputs, max_new_tokens=30)
print(tokenizer.decode(output[0], skip_special_tokens=True))