import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# Пути к файлам
model_fp32 = "/root/JiRackTernary1/new/model/jirack_1b.onnx"
model_int8 = "/root/JiRackTernary1/new/model/jirack_1b_int8.onnx"

print("🚀 Запуск квантования JiRack (BitNet) в INT8...")

try:
    # Оставляем только необходимые аргументы
    quantize_dynamic(
        model_input=model_fp32,
        model_output=model_int8,
        weight_type=QuantType.QInt8  # Используем знаковый Int8, он лучше для тернарных весов
    )
    print(f"✅ Успех! Квантованная модель сохранена: {model_int8}")
except Exception as e:
    print(f"❌ Ошибка при квантовании: {e}")