"""Weight-only block-wise quantization of the fp16 Core ML model (int4 or int8)."""
import sys

import coremltools as ct
import coremltools.optimize.coreml as cto

src, bits, out = sys.argv[1], sys.argv[2], sys.argv[3]
dtype = {"4": "int4", "8": "int8"}[bits]
cfg = cto.OptimizationConfig(
    global_config=cto.OpLinearQuantizerConfig(
        mode="linear_symmetric", dtype=dtype, granularity="per_block", block_size=32
    )
)
m = ct.models.MLModel(src, skip_model_load=True)
q = cto.linear_quantize_weights(m, cfg)
q.save(out)
print("saved", out)