# Copyright (c) 2026 Edison dos Santos
# Licensed under the Apache License 2.0 (see LICENSE file in root)
# Part of Ghost Assistant: https://github.com/Edison2ST/GhostAssistantONNXFiles

import onnx
from onnxruntime.quantization import matmul_nbits_quantizer, quant_utils
from pathlib import Path
import os

model_fp32 = "omni_v2_wsl_final.onnx"
model_q4 = "omni_v2_q4.onnx"

print(f"Loading {model_fp32}...")
# We load the model directly to avoid the symbolic shape inference crash
model = onnx.load(model_fp32)

# In ORT 1.23, we define the bits inside the algo_config
# block_size=32 is the best for Dimensity 6300 (ARM)
quant_config = matmul_nbits_quantizer.DefaultWeightOnlyQuantConfig(
    block_size=32,
    is_symmetric=True,
    accuracy_level=4
)
# Force 4-bit by setting bits explicitly if the default isn't 4
quant_config.bits = 4 

print("Quantizing to 4-bit... (This targets the MatMul/Transformer layers)")
quantizer = matmul_nbits_quantizer.MatMulNBitsQuantizer(
    model=model,
    algo_config=quant_config
)

quantizer.process()

print(f"Saving to {model_q4}...")
# Use the internal save method for better compatibility with 1.23
quantizer.model.save_model_to_file(model_q4, use_external_data_format=False)

# Size Check
size_fp32 = os.path.getsize(model_fp32) / (1024*1024)
size_q4 = os.path.getsize(model_q4) / (1024*1024)

print("-" * 30)
print(f"Success!")
print(f"Original Size: {size_fp32:.2f} MB")
print(f"Q4 Size:       {size_q4:.2f} MB")
print("-" * 30)