def __init__(
        self,
        input_name,
        output_name,
        number_of_dimensions,
        default_class_label,
        **kwargs
    ):
        """
        Create a KNearestNeighborsClassifierBuilder object.

        Parameters
        ----------
        input_name
                Name of the model input.

        output_name
                Name of the output.

        number_of_dimensions
                Number of dimensions of the input data.

        default_class_label
                The default class label to use for predictions. Must be either an
                int64 or a string.

        number_of_neighbors
                Number of neighbors to use for predictions. Default = 5 with allowed values
                between 1-1000.

        weighting_scheme
                Weight function used in prediction. One of ``'uniform'`` (default) or
                ``'inverse_distance'``.

        index_type
                Algorithm to compute nearest neighbors. One of ``'linear'`` (default), or
                ``'kd_tree'``.

        leaf_size
                Leaf size for the kd-tree. Ignored if index type is ``'linear'``. Default = 30.
        """
        self.spec = coremltools.proto.Model_pb2.Model()
        self.spec.specificationVersion = (
            coremltools._MINIMUM_NEAREST_NEIGHBORS_SPEC_VERSION
        )

        # the model is initially empty - assume it's updatable
        self.is_updatable = True

        if number_of_dimensions <= 0:
            raise ValueError("number_of_dimensions must be >= 0")
        self.spec.kNearestNeighborsClassifier.nearestNeighborsIndex.numberOfDimensions = (
            number_of_dimensions
        )

        input_type = kwargs.get(
            self._PARAMETER_KEY_INPUT_TYPE, self._PARAMETER_DEFAULT_INPUT_TYPE
        )
        input_feature_type = proto.FeatureTypes_pb2.ArrayFeatureType.FLOAT32
        if input_type == datatypes.Double:
            input_feature_type = proto.FeatureTypes_pb2.ArrayFeatureType.DOUBLE

        input_feature = self.spec.description.input.add()
        input_feature.name = input_name
        input_feature.type.multiArrayType.dataType = input_feature_type
        input_feature.type.multiArrayType.shape.extend([number_of_dimensions])

        training_features = self.spec.description.trainingInput.add()
        training_features.name = input_name
        training_features.type.multiArrayType.dataType = input_feature_type
        training_features.type.multiArrayType.shape.extend([number_of_dimensions])

        output_label = self.spec.description.output.add()
        output_label.name = output_name
        output_label_probs = self.spec.description.output.add()
        output_label_probs.name = output_name + "Probs"
        training_features = self.spec.description.trainingInput.add()
        training_features.name = output_name

        if self._is_valid_text_type(default_class_label):
            output_label.type.stringType.MergeFromString(b"")
            training_features.type.stringType.MergeFromString(b"")
            output_label_probs.type.dictionaryType.stringKeyType.MergeFromString(b"")
            self.spec.kNearestNeighborsClassifier.stringClassLabels.MergeFromString(b"")
            self.spec.kNearestNeighborsClassifier.defaultStringLabel = (
                default_class_label
            )
        elif self._is_valid_number_type(default_class_label):
            output_label.type.int64Type.MergeFromString(b"")
            training_features.type.int64Type.MergeFromString(b"")
            output_label_probs.type.dictionaryType.int64KeyType.MergeFromString(b"")
            self.spec.kNearestNeighborsClassifier.int64ClassLabels.MergeFromString(b"")
            self.spec.kNearestNeighborsClassifier.defaultInt64Label = (
                default_class_label
            )
        else:
            raise TypeError(
                "default_class_label type ({}) is invalid. Must be either string or int64".format(
                    type(default_class_label)
                )
            )

        self.spec.description.predictedFeatureName = output_label.name
        self.spec.description.predictedProbabilitiesName = output_label_probs.name

        number_of_neighbors = kwargs.get(
            self._PARAMETER_KEY_NUMBER_OF_NEIGHBORS,
            self._PARAMETER_DEFAULT_NUMBER_OF_NEIGHBORS,
        )
        self.set_number_of_neighbors_with_bounds(
            number_of_neighbors, allowed_range=(1, 1000)
        )  # Can we think of a more sensible default value?

        self.weighting_scheme = kwargs.get(
            self._PARAMETER_KEY_WEIGHTING_SCHEME,
            self._PARAMETER_DEFAULT_WEIGHTING_SCHEME,
        )

        index_type = kwargs.get(
            self._PARAMETER_KEY_INDEX_TYPE, self._PARAMETER_DEFAULT_INDEX_TYPE
        )
        leaf_size = kwargs.get(
            self._PARAMETER_KEY_LEAF_SIZE, self._PARAMETER_DEFAULT_LEAF_SIZE
        )
        self.set_index_type(index_type, leaf_size)

        # SED is currently the only supported distance metric
        self.spec.kNearestNeighborsClassifier.nearestNeighborsIndex.squaredEuclideanDistance.MergeFromString(
            b""
        )

        Create a KNearestNeighborsClassifierBuilder object.

        Parameters
        ----------
        input_name
                Name of the model input.

        output_name
                Name of the output.

        number_of_dimensions
                Number of dimensions of the input data.

        default_class_label
                The default class label to use for predictions. Must be either an
                int64 or a string.

        number_of_neighbors
                Number of neighbors to use for predictions. Default = 5 with allowed values
                between 1-1000.

        weighting_scheme
                Weight function used in prediction. One of ``'uniform'`` (default) or
                ``'inverse_distance'``.

        index_type
                Algorithm to compute nearest neighbors. One of ``'linear'`` (default), or
                ``'kd_tree'``.

        leaf_size
                Leaf size for the kd-tree. Ignored if index type is ``'linear'``. Default = 30.
        
coremltools/models/nearest_neighbors/builder.py