






"""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		(							UpperCAmelCase__       :					List[str]      )		->List[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		(							UpperCAmelCase__       :					List[str],						UpperCAmelCase__       :					Dict      )		->List[Any]:
						A__						:  List[Any]		  =  {}

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

												A__						:  int		  =  key.replace("""heads.cmd.mim_head.cls.predictions""",						"""mmm_image_head"""      )
												A__						:  str		  =  key.replace("""heads.cmd.mlm_head.cls.predictions""",						"""mmm_text_head"""      )
												A__						:  Dict		  =  key.replace("""heads.cmd.itm_head.cls""",						"""itm_head"""      )
												A__						:  List[Any]		  =  key.replace("""heads.cmd.itm_head.pooler""",						"""itm_head.pooler"""      )
												A__						:  List[Any]		  =  key.replace("""heads.cmd.clip_head.logit_scale""",						"""flava.logit_scale"""      )
												A__						:  Any		  =  key.replace("""heads.fairseq_mlm.cls.predictions""",						"""mlm_head"""      )
												A__						:  str		  =  key.replace("""heads.imagenet.mim_head.cls.predictions""",						"""mim_head"""      )
												A__						:  Optional[int]		  =  key.replace("""mm_text_projection""",						"""flava.text_to_mm_projection"""      )
												A__						:  int		  =  key.replace("""mm_image_projection""",						"""flava.image_to_mm_projection"""      )
												A__						:  Dict		  =  key.replace("""image_encoder.module""",						"""flava.image_model"""      )
												A__						:  Optional[int]		  =  key.replace("""text_encoder.module""",						"""flava.text_model"""      )
												A__						:  int		  =  key.replace("""mm_encoder.module.encoder.cls_token""",						"""flava.multimodal_model.cls_token"""      )
												A__						:  Optional[Any]		  =  key.replace("""mm_encoder.module""",						"""flava.multimodal_model"""      )
												A__						:  Optional[int]		  =  key.replace("""text_projection""",						"""flava.text_projection"""      )
												A__						:  List[Any]		  =  key.replace("""image_projection""",						"""flava.image_projection"""      )

												A__						:  Dict		  =  value.float()

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

						return upgrade




@torch.no_grad()
def      _lowerCAmelCase		(							UpperCAmelCase__       :					List[Any],						UpperCAmelCase__       :					List[Any],						UpperCAmelCase__       :					List[str],						UpperCAmelCase__       :					Optional[Any]=None      )		->Any:
						if config_path is not None:
												A__						:  Optional[int]		  =  FlavaConfig.from_pretrained(UpperCAmelCase__      )
						else:
												A__						:  List[str]		  =  FlavaConfig()

						A__						:  int		  =  FlavaForPreTraining(UpperCAmelCase__      ).eval()

						A__						:  Tuple		  =  convert_dalle_checkpoint(UpperCAmelCase__,						UpperCAmelCase__,						save_checkpoint=UpperCAmelCase__      )

						if os.path.exists(UpperCAmelCase__      ):
												A__						:  Dict		  =  torch.load(UpperCAmelCase__,						map_location="""cpu"""      )
						else:
												A__						:  int		  =  torch.hub.load_state_dict_from_url(UpperCAmelCase__,						map_location="""cpu"""      )

						A__						:  Dict		  =  upgrade_state_dict(UpperCAmelCase__,						UpperCAmelCase__      )
						hf_model.load_state_dict(UpperCAmelCase__      )
						A__						:  Optional[Any]		  =  hf_model.state_dict()
						A__						:  Tuple		  =  count_parameters(UpperCAmelCase__      )
						A__						:  List[Any]		  =  count_parameters(UpperCAmelCase__      ) + count_parameters(UpperCAmelCase__      )

						assert torch.allclose(UpperCAmelCase__,						UpperCAmelCase__,						atol=1e-3      )

						hf_model.save_pretrained(UpperCAmelCase__      )


if __name__ == "__main__":
						A_    =							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_    =							parser.parse_args()

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














"""simple docstring"""



def      _lowerCAmelCase		(							UpperCAmelCase__       :					int,						UpperCAmelCase__       :					int      )		->str:
						if a < 0 or b < 0:
												raise ValueError("""the value of both inputs must be positive"""      )

						A__						:  str		  =  str(bin(UpperCAmelCase__      )      )[2:]  # remove the leading "0b"
						A__						:  List[Any]		  =  str(bin(UpperCAmelCase__      )      )[2:]  # remove the leading "0b"

						A__						:  List[str]		  =  max(len(UpperCAmelCase__      ),						len(UpperCAmelCase__      )      )

						return "0b" + "".join(
						    str(int(char_a == """1""" and char_b == """1"""      )      )
						    for char_a, char_b in zip(a_binary.zfill(UpperCAmelCase__      ),						b_binary.zfill(UpperCAmelCase__      )      )      )


if __name__ == "__main__":
						import doctest

						doctest.testmod()






