import torch
from transformers import AutoTokenizer
from auto_round import AutoRound

# You need to use datasets 3.6.0 since >=4.0 doesn't support codeparrot/github-code-clean dataset
# You also need torchvision for some reason even if it's not used, otherwise it will throw an error "Unrecognized image processor ...".
# TORCHVISION_AVAILABLE = False is there because datasets 3.6.0 uses the old torchvision api, and isn't compatible w/ this torchvision.'(And it will throw ImportError: cannot import name 'VideoReader' from 'torchvision.io')
# Downgrading torchvision might solve the issue, but there's also a chance that doing that breaks other deps. Since I don't want to damage my brain anymore, I decided to just simply told datasets not to use torchvision.
# To replicate the setup:
#       uv venv --python 3.13 .venv
#       source .venv/bin/activate
#       uv pip install "torch==2.12.1" "torchvision==0.27.1" \
#                      --index-url https://download.pytorch.org/whl/cu130
#       uv pip install "transformers==5.12.1" "auto-round-nightly==0.14.0.dev20260625" \
#                      "accelerate==1.14.0" "datasets==3.6.0" "pillow"
# This is why I don't like Python, it's dependency hell.
import datasets
datasets.config.TORCHVISION_AVAILABLE = False

model_name_or_path = "."
output_dir = "./Ornith-1.0-35B-INT8-AutoRound/"

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)

# auto-round doesn't support quantizing linear_attn, but I wanted to be sure.
ignore_keywords = [
    "embed_tokens",
    "linear_attn",
    "shared_expert",
    "mlp.gate",
    "visual",
    "lm_head",
]

layer_config = {}
for keyword in ignore_keywords:
    layer_config[keyword] = {"bits": 16}

seqlen = 2048
nsamples = 1024

dataset = "NeelNanda/pile-10k:num=256,codeparrot/github-code-clean:num=768"

ar = AutoRound(
    model=model_name_or_path,
    tokenizer=tokenizer,
    scheme="W8A16",
    enable_torch_compile=False,
    group_size=-1,
    sym=True,
    layer_config=layer_config,

    dataset=dataset,

    device_map="0,1",
    batch_size=8,
    seqlen=seqlen,
    iters=1000,
    nsamples=nsamples,
    low_gpu_mem_usage=True,
)

ar.quantize_and_save(output_dir, format="auto_round")