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})