content: from .xbert import BertConfig, BertForMaskedLM, BertLMHeadModel, BertModel
  
  import logging
  logger = logging.getLogger(__name__)
  
  def build_bert(model_config, pretrain, checkpoint):
      """build text encoder.
  
      Args:
          model_config (dict): model config.
          pretrain (bool): Whether to do pretrain or finetuning.
          checkpoint (bool): whether to do gradient_checkpointing.
  
      Returns: TODO
  
      """
      bert_config = BertConfig.from_json_file(model_config.text_encoder.config)
      bert_config.encoder_width = model_config.vision_encoder.d_model
      bert_config.gradient_checkpointing = checkpoint
      bert_config.fusion_layer = model_config.text_encoder.fusion_layer
  
      if not model_config.multimodal.enable:
          bert_config.fusion_layer = bert_config.num_hidden_layers
  
      if pretrain:
          text_encoder, loading_info = BertForMaskedLM.from_pretrained(
              model_config.text_encoder.pretrained,
              config=bert_config,
              output_loading_info=True,
          )
      else:
          text_encoder, loading_info = BertModel.from_pretrained(
              model_config.text_encoder.pretrained,
              config=bert_config,
              add_pooling_layer=False,
              output_loading_info=True,
          )
  
      return text_encoder
  
  
  def build_bert_decoder(model_config, checkpoint):
      """build text decoder the same as the multimodal encoder.
  
      Args:
          model_config (dict): model config.
          pretrain (bool): Whether to do pretrain or finetuning.
          checkpoint (bool): whether to do gradient_checkpointing.
  
      Returns: TODO
  
      """
      bert_config = BertConfig.from_json_file(model_config.text_encoder.config)
      bert_config.encoder_width = model_config.vision_encoder.d_model
      bert_config.gradient_checkpointing = checkpoint
  
      bert_config.fusion_layer = 0
      bert_config.num_hidden_layers = (
          bert_config.num_hidden_layers - model_config.text_encoder.fusion_layer
      )
  
      text_decoder, loading_info = BertLMHeadModel.from_pretrained(
          model_config.text_encoder.pretrained,
          config=bert_config,
          output_loading_info=True,
      )
  
      return text_decoder
  