Utility function to convert a full precision (float) MLModel to a
    nbit quantized MLModel (float16).

    :param full_precision_model: MLModel
        Model which will be converted to half precision. Currently conversion
        for only neural network models is supported. If a pipeline model is
        passed in then all embedded neural network models embedded within
        will be converted.

    :param nbits: Int
        Number of bits per quantized weight. Only 8-bit and lower
        quantization is supported

    :param quantization_mode: str
        One of:
         "linear":
            Simple linear quantization with scale and bias

         "linear_lut":
            Simple linear quantization represented as a lookup table

         "kmeans_lut":
            LUT based quantization, where LUT is generated by K-Means clustering

         "custom_lut":
            LUT quantization where LUT and quantized weight params are
            calculated using a custom function. If this mode is selected then
            a custom function must be passed in kwargs with key lut_function.
            The function must have input params (nbits, wp) where nbits is the
            number of quantization bits and wp is the list of weights for a
            given layer. The function should return two parameters (lut, qw)
            where lut is an array of length (2^nbits)containing LUT values and
            qw is the list of quantized weight parameters. See
            _get_linear_lookup_table_and_weight for a sample implementation.

    :param sample_data: str | [dict]
        Data used to characterize performance of the quantized model in
        comparison to the full precision model. Either a list of sample input
        dictionaries or an absolute path to a directory containing images.
        Path to a directory containing images is only valid for models with
        one image input. For all other models a list of sample inputs must be
        provided.

    :param **kwargs:
        See below

    :Keyword Arguments:
        * *lut_function* (``callable function``) --
          A callable function provided when quantization mode is set to
          _QUANTIZATION_MODE_CUSTOM_LOOKUP_TABLE. See quantization_mode for
          more details

    Returns
    -------
    model: MLModel
        The quantized MLModel instance if running on macOS 10.14 or later,
        otherwise the quantized model specification is returned

    Examples
    --------
    .. sourcecode:: python
        >>> import coremltools
        >>> from coremltools.models.neural_network import quantization_utils
        >>> model = coremltools.models.MLModel('my_model.mlmodel')
        >>> quantized_model = quantization_utils.quantize_weights(model, 8, "linear")

def quantize_weights(full_precision_model,
                     nbits,
                     quantization_mode="linear",
                     sample_data=None,
                     **kwargs):
    """
    Utility function to convert a full precision (float) MLModel to a
    nbit quantized MLModel (float16).

    :param full_precision_model: MLModel
        Model which will be converted to half precision. Currently conversion
        for only neural network models is supported. If a pipeline model is
        passed in then all embedded neural network models embedded within
        will be converted.

    :param nbits: Int
        Number of bits per quantized weight. Only 8-bit and lower
        quantization is supported

    :param quantization_mode: str
        One of:
         "linear":
            Simple linear quantization with scale and bias

         "linear_lut":
            Simple linear quantization represented as a lookup table

         "kmeans_lut":
            LUT based quantization, where LUT is generated by K-Means clustering

         "custom_lut":
            LUT quantization where LUT and quantized weight params are
            calculated using a custom function. If this mode is selected then
            a custom function must be passed in kwargs with key lut_function.
            The function must have input params (nbits, wp) where nbits is the
            number of quantization bits and wp is the list of weights for a
            given layer. The function should return two parameters (lut, qw)
            where lut is an array of length (2^nbits)containing LUT values and
            qw is the list of quantized weight parameters. See
            _get_linear_lookup_table_and_weight for a sample implementation.

    :param sample_data: str | [dict]
        Data used to characterize performance of the quantized model in
        comparison to the full precision model. Either a list of sample input
        dictionaries or an absolute path to a directory containing images.
        Path to a directory containing images is only valid for models with
        one image input. For all other models a list of sample inputs must be
        provided.

    :param **kwargs:
        See below

    :Keyword Arguments:
        * *lut_function* (``callable function``) --
          A callable function provided when quantization mode is set to
          _QUANTIZATION_MODE_CUSTOM_LOOKUP_TABLE. See quantization_mode for
          more details

    Returns
    -------
    model: MLModel
        The quantized MLModel instance if running on macOS 10.14 or later,
        otherwise the quantized model specification is returned

    Examples
    --------
    .. sourcecode:: python
        >>> import coremltools
        >>> from coremltools.models.neural_network import quantization_utils
        >>> model = coremltools.models.MLModel('my_model.mlmodel')
        >>> quantized_model = quantization_utils.quantize_weights(model, 8, "linear")

    """
    qmode_mapping = {
        "linear": _QUANTIZATION_MODE_LINEAR_QUANTIZATION,
        "kmeans": _QUANTIZATION_MODE_LOOKUP_TABLE_KMEANS,
        "linear_lut": _QUANTIZATION_MODE_LOOKUP_TABLE_LINEAR,
        "custom_lut": _QUANTIZATION_MODE_CUSTOM_LOOKUP_TABLE,
        "dequantization": _QUANTIZATION_MODE_DEQUANTIZE
    }
    try:
        qmode = qmode_mapping[quantization_mode]
    except KeyError:
        raise Exception("Invalid quantization mode. Quantization mode must be "
                        "one of {}".format(qmode_mapping))

    print("Quantizing using {} quantization".format(quantization_mode))
    spec = full_precision_model.get_spec()
    qspec = quantize_spec_weights(spec, nbits, qmode, **kwargs)

    if macos_version() < (10, 14):
        print("WARNING! Unable to return a quantized MLModel instance since OS != macOS 10.14 or later")
        print("Returning quantized model specification instead")
        return qspec

    quantized_model = _get_model(qspec)
    if not sample_data:
        return quantized_model

    compare_models(full_precision_model, quantized_model, sample_data)
    return quantized_model