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_g128_4bit after upload
5quantized_model_id = "/home/jaro/models/quantized/Virtuoso-Medium-v2_gptq_g128_4bit" # or "JustJaro/Virtuoso-Medium-v2_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))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 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("arcee-ai/Virtuoso-Medium-v2",
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 upload_only: bool = typer.Option(False, help="Only upload the quantized model to the Hugging Face Hub."),
159):
160 # Prepare destination directory and model names.
161 model_name = source_model.split("/")[-1]
162 quantized_model_name = f"{model_name}_gptq_g128_4bit"
163 quantized_model_dir = os.path.expanduser(os.path.join("~/models/quantized", quantized_model_name))
164 if not os.path.exists(quantized_model_dir) or not upload_only:
165 os.makedirs(quantized_model_dir, exist_ok=True)
166
167 os.makedirs(quantized_model_dir, exist_ok=True)
168
169 typer.echo("Loading tokenizer from source model...")
170 tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)
171
172 typer.echo("Loading calibration dataset...")
173 typer.echo(f"Calibration dataset: {calibration_dataset}")
174 calibration_data = get_calibration_dataset(tokenizer_obj, nsamples, seq_len, calibration_dataset)
175 if not calibration_data:
176 typer.echo("Calibration dataset is empty. Aborting.", err=True)
177 raise typer.Exit(code=1)
178
179 quantize_config = QuantizeConfig(bits=4, group_size=128, damp_percent=0.01)
180 device = "cuda:0" if torch.cuda.is_available() else "cpu"
181 typer.echo(f"Loading model in {device} mode...")
182 model = GPTQModel.load(source_model, quantize_config)
183
184 typer.echo("Quantizing model...")
185 model.quantize(calibration_data, auto_gc=False, batch_size=int(nsamples*0.1))
186 # Retrieve Hugging Face user info for README generation.
187 package_versions = get_pinned_package_versions()
188 username = get_my_user(hf_token)
189
190 script_content = self_read_script()
191
192 typer.echo(f"Saving quantized model to {quantized_model_dir} using Transformers safe serialization...")
193 try:
194 model.save_pretrained(quantized_model_dir)
195 tokenizer_obj.save_pretrained(quantized_model_dir)
196 except Exception as ex:
197 typer.echo(f"Error during saving with safe_serialization: {ex}. Aborting.")
198 raise
199 typer.echo(f"Model uploaded to Hugging Face repo: {quantized_model_name}")
200 else:
201 tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)
202 package_versions = get_pinned_package_versions()
203 username = get_my_user(hf_token)
204 script_content = self_read_script()
205
206
207 device = "cuda:0" if torch.cuda.is_available() else "cpu"
208 model = GPTQModel.load(quantized_model_dir, device=device)
209 avg_ppl = calculate_avg_ppl(model, tokenizer_obj)
210 typer.echo(f"Average perplexity (PPL) on wikitext v2 dataset: {avg_ppl}")
211 deps = Path("./pyproject.toml")
212 shutil.copy(deps, quantized_model_dir)
213 generate_readme(calibration_dataset, nsamples, quantized_model_dir,
214 quantized_model_name, script_content, seq_len, source_model, username, avg_ppl)
215 GPTQModel.push_to_hub(quantized_path=quantized_model_dir, private=False, repo_id=quantized_model_name,
216 token=HF_TOKEN)
217 typer.echo(f"Model uploaded to Hugging Face repo: {quantized_model_name}")
218 demo_input = tokenizer_obj("test is", return_tensors="pt").to(device)
219 generated_ids = model.generate(**demo_input)
220 output_text = tokenizer_obj.decode(generated_ids[0])
221 typer.echo(f"Inference demo output: {output_text}")
222 typer.echo(f"Average perplexity (PPL) on calibration dataset: {avg_ppl}")
223
224
225def self_read_script():
226 try:
227 script_path = os.path.abspath(__file__)
228 with open(script_path, "r") as f:
229 script_content = f.read()
230 except Exception as e:
231 script_content = "Error reading script content: " + str(e)
232 return script_content
233
234
235def get_my_user(hf_token):
236 api = HfApi(token=hf_token)
237 user_info = api.whoami()
238 try:
239 username = user_info.get("name") or user_info.get("username")
240 except Exception as e:
241 typer.echo(f"Error retrieving username from Hugging Face API: {e}. Using default username.")
242 username = api.whoami()
243 if not username:
244 typer.echo("Could not determine your Hugging Face username from the token, defaulting to hard coded username.",
245 err=True)
246 username = "JustJaro"
247 return username
248
249
250def generate_readme(calibration_dataset, nsamples, quantized_model_dir,
251 quantized_model_name, script_content, seq_len, source_model, username, avg_ppl):
252 readme_content = f"""---
253tags:
254- gptq
255- quantization
256- 4bit
257- confidentialmind
258- text-generation
259- apache2.0
260- mistral-small-24b
261---
262# 🔥 Quantized Model: {quantized_model_name} 🔥
263
264This is a 4-bit quantized version of [{source_model}](https://huggingface.co/{source_model}) model, quantized by [ConfidentialMind.com](https://www.confidentialmind.com) 🤖✨
265It leverages the open-source GPTQModel quantization to achieve 4-bit precision with a group size of 128 resulting in a
266smaller,
267faster model with minimal performance degradation.
268
269Ran on a single NVIDIA A100 GPU with 80GB of VRAM.
270
271*Note* `batch_size` is set quite high as the model is small, you may need to adjust this to your GPU VRAM.
272
273## Model Details
274- **Original Model:** [{source_model}](https://huggingface.co/{source_model})
275- **Quantized Model:** {quantized_model_name} (this repository)
276- **Quantization Method:** GPTQ (4-bit, group size 128)
277- **Quantization Library:** [GPTQModel](https://github.com/ModelCloud/GPTQModel/tree/main)
278- **Calibration Dataset:** {calibration_dataset} (using {nsamples} samples with seq len {seq_len})
279- **Quantized by:** [ConfidentialMind.com](https://www.confidentialmind.com)
280
281## Usage
282
283```python
284from gptqmodel import GPTQModel
285from transformers import AutoTokenizer
286
287# Use the local directory or {username}/{quantized_model_name} after upload
288quantized_model_id = "{quantized_model_dir}" # or "{username}/{quantized_model_name}"
289tokenizer = AutoTokenizer.from_pretrained(quantized_model_id)
290model = GPTQModel.load(quantized_model_id, device="cuda:0") # or "cpu"
291
292input_text = "This is a test prompt"
293inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")
294outputs = model.generate(**inputs)
295print(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.
## Quantization Performance
Average perplexity (PPL) on wikitext v2 dataset: 6.455169136972343
## Disclaimer
This model is for research purposes only. It may inherit limitations and biases from the original model and the quantization process. Please use responsibly and refer to the original model card for more details.
## Contact
For any questions or support, please visit [ConfidentialMind.com](https://www.confidentialmind.com) or contact us directly.
## License
This model inherits the license from the original model. Please refer to the original model card for more details.
Original model card: `arcee-ai/Virtuoso-Medium-v2`
## Author
This model was quantized by [Jaro](https://www.linkedin.com/in/jaroai/)
## Acknowledgements
Quantization performed using the GPTQModel pipeline.
TODO: Add `gptqmodel.utils.eval` integration and auto-generation of eval table.
---
*Generated and quantized using GPTQModel.*