import onnx, os, traceback
from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer

os.makedirs("/tmp/nemo/en-int4", exist_ok=True)
print("loading fp16 EN encoder (1.2GB)...", flush=True)
model = onnx.load("/tmp/nemo/sh/encoder_model.onnx")  # resolves encoder_model.onnx.data
print("quantizing MatMul -> int4 MatMulNBits (block_size=128)...", flush=True)
quant = MatMulNBitsQuantizer(model, bits=4, block_size=128, is_symmetric=False)
quant.process()
out = "/tmp/nemo/en-int4/encoder_model.onnx"
m = quant.model
print("saving int4 encoder...", flush=True)
saved = False
try:
    m.save_model_to_file(out, use_external_data_format=True)
    saved = True
except Exception as e:
    print("save_model_to_file failed:", repr(e), flush=True)
    try:
        mp = m.model if hasattr(m, "model") and isinstance(getattr(m, "model"), onnx.ModelProto) else m
        if not isinstance(mp, onnx.ModelProto):
            import onnx_ir as ir
            mp = ir.serde.serialize_model(m) if hasattr(ir, "serde") else ir.to_proto(m)
        onnx.save(mp, out, save_as_external_data=True, all_tensors_to_one_file=True,
                  location="encoder_model.onnx.data", convert_attribute=False)
        saved = True
    except Exception as e2:
        print("fallback save failed:", repr(e2)); traceback.print_exc()

if saved:
    g = os.path.getsize(out)
    d = os.path.getsize(out + ".data") if os.path.exists(out + ".data") else 0
    print(f"DONE int4 encoder: graph {g} B + data {d} B = {(g+d)/1e6:.0f} MB (fp16 was 1249 MB)")
    import collections
    mm = onnx.load(out, load_external_data=False)
    ops = collections.Counter(n.op_type for n in mm.graph.node)
    print("MatMulNBits:", ops.get("MatMulNBits"), "| ConvInteger:", ops.get("ConvInteger"), "| Conv:", ops.get("Conv"), "| MatMul:", ops.get("MatMul"))