
"""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		):
	# 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	,  _UpperCamelCase		):
	__lowerCAmelCase      :						List[str]						      =				{}

	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      :						Any						      =				key.replace('heads.cmd.mlm_head.cls.predictions'	,  'mmm_text_head'		)
		__lowerCAmelCase      :						Optional[int]						      =				key.replace('heads.cmd.itm_head.cls'	,  'itm_head'		)
		__lowerCAmelCase      :						int						      =				key.replace('heads.cmd.itm_head.pooler'	,  'itm_head.pooler'		)
		__lowerCAmelCase      :						Optional[Any]						      =				key.replace('heads.cmd.clip_head.logit_scale'	,  'flava.logit_scale'		)
		__lowerCAmelCase      :						Dict						      =				key.replace('heads.fairseq_mlm.cls.predictions'	,  'mlm_head'		)
		__lowerCAmelCase      :						Tuple						      =				key.replace('heads.imagenet.mim_head.cls.predictions'	,  'mim_head'		)
		__lowerCAmelCase      :						List[str]						      =				key.replace('mm_text_projection'	,  'flava.text_to_mm_projection'		)
		__lowerCAmelCase      :						Any						      =				key.replace('mm_image_projection'	,  'flava.image_to_mm_projection'		)
		__lowerCAmelCase      :						str						      =				key.replace('image_encoder.module'	,  'flava.image_model'		)
		__lowerCAmelCase      :						str						      =				key.replace('text_encoder.module'	,  'flava.text_model'		)
		__lowerCAmelCase      :						int						      =				key.replace('mm_encoder.module.encoder.cls_token'	,  'flava.multimodal_model.cls_token'		)
		__lowerCAmelCase      :						Dict						      =				key.replace('mm_encoder.module'	,  'flava.multimodal_model'		)
		__lowerCAmelCase      :						Tuple						      =				key.replace('text_projection'	,  'flava.text_projection'		)
		__lowerCAmelCase      :						str						      =				key.replace('image_projection'	,  'flava.image_projection'		)

		__lowerCAmelCase      :						Tuple						      =				value.float()

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

	return upgrade
@torch.no_grad()
def    __lowerCAmelCase    (_UpperCamelCase	,  _UpperCamelCase	,  _UpperCamelCase	,  _UpperCamelCase=None		):
	if config_path is not None:
		__lowerCAmelCase      :						Union[str, Any]						      =				FlavaConfig.from_pretrained(_UpperCamelCase		)
	else:
		__lowerCAmelCase      :						Any						      =				FlavaConfig()

	__lowerCAmelCase      :						int						      =				FlavaForPreTraining(_UpperCamelCase		).eval()

	__lowerCAmelCase      :						List[Any]						      =				convert_dalle_checkpoint(_UpperCamelCase	,  _UpperCamelCase	,  save_checkpoint=_UpperCamelCase		)

	if os.path.exists(_UpperCamelCase		):
		__lowerCAmelCase      :						Any						      =				torch.load(_UpperCamelCase	,  map_location='cpu'		)
	else:
		__lowerCAmelCase      :						Optional[int]						      =				torch.hub.load_state_dict_from_url(_UpperCamelCase	,  map_location='cpu'		)

	__lowerCAmelCase      :						Tuple						      =				upgrade_state_dict(_UpperCamelCase	,  _UpperCamelCase		)
	hf_model.load_state_dict(_UpperCamelCase		)
	__lowerCAmelCase      :						Tuple						      =				hf_model.state_dict()
	__lowerCAmelCase      :						List[Any]						      =				count_parameters(_UpperCamelCase		)
	__lowerCAmelCase      :						Any						      =				count_parameters(_UpperCamelCase		) + count_parameters(_UpperCamelCase		)

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

	hf_model.save_pretrained(_UpperCamelCase		)


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

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

"""simple docstring"""







from __future__ import annotations
def    __lowerCAmelCase    (_UpperCamelCase		):
	__lowerCAmelCase      :						List[str]						      =				str(_UpperCamelCase		)
	return len(_UpperCamelCase		) == 9 and set(_UpperCamelCase		) == set('123456789'		)
def    __lowerCAmelCase    ():
	for base_num in range(9999	,  4999	,  -1		):
		__lowerCAmelCase      :						Union[str, Any]						      =				10_0002 * base_num
		if is_9_pandigital(_UpperCamelCase		):
			return candidate

	for base_num in range(333	,  99	,  -1		):
		__lowerCAmelCase      :						Dict						      =				100_2003 * base_num
		if is_9_pandigital(_UpperCamelCase		):
			return candidate

	return None


if __name__ == "__main__":
							print(f'{solution() = }')