# SPDX-License-Identifier: Apache-2.0

from ...common._registration import register_converter


def convert_dictionary_vectorizer(scope, operator, container):
    op_type = 'DictVectorizer'
    attrs = {'name': operator.full_name}
    raw_model = operator.raw_operator.dictVectorizer
    if raw_model.HasField('stringToIndex'):
        attrs['string_vocabulary'] = raw_model.stringToIndex.vector
    else:
        attrs['int64_vocabulary'] = raw_model.int64ToIndex.vector

    container.add_node(op_type, [operator.inputs[0].full_name], [operator.outputs[0].full_name],
                       op_domain='ai.onnx.ml', **attrs)


register_converter('dictVectorizer', convert_dictionary_vectorizer)
