import argparse
import os
from transformers import AutoTokenizer, Qwen3VLMoeForConditionalGeneration
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
def quant(model_path, offload_folder=None, output_path=None):
    if offload_folder is None:
        offload_folder = "./offload"
    os.makedirs(offload_folder, exist_ok=True)
    model = Qwen3VLMoeForConditionalGeneration.from_pretrained(
        model_path, 
        device_map="auto",
        torch_dtype="auto",
        offload_folder=offload_folder,
        offload_state_dict=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    # Configure the simple PTQ quantization
    recipe = QuantizationModifier(targets="Linear",
                                  scheme="FP8_DYNAMIC",
                                  ignore=["lm_head", "re:.*mlp.gate$"])
    # Apply the quantization algorithm.
    oneshot(model=model, recipe=recipe)
    # Save the model.
    # SAVE_DIR = model_path.split("/")[1] + "-FP8-Dynamic"
    if output_path is None:
        SAVE_DIR = model_path + "-FP8-Dynamic"
    else:
        SAVE_DIR = output_path + "-FP8-Dynamic"
    model.save_pretrained(SAVE_DIR)
    tokenizer.save_pretrained(SAVE_DIR)
if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog='fp8_quant', description='fp8 quant')
    parser.add_argument("--model_path", type=str, required=True)
    parser.add_argument("--output_path", type=str, required=True)
    parser.add_argument("--offload_folder", type=str, default="./offload")
    args = parser.parse_args()
    print(f"run fp8 quant: {args.model_path}")
    print(f"offload folder: {args.offload_folder}")
    print(f"output_path folder: {args.output_path}")
    quant(args.model_path, args.offload_folder, args.output_path)