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/SmolLM-135M_gptq_g32_4bit after upload
5quantized_model_id = "/home/jarouljanov/models/quantized/SmolLM-135M_gptq_g32_4bit" # or "JustJaro/SmolLM-135M_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
18import math
19from enum import Enum
20from pathlib import Path
21from typing import List, Union
22
23import torch
24import typer
25from datasets import load_dataset
26from dotenv import load_dotenv, find_dotenv
27from gptqmodel import GPTQModel, QuantizeConfig
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
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 except Exception as e:
87 print(f"Error loading dataset '{calibration_dataset}' via load_dataset: {e}")
88 # Fallback: if the supplied calibration_dataset is a local path, try to load it as JSON-lines.
89 if os.path.exists(calibration_dataset):
90 try:
91 ds = load_dataset("json", data_files=calibration_dataset, split="train")
92 print(f"Loaded calibration dataset from local file {calibration_dataset}.")
93 except Exception as e2:
94 print(f"Error loading local json dataset from '{calibration_dataset}': {e2}")
95 return []
96 else:
97 return []
98
99 print(f"Dataset features: {ds.features}")
100
101 # Filter examples that have at least 80% 'seqlen' of extracted text (wikitext-2-raw-v1 dataset has short examples).
102 ds = ds.filter(lambda x: len(get_text_from_example(x)) <= int(seqlen * 0.8))
103 sample_range = min(nsamples, len(ds))
104 calibration_data = []
105 for i in range(sample_range):
106 example = ds[i]
107 text = get_text_from_example(example)
108 tokenized = tokenizer(text, truncation=True, max_length=seqlen, return_tensors="pt")
109 tokenized = {k: v.squeeze(0) for k, v in tokenized.items()}
110 calibration_data.append(tokenized)
111 return calibration_data
112
113
114def calculate_perplexity_manual(model, tokenizer, dataset_name="wikitext", dataset_config="wikitext-2-raw-v1",
115 split="test", max_samples=100, max_length=512) -> Union[float, str]:
116 """
117 Calculate perplexity manually using a dataset.
118 Based on the research from GPTQModel documentation.
119 """
120 try:
121 # Load test dataset
122 if "/" in dataset_name:
123 dataset = load_dataset(dataset_name, split=split)
124 else:
125 dataset = load_dataset(dataset_name, dataset_config, split=split)
126
127 # Filter out empty texts
128 texts = [text for text in dataset["text"] if text.strip()]
129
130 # Limit samples for efficiency
131 texts = texts[:max_samples]
132
133 typer.echo(f"Calculating perplexity on {len(texts)} samples from {dataset_name}...")
134
135 model.model.eval()
136 total_loss = 0.0
137 total_tokens = 0
138
139 with torch.no_grad():
140 for i, text in enumerate(texts):
141 if i % 20 == 0:
142 typer.echo(f"Processing sample {i+1}/{len(texts)}")
143
144 # Tokenize the text
145 inputs = tokenizer(
146 text,
147 return_tensors="pt",
148 truncation=True,
149 max_length=max_length,
150 padding=False
151 )
152
153 input_ids = inputs.input_ids.to(model.model.device)
154
155 # Skip if too short
156 if input_ids.size(1) < 2:
157 continue
158
159 # Calculate loss
160 outputs = model.model(input_ids, labels=input_ids)
161 loss = outputs.loss
162
163 # Accumulate loss and token count
164 total_loss += loss.item() * input_ids.size(1)
165 total_tokens += input_ids.size(1)
166
167 if total_tokens == 0:
168 return "N/A (No valid tokens processed)"
169
170 # Calculate perplexity
171 avg_loss = total_loss / total_tokens
172 perplexity = math.exp(avg_loss)
173
174 return perplexity
175
176 except Exception as e:
177 typer.echo(f"Error calculating perplexity manually: {e}")
178 return f"N/A (Error: {str(e)})"
179
180
181def calculate_perplexity_lm_eval(model, tokenizer) -> Union[float, str]:
182 """
183 Calculate perplexity using lm-eval framework if available.
184 Based on GPTQModel documentation research.
185 """
186 try:
187 from gptqmodel.utils.eval import EVAL
188
189 # Try to use GPTQModel's built-in evaluation
190 typer.echo("Attempting to calculate perplexity using GPTQModel.eval...")
191
192 # Create a temporary directory to save the model for evaluation
193 temp_model_path = "/tmp/temp_gptq_model"
194 os.makedirs(temp_model_path, exist_ok=True)
195
196 model.save_pretrained(temp_model_path)
197 tokenizer.save_pretrained(temp_model_path)
198
199 # Use GPTQModel.eval with lm-eval framework
200 results = GPTQModel.eval(
201 temp_model_path,
202 framework=EVAL.LM_EVAL,
203 tasks=["wikitext"],
204 output_file=None
205 )
206
207 # Clean up temporary directory
208 shutil.rmtree(temp_model_path, ignore_errors=True)
209
210 # Extract perplexity from results
211 if "wikitext" in results.get("results", {}):
212 wikitext_results = results["results"]["wikitext"]
213 if "perplexity" in wikitext_results:
214 return wikitext_results["perplexity"]
215
216 return "N/A (Perplexity not found in lm-eval results)"
217
218 except ImportError:
219 typer.echo("lm-eval framework not available, falling back to manual calculation")
220 return None
221 except Exception as e:
222 typer.echo(f"Error using lm-eval: {e}, falling back to manual calculation")
223 return None
224
225
226def calculate_avg_ppl(model, tokenizer):
227 """
228 Computes the average perplexity using multiple methods.
229 First tries lm-eval framework, then falls back to manual calculation.
230 """
231 typer.echo("Starting perplexity calculation...")
232
233 # Method 1: Try lm-eval framework
234 ppl_result = calculate_perplexity_lm_eval(model, tokenizer)
235 if ppl_result is not None and not isinstance(ppl_result, str):
236 typer.echo(f"✓ Perplexity calculated using lm-eval: {ppl_result:.4f}")
237 return ppl_result
238
239 # Method 2: Manual calculation
240 typer.echo("Using manual perplexity calculation...")
241 ppl_result = calculate_perplexity_manual(model, tokenizer)
242
243 if isinstance(ppl_result, float):
244 typer.echo(f"✓ Perplexity calculated manually: {ppl_result:.4f}")
245 return ppl_result
246 else:
247 typer.echo(f"⚠ Perplexity calculation failed: {ppl_result}")
248 return ppl_result
249
250
251def get_pinned_package_versions():
252 """
253 Retrieves pinned package versions using 'uv pip freeze'.
254 Returns a dictionary mapping lowercased package names to their versions.
255 """
256 try:
257 result = subprocess.run(["uv", "pip", "freeze"], capture_output=True, text=True, check=True)
258 packages_output = result.stdout.strip()
259 versions = {}
260 for line in packages_output.splitlines():
261 if "==" in line:
262 package_name, package_version = line.split("==", 1)
263 versions[package_name.lower()] = package_version
264 return versions
265 except subprocess.CalledProcessError as e:
266 typer.echo(f"Error running 'uv pip freeze': {e}", err=True)
267 return {}
268 except FileNotFoundError:
269 typer.echo("uv command not found. Make sure uv is installed and in your PATH.", err=True)
270 return {}
271
272
273def self_read_script():
274 """
275 Reads the current script file content for inclusion in README.
276 """
277 try:
278 script_path = os.path.abspath(__file__)
279 with open(script_path, "r") as f:
280 script_content = f.read()
281 except Exception as e:
282 script_content = "Error reading script content: " + str(e)
283 return script_content
284
285
286def get_my_user(hf_token):
287 """
288 Gets the Hugging Face username from the provided token.
289 """
290 api = HfApi(token=hf_token)
291 user_info = api.whoami()
292 try:
293 username = user_info.get("name") or user_info.get("username")
294 except Exception as e:
295 typer.echo(f"Error retrieving username from Hugging Face API: {e}. Using default username.")
296 username = api.whoami()
297 if not username:
298 typer.echo("Could not determine your Hugging Face username from the token, defaulting to hard coded username.",
299 err=True)
300 username = "JustJaro"
301 return username
302
303
304def generate_readme(calibration_dataset, nsamples, quantized_model_dir,
305 quantized_model_name, script_content, seq_len, source_model, username, avg_ppl):
306 """
307 Generates a comprehensive README.md file for the quantized model.
308 """
309 # Format perplexity value for display
310 if isinstance(avg_ppl, float):
311 ppl_display = f"{avg_ppl:.4f}"
312 else:
313 ppl_display = str(avg_ppl)
314
315 readme_content = f"""---
316tags:
317- gptq
318- quantization
319- 4bit
320- confidentialmind
321- text-generation
322- apache2.0
323- mistral-small-24b
324---
325# 🔥 Quantized Model: {quantized_model_name} 🔥
326
327This is a 4-bit quantized version of [{source_model}](https://huggingface.co/{source_model}) model, quantized by [ConfidentialMind.com](https://www.confidentialmind.com) 🤖✨
328It leverages the open-source GPTQModel quantization to achieve 4-bit precision with a group size of 128 resulting in a
329smaller,
330faster model with minimal performance degradation.
331
332Ran on a single NVIDIA A100 GPU with 80GB of VRAM.
333
334*Note* `batch_size` is set quite high as the model is small, you may need to adjust this to your GPU VRAM.
335
336## Model Details
337- **Original Model:** [{source_model}](https://huggingface.co/{source_model})
338- **Quantized Model:** {quantized_model_name} (this repository)
339- **Quantization Method:** GPTQ (4-bit, group size 128)
340- **Quantization Library:** [GPTQModel](https://github.com/ModelCloud/GPTQModel/tree/main)
341- **Calibration Dataset:** {calibration_dataset} (using {nsamples} samples with seq len {seq_len})
342- **Quantized by:** [ConfidentialMind.com](https://www.confidentialmind.com)
343
344## Usage
345
346```python
347from gptqmodel import GPTQModel
348from transformers import AutoTokenizer
349
350# Use the local directory or {username}/{quantized_model_name} after upload
351quantized_model_id = "{quantized_model_dir}" # or "{username}/{quantized_model_name}"
352tokenizer = AutoTokenizer.from_pretrained(quantized_model_id)
353model = GPTQModel.load(quantized_model_id, device="cuda:0") # or "cpu"
354
355input_text = "This is a test prompt"
356inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")
357outputs = model.generate(**inputs)
358print(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.