kplug / draft /register.py
xusong28
update
c10350f
raw
history blame
958 Bytes
# 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"