File size: 1,376 Bytes
413bf7f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
from transformers import AutoConfig, AutoModel, AutoTokenizer
from transformers.models.auto.modeling_auto import MODEL_MAPPING
from transformers.models.auto.configuration_auto import CONFIG_MAPPING
from .config import BilingualConfig
from .modeling import BilingualModel, BilingualForMaskedLM, BilingualForMultipleChoice, BilingualForQuestionAnswering, BilingualForSequenceClassification, BilingualForTokenClassification
AutoConfig.register("bilingual", BilingualConfig)
AutoModel.register(BilingualConfig, BilingualModel)
AutoModel.register(BilingualConfig, BilingualForMaskedLM)
AutoModel.register(BilingualConfig, BilingualForMultipleChoice)
AutoModel.register(BilingualConfig, BilingualForQuestionAnswering)
AutoModel.register(BilingualConfig, BilingualForSequenceClassification)
AutoModel.register(BilingualConfig, BilingualForTokenClassification)
AutoTokenizer.register(BilingualConfig, AutoTokenizer)
CONFIG_MAPPING.update({"bilingual": BilingualConfig})
MODEL_MAPPING.update({BilingualConfig: BilingualModel})
MODEL_MAPPING.update({BilingualConfig: BilingualForMaskedLM})
MODEL_MAPPING.update({BilingualConfig: BilingualForMultipleChoice})
MODEL_MAPPING.update({BilingualConfig: BilingualForQuestionAnswering})
MODEL_MAPPING.update({BilingualConfig: BilingualForSequenceClassification})
MODEL_MAPPING.update({BilingualConfig: BilingualForTokenClassification}) |