





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       __UpperCamelCase  (_SCREAMING_SNAKE_CASE		)    ->       List[Any]:
  # 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       __UpperCamelCase  (_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE		)    ->       Dict:
  lowercase__           =  {}

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

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

    lowercase__           =  value.float()

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

  return upgrade




@torch.no_grad()
def       __UpperCamelCase  (_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE=None		)    ->       List[Any]:
  if config_path is not None:
    lowercase__           =  FlavaConfig.from_pretrained(_SCREAMING_SNAKE_CASE		)
  else:
    lowercase__           =  FlavaConfig()

  lowercase__           =  FlavaForPreTraining(_SCREAMING_SNAKE_CASE		).eval()

  lowercase__           =  convert_dalle_checkpoint(_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE		,						save_checkpoint=_SCREAMING_SNAKE_CASE		)

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

  lowercase__           =  upgrade_state_dict(_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE		)
  hf_model.load_state_dict(_SCREAMING_SNAKE_CASE		)
  lowercase__           =  hf_model.state_dict()
  lowercase__           =  count_parameters(_SCREAMING_SNAKE_CASE		)
  lowercase__           =  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__":
       lowercase_          =		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""")
       lowercase_          =		parser.parse_args()

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









from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase_          =		logging.get_logger(__name__)

lowercase_          =		{
    """transfo-xl-wt103""": """https://huggingface.co/transfo-xl-wt103/resolve/main/config.json""",
}
class     SCREAMING_SNAKE_CASE       (UpperCAmelCase				):
   _UpperCamelCase   :		Optional[Any]			 =	'transfo-xl'
   _UpperCamelCase   :		Any			 =	['mems']
   _UpperCamelCase   :		Any			 =	{
       'n_token': 'vocab_size',
       'hidden_size': 'd_model',
       'num_attention_heads': 'n_head',
       'num_hidden_layers': 'n_layer',
   }


   def __init__(   self      :    Optional[Any]      ,					a      :    Optional[int]=267_735      ,					a      :    str=[20_000, 40_000, 200_000]      ,					a      :    str=1_024      ,					a      :    str=1_024      ,					a      :    int=16      ,					a      :    Optional[int]=64      ,					a      :    Optional[int]=4_096      ,					a      :    int=4      ,					a      :    Tuple=False      ,					a      :    Any=18      ,					a      :    Tuple=1_600      ,					a      :    Union[str, Any]=1_000      ,					a      :    str=True      ,					a      :    Dict=True      ,					a      :    Any=0      ,					a      :    List[Any]=-1      ,					a      :    List[Any]=True      ,					a      :    Tuple=0.1      ,					a      :    List[Any]=0.0      ,					a      :    Optional[Any]=True      ,					a      :    int="normal"      ,					a      :    Optional[Any]=0.01      ,					a      :    str=0.01      ,					a      :    List[Any]=0.02      ,					a      :    List[Any]=1E-5      ,					a      :    Optional[Any]=0      ,					**a      :    Optional[int]      ,					)->       Optional[int]:





     """simple docstring"""







     lowercase__           =  vocab_size
     lowercase__           =  []
     self.cutoffs.extend(a	)
     if proj_share_all_but_first:
       lowercase__           =  [False] + [True] * len(self.cutoffs	)
     else:
       lowercase__           =  [False] + [False] * len(self.cutoffs	)
     lowercase__           =  d_model
     lowercase__           =  d_embed
     lowercase__           =  d_head
     lowercase__           =  d_inner
     lowercase__           =  div_val
     lowercase__           =  pre_lnorm
     lowercase__           =  n_layer
     lowercase__           =  n_head
     lowercase__           =  mem_len
     lowercase__           =  same_length
     lowercase__           =  attn_type
     lowercase__           =  clamp_len
     lowercase__           =  sample_softmax
     lowercase__           =  adaptive
     lowercase__           =  dropout
     lowercase__           =  dropatt
     lowercase__           =  untie_r
     lowercase__           =  init
     lowercase__           =  init_range
     lowercase__           =  proj_init_std
     lowercase__           =  init_std
     lowercase__           =  layer_norm_epsilon
     super().__init__(eos_token_id=a      ,					**a	)


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]	)->       Union[str, Any]:





     """simple docstring"""







     logger.info(f"""The model {self.model_type} is one of the few models that has no sequence length limit."""	)
     return -1


   @max_position_embeddings.setter
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any      ,					a      :    Optional[int]	)->       Optional[int]:





     """simple docstring"""







     raise NotImplementedError(
         f"""The model {self.model_type} is one of the few models that has no sequence length limit."""	)


