from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer, AwqConfig

model_path = "facebook/opt-125m"
quant_path = "test-quant/opt-125m-awq-no-kproj"
modules_to_not_convert = ["k_proj"]

quant_config = {"zero_point": True, "q_group_size": 128, "w_bit": 4, "version":"GEMM", "modules_to_not_convert": modules_to_not_convert}

# Load model
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize
model.quantize(tokenizer, quant_config=quant_config, modules_to_not_convert=modules_to_not_convert)