from transformers import AutoModelForMaskedLM
from .modeling_custom_llama import CustomLlamaConfig, CustomLlamaForCausalLM, CustomLlamaForMaskedLM
from transformers import CONFIG_MAPPING, MODEL_MAPPING
# Assuming CustomLlamaConfig is your config class
AutoModelForMaskedLM.register(CustomLlamaConfig, CustomLlamaForMaskedLM)
CONFIG_MAPPING.update({"custom_llama": CustomLlamaConfig})
# MODEL_MAPPING.update({"custom_llama": CustomLlamaForMaskedLM})