#!/usr/bin/env python
"""W4A16 compressed-tensors export - runs in venv_quant. int4 / group 32 / symmetric,
matching google/gemma-4-12B-it-qat-w4a16-ct. Embeddings + lm_head stay high precision
(Google's ignore list) unless --quant-embeddings."""
import os, sys, json, argparse
os.environ.setdefault("HF_HOME","/workspace/.hf_home")
ap=argparse.ArgumentParser()
ap.add_argument("--src",required=True); ap.add_argument("--out",required=True)
ap.add_argument("--base",default="/workspace/models/gemma4-qat-unquantized")
A=ap.parse_args()
from transformers import AutoTokenizer, AutoModelForCausalLM
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
IGNORE=["lm_head","re:.*embed_vision.*","re:.*vision_embedder.*","re:.*embed_audio.*","re:.*audio_tower.*","re:.*vision_tower.*","re:.*multimodal_embedder.*","re:.*patch_dense.*","re:.*patch_ln.*","re:.*pos_norm.*","re:.*embedding_projection.*"]
# Match google/gemma-4-12B-it-qat-w4a16-ct EXACTLY: int4, group 32, symmetric.
# The QAT weights are calibrated for a group-32 grid; the llm-compressor default of
# group 128 is a coarser grid and forfeits part of that calibration.
from compressed_tensors.quantization import (QuantizationArgs, QuantizationScheme,
                                             QuantizationStrategy, QuantizationType)
_w=QuantizationArgs(num_bits=4, type=QuantizationType.INT, symmetric=True,
                    strategy=QuantizationStrategy.GROUP, group_size=32)
recipe=QuantizationModifier(
    config_groups={"group_0":QuantizationScheme(targets=["Linear"], weights=_w)},
    ignore=IGNORE)
tok=AutoTokenizer.from_pretrained(A.src)
m=AutoModelForCausalLM.from_pretrained(A.src, dtype="bfloat16", device_map="cpu")
oneshot(model=m, recipe=recipe, output_dir=A.out, tokenizer=tok)
# re-assert the three keys AFTER save_pretrained (handover 6.7)
cfg=os.path.join(A.out,"config.json"); c=json.load(open(cfg))
b=json.load(open(os.path.join(A.base,"config.json")))
for sec,k in (("vision_config","model_patch_size"),("vision_config","num_soft_tokens"),
              ("audio_config","output_proj_dims")):
    c.setdefault(sec,{})
    if c[sec].get(k) is None: c[sec][k]=b[sec][k]
json.dump(c,open(cfg,"w"),indent=2)
c=json.load(open(cfg))
for sec,k in (("vision_config","model_patch_size"),("vision_config","num_soft_tokens"),
              ("audio_config","output_proj_dims")):
    assert c[sec].get(k) is not None, f"MISSING {sec}.{k}"
sz=sum(os.path.getsize(os.path.join(A.out,f)) for f in os.listdir(A.out) if f.endswith('.safetensors'))
print(f"[quant] W4A16 written -> {A.out}  {sz/1e9:.2f} GB  (config keys asserted)")