# fp32 encoder -> 4-bit MatMulNBits (block 32, uint8 zero points: what onnxruntime-web's WebGPU kernel runs)
import sys, time, onnx, onnxruntime
from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer, DefaultWeightOnlyQuantConfig
from onnxruntime.quantization import QuantFormat
src, out = sys.argv[1], sys.argv[2]
print('onnxruntime', onnxruntime.__version__, flush=True)
t = time.time()
m = onnx.load(src)
cfg = DefaultWeightOnlyQuantConfig(block_size=32, is_symmetric=False, quant_format=QuantFormat.QOperator, op_types_to_quantize=('MatMul',))
q = MatMulNBitsQuantizer(m, algo_config=cfg)
q.process()
q.model.save_model_to_file(out, use_external_data_format=False)
print('done', round(time.time() - t), 's', flush=True)