import sys, torch
from transformers import AutoModelForImageTextToText, AutoTokenizer
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.modeling.moe.linearize import load_quantizable_moe

MODEL, OUT = sys.argv[1], sys.argv[2]
NSAMPLES = int(sys.argv[3]) if len(sys.argv) > 3 else 64

# Mirror Mistral official NVFP4: quantize expert FFN (routed via MoE-linearize + shared),
# keep bf16: vision, all MLA attention, MoE router gate, embeddings, lm_head
IGNORE = [
    "re:.*lm_head.*", "re:.*embed_tokens.*",
    "re:.*vision_tower.*", "re:.*multi_modal_projector.*",
    "re:.*self_attn.*", r"re:.*mlp\.gate$",
]
recipe = QuantizationModifier(targets="Linear", scheme="NVFP4", ignore=IGNORE)

with load_quantizable_moe(AutoModelForImageTextToText):
    model = AutoModelForImageTextToText.from_pretrained(MODEL, dtype="auto", device_map="cpu")

oneshot(model=model, recipe=recipe, processor=AutoTokenizer.from_pretrained(MODEL), dataset="ultrachat_200k",
        splits={"calibration": f"train_gen[:{NSAMPLES}]"},
        num_calibration_samples=NSAMPLES, max_seq_length=512)
model.save_pretrained(OUT, save_compressed=True)
print("SAVED", OUT)