bilingual_impl / __init__.py
dangvantuan's picture
Upload 3 files
413bf7f verified
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})