import torchimport argparsefrom datasets import load_datasetfrom transformers import WhisperProcessorfrom llmcompressor import oneshotfrom llmcompressor.modifiers.quantization import GPTQModifierfrom llmcompressor.transformers.tracing import TraceableWhisperForConditionalGenerationimport osfrom compressed_tensors.quantization import QuantizationArgs, QuantizationType, QuantizationStrategy, ActivationOrdering, QuantizationSchemefrom llmcompressor.modifiers.smoothquant import SmoothQuantModifier
parser = argparse.ArgumentParser()parser.add_argument('--model_path', type=str)parser.add_argument('--quant_path', type=str)parser.add_argument('--calib_size', type=int, default=256)parser.add_argument('--dampening_frac', type=float, default=0.1) parser.add_argument('--observer', type=str, default="minmax")parser.add_argument('--save_dir', type=str, required=True)
args = parser.parse_args()model_id = args.model_path
model = TraceableWhisperForConditionalGeneration.from_pretrained( model_id, device_map="auto", torch_dtype="auto",)model.config.forced_decoder_ids = Noneprocessor = WhisperProcessor.from_pretrained(model_id)
# Configure processor the dataset task.processor.tokenizer.set_prefix_tokens(language="en", task="transcribe")
# Select calibration dataset.DATASET_ID = "MLCommons/peoples_speech"DATASET_SUBSET = "test"DATASET_SPLIT = "test"
# Select number of samples for calibration. 512 samples is a good place to start.# Increasing the number of samples can improve accuracy.
NUM_CALIBRATION_SAMPLES = args.calib_sizeMAX_SEQUENCE_LENGTH = 2048dampening_frac=args.dampening_fracactorder_arg=args.actordergroup_size=args.group_size
# Load dataset and preprocess.ds = load_dataset( DATASET_ID, DATASET_SUBSET, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]", trust_remote_code=True,)
def preprocess(example): return { "array": example["audio"]["array"], "sampling_rate": example["audio"]["sampling_rate"], "text": " " + example["text"].capitalize(), }
ds = ds.map(preprocess, remove_columns=ds.column_names)
# Process inputs.def process(sample): inputs = processor( audio=sample["array"], sampling_rate=sample["sampling_rate"], text=sample["text"], add_special_tokens=True, return_tensors="pt", )
inputs["input_features"] = inputs["input_features"].to(dtype=model.dtype) inputs["decoder_input_ids"] = inputs["labels"] del inputs["labels"]
return inputs
ds = ds.map(process, remove_columns=ds.column_names)
# Define a oneshot data collator for multimodal inputs.def data_collator(batch): assert len(batch) == 1 return {key: torch.tensor(value) for key, value in batch[0].items()}
ignore=["lm_head"]
#Reciperecipe = [ GPTQModifier( targets="Linear", scheme="W8A8", sequential_targets=["WhisperEncoderLayer", "WhisperDecoderLayer"], ignore=ignore, )]
# Apply algorithms.oneshot( model=model, dataset=ds, recipe=recipe, max_seq_length=MAX_SEQUENCE_LENGTH, num_calibration_samples=NUM_CALIBRATION_SAMPLES, data_collator=data_collator,)
# Save to disk compressed.save_name = f"{model_id.split('/')[-1]}-quantized.w8a8"save_path = os.path.join(args.save_dir, save_name)print("Saving model:", save_path)model.save_pretrained(save_path, save_compressed=True)processor.save_pretrained(save_path)