Upload quantize.py with huggingface_hub

quantize.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+"""Quantize the exported fp32 ONNX graphs, matching onnx-community/convert-to-onnx recipes.
+
+  q4     : MatMulNBitsQuantizer(bits=4, block_size=32, is_symmetric=True)
+  q4f16  : float16(keep_io_types=True) applied to the q4 graph
+  fp16   : float16(keep_io_types=True) applied to the fp32 graph
+
+Run: python quantize.py q4            # or: fp16 q4f16
+Large graphs (encoder) are saved with external data.
+"""
+import sys
+from pathlib import Path
+
+import onnx
+from onnxconverter_common import float16 as onnx_float16
+from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer
+
+OUT = Path("onnx")
+STEMS = ["encoder_model", "decoder_model_merged"]
+
+
+def _save(model, path: Path):
+    onnx.save(model, str(path), save_as_external_data=True, all_tensors_to_one_file=True,
+              location=path.name + "_data", convert_attribute=True)
+    print(f"  wrote {path.name} ({(path.stat().st_size)/1e6:.1f} MB graph + external data)")
+
+
+def make_q4(stem: str):
+    m = onnx.load(str(OUT / f"{stem}.onnx"), load_external_data=True)
+    q = MatMulNBitsQuantizer(m, bits=4, block_size=32, is_symmetric=True)
+    q.process()
+    _save(q.model.model, OUT / f"{stem}_q4.onnx")
+
+
+def make_fp16(stem: str):
+    m = onnx.load(str(OUT / f"{stem}.onnx"), load_external_data=True)
+    # disable_shape_infer: the >2GB encoder overflows the single-proto serialize in infer_shapes.
+    fp16 = onnx_float16.convert_float_to_float16(m, keep_io_types=True, disable_shape_infer=True)
+    _save(fp16, OUT / f"{stem}_fp16.onnx")
+
+
+def make_q4f16(stem: str):
+    m = onnx.load(str(OUT / f"{stem}_q4.onnx"), load_external_data=True)
+    fp16 = onnx_float16.convert_float_to_float16(m, keep_io_types=True, disable_shape_infer=True)
+    _save(fp16, OUT / f"{stem}_q4f16.onnx")
+
+
+def main():
+    targets = sys.argv[1:] or ["q4"]
+    for stem in STEMS:
+        for t in targets:
+            print(f"[{t}] {stem}")
+            {"q4": make_q4, "fp16": make_fp16, "q4f16": make_q4f16}[t](stem)
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()