import argparseimport torchimport osfrom datasets import load_datasetfrom transformers import WhisperProcessorfrom llmcompressor import oneshotfrom llmcompressor.modifiers.quantization import QuantizationModifierfrom llmcompressor.transformers.tracing import TraceableWhisperForConditionalGenerationfrom compressed_tensors.quantization import QuantizationType # --- Args ---parser = argparse.ArgumentParser()parser.add_argument('--model_path', type=str, required=True)parser.add_argument('--quant_path', type=str, required=True)parser.add_argument('--observer', type=str, default="minmax")args = parser.parse_args() # --- Load Model ---model = TraceableWhisperForConditionalGeneration.from_pretrained( args.model_path, device_map="auto", torch_dtype="auto",)model.config.forced_decoder_ids = Noneprocessor = WhisperProcessor.from_pretrained(args.model_path) # --- Recipe (FP8 Dynamic) ---recipe = [ QuantizationModifier( targets="Linear", scheme="FP8_DYNAMIC", sequential_targets=["WhisperEncoderLayer", "WhisperDecoderLayer"], ignore=["re:.*lm_head"], )] # --- Run oneshot ---oneshot( model=model, recipe=recipe, trust_remote_code_model=True,) # --- Save ---os.makedirs(args.quant_path, exist_ok=True)model.save_pretrained(args.quant_path, save_compressed=True)processor.save_pretrained(args.quant_path)