Views
No views yet
1from datasets import load_dataset
2from gptqmodel import GPTQModel, QuantizeConfig
3from huggingface_hub import constants
4
5model_id = "Qwen/Qwen3-32B"
6# Save the quantized model in the HF cache directory
7cache_dir = constants.HF_HUB_CACHE
8quant_path = os.path.join(cache_dir, "models--quantized--" + model_id.replace("/", "--") + "mixed--calibration")
9os.makedirs(quant_path, exist_ok=True)
10
11# Load calibration data (512 samples from C4)
12calibration_dataset = load_dataset(
13 "allenai/c4",
14 data_files="en/c4-train.00001-of-01024.json.gz",
15 split="train"
16 ).select(range(512))["text"]
17
18# Add custom dataset
19custom_calibration_dataset = []
20with open("./data/custom_calibration_dataset.jsonl", "r") as f:
21 for line in f:
22 if line.strip(): # Skip empty lines
23 item = json.loads(line)
24 custom_calibration_dataset.append(item["text"])
25
26# randomly choose 512 samples from custom_calibration_dataset, and add the last 6 samples
27selected_samples = random.sample(custom_calibration_dataset, 512)
28selected_samples.extend(custom_calibration_dataset[-6:])
29selected_samples = list(set(selected_samples))
30calibration_dataset.extend(selected_samples)
31
32# shuffle calibration_dataset
33random.shuffle(calibration_dataset)
34
35# Configure and run quantization
36quant_config = QuantizeConfig(bits=4, group_size=128)
37model = GPTQModel.load(model_id, quant_config)
38model.quantize(calibration_dataset, batch_size=2)
39model.save(quant_path)
40