

'''simple docstring'''




import argparse

from transformers import BigBirdConfig, BigBirdForPreTraining, BigBirdForQuestionAnswering, load_tf_weights_in_big_bird
from transformers.utils import logging


logging.set_verbosity_info()


def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : Optional[Any]							, SCREAMING_SNAKE_CASE   : int							, SCREAMING_SNAKE_CASE   : List[Any]							, SCREAMING_SNAKE_CASE   : str    ):
  # Initialise PyTorch model
  UpperCAmelCase    				=						BigBirdConfig.from_json_file(SCREAMING_SNAKE_CASE    )
  print(f'''Building PyTorch model from configuration: {config}'''    )

  if is_trivia_qa:
    UpperCAmelCase    				=						BigBirdForQuestionAnswering(SCREAMING_SNAKE_CASE    )
  else:
    UpperCAmelCase    				=						BigBirdForPreTraining(SCREAMING_SNAKE_CASE    )

  # Load weights from tf checkpoint
  load_tf_weights_in_big_bird(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, is_trivia_qa=SCREAMING_SNAKE_CASE    )

  # Save pytorch-model
  print(f'''Save PyTorch model to {pytorch_dump_path}'''    )
  model.save_pretrained(SCREAMING_SNAKE_CASE    )


if __name__ == "__main__":
  _a					:			List[str]							=    argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
  )
  parser.add_argument(
      '--big_bird_config_file',
      default=None,
      type=str,
      required=True,
      help=(
          'The config json file corresponding to the pre-trained BERT model. \n'
          'This specifies the model architecture.'
      ),
  )
  parser.add_argument(
      '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
  )
  parser.add_argument(
      '--is_trivia_qa', action='store_true', help='Whether to convert a model with a trivia_qa head.'
  )
  _a					:			List[str]							=    parser.parse_args()
  convert_tf_checkpoint_to_pytorch(
      args.tf_checkpoint_path, args.big_bird_config_file, args.pytorch_dump_path, args.is_trivia_qa
  )





'''simple docstring'''




import argparse
import os

import torch

from transformers import FlavaConfig, FlavaForPreTraining
from transformers.models.flava.convert_dalle_to_flava_codebook import convert_dalle_checkpoint


def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : str    ):
  # encoder.embeddings are double copied in original FLAVA
  return sum(param.float().sum() if 'encoder.embeddings' not in key else 0 for key, param in state_dict.items()    )


def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : int							, SCREAMING_SNAKE_CASE   : Optional[int]    ):
  UpperCAmelCase    				=						{}

  for key, value in state_dict.items():
    if "text_encoder.embeddings" in key or "image_encoder.embeddings" in key:
      continue

    UpperCAmelCase    				=						key.replace('heads.cmd.mim_head.cls.predictions'							, 'mmm_image_head'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.mlm_head.cls.predictions'							, 'mmm_text_head'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.itm_head.cls'							, 'itm_head'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.itm_head.pooler'							, 'itm_head.pooler'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.clip_head.logit_scale'							, 'flava.logit_scale'    )
    UpperCAmelCase    				=						key.replace('heads.fairseq_mlm.cls.predictions'							, 'mlm_head'    )
    UpperCAmelCase    				=						key.replace('heads.imagenet.mim_head.cls.predictions'							, 'mim_head'    )
    UpperCAmelCase    				=						key.replace('mm_text_projection'							, 'flava.text_to_mm_projection'    )
    UpperCAmelCase    				=						key.replace('mm_image_projection'							, 'flava.image_to_mm_projection'    )
    UpperCAmelCase    				=						key.replace('image_encoder.module'							, 'flava.image_model'    )
    UpperCAmelCase    				=						key.replace('text_encoder.module'							, 'flava.text_model'    )
    UpperCAmelCase    				=						key.replace('mm_encoder.module.encoder.cls_token'							, 'flava.multimodal_model.cls_token'    )
    UpperCAmelCase    				=						key.replace('mm_encoder.module'							, 'flava.multimodal_model'    )
    UpperCAmelCase    				=						key.replace('text_projection'							, 'flava.text_projection'    )
    UpperCAmelCase    				=						key.replace('image_projection'							, 'flava.image_projection'    )

    UpperCAmelCase    				=						value.float()

  for key, value in codebook_state_dict.items():
    UpperCAmelCase    				=						value

  return upgrade


@torch.no_grad()
def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : List[str]							, SCREAMING_SNAKE_CASE   : Optional[Any]							, SCREAMING_SNAKE_CASE   : Optional[int]							, SCREAMING_SNAKE_CASE   : Any=None    ):
  if config_path is not None:
    UpperCAmelCase    				=						FlavaConfig.from_pretrained(SCREAMING_SNAKE_CASE    )
  else:
    UpperCAmelCase    				=						FlavaConfig()

  UpperCAmelCase    				=						FlavaForPreTraining(SCREAMING_SNAKE_CASE    ).eval()

  UpperCAmelCase    				=						convert_dalle_checkpoint(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, save_checkpoint=SCREAMING_SNAKE_CASE    )

  if os.path.exists(SCREAMING_SNAKE_CASE    ):
    UpperCAmelCase    				=						torch.load(SCREAMING_SNAKE_CASE							, map_location='cpu'    )
  else:
    UpperCAmelCase    				=						torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE							, map_location='cpu'    )

  UpperCAmelCase    				=						upgrade_state_dict(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE    )
  hf_model.load_state_dict(SCREAMING_SNAKE_CASE    )
  UpperCAmelCase    				=						hf_model.state_dict()
  UpperCAmelCase    				=						count_parameters(SCREAMING_SNAKE_CASE    )
  UpperCAmelCase    				=						count_parameters(SCREAMING_SNAKE_CASE    ) + count_parameters(SCREAMING_SNAKE_CASE    )

  assert torch.allclose(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, atol=1E-3    )

  hf_model.save_pretrained(SCREAMING_SNAKE_CASE    )


if __name__ == "__main__":
  _a					:			str							=    argparse.ArgumentParser()
  parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
  parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to flava checkpoint')
  parser.add_argument('--codebook_path', default=None, type=str, help='Path to flava codebook checkpoint')
  parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
  _a					:			str							=    parser.parse_args()

  convert_flava_checkpoint(args.checkpoint_path, args.codebook_path, args.pytorch_dump_folder_path, args.config_path)


