




import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
					import jax.numpy as jnp

if is_tf_available():
					import tensorflow as tf

if is_torch_available():
					import torch
class       a__   ( unittest.TestCase				):






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	{
											    "task_specific_params": {
											        "summarization": {"length_penalty": 1.0, "max_length": 1_2_8, "min_length": 1_2, "num_beams": 4},
											        "summarization_cnn": {"length_penalty": 2.0, "max_length": 1_4_2, "min_length": 5_6, "num_beams": 4},
											        "summarization_xsum": {"length_penalty": 1.0, "max_length": 6_2, "min_length": 1_1, "num_beams": 6},
											    }
											}
											__lowerCAmelCase           =	{
											    "task_specific_params.summarization.length_penalty": 1.0,
											    "task_specific_params.summarization.max_length": 1_2_8,
											    "task_specific_params.summarization.min_length": 1_2,
											    "task_specific_params.summarization.num_beams": 4,
											    "task_specific_params.summarization_cnn.length_penalty": 2.0,
											    "task_specific_params.summarization_cnn.max_length": 1_4_2,
											    "task_specific_params.summarization_cnn.min_length": 5_6,
											    "task_specific_params.summarization_cnn.num_beams": 4,
											    "task_specific_params.summarization_xsum.length_penalty": 1.0,
											    "task_specific_params.summarization_xsum.max_length": 6_2,
											    "task_specific_params.summarization_xsum.min_length": 1_1,
											    "task_specific_params.summarization_xsum.num_beams": 6,
											}

											self.assertEqual(flatten_dict(_A      )		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											self.assertTrue(np.allclose(transpose(_A      )		,      x.transpose()      )      )

											__lowerCAmelCase           =	np.random.randn(3		,      4		,      5      )
											self.assertTrue(np.allclose(transpose(_A		,      axes=(1, 2, 0)      )		,      x.transpose((1, 2, 0)      )      )      )






						@require_torch
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											__lowerCAmelCase           =	torch.tensor(_A      )
											self.assertTrue(np.allclose(transpose(_A      )		,      transpose(_A      ).numpy()      )      )

											__lowerCAmelCase           =	np.random.randn(3		,      4		,      5      )
											__lowerCAmelCase           =	torch.tensor(_A      )
											self.assertTrue(np.allclose(transpose(_A		,      axes=(1, 2, 0)      )		,      transpose(_A		,      axes=(1, 2, 0)      ).numpy()      )      )






						@require_tf
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											__lowerCAmelCase           =	tf.constant(_A      )
											self.assertTrue(np.allclose(transpose(_A      )		,      transpose(_A      ).numpy()      )      )

											__lowerCAmelCase           =	np.random.randn(3		,      4		,      5      )
											__lowerCAmelCase           =	tf.constant(_A      )
											self.assertTrue(np.allclose(transpose(_A		,      axes=(1, 2, 0)      )		,      transpose(_A		,      axes=(1, 2, 0)      ).numpy()      )      )






						@require_flax
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											__lowerCAmelCase           =	jnp.array(_A      )
											self.assertTrue(np.allclose(transpose(_A      )		,      np.asarray(transpose(_A      )      )      )      )

											__lowerCAmelCase           =	np.random.randn(3		,      4		,      5      )
											__lowerCAmelCase           =	jnp.array(_A      )
											self.assertTrue(np.allclose(transpose(_A		,      axes=(1, 2, 0)      )		,      np.asarray(transpose(_A		,      axes=(1, 2, 0)      )      )      )      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											self.assertTrue(np.allclose(reshape(_A		,      (4, 3)      )		,      np.reshape(_A		,      (4, 3)      )      )      )

											__lowerCAmelCase           =	np.random.randn(3		,      4		,      5      )
											self.assertTrue(np.allclose(reshape(_A		,      (1_2, 5)      )		,      np.reshape(_A		,      (1_2, 5)      )      )      )






						@require_torch
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											__lowerCAmelCase           =	torch.tensor(_A      )
											self.assertTrue(np.allclose(reshape(_A		,      (4, 3)      )		,      reshape(_A		,      (4, 3)      ).numpy()      )      )

											__lowerCAmelCase           =	np.random.randn(3		,      4		,      5      )
											__lowerCAmelCase           =	torch.tensor(_A      )
											self.assertTrue(np.allclose(reshape(_A		,      (1_2, 5)      )		,      reshape(_A		,      (1_2, 5)      ).numpy()      )      )






						@require_tf
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											__lowerCAmelCase           =	tf.constant(_A      )
											self.assertTrue(np.allclose(reshape(_A		,      (4, 3)      )		,      reshape(_A		,      (4, 3)      ).numpy()      )      )

											__lowerCAmelCase           =	np.random.randn(3		,      4		,      5      )
											__lowerCAmelCase           =	tf.constant(_A      )
											self.assertTrue(np.allclose(reshape(_A		,      (1_2, 5)      )		,      reshape(_A		,      (1_2, 5)      ).numpy()      )      )






						@require_flax
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											__lowerCAmelCase           =	jnp.array(_A      )
											self.assertTrue(np.allclose(reshape(_A		,      (4, 3)      )		,      np.asarray(reshape(_A		,      (4, 3)      )      )      )      )

											__lowerCAmelCase           =	np.random.randn(3		,      4		,      5      )
											__lowerCAmelCase           =	jnp.array(_A      )
											self.assertTrue(np.allclose(reshape(_A		,      (1_2, 5)      )		,      np.asarray(reshape(_A		,      (1_2, 5)      )      )      )      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(1		,      3		,      4      )
											self.assertTrue(np.allclose(squeeze(_A      )		,      np.squeeze(_A      )      )      )

											__lowerCAmelCase           =	np.random.randn(1		,      4		,      1		,      5      )
											self.assertTrue(np.allclose(squeeze(_A		,      axis=2      )		,      np.squeeze(_A		,      axis=2      )      )      )






						@require_torch
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(1		,      3		,      4      )
											__lowerCAmelCase           =	torch.tensor(_A      )
											self.assertTrue(np.allclose(squeeze(_A      )		,      squeeze(_A      ).numpy()      )      )

											__lowerCAmelCase           =	np.random.randn(1		,      4		,      1		,      5      )
											__lowerCAmelCase           =	torch.tensor(_A      )
											self.assertTrue(np.allclose(squeeze(_A		,      axis=2      )		,      squeeze(_A		,      axis=2      ).numpy()      )      )






						@require_tf
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(1		,      3		,      4      )
											__lowerCAmelCase           =	tf.constant(_A      )
											self.assertTrue(np.allclose(squeeze(_A      )		,      squeeze(_A      ).numpy()      )      )

											__lowerCAmelCase           =	np.random.randn(1		,      4		,      1		,      5      )
											__lowerCAmelCase           =	tf.constant(_A      )
											self.assertTrue(np.allclose(squeeze(_A		,      axis=2      )		,      squeeze(_A		,      axis=2      ).numpy()      )      )






						@require_flax
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(1		,      3		,      4      )
											__lowerCAmelCase           =	jnp.array(_A      )
											self.assertTrue(np.allclose(squeeze(_A      )		,      np.asarray(squeeze(_A      )      )      )      )

											__lowerCAmelCase           =	np.random.randn(1		,      4		,      1		,      5      )
											__lowerCAmelCase           =	jnp.array(_A      )
											self.assertTrue(np.allclose(squeeze(_A		,      axis=2      )		,      np.asarray(squeeze(_A		,      axis=2      )      )      )      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											self.assertTrue(np.allclose(expand_dims(_A		,      axis=1      )		,      np.expand_dims(_A		,      axis=1      )      )      )






						@require_torch
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											__lowerCAmelCase           =	torch.tensor(_A      )
											self.assertTrue(np.allclose(expand_dims(_A		,      axis=1      )		,      expand_dims(_A		,      axis=1      ).numpy()      )      )






						@require_tf
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											__lowerCAmelCase           =	tf.constant(_A      )
											self.assertTrue(np.allclose(expand_dims(_A		,      axis=1      )		,      expand_dims(_A		,      axis=1      ).numpy()      )      )






						@require_flax
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	np.random.randn(3		,      4      )
											__lowerCAmelCase           =	jnp.array(_A      )
											self.assertTrue(np.allclose(expand_dims(_A		,      axis=1      )		,      np.asarray(expand_dims(_A		,      axis=1      )      )      )      )










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)




