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.1from gptqmodel import GPTQModel
2from transformers import AutoTokenizer
3
4# Use the local directory or JustJaro/Virtuoso-Medium-v2_gptq_g32_4bit after upload
5quantized_model_id = "/home/jaro/models/quantized/Virtuoso-Medium-v2_gptq_g32_4bit" # or "JustJaro/Virtuoso-Medium-v2_gptq_g32_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))1uv venv
2source venv/bin/activate
3uv sync1HF_TOKEN=<YOUR_HF_TOKEN>
2TOKENIZERS_PARALLELISM="true"
3PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True1#!/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 enum import Enum
19from pathlib import Path
20from typing import List
21
22import torch
23import typer
24from datasets import load_dataset
25from dotenv import load_dotenv, find_dotenv
26from gptqmodel import GPTQModel, QuantizeConfig
27from gptqmodel.utils import Perplexity
28# For later pushing to the model hub
29from huggingface_hub import HfApi
30from transformers import AutoTokenizer, PreTrainedTokenizerBase
31
32load_dotenv(find_dotenv())
33HF_TOKEN = os.getenv("HF_TOKEN")
34
35app = typer.Typer()
36
37class GroupSize(str, Enum):
38 accurate:int = 32
39 balanced:int = 64
40 fast:int = 128
41
42
43def get_text_from_example(example: dict) -> str:
44 """
45 Returns text from a dataset example.
46 If the example contains a "text" field, and it is nonempty, that text is used.
47 Otherwise, if it has a "messages" field (a list of dicts with a "content" key),
48 the function returns the concatenation of all non-empty message contents.
49 """
50 if "text" in example and example["text"]:
51 return example["text"]
52 elif "messages" in example:
53 contents = [msg.get("content", "").strip() for msg in example["messages"]]
54 return " ".join([s for s in contents if s])
55 else:
56 return ""
57
58
59def get_calibration_dataset(
60 tokenizer: PreTrainedTokenizerBase,
61 nsamples: int,
62 seqlen: int,
63 calibration_dataset: str
64 ) -> List[dict]:
65 """
66 Loads a calibration dataset from the Hugging Face Hub (or from a local file).
67 It accepts datasets with a single "text" field (like wikitext)
68 or with a "messages" field (as in the Neural Magic LLM Compression Calibration dataset).
69 Only examples whose extracted text length is at least 'seqlen' are kept.
70 Each chosen example is tokenized (with truncation up to 'seqlen') and returned as a dict.
71 """
72 ds = None
73 try:
74 # Attempt to load from HF Hub.
75 try:
76 if "/" in calibration_dataset:
77 parts = calibration_dataset.split("/", 1)
78 ds = load_dataset(parts[0], parts[1], split="train")
79 else:
80 ds = load_dataset(calibration_dataset, split="train")
81 except Exception as e:
82 print(f"Error loading dataset '{calibration_dataset}' via load_dataset: {e}")
83 ds = load_dataset(calibration_dataset, split="train")
84 print(f"Loaded calibration dataset from full remote path {calibration_dataset}.")
85
86
87 except Exception as e:
88 print(f"Error loading dataset '{calibration_dataset}' via load_dataset: {e}")
89 # Fallback: if the supplied calibration_dataset is a local path, try to load it as JSON-lines.
90 if os.path.exists(calibration_dataset):
91 try:
92 ds = load_dataset("json", data_files=calibration_dataset, split="train")
93 print(f"Loaded calibration dataset from local file {calibration_dataset}.")
94 except Exception as e2:
95 print(f"Error loading local json dataset from '{calibration_dataset}': {e2}")
96 return []
97 else:
98 return []
99
100 print(f"Dataset features: {ds.features}")
101
102 # Filter examples that have at least 80% 'seqlen' of extracted text (wikitext-2-raw-v1 dataset has short examples).
103 ds = ds.filter(lambda x: len(get_text_from_example(x)) <= int(seqlen*0.8))
104 sample_range = min(nsamples, len(ds))
105 calibration_data = []
106 for i in range(sample_range):
107 example = ds[i]
108 text = get_text_from_example(example)
109 tokenized = tokenizer(text, truncation=True, max_length=seqlen, return_tensors="pt")
110 tokenized = {k: v.squeeze(0) for k, v in tokenized.items()}
111 calibration_data.append(tokenized)
112 return calibration_data
113
114
115def calculate_avg_ppl(model, tokenizer):
116 """
117 Computes the average perplexity on the wikitext-2-raw-v1 train split using GPTQModel's Perplexity utility.
118 """
119 ppl = Perplexity(
120 model=model,
121 tokenizer=tokenizer,
122 dataset_path="wikitext",
123 dataset_name="wikitext-2-raw-v1",
124 split="train",
125 text_column="text",
126 )
127 ppl_values = ppl.calculate(n_ctx=512, n_batch=512)
128 avg = sum(ppl_values) / len(ppl_values)
129 return avg
130
131
132def get_pinned_package_versions():
133 """
134 Retrieves pinned package versions using 'uv pip freeze'.
135 Returns a dictionary mapping lowercased package names to their versions.
136 """
137 try:
138 result = subprocess.run(["uv", "pip", "freeze"], capture_output=True, text=True, check=True)
139 packages_output = result.stdout.strip()
140 versions = {}
141 for line in packages_output.splitlines():
142 if "==" in line:
143 package_name, package_version = line.split("==", 1)
144 versions[package_name.lower()] = package_version
145 return versions
146 except subprocess.CalledProcessError as e:
147 typer.echo(f"Error running 'uv pip freeze': {e}", err=True)
148 return {}
149 except FileNotFoundError:
150 typer.echo("uv command not found. Make sure uv is installed and in your PATH.", err=True)
151 return {}
152
153
154@app.command()
155def main(
156 seq_len: int = typer.Option(4096, help="Sequence length for tokenization and calibration."),
157 nsamples: int = typer.Option(256, help="Number of samples to use for calibration."),
158 source_model: str = typer.Option("arcee-ai/Virtuoso-Medium-v2",
159 help="Source model HF repository identifier."),
160 calibration_dataset: str = typer.Option("wikitext/wikitext-2-raw-v1",
161 help="Calibration dataset identifier (in 'dataset/config' format) or local file path."),
162 hf_token: str = typer.Option(HF_TOKEN,
163 help="Hugging Face token for creating/updating your repo."),
164 upload_only: bool = typer.Option(False, help="Only upload the quantized model to the Hugging Face Hub."),
165 # Allow for 32, 64, 128 only using typer:
166 group_size: GroupSize = typer.Option(GroupSize.accurate, help="Group size for quantization accurate: 32, "
167 "balanced: 64, fast: 128. Default: accurate."),
168):
169 # Prepare destination directory and model names.
170 model_name = source_model.split("/")[-1]
171 quantized_model_name = f"{model_name}_gptq_g{int(group_size.value)}_4bit"
172 quantized_model_dir = os.path.expanduser(os.path.join("~/models/quantized", quantized_model_name))
173 if not os.path.exists(quantized_model_dir) or not upload_only:
174 os.makedirs(quantized_model_dir, exist_ok=True)
175
176 os.makedirs(quantized_model_dir, exist_ok=True)
177
178 typer.echo("Loading tokenizer from source model...")
179 tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)
180
181 typer.echo("Loading calibration dataset...")
182 typer.echo(f"Calibration dataset: {calibration_dataset}")
183 calibration_data = get_calibration_dataset(tokenizer_obj, nsamples, seq_len, calibration_dataset)
184 if not calibration_data:
185 typer.echo("Calibration dataset is empty. Aborting.", err=True)
186 raise typer.Exit(code=1)
187
188 quantize_config = QuantizeConfig(bits=4, group_size=int(group_size.value), damp_percent=0.01)
189 device = "cuda:0" if torch.cuda.is_available() else "cpu"
190 typer.echo(f"Loading model in {device} mode...")
191 model = GPTQModel.load(source_model, quantize_config)
192
193 typer.echo("Quantizing model...")
194 group_size_factor = int(128 / int(group_size.value))
195 model.quantize(calibration_data, auto_gc=False, batch_size=int((nsamples*0.1)/group_size_factor))
196 # Retrieve Hugging Face user info for README generation.
197 package_versions = get_pinned_package_versions()
198 username = get_my_user(hf_token)
199
200 script_content = self_read_script()
201
202 typer.echo(f"Saving quantized model to {quantized_model_dir} using Transformers safe serialization...")
203 try:
204 model.save_pretrained(quantized_model_dir)
205 tokenizer_obj.save_pretrained(quantized_model_dir)
206 except Exception as ex:
207 typer.echo(f"Error during saving with safe_serialization: {ex}. Aborting.")
208 raise
209 typer.echo(f"Model uploaded to Hugging Face repo: {quantized_model_name}")
210 else:
211 tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)
212 package_versions = get_pinned_package_versions()
213 username = get_my_user(hf_token)
214 script_content = self_read_script()
215
216
217 device = "cuda:0" if torch.cuda.is_available() else "cpu"
218 model = GPTQModel.load(quantized_model_dir, device=device)
219 avg_ppl = calculate_avg_ppl(model, tokenizer_obj)
220 typer.echo(f"Average perplexity (PPL) on wikitext v2 dataset: {avg_ppl}")
221 deps = Path("./pyproject.toml")
222 shutil.copy(deps, quantized_model_dir)
223 generate_readme(calibration_dataset, nsamples, quantized_model_dir,
224 quantized_model_name, script_content, seq_len, source_model, username, avg_ppl)
225 GPTQModel.push_to_hub(quantized_path=quantized_model_dir, private=False, repo_id=quantized_model_name,
226 token=HF_TOKEN)
227 typer.echo(f"Model uploaded to Hugging Face repo: {quantized_model_name}")
228 demo_input = tokenizer_obj("test is", return_tensors="pt").to(device)
229 generated_ids = model.generate(**demo_input)
230 output_text = tokenizer_obj.decode(generated_ids[0])
231 typer.echo(f"Inference demo output: {output_text}")
232 typer.echo(f"Average perplexity (PPL) on calibration dataset: {avg_ppl}")
233
234
235def self_read_script():
236 try:
237 script_path = os.path.abspath(__file__)
238 with open(script_path, "r") as f:
239 script_content = f.read()
240 except Exception as e:
241 script_content = "Error reading script content: " + str(e)
242 return script_content
243
244
245def get_my_user(hf_token):
246 api = HfApi(token=hf_token)
247 user_info = api.whoami()
248 try:
249 username = user_info.get("name") or user_info.get("username")
250 except Exception as e:
251 typer.echo(f"Error retrieving username from Hugging Face API: {e}. Using default username.")
252 username = api.whoami()
253 if not username:
254 typer.echo("Could not determine your Hugging Face username from the token, defaulting to hard coded username.",
255 err=True)
256 username = "JustJaro"
257 return username
258
259
260def generate_readme(calibration_dataset, nsamples, quantized_model_dir,
261 quantized_model_name, script_content, seq_len, source_model, username, avg_ppl):
262 readme_content = f"""---
263tags:
264- gptq
265- quantization
266- 4bit
267- confidentialmind
268- text-generation
269- apache2.0
270- mistral-small-24b
271---
272# 🔥 Quantized Model: {quantized_model_name} 🔥
273
274This is a 4-bit quantized version of [{source_model}](https://huggingface.co/{source_model}) model, quantized by [ConfidentialMind.com](https://www.confidentialmind.com) 🤖✨
275It leverages the open-source GPTQModel quantization to achieve 4-bit precision with a group size of 128 resulting in a
276smaller,
277faster model with minimal performance degradation.
278
279Ran on a single NVIDIA A100 GPU with 80GB of VRAM.
280
281*Note* `batch_size` is set quite high as the model is small, you may need to adjust this to your GPU VRAM.
282
283## Model Details
284- **Original Model:** [{source_model}](https://huggingface.co/{source_model})
285- **Quantized Model:** {quantized_model_name} (this repository)
286- **Quantization Method:** GPTQ (4-bit, group size 128)
287- **Quantization Library:** [GPTQModel](https://github.com/ModelCloud/GPTQModel/tree/main)
288- **Calibration Dataset:** {calibration_dataset} (using {nsamples} samples with seq len {seq_len})
289- **Quantized by:** [ConfidentialMind.com](https://www.confidentialmind.com)
290
291## Usage
292
293```python
294from gptqmodel import GPTQModel
295from transformers import AutoTokenizer
296
297# Use the local directory or {username}/{quantized_model_name} after upload
298quantized_model_id = "{quantized_model_dir}" # or "{username}/{quantized_model_name}"
299tokenizer = AutoTokenizer.from_pretrained(quantized_model_id)
300model = GPTQModel.load(quantized_model_id, device="cuda:0") # or "cpu"
301
302input_text = "This is a test prompt"
303inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")
304outputs = model.generate(**inputs)
305print(tokenizer.decode(outputs[0], skip_special_tokens=True))1uv venv
2source venv/bin/activate
3uv sync1HF_TOKEN=<YOUR_HF_TOKEN>
2TOKENIZERS_PARALLELISM="true"
3PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True{script_content}{source_model}gptqmodel.utils.eval integration and auto-generation of eval table.