Arcee-Blitz-GPTQ-G32-W4A16

ConfidentialMind 🚀🧠

Generative AI Software Infrastructure Simplified 🎉

---

🔥 Quantized Model: Arcee-Blitz-GPTQ-G32-W4A16 🦾 🔥

<details> <summary><strong>Model Details</strong></summary>

• Original Model: arcee-ai/Arcee-Blitz
• Quantized Model: Arcee-Blitz-GPTQ-G32-W4A16 (this repository)
• Quantization Method: GPTQ (4-bit, group size 32)
• Quantization Library: GPTQModel
• Calibration Dataset: neuralmagic/LLM_compression_calibration (using 1638 samples with seq len 6553)
• Quantized by: ConfidentialMind.com

</details> <details> <summary><strong>Usage</strong></summary>

from gptqmodel import GPTQModel
from transformers import AutoTokenizer

# Use the local directory or JustJaro/Arcee-Blitz-GPTQ-G32-W4A16 after upload
quantized_model_id = "/home/jaro/models/quantized/Arcee-Blitz-GPTQ-G32-W4A16"  # or "JustJaro/Arcee-Blitz-GPTQ-G32-W4A16"
tokenizer = AutoTokenizer.from_pretrained(quantized_model_id)
model = GPTQModel.load(quantized_model_id, device="cuda:0")  # or "cpu"

input_text = "This is a test prompt"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

</details> <details> <summary><strong>Package Versions and Installation Instructions</strong></summary>

See pyproject.toml for the exact UV project file. See the GPTQModel repo for more details on how to install the package.

Use the provided pyproject.toml:

uv venv
source venv/bin/activate
uv sync

</details> <details> <summary><strong>Quantization Script</strong></summary>

Below is the exact quantize.py script used to generate this model:

#!/usr/bin/env python3
"""
This script loads a source Hugging Face model and a calibration dataset,
quantizes the model using GPTQModel (with 4-bit precision and a dynamic group size),
saves the quantized model with Transformers’ safe serialization under ~/models/quantized/,
and then creates/updates a Hugging Face repository by uploading the model, tokenizer,
and an auto–generated README.md that includes proper foldable sections, badges, and warnings.

Usage example:
    python quantize.py --source-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
                       --calibration-dataset wikitext/wikitext-2-raw-v1 \
                       --seq-len 1024 --nsamples 256 --hf-token <YOUR_HF_TOKEN>
"""

import os
import shutil
import subprocess
from enum import Enum
from pathlib import Path
from typing import List

import torch
import typer
from datasets import load_dataset
from dotenv import load_dotenv, find_dotenv
from gptqmodel import GPTQModel, QuantizeConfig
from gptqmodel.utils import Perplexity
# For later pushing to the model hub
from huggingface_hub import HfApi
from transformers import AutoTokenizer, PreTrainedTokenizerBase

load_dotenv(find_dotenv())
HF_TOKEN = os.getenv("HF_TOKEN")

app = typer.Typer()


class GroupSize(str, Enum):
    accurate: int = 32
    balanced: int = 64
    fast: int = 128


def get_text_from_example(example: dict) -> str:
    """
    Returns text from a dataset example.
    If the example contains a "text" field, that text is used.
    Otherwise, if it has a "messages" field (a list of dicts with a "content" key),
    the contents of all messages are concatenated.
    """
    if "text" in example and example["text"]:
        return example["text"]
    elif "messages" in example:
        contents = [msg.get("content", "").strip() for msg in example["messages"]]
        return " ".join([s for s in contents if s])
    else:
        return ""


def get_calibration_dataset(
    tokenizer: PreTrainedTokenizerBase,
    nsamples: int,
    seqlen: int,
    calibration_dataset: str
) -> List[dict]:
    """
    Loads and tokenizes a calibration dataset from the HF Hub (or a local file).
    Only examples with at least 80% of seqlen characters (after extraction) are kept.
    """
    ds = None
    try:
        try:
            if "/" in calibration_dataset:
                parts = calibration_dataset.split("/", 1)
                ds = load_dataset(parts[0], parts[1], split="train")
            else:
                ds = load_dataset(calibration_dataset, split="train")
        except Exception as e:
            print(f"Error loading dataset '{calibration_dataset}' via load_dataset: {e}")
            ds = load_dataset(calibration_dataset, split="train")
            print(f"Loaded calibration dataset from full remote path {calibration_dataset}.")
    except Exception as e:
        print(f"Error loading dataset '{calibration_dataset}' via load_dataset: {e}")
        if os.path.exists(calibration_dataset):
            try:
                ds = load_dataset("json", data_files=calibration_dataset, split="train")
                print(f"Loaded calibration dataset from local file {calibration_dataset}.")
            except Exception as e2:
                print(f"Error loading local json dataset from '{calibration_dataset}': {e2}")
                return []
        else:
            return []

    print(f"Dataset features: {ds.features}")

    ds = ds.filter(lambda x: len(get_text_from_example(x)) <= int(seqlen * 0.8))
    sample_range = min(nsamples, len(ds))
    calibration_data = []
    for i in range(sample_range):
        example = ds[i]
        text = get_text_from_example(example)
        tokenized = tokenizer(text, truncation=True, max_length=seqlen, return_tensors="pt")
        tokenized = {k: v.squeeze(0) for k, v in tokenized.items()}
        calibration_data.append(tokenized)
    return calibration_data


def calculate_avg_ppl(model, tokenizer, dataset_name="wikitext-2-raw-v1"):
    """
    Computes the average perplexity on the wikitext-2-raw-v1 training split.
    """
    ppl = Perplexity(
        model=model,
        tokenizer=tokenizer,
        dataset_path="wikitext",
        dataset_name=dataset_name,
        split="train",
        text_column="text",
    )
    ppl_values = ppl.calculate(n_ctx=512, n_batch=512)
    avg = sum(ppl_values) / len(ppl_values)
    return avg, dataset_name


def get_pinned_package_versions():
    """
    Retrieves pinned package versions via 'uv pip freeze'.
    """
    try:
        result = subprocess.run(["uv", "pip", "freeze"], capture_output=True, text=True, check=True)
        packages_output = result.stdout.strip()
        versions = {}
        for line in packages_output.splitlines():
            if "==" in line:
                package_name, package_version = line.split("==", 1)
                versions[package_name.lower()] = package_version
        return versions
    except subprocess.CalledProcessError as e:
        typer.echo(f"Error running 'uv pip freeze': {e}", err=True)
        return {}
    except FileNotFoundError:
        typer.echo("uv command not found. Make sure uv is installed and in your PATH.", err=True)
        return {}


def prepare_model_dir(model_dir: str):
    """Removes the given directory if it exists and creates a new one."""
    if os.path.exists(model_dir):
        shutil.rmtree(model_dir)
    os.makedirs(model_dir, exist_ok=True)


def self_read_script():
    """Returns the full text of this script."""
    try:
        script_path = os.path.abspath(__file__)
        with open(script_path, "r") as f:
            script_content = f.read()
    except Exception as e:
        script_content = "Error reading script content: " + str(e)
    return script_content


def get_my_user(hf_token):
    """Retrieves your Hugging Face username from your token."""
    api = HfApi(token=hf_token)
    user_info = api.whoami()
    try:
        username = user_info.get("name") or user_info.get("username")
    except Exception as e:
        typer.echo(f"Error retrieving username from Hugging Face API: {e}. Using default username.")
        username = api.whoami()
    if not username:
        typer.echo("Could not determine your Hugging Face username from the token. Using default username.", err=True)
        username = "JustJaro"
    return username


def make_details_section(title: str, content: str) -> str:
    """
    Returns a markdown string for a collapsible section.
    The format is:
      <details>
        <summary><strong>{title}</strong></summary>

      {content}

      </details>
    """
    return f"<details>\n  <summary><strong>{title}</strong></summary>\n\n{content}\n\n</details>\n"


def generate_readme(
        calibration_dataset: str,
        nsamples: int,
        quantized_model_dir: str,
        quantized_model_name: str,
        script_content: str,
        seq_len: int,
        source_model: str,
        username: str,
        avg_ppl: float,
        group_size_int: int,
        ppl_dataset: str,
) -> None:
    """
    Creates a README.md with a YAML front matter, title (with a warning if perplexity is high),
    and a series of foldable sections.
    """
    import random

    # Pick a random emoji for the title
    chosen_emoji = random.choice(["⚡️", "🐣", "🦾", "🤖", "🧠", "🧐", "🚀"])

    # Warning if average perplexity is above 30
    warning_text = ""
    if avg_ppl > 30:
        warning_text = f"\n**⚠️ WARNING: High Perplexity Detected!** The average perplexity is {avg_ppl:.2f}, which exceeds the recommended threshold.\n"

    # YAML front matter and top header
    front_matter = (
        "---\n"
        'company: "ConfidentialMind"\n'
        'emoji: "🧠"\n'
        'colorFrom: "blue"\n'
        'colorTo: "purple"\n'
        'pinned: true\n'
        'authors: "JustJaro"\n'
        "---\n\n"
        "# ConfidentialMind 🚀🧠\n\n"
        "Generative AI Software Infrastructure Simplified 🎉\n\n"
        "[![Website](https://img.shields.io/badge/Website-confidentialmind.com-blue)](https://confidentialmind.com)  \n"
        "[![Email](https://img.shields.io/badge/Email-info%40confidentialmind.com-orange)](mailto:[email protected])\n\n"
    )

    # Main title block for the quantized model
    title = f"# 🔥 Quantized Model: {quantized_model_name} {chosen_emoji} 🔥\n{warning_text}\n"

    # Build each collapsible section using the helper:

    model_details_content = (
        f"- **Original Model:** [{source_model}](https://huggingface.co/{source_model})\n"
        f"- **Quantized Model:** {quantized_model_name} (this repository)\n"
        f"- **Quantization Method:** GPTQ (4-bit, group size {group_size_int})\n"
        f"- **Quantization Library:** [GPTQModel](https://github.com/ModelCloud/GPTQModel/tree/main)\n"
        f"- **Calibration Dataset:** {calibration_dataset} (using {nsamples} samples with seq len {seq_len})\n"
        f"- **Quantized by:** [ConfidentialMind.com](https://www.confidentialmind.com)"
    )
    model_details_section = make_details_section("Model Details", model_details_content)

    usage_content = (
        f"```python\n"
        f"from gptqmodel import GPTQModel\n"
        f"from transformers import AutoTokenizer\n\n"
        f"# Use the local directory or {username}/{quantized_model_name} after upload\n"
        f'quantized_model_id = "{quantized_model_dir}"  # or "{username}/{quantized_model_name}"\n'
        f"tokenizer = AutoTokenizer.from_pretrained(quantized_model_id)\n"
        f'model = GPTQModel.load(quantized_model_id, device="cuda:0")  # or "cpu"\n\n'
        f'input_text = "This is a test prompt"\n'
        f'inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")\n'
        f"outputs = model.generate(**inputs)\n"
        f"print(tokenizer.decode(outputs[0], skip_special_tokens=True))\n"
        f"```"
    )
    usage_section = make_details_section("Usage", usage_content)

    package_content = (
        "See `pyproject.toml` for the exact UV project file. See the "
        "[GPTQModel](https://github.com/ModelCloud/GPTQModel/tree/main) repo for more details on how to install the package.\n\n"
        "Use the provided `pyproject.toml`:\n\n"
        "```bash\n"
        "uv venv\n"
        "source venv/bin/activate\n"
        "uv sync\n"
        "```"
    )
    package_section = make_details_section("Package Versions and Installation Instructions", package_content)

    script_content_md = (
        "Below is the exact `quantize.py` script used to generate this model:\n\n"
        "```python\n"
        f"{script_content}\n"
        "```"
    )
    script_section = make_details_section("Quantization Script", script_content_md)

    performance_content = f"**Average perplexity (PPL) on {ppl_dataset} dataset:** {avg_ppl:.2f}"
    performance_section = make_details_section("Quantization Performance", performance_content)

    disclaimer_content = (
        "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."
    )
    disclaimer_section = make_details_section("Disclaimer", disclaimer_content)

    contact_content = (
        "For any questions or support, please visit [ConfidentialMind](https://www.confidentialmind.com) or contact us directly.\n\n"
        "[![LinkedIn](https://img.shields.io/badge/LinkedIn-ConfidentialMind-blue)](https://www.linkedin.com/company/confidentialmind/)"
    )
    contact_section = make_details_section("Contact", contact_content)

    license_content = (
        "This model inherits the license from the original model. Please refer to the original model card for more details.\n\n"
        f"Original model card: `{source_model}`"
    )
    license_section = make_details_section("License", license_content)

    author_content = (
        "This model was quantized by [![LinkedIn](https://img.shields.io/badge/LinkedIn-Jaro-blue)](https://www.linkedin.com/in/jaroai/)"
    )
    author_section = make_details_section("Author", author_content)

    ack_content = (
        "Quantization performed using the GPTQModel pipeline.\n\n"
        "**TODO:**\n"
        "- HELMET\n"
        "- Eluther evaluation harness"
    )
    ack_section = make_details_section("Acknowledgements", ack_content)

    # Combine everything into one README content string.
    readme_content = (
            front_matter +
            title + "\n" +
            model_details_section +
            usage_section +
            package_section +
            script_section +
            performance_section +
            disclaimer_section +
            contact_section +
            license_section +
            author_section +
            ack_section
    )

    readme_path = os.path.join(quantized_model_dir, "README.md")
    with open(readme_path, "w") as f:
        f.write(readme_content)
    typer.echo("README.md created with detailed information.")
    typer.echo(f"README.md saved to {readme_path}")


@app.command()
def main(
    seq_len: int = typer.Option(4096, help="Sequence length for tokenization and calibration."),
    nsamples: int = typer.Option(512, help="Number of samples to use for calibration."),
    source_model: str = typer.Option("rombodawg/Rombos-LLM-V2.6-Qwen-14b",
                                     help="Source model HF repository identifier."),
    calibration_dataset: str = typer.Option("wikitext/wikitext-2-raw-v1",
                                              help="Calibration dataset identifier (in 'dataset/config' format) or local file path."),
    hf_token: str = typer.Option(HF_TOKEN, help="Hugging Face token for creating/updating your repo."),
    upload_only: bool = typer.Option(False, help="Only upload the quantized model to the Hugging Face Hub."),
    # Allow for 32, 64, 128 only using typer:
    group_size: GroupSize = typer.Option(GroupSize.accurate, help="Group size for quantization: accurate (32), balanced (64), fast (128)."),
    mse: bool = typer.Option(False, help="Use MSE instead of MAE for the loss function."),
    size_multi: float = typer.Option(3.5, help="Model size multiplier depends on the source model. Default: 1."),
):
    # Prepare destination directory and model names.
    model_name = source_model.split("/")[-1]
    if size_multi != 1:
        size_multiplier = size_multi
        size_multiplier_len = size_multiplier / 2
    else:
        size_multiplier = 1
        size_multiplier_len = 1

    nsamples = int(nsamples * size_multiplier)
    seq_len = int(seq_len * size_multiplier_len)
    quantized_model_name = f"{model_name}-GPTQ-G{int(group_size.value)}-W4A16"
    quantized_model_dir = os.path.expanduser(os.path.join("~/models/quantized", quantized_model_name))

    if not upload_only:
        prepare_model_dir(quantized_model_dir)

        typer.echo("Loading tokenizer from source model...")
        tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)

        typer.echo("Loading calibration dataset...")
        typer.echo(f"Calibration dataset: {calibration_dataset}")
        calibration_data = get_calibration_dataset(tokenizer_obj, nsamples, seq_len, calibration_dataset)
        if not calibration_data:
            typer.echo("Calibration dataset is empty. Aborting.", err=True)
            raise typer.Exit(code=1)

        if mse:
            mse_val = 0.01
            quantize_config = QuantizeConfig(bits=4, group_size=int(group_size.value), damp_percent=0.015, mse=mse_val)
        else:
            quantize_config = QuantizeConfig(bits=4, group_size=int(group_size.value), damp_percent=0.01)

        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        typer.echo(f"Loading model in {device} mode...")
        model = GPTQModel.load(source_model, quantize_config)

        typer.echo("Quantizing model...")
        group_size_factor = int(128 / int(group_size.value))
        batch_size = max(
            1, int(int((nsamples * 0.1) / group_size_factor) * int(size_multiplier_len))
        )
        model.quantize(calibration_data, auto_gc=False, batch_size=batch_size)

        package_versions = get_pinned_package_versions()
        username = get_my_user(hf_token)
        script_content = self_read_script()

        typer.echo(f"Saving quantized model to {quantized_model_dir} using Transformers safe serialization...")
        try:
            model.save_pretrained(quantized_model_dir)
            tokenizer_obj.save_pretrained(quantized_model_dir)
        except Exception as ex:
            typer.echo(f"Error during saving: {ex}. Aborting.")
            raise
        typer.echo(f"Model saved successfully to {quantized_model_dir}.")
    else:
        tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)
        package_versions = get_pinned_package_versions()
        username = get_my_user(hf_token)
        script_content = self_read_script()
        device = "cuda:0" if torch.cuda.is_available() else "cpu"

    # Load the (possibly quantized) model for evaluation.
    model = GPTQModel.load(quantized_model_dir, device=device)
    avg_ppl, ppl_dataset = calculate_avg_ppl(model, tokenizer_obj)
    typer.echo(f"Average perplexity (PPL) on wikitext-2-raw-v1 dataset: {avg_ppl:.2f}")

    deps = Path("./pyproject.toml")
    shutil.copy(deps, quantized_model_dir)

    # Note: pass the dynamic group size as an integer.
    generate_readme(calibration_dataset, nsamples, quantized_model_dir,
                    quantized_model_name, script_content, seq_len,
                    source_model, username, avg_ppl, int(group_size.value), ppl_dataset)
    GPTQModel.push_to_hub(quantized_path=quantized_model_dir, private=False,
                           repo_id=quantized_model_name, token=HF_TOKEN)
    typer.echo(f"Model pushed to Hugging Face repo: {quantized_model_name}")

    demo_input = tokenizer_obj("test is", return_tensors="pt").to(device)
    generated_ids = model.generate(**demo_input)
    output_text = tokenizer_obj.decode(generated_ids[0])
    typer.echo(f"Inference demo output: {output_text}")
    typer.echo(f"Average perplexity (PPL) on calibration dataset: {avg_ppl:.2f}")


if __name__ == "__main__":
    app()

</details> <details> <summary><strong>Quantization Performance</strong></summary>

Average perplexity (PPL) on wikitext-2-raw-v1 dataset: 7.86

</details> <details> <summary><strong>Disclaimer</strong></summary>

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.

</details> <details> <summary><strong>Contact</strong></summary>

For any questions or support, please visit ConfidentialMind or contact us directly.

</details> <details> <summary><strong>License</strong></summary>

This model inherits the license from the original model. Please refer to the original model card for more details.

Original model card: arcee-ai/Arcee-Blitz

</details> <details> <summary><strong>Author</strong></summary>

This model was quantized by

</details> <details> <summary><strong>Acknowledgements</strong></summary>

Quantization performed using the GPTQModel pipeline.

TODO:

• HELMET
• Eluther evaluation harness

</details>