Views
No views yet
batch_size is set quite high as the model is small, you may need to adjust this to your GPU VRAM.
Note2 Due to the "packed" nature of mistral-small weights, MSE was used agressively along with a higher damping factor - this resulted in lesser loss and perplexity, however G32 is more advised1from gptqmodel import GPTQModel
2from transformers import AutoTokenizer
3
4# Use the local directory or JustJaro/Mistral-Small-24B-Instruct-2501_gptq_g128_4bit after upload
5quantized_model_id = "/home/jaro/models/quantized/Mistral-Small-24B-Instruct-2501_gptq_g128_4bit" # or "JustJaro/Mistral-Small-24B-Instruct-2501_gptq_g128_4bit"
6tokenizer = AutoTokenizer.from_pretrained(quantized_model_id)
7model = GPTQModel.load(quantized_model_id, device="cuda:0") # or "cpu"
8
9input_text = "This is a test prompt"
10inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")
11outputs = model.generate(**inputs)
12print(tokenizer.decode(outputs[0], skip_special_tokens=True))1pip install \
2 gptqmodel==1.9.0 \
3 typer==0.15.1 \
4 huggingface_hub==<version> \
5 datasets==3.3.0 \
6 transformers==4.48.3 \
7 safetensors==0.5.2 \
8 torch==2.6.01uv venv
2source venv/bin/activate
3uv sync1HF_TOKEN=<YOUR_HF_TOKEN>
2TOKENIZERS_PARALLELISM="true"
3PYTORCH_CUDA_ALLOC_CONF=expandable_segments:Truereadme_content = f"""{MakeYourown}"" and using the variables passed to the function.1#!/usr/bin/env python3
2"""
3This script loads a source Hugging Face model and a calibration dataset,
4quantizes the model using GPTQModel (with 4-bit precision and group size 128),
5saves the quantized model using the Transformers API with safetensors (safe serialization)
6under ~/models/quantized/, and then creates/updates a Hugging Face repository (with the
7_gptq_g128_4bit suffix) by uploading the model, tokenizer, and an auto-generated README.md.
8
9Usage example:
10 python quantize.py --source-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
11 --calibration-dataset wikitext/wikitext-2-raw-v1 \
12 --seq-len 1024 --nsamples 256 --hf-token <YOUR_HF_TOKEN>
13"""
14
15import os
16import shutil
17import subprocess
18from pathlib import Path
19from typing import List
20
21import torch
22import typer
23from datasets import load_dataset
24from dotenv import load_dotenv, find_dotenv
25from gptqmodel import GPTQModel, QuantizeConfig
26from gptqmodel.utils import Perplexity
27# For later pushing to the model hub
28from huggingface_hub import HfApi
29from transformers import AutoTokenizer, PreTrainedTokenizerBase
30
31load_dotenv(find_dotenv())
32HF_TOKEN = os.getenv("HF_TOKEN")
33
34app = typer.Typer()
35
36
37def get_text_from_example(example: dict) -> str:
38 """
39 Returns text from a dataset example.
40 If the example contains a "text" field, and it is nonempty, that text is used.
41 Otherwise, if it has a "messages" field (a list of dicts with a "content" key),
42 the function returns the concatenation of all non-empty message contents.
43 """
44 if "text" in example and example["text"]:
45 return example["text"]
46 elif "messages" in example:
47 contents = [msg.get("content", "").strip() for msg in example["messages"]]
48 return " ".join([s for s in contents if s])
49 else:
50 return ""
51
52
53def get_calibration_dataset(
54 tokenizer: PreTrainedTokenizerBase,
55 nsamples: int,
56 seqlen: int,
57 calibration_dataset: str
58 ) -> List[dict]:
59 """
60 Loads a calibration dataset from the Hugging Face Hub (or from a local file).
61 It accepts datasets with a single "text" field (like wikitext)
62 or with a "messages" field (as in the Neural Magic LLM Compression Calibration dataset).
63 Only examples whose extracted text length is at least 'seqlen' are kept.
64 Each chosen example is tokenized (with truncation up to 'seqlen') and returned as a dict.
65 """
66 ds = None
67 try:
68 # Attempt to load from HF Hub.
69 try:
70 if "/" in calibration_dataset:
71 parts = calibration_dataset.split("/", 1)
72 ds = load_dataset(parts[0], parts[1], split="train")
73 else:
74 ds = load_dataset(calibration_dataset, split="train")
75 except Exception as e:
76 print(f"Error loading dataset '{calibration_dataset}' via load_dataset: {e}")
77 ds = load_dataset(calibration_dataset, split="train")
78 print(f"Loaded calibration dataset from full remote path {calibration_dataset}.")
79
80
81 except Exception as e:
82 print(f"Error loading dataset '{calibration_dataset}' via load_dataset: {e}")
83 # Fallback: if the supplied calibration_dataset is a local path, try to load it as JSON-lines.
84 if os.path.exists(calibration_dataset):
85 try:
86 ds = load_dataset("json", data_files=calibration_dataset, split="train")
87 print(f"Loaded calibration dataset from local file {calibration_dataset}.")
88 except Exception as e2:
89 print(f"Error loading local json dataset from '{calibration_dataset}': {e2}")
90 return []
91 else:
92 return []
93
94 print(f"Dataset features: {ds.features}")
95
96 # Filter examples that have at least 80% 'seqlen' of extracted text.
97 ds = ds.filter(lambda x: len(get_text_from_example(x)) >= int(seqlen*0.8))
98 sample_range = min(nsamples, len(ds))
99 calibration_data = []
100 for i in range(sample_range):
101 example = ds[i]
102 text = get_text_from_example(example)
103 tokenized = tokenizer(text, truncation=True, max_length=seqlen, return_tensors="pt")
104 tokenized = {k: v.squeeze(0) for k, v in tokenized.items()}
105 calibration_data.append(tokenized)
106 return calibration_data
107
108
109def calculate_avg_ppl(model, tokenizer):
110 """
111 Computes the average perplexity on the wikitext-2-raw-v1 train split using GPTQModel's Perplexity utility.
112 """
113 ppl = Perplexity(
114 model=model,
115 tokenizer=tokenizer,
116 dataset_path="wikitext",
117 dataset_name="wikitext-2-raw-v1",
118 split="train",
119 text_column="text",
120 )
121 ppl_values = ppl.calculate(n_ctx=512, n_batch=512)
122 avg = sum(ppl_values) / len(ppl_values)
123 return avg
124
125
126def get_pinned_package_versions():
127 """
128 Retrieves pinned package versions using 'uv pip freeze'.
129 Returns a dictionary mapping lowercased package names to their versions.
130 """
131 try:
132 result = subprocess.run(["uv", "pip", "freeze"], capture_output=True, text=True, check=True)
133 packages_output = result.stdout.strip()
134 versions = {}
135 for line in packages_output.splitlines():
136 if "==" in line:
137 package_name, package_version = line.split("==", 1)
138 versions[package_name.lower()] = package_version
139 return versions
140 except subprocess.CalledProcessError as e:
141 typer.echo(f"Error running 'uv pip freeze': {e}", err=True)
142 return {}
143 except FileNotFoundError:
144 typer.echo("uv command not found. Make sure uv is installed and in your PATH.", err=True)
145 return {}
146
147
148@app.command()
149def main(
150 seq_len: int = typer.Option(4096, help="Sequence length for tokenization and calibration."),
151 nsamples: int = typer.Option(512, help="Number of samples to use for calibration."),
152 source_model: str = typer.Option("mistralai/Mistral-Small-24B-Instruct-2501",
153 help="Source model HF repository identifier."),
154 calibration_dataset: str = typer.Option("wikitext/wikitext-2-raw-v1",
155 help="Calibration dataset identifier (in 'dataset/config' format) or local file path."),
156 hf_token: str = typer.Option(HF_TOKEN,
157 help="Hugging Face token for creating/updating your repo."),
158):
159 # Prepare destination directory and model names.
160 model_name = source_model.split("/")[-1]
161 quantized_model_name = f"{model_name}_gptq_g128_4bit"
162 quantized_model_dir = os.path.expanduser(os.path.join("~/models/quantized", quantized_model_name))
163 if not os.path.exists(quantized_model_dir):
164 os.makedirs(quantized_model_dir, exist_ok=True)
165
166 os.makedirs(quantized_model_dir, exist_ok=True)
167
168 typer.echo("Loading tokenizer from source model...")
169 tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)
170
171 typer.echo("Loading calibration dataset...")
172 typer.echo(f"Calibration dataset: {calibration_dataset}")
173 calibration_data = get_calibration_dataset(tokenizer_obj, nsamples, seq_len, calibration_dataset)
174 if not calibration_data:
175 typer.echo("Calibration dataset is empty. Aborting.", err=True)
176 raise typer.Exit(code=1)
177
178 quantize_config = QuantizeConfig(bits=4, group_size=128, mse=0.01, damp_percent=0.015)
179 device = "cuda:0" if torch.cuda.is_available() else "cpu"
180 typer.echo(f"Loading model in {device} mode...")
181 model = GPTQModel.load(source_model, quantize_config)
182
183 typer.echo("Quantizing model...")
184 model.quantize(calibration_data, auto_gc=False, batch_size=int(nsamples*0.1))
185 # Retrieve Hugging Face user info for README generation.
186 package_versions = get_pinned_package_versions()
187 username = get_my_user(hf_token)
188
189 script_content = self_read_script()
190
191 typer.echo(f"Saving quantized model to {quantized_model_dir} using Transformers safe serialization...")
192 try:
193 model.save_pretrained(quantized_model_dir)
194 tokenizer_obj.save_pretrained(quantized_model_dir)
195 except Exception as ex:
196 typer.echo(f"Error during saving with safe_serialization: {ex}. Aborting.")
197 raise
198 typer.echo(f"Model uploaded to Hugging Face repo: {quantized_model_name}")
199 else:
200 tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)
201 package_versions = get_pinned_package_versions()
202 username = get_my_user(hf_token)
203 script_content = self_read_script()
204
205
206 device = "cuda:0" if torch.cuda.is_available() else "cpu"
207 model = GPTQModel.load(quantized_model_dir, device=device)
208 avg_ppl = calculate_avg_ppl(model, tokenizer_obj)
209 typer.echo(f"Average perplexity (PPL) on wikitext v2 dataset: {avg_ppl}")
210 deps = Path("./pyproject.toml")
211 shutil.copy(deps, quantized_model_dir)
212 generate_readme(calibration_dataset, nsamples, package_versions, quantized_model_dir,
213 quantized_model_name, script_content, seq_len, source_model, username, avg_ppl)
214 GPTQModel.push_to_hub(quantized_path=quantized_model_dir, private=False, repo_id=quantized_model_name,
215 token=HF_TOKEN)
216 typer.echo(f"Model uploaded to Hugging Face repo: {quantized_model_name}")
217 demo_input = tokenizer_obj("test is", return_tensors="pt").to(device)
218 generated_ids = model.generate(**demo_input)
219 output_text = tokenizer_obj.decode(generated_ids[0])
220 typer.echo(f"Inference demo output: {output_text}")
221 typer.echo(f"Average perplexity (PPL) on calibration dataset: {avg_ppl}")
222
223
224def self_read_script():
225 try:
226 script_path = os.path.abspath(__file__)
227 with open(script_path, "r") as f:
228 script_content = f.read()
229 except Exception as e:
230 script_content = "Error reading script content: " + str(e)
231 return script_content
232
233
234def get_my_user(hf_token):
235 api = HfApi(token=hf_token)
236 user_info = api.whoami()
237 try:
238 username = user_info.get("name") or user_info.get("username")
239 except Exception as e:
240 typer.echo(f"Error retrieving username from Hugging Face API: {e}. Using default username.")
241 username = api.whoami()
242 if not username:
243 typer.echo("Could not determine your Hugging Face username from the token, defaulting to hard coded username.",
244 err=True)
245 username = "JustJaro"
246 return username
247
248
249def generate_readme(calibration_dataset, nsamples, package_versions, quantized_model_dir,
250 quantized_model_name, script_content, seq_len, source_model, username, avg_ppl):
251 readme_content = f"""{MakeYourown}""1uv venv
2source venv/bin/activate
3uv sync1HF_TOKEN=<YOUR_HF_TOKEN>
2TOKENIZERS_PARALLELISM="true"
3PYTORCH_CUDA_ALLOC_CONF=expandable_segments:Truegptqmodel.utils.eval integration and auto-generation of eval table, fix README.md generation.