#!/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()