




SCREAMING_SNAKE_CASE__       :      Optional[Any]  =   'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'

def      A      (      _SCREAMING_SNAKE_CASE    )	-> bytes:
       # Make sure the supplied data is a bytes-like object
       if not isinstance(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    ):
              lowerCamelCase   :      int			 =					f'''a bytes-like object is required, not \'{data.__class__.__name__}\''''
              raise TypeError(_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      Optional[int]			 =					"".join(bin(_SCREAMING_SNAKE_CASE    )[2:].zfill(8    ) for byte in data    )

       lowerCamelCase   :      List[str]			 =					len(_SCREAMING_SNAKE_CASE    ) % 6 != 0

       if padding_needed:
              # The padding that will be added later
              lowerCamelCase   :      Optional[int]			 =					B"=" * ((6 - len(_SCREAMING_SNAKE_CASE    ) % 6) // 2)

              # Append binary_stream with arbitrary binary digits (0's by default) to make its
              # length a multiple of 6.
              binary_stream += "0" * (6 - len(_SCREAMING_SNAKE_CASE    ) % 6)
       else:
              lowerCamelCase   :      List[Any]			 =					B""

       # Encode every 6 binary digits to their corresponding Base64 character
       return (
           "".join(
               B64_CHARSET[int(binary_stream[index : index + 6] ,2    )]
               for index in range(0 ,len(_SCREAMING_SNAKE_CASE    ) ,6    )    ).encode()
           + padding
       )

def      A      (      _SCREAMING_SNAKE_CASE    )	-> bytes:
       # Make sure encoded_data is either a string or a bytes-like object
       if not isinstance(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    ) and not isinstance(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    ):
              lowerCamelCase   :      List[Any]			 =					(
                  "argument should be a bytes-like object or ASCII string, "
                  f'''not \'{encoded_data.__class__.__name__}\''''
              )
              raise TypeError(_SCREAMING_SNAKE_CASE    )

       # In case encoded_data is a bytes-like object, make sure it contains only
       # ASCII characters so we convert it to a string object
       if isinstance(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    ):
              try:
                     lowerCamelCase   :      Optional[Any]			 =					encoded_data.decode("utf-8"    )
              except UnicodeDecodeError:
                     raise ValueError("base64 encoded data should only contain ASCII characters"    )

       lowerCamelCase   :      Optional[Any]			 =					encoded_data.count("="    )

       # Check if the encoded string contains non base64 characters
       if padding:
              assert all(
                  char in B64_CHARSET for char in encoded_data[:-padding]    ), "Invalid base64 character(s) found."
       else:
              assert all(
                  char in B64_CHARSET for char in encoded_data    ), "Invalid base64 character(s) found."

       # Check the padding
       assert len(_SCREAMING_SNAKE_CASE    ) % 4 == 0 and padding < 3, "Incorrect padding"

       if padding:
              # Remove padding if there is one
              lowerCamelCase   :      Optional[Any]			 =					encoded_data[:-padding]

              lowerCamelCase   :      Tuple			 =					"".join(
                  bin(B64_CHARSET.index(_SCREAMING_SNAKE_CASE    )    )[2:].zfill(6    ) for char in encoded_data    )[: -padding * 2]
       else:
              lowerCamelCase   :      Dict			 =					"".join(
                  bin(B64_CHARSET.index(_SCREAMING_SNAKE_CASE    )    )[2:].zfill(6    ) for char in encoded_data    )

       lowerCamelCase   :      int			 =					[
           int(binary_stream[index : index + 8] ,2    )
           for index in range(0 ,len(_SCREAMING_SNAKE_CASE    ) ,8    )
       ]

       return bytes(_SCREAMING_SNAKE_CASE    )


if __name__ == "__main__":
       import doctest

       doctest.testmod()








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			(	a_						)		->	List[str]:







				"""simple docstring"""



				return sum(param.float().sum() if "encoder.embeddings" not in key else 0 for key, param in state_dict.items()						)


def 						UpperCAmelCase			(	a_		,   a_						)		->	Tuple:







				"""simple docstring"""



				__A   			= {}

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

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

								__A   			= value.float()

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

				return upgrade


@torch.no_grad()
def 						UpperCAmelCase			(	a_		,   a_		,   a_		,   a_=None						)		->	Tuple:







				"""simple docstring"""



				if config_path is not None:
								__A   			= FlavaConfig.from_pretrained(a_						)
				else:
								__A   			= FlavaConfig()

				__A   			= FlavaForPreTraining(a_						).eval()

				__A   			= convert_dalle_checkpoint(a_		,   a_		,   save_checkpoint=a_						)

				if os.path.exists(a_						):
								__A   			= torch.load(a_		,   map_location="cpu"						)
				else:
								__A   			= torch.hub.load_state_dict_from_url(a_		,   map_location="cpu"						)

				__A   			= upgrade_state_dict(a_		,   a_						)
				hf_model.load_state_dict(a_						)
				__A   			= hf_model.state_dict()
				__A   			= count_parameters(a_						)
				__A   			= count_parameters(a_						) + count_parameters(a_						)

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

				hf_model.save_pretrained(a_						)


if __name__ == "__main__":
					SCREAMING_SNAKE_CASE							:Any							 =				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')
					SCREAMING_SNAKE_CASE							:Optional[int]							 =				parser.parse_args()

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