




from datasets.utils.patching import _PatchedModuleObj, patch_submodule

from . import _test_patching







def      UpperCAmelCase_				(						):
					import os as original_os
					from os import path as original_path
					from os import rename as original_rename
					from os.path import dirname as original_dirname
					from os.path import join as original_join

					assert _test_patching.os is original_os
					assert _test_patching.path is original_path
					assert _test_patching.join is original_join

					assert _test_patching.renamed_os is original_os
					assert _test_patching.renamed_path is original_path
					assert _test_patching.renamed_join is original_join

					__lowerCAmelCase           =	"__test_patch_submodule_mock__"
					with patch_submodule(_test_patching				,				"os.path.join"				,				UpperCamelCase__					):
										# Every way to access os.path.join must be patched, and the rest must stay untouched

										# check os.path.join
										assert isinstance(_test_patching.os				,				_PatchedModuleObj					)
										assert isinstance(_test_patching.os.path				,				_PatchedModuleObj					)
										assert _test_patching.os.path.join is mock

										# check path.join
										assert isinstance(_test_patching.path				,				_PatchedModuleObj					)
										assert _test_patching.path.join is mock

										# check join
										assert _test_patching.join is mock

										# check that the other attributes are untouched
										assert _test_patching.os.rename is original_rename
										assert _test_patching.path.dirname is original_dirname
										assert _test_patching.os.path.dirname is original_dirname

										# Even renamed modules or objects must be patched

										# check renamed_os.path.join
										assert isinstance(_test_patching.renamed_os				,				_PatchedModuleObj					)
										assert isinstance(_test_patching.renamed_os.path				,				_PatchedModuleObj					)
										assert _test_patching.renamed_os.path.join is mock

										# check renamed_path.join
										assert isinstance(_test_patching.renamed_path				,				_PatchedModuleObj					)
										assert _test_patching.renamed_path.join is mock

										# check renamed_join
										assert _test_patching.renamed_join is mock

										# check that the other attributes are untouched
										assert _test_patching.renamed_os.rename is original_rename
										assert _test_patching.renamed_path.dirname is original_dirname
										assert _test_patching.renamed_os.path.dirname is original_dirname

					# check that everthing is back to normal when the patch is over

					assert _test_patching.os is original_os
					assert _test_patching.path is original_path
					assert _test_patching.join is original_join

					assert _test_patching.renamed_os is original_os
					assert _test_patching.renamed_path is original_path
					assert _test_patching.renamed_join is original_join







def      UpperCAmelCase_				(						):
					assert _test_patching.open is open

					__lowerCAmelCase           =	"__test_patch_submodule_builtin_mock__"
					# _test_patching has "open" in its globals
					assert _test_patching.open is open
					with patch_submodule(_test_patching				,				"open"				,				UpperCamelCase__					):
										assert _test_patching.open is mock

					# check that everthing is back to normal when the patch is over

					assert _test_patching.open is open







def      UpperCAmelCase_				(						):
					# pandas.read_csv is not present in _test_patching
					__lowerCAmelCase           =	"__test_patch_submodule_missing_mock__"
					with patch_submodule(_test_patching				,				"pandas.read_csv"				,				UpperCamelCase__					):
										pass







def      UpperCAmelCase_				(						):
					# builtin should always be mocked even if they're not in the globals
					# in case they're loaded at one point
					__lowerCAmelCase           =	"__test_patch_submodule_missing_builtin_mock__"
					# _test_patching doesn't have "len" in its globals
					assert getattr(_test_patching				,				"len"				,				UpperCamelCase__					) is None
					with patch_submodule(_test_patching				,				"len"				,				UpperCamelCase__					):
										assert _test_patching.len is mock
					assert _test_patching.len is len







def      UpperCAmelCase_				(						):
					__lowerCAmelCase           =	"__test_patch_submodule_start_and_stop_mock__"
					__lowerCAmelCase           =	patch_submodule(_test_patching				,				"open"				,				UpperCamelCase__					)
					assert _test_patching.open is open
					patch.start()
					assert _test_patching.open is mock
					patch.stop()
					assert _test_patching.open is open







def      UpperCAmelCase_				(						):
					from os import rename as original_rename
					from os.path import dirname as original_dirname
					from os.path import join as original_join

					__lowerCAmelCase           =	"__test_patch_submodule_successive_join__"
					__lowerCAmelCase           =	"__test_patch_submodule_successive_dirname__"
					__lowerCAmelCase           =	"__test_patch_submodule_successive_rename__"
					assert _test_patching.os.path.join is original_join
					assert _test_patching.os.path.dirname is original_dirname
					assert _test_patching.os.rename is original_rename

					with patch_submodule(_test_patching				,				"os.path.join"				,				UpperCamelCase__					):
										with patch_submodule(_test_patching				,				"os.rename"				,				UpperCamelCase__					):
															with patch_submodule(_test_patching				,				"os.path.dirname"				,				UpperCamelCase__					):
																				assert _test_patching.os.path.join is mock_join
																				assert _test_patching.os.path.dirname is mock_dirname
																				assert _test_patching.os.rename is mock_rename

    # try another order
					with patch_submodule(_test_patching				,				"os.rename"				,				UpperCamelCase__					):
										with patch_submodule(_test_patching				,				"os.path.join"				,				UpperCamelCase__					):
															with patch_submodule(_test_patching				,				"os.path.dirname"				,				UpperCamelCase__					):
																				assert _test_patching.os.path.join is mock_join
																				assert _test_patching.os.path.dirname is mock_dirname
																				assert _test_patching.os.rename is mock_rename

					assert _test_patching.os.path.join is original_join
					assert _test_patching.os.path.dirname is original_dirname
					assert _test_patching.os.rename is original_rename







def      UpperCAmelCase_				(						):
					__lowerCAmelCase           =	"__test_patch_submodule_doesnt_exist_mock__"
					with patch_submodule(_test_patching				,				"__module_that_doesn_exist__.__attribute_that_doesn_exist__"				,				UpperCamelCase__					):
										pass
					with patch_submodule(_test_patching				,				"os.__attribute_that_doesn_exist__"				,				UpperCamelCase__					):
										pass










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      _a				(						SCREAMING_SNAKE_CASE_							:    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      _a				(						SCREAMING_SNAKE_CASE_							:    Dict				,				SCREAMING_SNAKE_CASE_							:    Dict					):
					__lowerCAmelCase           =	{}

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

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

										__lowerCAmelCase           =	value.float()

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

					return upgrade







@torch.no_grad()
def      _a				(						SCREAMING_SNAKE_CASE_							:    List[Any]				,				SCREAMING_SNAKE_CASE_							:    Optional[int]				,				SCREAMING_SNAKE_CASE_							:    List[str]				,				SCREAMING_SNAKE_CASE_							:    int=None					):
					if config_path is not None:
										__lowerCAmelCase           =	FlavaConfig.from_pretrained(SCREAMING_SNAKE_CASE_					)
					else:
										__lowerCAmelCase           =	FlavaConfig()

					__lowerCAmelCase           =	FlavaForPreTraining(SCREAMING_SNAKE_CASE_					).eval()

					__lowerCAmelCase           =	convert_dalle_checkpoint(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				save_checkpoint=SCREAMING_SNAKE_CASE_					)

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

					__lowerCAmelCase           =	upgrade_state_dict(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_					)
					hf_model.load_state_dict(SCREAMING_SNAKE_CASE_					)
					__lowerCAmelCase           =	hf_model.state_dict()
					__lowerCAmelCase           =	count_parameters(SCREAMING_SNAKE_CASE_					)
					__lowerCAmelCase           =	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__":
					UpperCamelCase__         =      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""")
					UpperCamelCase__         =      parser.parse_args()

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




