"""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    :			Optional[Any]   )   ->			Union[str, Any]:





      """simple docstring"""






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





def 				lowerCAmelCase__      (     _UpperCamelCase    :			str	,     _UpperCamelCase    :			Optional[int]   )   ->			List[str]:





      """simple docstring"""






      snake_case        =     {}

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

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

            snake_case        =     value.float()

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

      return upgrade





@torch.no_grad()
def 				lowerCAmelCase__      (     _UpperCamelCase    :			str	,     _UpperCamelCase    :			Optional[Any]	,     _UpperCamelCase    :			Any	,     _UpperCamelCase    :			List[str]=None   )   ->			Union[str, Any]:





      """simple docstring"""






      if config_path is not None:
            snake_case        =     FlavaConfig.from_pretrained(_UpperCamelCase   )
      else:
            snake_case        =     FlavaConfig()

      snake_case        =     FlavaForPreTraining(_UpperCamelCase   ).eval()

      snake_case        =     convert_dalle_checkpoint(_UpperCamelCase	,     _UpperCamelCase	,     save_checkpoint=_UpperCamelCase   )

      if os.path.exists(_UpperCamelCase   ):
            snake_case        =     torch.load(_UpperCamelCase	,     map_location='cpu'   )
      else:
            snake_case        =     torch.hub.load_state_dict_from_url(_UpperCamelCase	,     map_location='cpu'   )

      snake_case        =     upgrade_state_dict(_UpperCamelCase	,     _UpperCamelCase   )
      hf_model.load_state_dict(_UpperCamelCase   )
      snake_case        =     hf_model.state_dict()
      snake_case        =     count_parameters(_UpperCamelCase   )
      snake_case        =     count_parameters(_UpperCamelCase   ) + count_parameters(_UpperCamelCase   )

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

      hf_model.save_pretrained(_UpperCamelCase   )


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

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



"""simple docstring"""
import os
import sys
import unittest


SCREAMING_SNAKE_CASE__   =					os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
SCREAMING_SNAKE_CASE__   =					os.path.join(git_repo_path, "src", "transformers")

SCREAMING_SNAKE_CASE__   =					"\n{0} = None\n"

SCREAMING_SNAKE_CASE__   =					"\nclass {0}(metaclass=DummyObject):\n    _backends = {1}\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, {1})\n"


SCREAMING_SNAKE_CASE__   =					"\ndef {0}(*args, **kwargs):\n    requires_backends({0}, {1})\n"



class 					lowerCAmelCase_      (     unittest.TestCase   ):



				"""simple docstring"""

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     find_backend('    _import_structure["models.albert"].append("AlbertTokenizerFast")'      )
										self.assertIsNone(lowerCAmelCase      )

										snake_case        =     find_backend('    if not is_tokenizers_available():'      )
										self.assertEqual(lowerCAmelCase						,      'tokenizers'      )

										snake_case        =     find_backend('    if not is_tensorflow_text_available():'      )
										self.assertEqual(lowerCAmelCase						,      'tensorflow_text'      )

										snake_case        =     find_backend('    if not (is_sentencepiece_available() and is_tokenizers_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tokenizers'      )

										snake_case        =     find_backend(
										    '    if not (is_sentencepiece_available() and is_tensorflow_text_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tensorflow_text'      )

										snake_case        =     find_backend(
										    '    if not (is_sentencepiece_available() and is_tokenizers_available() and is_vision_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tokenizers_and_vision'      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     read_init()
										# We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
										self.assertIn('torch'						,      lowerCAmelCase      )
										self.assertIn('tensorflow_text'						,      lowerCAmelCase      )
										self.assertIn('sentencepiece_and_tokenizers'						,      lowerCAmelCase      )

										# Likewise, we can't assert on the exact content of a key
										self.assertIn('BertModel'						,      objects['torch']      )
										self.assertIn('TFBertModel'						,      objects['tf']      )
										self.assertIn('FlaxBertModel'						,      objects['flax']      )
										self.assertIn('BertModel'						,      objects['torch']      )
										self.assertIn('TFBertTokenizer'						,      objects['tensorflow_text']      )
										self.assertIn('convert_slow_tokenizer'						,      objects['sentencepiece_and_tokenizers']      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     create_dummy_object('CONSTANT'						,      '\'torch\''      )
										self.assertEqual(lowerCAmelCase						,      '\nCONSTANT = None\n'      )

										snake_case        =     create_dummy_object('function'						,      '\'torch\''      )
										self.assertEqual(
										    lowerCAmelCase						,      '\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n'      )

										snake_case        =     '\nclass FakeClass(metaclass=DummyObject):\n    _backends = \'torch\'\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, \'torch\')\n'
										snake_case        =     create_dummy_object('FakeClass'						,      '\'torch\''      )
										self.assertEqual(lowerCAmelCase						,      lowerCAmelCase      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     '# This file is autogenerated by the command `make fix-copies`, do not edit.\nfrom ..utils import DummyObject, requires_backends\n\n\nCONSTANT = None\n\n\ndef function(*args, **kwargs):\n    requires_backends(function, ["torch"])\n\n\nclass FakeClass(metaclass=DummyObject):\n    _backends = ["torch"]\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, ["torch"])\n'
										snake_case        =     create_dummy_files({'torch': ['CONSTANT', 'function', 'FakeClass']}      )
										self.assertEqual(dummy_files['torch']						,      lowerCAmelCase      )


