import onnx, os, time, numpy as np, sys
from onnxconverter_common import float16
from onnxruntime.quantization import quantize_dynamic, QuantType
from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer
D = "/Users/shreyas/work/rnd/gliner2/work/export/onnx"
model = onnx.load(D + "/model.onnx")
print("loaded fp32; nodes", len(model.graph.node), flush=True)

# fp16: keep gather/index ops in fp32 via op_block_list defaults; blocks handle the int paths
t0 = time.time()
m16 = float16.convert_float_to_float16(model, keep_io_types=True, disable_shape_infer=True)
onnx.save_model(m16, D + "/model_fp16.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="model_fp16.onnx.data", size_threshold=1024)
print("fp16 done", round(time.time()-t0,1), "s", flush=True)
del m16

# q4: MatMulNBits (what onnx-community ships as model_q4.onnx); embeddings stay fp32
t0 = time.time()
q = MatMulNBitsQuantizer(onnx.load(D + "/model.onnx"), block_size=32, is_symmetric=True, accuracy_level=4)
q.process()
q.model.save_model_to_file(D + "/model_q4.onnx", use_external_data_format=True)
print("q4 done", round(time.time()-t0,1), "s", flush=True)

# q4f16: q4 matmuls with fp16 everything else
t0 = time.time()
m = onnx.load(D + "/model_q4.onnx")
m = float16.convert_float_to_float16(m, keep_io_types=True, disable_shape_infer=True, op_block_list=float16.DEFAULT_OP_BLOCK_LIST + ["MatMulNBits"])
onnx.save_model(m, D + "/model_q4f16.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="model_q4f16.onnx.data", size_threshold=1024)
print("q4f16 done", round(time.time()-t0,1), "s", flush=True)