Views
No views yet
1from datasets import load_dataset
2from gptqmodel import GPTQModel, QuantizeConfig
3import sys
4
5model_id = sys.argv[1]
6quant_path = "quantized_model"
7
8# Load calibration data (1024 samples from C4)
9calibration_dataset = load_dataset(
10 "allenai/c4",
11 data_files="en/c4-train.00001-of-01024.json.gz",
12 split="train"
13 ).select(range(1024))["text"]
14
15# Configure and run quantization
16quant_config = QuantizeConfig(bits=4, group_size=128)
17model = GPTQModel.load(model_id, quant_config)
18model.quantize(calibration_dataset, batch_size=2)
19model.save(quant_path)