




"""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   _snake_case    (       lowercase__	):
				# 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   _snake_case    (       lowercase__							,					lowercase__	):
				_lowerCamelCase  :    Tuple					=					{}

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

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

								_lowerCamelCase  :    List[Any]					=					value.float()

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

				return upgrade







@torch.no_grad()
def   _snake_case    (       lowercase__							,					lowercase__							,					lowercase__							,					lowercase__=None	):
				if config_path is not None:
								_lowerCamelCase  :    int					=					FlavaConfig.from_pretrained(lowercase__	)
				else:
								_lowerCamelCase  :    int					=					FlavaConfig()

				_lowerCamelCase  :    str					=					FlavaForPreTraining(lowercase__	).eval()

				_lowerCamelCase  :    Dict					=					convert_dalle_checkpoint(lowercase__							,					lowercase__							,					save_checkpoint=lowercase__	)

				if os.path.exists(lowercase__	):
								_lowerCamelCase  :    Optional[int]					=					torch.load(lowercase__							,					map_location='cpu'	)
				else:
								_lowerCamelCase  :    str					=					torch.hub.load_state_dict_from_url(lowercase__							,					map_location='cpu'	)

				_lowerCamelCase  :    Dict					=					upgrade_state_dict(lowercase__							,					lowercase__	)
				hf_model.load_state_dict(lowercase__	)
				_lowerCamelCase  :    List[str]					=					hf_model.state_dict()
				_lowerCamelCase  :    int					=					count_parameters(lowercase__	)
				_lowerCamelCase  :    Any					=					count_parameters(lowercase__	) + count_parameters(lowercase__	)

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

				hf_model.save_pretrained(lowercase__	)


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)





"""simple docstring"""






def   _snake_case    (       lowercase__	):
				_lowerCamelCase  :    List[str]					=					abs(lowercase__	)
				_lowerCamelCase  :    Optional[int]					=					0
				while n > 0:
								res += n % 10
								n //= 10
				return res







def   _snake_case    (       lowercase__	):
				_lowerCamelCase  :    List[str]					=					abs(lowercase__	)
				return n if n < 10 else n % 10 + sum_of_digits(n // 10	)







def   _snake_case    (       lowercase__	):
				return sum(int(lowercase__	) for c in str(abs(lowercase__	)	)	)







def   _snake_case    (       ):
				from collections.abc import Callable
				from timeit import timeit

				def benchmark_a_function(lowercase__							,					lowercase__	) -> None:
								_lowerCamelCase  :    int					=					f'''{func.__name__}({value})'''
								_lowerCamelCase  :    Optional[Any]					=					timeit(f'''__main__.{call}'''							,					setup='import __main__'	)
								print(f'''{call:56} = {func(lowercase__	)} -- {timing:.4f} seconds'''	)

				for value in (262144, 1125899906842624, 1267650600228229401496703205376):
								for func in (sum_of_digits, sum_of_digits_recursion, sum_of_digits_compact):
												benchmark_a_function(lowercase__							,					lowercase__	)
								print()


if __name__ == "__main__":
			import doctest

			doctest.testmod()
			benchmark()