# coding=utf-8
# author: xusong <xusong28@jd.com>
# time: 2022/8/25 10:49

"""

注册后，才能被AutoConfig识别
https://github.com/huggingface/transformers/blob/main/src/transformers/models/auto/configuration_auto.py
https://github.com/huggingface/transformers/blob/main/src/transformers/models/auto/modeling_auto.py
"""

import transformers
# import kplug
#
# transformers.models.kplug = kplug

from transformers.models.auto.modeling_auto import MODEL_MAPPING_NAMES, MODEL_WITH_LM_HEAD_MAPPING_NAMES, MODEL_FOR_MASKED_LM_MAPPING_NAMES, CONFIG_MAPPING_NAMES


CONFIG_MAPPING_NAMES["kplug"] = "BertConfig"
MODEL_MAPPING_NAMES["kplug"] = "KplugModel"
MODEL_WITH_LM_HEAD_MAPPING_NAMES["kplug"] = "KplugForMaskedLM"
MODEL_FOR_MASKED_LM_MAPPING_NAMES["kplug"] = "KplugForMaskedLM"

# CONFIG_MAPPING = _LazyConfigMapping(CONFIG_MAPPING_NAMES)

from transformers.models.auto.configuration_auto import CONFIG_MAPPING


CONFIG_MAPPING["kplug"] = "BertConfig"