


import tempfile
import unittest

from transformers import DPRConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import DPRContextEncoder, DPRQuestionEncoder, DPRReader, DPRReaderTokenizer
  from transformers.models.dpr.modeling_dpr import (
      DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
      DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
      DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
  )







class 				__SCREAMING_SNAKE_CASE   :


      def __init__(							self							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__=13							,						SCREAMING_SNAKE_CASE__=7							,						SCREAMING_SNAKE_CASE__=False							,						SCREAMING_SNAKE_CASE__=True							,						SCREAMING_SNAKE_CASE__=True							,						SCREAMING_SNAKE_CASE__=True							,						SCREAMING_SNAKE_CASE__=99							,						SCREAMING_SNAKE_CASE__=32							,						SCREAMING_SNAKE_CASE__=5							,						SCREAMING_SNAKE_CASE__=4							,						SCREAMING_SNAKE_CASE__=37							,						SCREAMING_SNAKE_CASE__="gelu"							,						SCREAMING_SNAKE_CASE__=0.1							,						SCREAMING_SNAKE_CASE__=0.1							,						SCREAMING_SNAKE_CASE__=512							,						SCREAMING_SNAKE_CASE__=16							,						SCREAMING_SNAKE_CASE__=2							,						SCREAMING_SNAKE_CASE__=0.02							,						SCREAMING_SNAKE_CASE__=3							,						SCREAMING_SNAKE_CASE__=4							,						SCREAMING_SNAKE_CASE__=None							,						SCREAMING_SNAKE_CASE__=0							,						):
            lowercase			:   Optional[int]					   =		parent
            lowercase			:   Tuple					   =		batch_size
            lowercase			:   Dict					   =		seq_length
            lowercase			:   Dict					   =		is_training
            lowercase			:   Any					   =		use_input_mask
            lowercase			:   Union[str, Any]					   =		use_token_type_ids
            lowercase			:   Optional[Any]					   =		use_labels
            lowercase			:   Any					   =		vocab_size
            lowercase			:   int					   =		hidden_size
            lowercase			:   Tuple					   =		num_hidden_layers
            lowercase			:   List[str]					   =		num_attention_heads
            lowercase			:   Any					   =		intermediate_size
            lowercase			:   Tuple					   =		hidden_act
            lowercase			:   str					   =		hidden_dropout_prob
            lowercase			:   str					   =		attention_probs_dropout_prob
            lowercase			:   List[str]					   =		max_position_embeddings
            lowercase			:   int					   =		type_vocab_size
            lowercase			:   Dict					   =		type_sequence_label_size
            lowercase			:   List[str]					   =		initializer_range
            lowercase			:   Dict					   =		num_labels
            lowercase			:   List[str]					   =		num_choices
            lowercase			:   Dict					   =		scope
            lowercase			:   List[Any]					   =		projection_dim


      def        __lowerCamelCase			(							self		):
            lowercase			:   Union[str, Any]					   =		ids_tensor([self.batch_size, self.seq_length]							,						self.vocab_size		)

            lowercase			:   Any					   =		None
            if self.use_input_mask:
                  lowercase			:   str					   =		random_attention_mask([self.batch_size, self.seq_length]		)

            lowercase			:   Tuple					   =		None
            if self.use_token_type_ids:
                  lowercase			:   Optional[int]					   =		ids_tensor([self.batch_size, self.seq_length]							,						self.type_vocab_size		)

            lowercase			:   str					   =		None
            lowercase			:   int					   =		None
            lowercase			:   Dict					   =		None
            if self.use_labels:
                  lowercase			:   Optional[int]					   =		ids_tensor([self.batch_size]							,						self.type_sequence_label_size		)
                  lowercase			:   Optional[Any]					   =		ids_tensor([self.batch_size, self.seq_length]							,						self.num_labels		)
                  lowercase			:   int					   =		ids_tensor([self.batch_size]							,						self.num_choices		)

            lowercase			:   str					   =		self.get_config()

            return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


      def        __lowerCamelCase			(							self		):
            return DPRConfig(
                projection_dim=self.projection_dim							,						vocab_size=self.vocab_size							,						hidden_size=self.hidden_size							,						num_hidden_layers=self.num_hidden_layers							,						num_attention_heads=self.num_attention_heads							,						intermediate_size=self.intermediate_size							,						hidden_act=self.hidden_act							,						hidden_dropout_prob=self.hidden_dropout_prob							,						attention_probs_dropout_prob=self.attention_probs_dropout_prob							,						max_position_embeddings=self.max_position_embeddings							,						type_vocab_size=self.type_vocab_size							,						initializer_range=self.initializer_range							,						)


      def        __lowerCamelCase			(							self							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__		):
            lowercase			:   List[Any]					   =		DPRContextEncoder(config=SCREAMING_SNAKE_CASE__		)
            model.to(SCREAMING_SNAKE_CASE__		)
            model.eval()
            lowercase			:   Tuple					   =		model(SCREAMING_SNAKE_CASE__							,						attention_mask=SCREAMING_SNAKE_CASE__							,						token_type_ids=SCREAMING_SNAKE_CASE__		)
            lowercase			:   Union[str, Any]					   =		model(SCREAMING_SNAKE_CASE__							,						token_type_ids=SCREAMING_SNAKE_CASE__		)
            lowercase			:   List[Any]					   =		model(SCREAMING_SNAKE_CASE__		)
            self.parent.assertEqual(result.pooler_output.shape							,						(self.batch_size, self.projection_dim or self.hidden_size)		)


      def        __lowerCamelCase			(							self							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__		):
            lowercase			:   Optional[Any]					   =		DPRQuestionEncoder(config=SCREAMING_SNAKE_CASE__		)
            model.to(SCREAMING_SNAKE_CASE__		)
            model.eval()
            lowercase			:   List[Any]					   =		model(SCREAMING_SNAKE_CASE__							,						attention_mask=SCREAMING_SNAKE_CASE__							,						token_type_ids=SCREAMING_SNAKE_CASE__		)
            lowercase			:   List[str]					   =		model(SCREAMING_SNAKE_CASE__							,						token_type_ids=SCREAMING_SNAKE_CASE__		)
            lowercase			:   List[Any]					   =		model(SCREAMING_SNAKE_CASE__		)
            self.parent.assertEqual(result.pooler_output.shape							,						(self.batch_size, self.projection_dim or self.hidden_size)		)


      def        __lowerCamelCase			(							self							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__		):
            lowercase			:   Dict					   =		DPRReader(config=SCREAMING_SNAKE_CASE__		)
            model.to(SCREAMING_SNAKE_CASE__		)
            model.eval()
            lowercase			:   int					   =		model(
                SCREAMING_SNAKE_CASE__							,						attention_mask=SCREAMING_SNAKE_CASE__							,						)

            self.parent.assertEqual(result.start_logits.shape							,						(self.batch_size, self.seq_length)		)
            self.parent.assertEqual(result.end_logits.shape							,						(self.batch_size, self.seq_length)		)
            self.parent.assertEqual(result.relevance_logits.shape							,						(self.batch_size,)		)






      def        __lowerCamelCase			(							self		):
            lowercase			:   Optional[int]					   =		self.prepare_config_and_inputs()
            (
                (
                lowercase
            )   ,			(
                lowercase
            )   ,			(
                lowercase
            )   ,			(
                lowercase
            )   ,			(
                lowercase
            )   ,			(
                lowercase
            )   ,			(
                lowercase
            )   ,			
            )			:   Optional[int]					   =		config_and_inputs
            lowercase			:   Any					   =		{'''input_ids''': input_ids}
            return config, inputs_dict







@require_torch
class 				__SCREAMING_SNAKE_CASE   (			A__  ,						A__  ,						unittest.TestCase   ):
      A				:							Union[str, Any]    							=     (
          (
              DPRContextEncoder,
              DPRQuestionEncoder,
              DPRReader,
          )
          if is_torch_available()
          else ()
      )
      A				:							Tuple    							=     {'feature-extraction': DPRQuestionEncoder} if is_torch_available() else {}

      A				:							List[str]    							=     False
      A				:							int    							=     False  # why?
      A				:							Any    							=     False
      A				:							Any    							=     False


      def        __lowerCamelCase			(							self		):
            lowercase			:   str					   =		DPRModelTester(self		)
            lowercase			:   Optional[Any]					   =		ConfigTester(self							,						config_class=SCREAMING_SNAKE_CASE__							,						hidden_size=37		)


      def        __lowerCamelCase			(							self		):
            self.config_tester.run_common_tests()


      def        __lowerCamelCase			(							self		):
            lowercase			:   Tuple					   =		self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_context_encoder(*SCREAMING_SNAKE_CASE__		)


      def        __lowerCamelCase			(							self		):
            lowercase			:   Any					   =		self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_question_encoder(*SCREAMING_SNAKE_CASE__		)


      def        __lowerCamelCase			(							self		):
            lowercase			:   List[Any]					   =		self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_reader(*SCREAMING_SNAKE_CASE__		)


      def        __lowerCamelCase			(							self		):
            lowercase			:   Any					   =		self.model_tester.prepare_config_and_inputs()[0]

            lowercase			:   str					   =		DPRQuestionEncoder(config=SCREAMING_SNAKE_CASE__		)
            model.to(SCREAMING_SNAKE_CASE__		)
            model.eval()

            with tempfile.TemporaryDirectory() as tmp_dirname:
                  model.save_pretrained(SCREAMING_SNAKE_CASE__		)
                  lowercase			:   List[str]					   =		DPRQuestionEncoder.from_pretrained(SCREAMING_SNAKE_CASE__							,						projection_dim=512		)

            self.assertIsNotNone(SCREAMING_SNAKE_CASE__		)






      @slow
      def        __lowerCamelCase			(							self		):
            for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowercase			:   List[Any]					   =		DPRContextEncoder.from_pretrained(SCREAMING_SNAKE_CASE__		)
                  self.assertIsNotNone(SCREAMING_SNAKE_CASE__		)

            for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowercase			:   Union[str, Any]					   =		DPRContextEncoder.from_pretrained(SCREAMING_SNAKE_CASE__		)
                  self.assertIsNotNone(SCREAMING_SNAKE_CASE__		)

            for model_name in DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowercase			:   str					   =		DPRQuestionEncoder.from_pretrained(SCREAMING_SNAKE_CASE__		)
                  self.assertIsNotNone(SCREAMING_SNAKE_CASE__		)

            for model_name in DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowercase			:   Optional[int]					   =		DPRReader.from_pretrained(SCREAMING_SNAKE_CASE__		)
                  self.assertIsNotNone(SCREAMING_SNAKE_CASE__		)







@require_torch
class 				__SCREAMING_SNAKE_CASE   (			unittest.TestCase   ):


      @slow
      def        __lowerCamelCase			(							self		):
            lowercase			:   Dict					   =		DPRQuestionEncoder.from_pretrained('''facebook/dpr-question_encoder-single-nq-base'''							,						return_dict=SCREAMING_SNAKE_CASE__		)
            model.to(SCREAMING_SNAKE_CASE__		)

            lowercase			:   Optional[Any]					   =		torch.tensor(
                [[101, 7592, 1010, 2003, 2026, 3899, 10140, 1029, 102]]							,						dtype=torch.long							,						device=SCREAMING_SNAKE_CASE__		)  # [CLS] hello, is my dog cute? [SEP]
            lowercase			:   int					   =		model(SCREAMING_SNAKE_CASE__		)[0]  # embedding shape = (1, 768)
            # compare the actual values for a slice.
            lowercase			:   int					   =		torch.tensor(
                [
                    [
                        0.03236253,
                        0.12753335,
                        0.16818509,
                        0.00279786,
                        0.3896933,
                        0.24264945,
                        0.2178971,
                        -0.02335227,
                        -0.08481959,
                        -0.14324117,
                    ]
                ]							,						dtype=torch.float							,						device=SCREAMING_SNAKE_CASE__							,						)
            self.assertTrue(torch.allclose(output[:, :10]							,						SCREAMING_SNAKE_CASE__							,						atol=1E-4		)		)






      @slow
      def        __lowerCamelCase			(							self		):
            lowercase			:   Union[str, Any]					   =		DPRReaderTokenizer.from_pretrained('''facebook/dpr-reader-single-nq-base'''		)
            lowercase			:   List[Any]					   =		DPRReader.from_pretrained('''facebook/dpr-reader-single-nq-base'''		)
            model.to(SCREAMING_SNAKE_CASE__		)

            lowercase			:   Optional[Any]					   =		tokenizer(
                questions='''What is love ?'''							,						titles='''Haddaway'''							,						texts='''What Is Love is a song recorded by the artist Haddaway'''							,						padding=SCREAMING_SNAKE_CASE__							,						return_tensors='''pt'''							,						)
            encoded_inputs.to(SCREAMING_SNAKE_CASE__		)

            lowercase			:   List[Any]					   =		model(**SCREAMING_SNAKE_CASE__		)

            # compare the actual values for a slice.
            lowercase			:   Optional[int]					   =		torch.tensor(
                [[-10.3005, -10.7765, -11.4872, -11.6841, -11.9312, -10.3002, -9.8544, -11.7378, -12.0821, -10.2975]]							,						dtype=torch.float							,						device=SCREAMING_SNAKE_CASE__							,						)

            lowercase			:   int					   =		torch.tensor(
                [[-11.0684, -11.7041, -11.5397, -10.3465, -10.8791, -6.8443, -11.9959, -11.0364, -10.0096, -6.8405]]							,						dtype=torch.float							,						device=SCREAMING_SNAKE_CASE__							,						)
            self.assertTrue(torch.allclose(outputs.start_logits[:, :10]							,						SCREAMING_SNAKE_CASE__							,						atol=1E-4		)		)
            self.assertTrue(torch.allclose(outputs.end_logits[:, :10]							,						SCREAMING_SNAKE_CASE__							,						atol=1E-4		)		)





import argparse
import os

import align
import numpy as np
import requests
import tensorflow as tf
import torch
from PIL import Image
from tokenizer import Tokenizer

from transformers import (
    AlignConfig,
    AlignModel,
    AlignProcessor,
    BertConfig,
    BertTokenizer,
    EfficientNetConfig,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__a   				=      logging.get_logger(__name__)




def 			__lowercase	(						_UpperCamelCase     )    ->Tuple:







      """simple docstring"""

      lowercase			:   Optional[int]					   =		tf.image.resize(_UpperCamelCase,    (346, 346)     )
      lowercase			:   Any					   =		tf.image.crop_to_bounding_box(_UpperCamelCase,    (346 - 289) // 2,    (346 - 289) // 2,    289,    289     )
      return image




def 			__lowercase	(						)    ->List[str]:







      """simple docstring"""

      lowercase			:   Tuple					   =		EfficientNetConfig.from_pretrained('''google/efficientnet-b7'''     )
      lowercase			:   Optional[int]					   =		289
      lowercase			:   Optional[Any]					   =		640
      lowercase			:   Any					   =		{'''0''': '''LABEL_0''', '''1''': '''LABEL_1'''}
      lowercase			:   Optional[Any]					   =		{'''LABEL_0''': 0, '''LABEL_1''': 1}
      lowercase			:   Union[str, Any]					   =		[]

      lowercase			:   Union[str, Any]					   =		BertConfig()
      lowercase			:   Any					   =		AlignConfig.from_text_vision_configs(
          text_config=_UpperCamelCase,    vision_config=_UpperCamelCase,    projection_dim=640     )
      return config




def 			__lowercase	(						)    ->Union[str, Any]:







      """simple docstring"""

      lowercase			:   int					   =		'''http://images.cocodataset.org/val2017/000000039769.jpg'''
      lowercase			:   int					   =		Image.open(requests.get(_UpperCamelCase,    stream=_UpperCamelCase     ).raw     )
      return im




def 			__lowercase	(						)    ->Optional[int]:







      """simple docstring"""

      lowercase			:   Any					   =		EfficientNetImageProcessor(
          do_center_crop=_UpperCamelCase,    rescale_factor=1 / 1_2_7.5,    rescale_offset=_UpperCamelCase,    do_normalize=_UpperCamelCase,    include_top=_UpperCamelCase,    resample=Image.BILINEAR,    )
      lowercase			:   Any					   =		BertTokenizer.from_pretrained('''bert-base-uncased'''     )
      lowercase			:   int					   =		64
      lowercase			:   List[Any]					   =		AlignProcessor(image_processor=_UpperCamelCase,    tokenizer=_UpperCamelCase     )
      return processor




def 			__lowercase	(						_UpperCamelCase     )    ->int:







      """simple docstring"""

      lowercase			:   str					   =		[v.split('''_'''     )[0].split('''block'''     )[1] for v in original_param_names if v.startswith('''block'''     )]
      lowercase			:   str					   =		list(set(_UpperCamelCase     )     )
      lowercase			:   List[str]					   =		sorted(_UpperCamelCase     )
      lowercase			:   str					   =		len(_UpperCamelCase     )
      lowercase			:   Dict					   =		{b: str(_UpperCamelCase     ) for b, i in zip(_UpperCamelCase,    range(_UpperCamelCase     )     )}

      lowercase			:   int					   =		[]
      rename_keys.append(('''stem_conv/kernel:0''', '''embeddings.convolution.weight''')     )
      rename_keys.append(('''stem_bn/gamma:0''', '''embeddings.batchnorm.weight''')     )
      rename_keys.append(('''stem_bn/beta:0''', '''embeddings.batchnorm.bias''')     )
      rename_keys.append(('''stem_bn/moving_mean:0''', '''embeddings.batchnorm.running_mean''')     )
      rename_keys.append(('''stem_bn/moving_variance:0''', '''embeddings.batchnorm.running_var''')     )

      for b in block_names:
            lowercase			:   Dict					   =		block_name_mapping[b]
            rename_keys.append((f"""block{b}_expand_conv/kernel:0""", f"""encoder.blocks.{hf_b}.expansion.expand_conv.weight""")     )
            rename_keys.append((f"""block{b}_expand_bn/gamma:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.weight""")     )
            rename_keys.append((f"""block{b}_expand_bn/beta:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.bias""")     )
            rename_keys.append(
                (f"""block{b}_expand_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.running_mean""")     )
            rename_keys.append(
                (f"""block{b}_expand_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.running_var""")     )
            rename_keys.append(
                (f"""block{b}_dwconv/depthwise_kernel:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight""")     )
            rename_keys.append((f"""block{b}_bn/gamma:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight""")     )
            rename_keys.append((f"""block{b}_bn/beta:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias""")     )
            rename_keys.append(
                (f"""block{b}_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean""")     )
            rename_keys.append(
                (f"""block{b}_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var""")     )

            rename_keys.append((f"""block{b}_se_reduce/kernel:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.reduce.weight""")     )
            rename_keys.append((f"""block{b}_se_reduce/bias:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.reduce.bias""")     )
            rename_keys.append((f"""block{b}_se_expand/kernel:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.expand.weight""")     )
            rename_keys.append((f"""block{b}_se_expand/bias:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.expand.bias""")     )
            rename_keys.append(
                (f"""block{b}_project_conv/kernel:0""", f"""encoder.blocks.{hf_b}.projection.project_conv.weight""")     )
            rename_keys.append((f"""block{b}_project_bn/gamma:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.weight""")     )
            rename_keys.append((f"""block{b}_project_bn/beta:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.bias""")     )
            rename_keys.append(
                (f"""block{b}_project_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.running_mean""")     )
            rename_keys.append(
                (f"""block{b}_project_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.running_var""")     )

      lowercase			:   Union[str, Any]					   =		{}
      for item in rename_keys:
            if item[0] in original_param_names:
                  lowercase			:   Dict					   =		'''vision_model.''' + item[1]

    # BERT text encoder
      lowercase			:   str					   =		[]
      lowercase			:   Optional[Any]					   =		'''tf_bert_model/bert'''
      lowercase			:   str					   =		'''text_model'''
      for i in range(12     ):
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/self/query/kernel:0""",
                    f"""{new}.encoder.layer.{i}.attention.self.query.weight""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/self/query/bias:0""",
                    f"""{new}.encoder.layer.{i}.attention.self.query.bias""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/self/key/kernel:0""",
                    f"""{new}.encoder.layer.{i}.attention.self.key.weight""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/self/key/bias:0""",
                    f"""{new}.encoder.layer.{i}.attention.self.key.bias""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/self/value/kernel:0""",
                    f"""{new}.encoder.layer.{i}.attention.self.value.weight""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/self/value/bias:0""",
                    f"""{new}.encoder.layer.{i}.attention.self.value.bias""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/output/dense/kernel:0""",
                    f"""{new}.encoder.layer.{i}.attention.output.dense.weight""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/output/dense/bias:0""",
                    f"""{new}.encoder.layer.{i}.attention.output.dense.bias""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/output/LayerNorm/gamma:0""",
                    f"""{new}.encoder.layer.{i}.attention.output.LayerNorm.weight""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/attention/output/LayerNorm/beta:0""",
                    f"""{new}.encoder.layer.{i}.attention.output.LayerNorm.bias""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/intermediate/dense/kernel:0""",
                    f"""{new}.encoder.layer.{i}.intermediate.dense.weight""",
                )     )
            rename_keys.append(
                (
                    f"""{old}/encoder/layer_._{i}/intermediate/dense/bias:0""",
                    f"""{new}.encoder.layer.{i}.intermediate.dense.bias""",
                )     )
            rename_keys.append(
                (f"""{old}/encoder/layer_._{i}/output/dense/kernel:0""", f"""{new}.encoder.layer.{i}.output.dense.weight""")     )
            rename_keys.append(
                (f"""{old}/encoder/layer_._{i}/output/dense/bias:0""", f"""{new}.encoder.layer.{i}.output.dense.bias""")     )
            rename_keys.append(
                (f"""{old}/encoder/layer_._{i}/output/LayerNorm/gamma:0""", f"""{new}.encoder.layer.{i}.output.LayerNorm.weight""")     )
            rename_keys.append(
                (f"""{old}/encoder/layer_._{i}/output/LayerNorm/beta:0""", f"""{new}.encoder.layer.{i}.output.LayerNorm.bias""")     )

      rename_keys.append((f"""{old}/embeddings/word_embeddings/weight:0""", f"""{new}.embeddings.word_embeddings.weight""")     )
      rename_keys.append(
          (f"""{old}/embeddings/position_embeddings/embeddings:0""", f"""{new}.embeddings.position_embeddings.weight""")     )
      rename_keys.append(
          (f"""{old}/embeddings/token_type_embeddings/embeddings:0""", f"""{new}.embeddings.token_type_embeddings.weight""")     )
      rename_keys.append((f"""{old}/embeddings/LayerNorm/gamma:0""", f"""{new}.embeddings.LayerNorm.weight""")     )
      rename_keys.append((f"""{old}/embeddings/LayerNorm/beta:0""", f"""{new}.embeddings.LayerNorm.bias""")     )

      rename_keys.append((f"""{old}/pooler/dense/kernel:0""", f"""{new}.pooler.dense.weight""")     )
      rename_keys.append((f"""{old}/pooler/dense/bias:0""", f"""{new}.pooler.dense.bias""")     )
      rename_keys.append(('''dense/kernel:0''', '''text_projection.weight''')     )
      rename_keys.append(('''dense/bias:0''', '''text_projection.bias''')     )
      rename_keys.append(('''dense/bias:0''', '''text_projection.bias''')     )
      rename_keys.append(('''temperature:0''', '''temperature''')     )

      for item in rename_keys:
            if item[0] in original_param_names:
                  lowercase			:   Tuple					   =		item[1]
      return key_mapping




def 			__lowercase	(						_UpperCamelCase,    _UpperCamelCase,    _UpperCamelCase     )    ->Any:







      """simple docstring"""

      list(hf_params.keys()     )

      for key, value in tf_params.items():
            if key not in key_mapping:
                  continue

            lowercase			:   List[str]					   =		key_mapping[key]
            if "_conv" in key and "kernel" in key:
                  lowercase			:   List[Any]					   =		torch.from_numpy(_UpperCamelCase     ).permute(3,    2,    0,    1     )
            elif "embeddings" in key:
                  lowercase			:   int					   =		torch.from_numpy(_UpperCamelCase     )
            elif "depthwise_kernel" in key:
                  lowercase			:   Union[str, Any]					   =		torch.from_numpy(_UpperCamelCase     ).permute(2,    3,    0,    1     )
            elif "kernel" in key:
                  lowercase			:   int					   =		torch.from_numpy(np.transpose(_UpperCamelCase     )     )
            elif "temperature" in key:
                  lowercase			:   Optional[int]					   =		value
            elif "bn/gamma" or "bn/beta" in key:
                  lowercase			:   Dict					   =		torch.from_numpy(np.transpose(_UpperCamelCase     )     ).squeeze()
            else:
                  lowercase			:   Union[str, Any]					   =		torch.from_numpy(_UpperCamelCase     )

            # Replace HF parameters with original TF model parameters
            hf_params[hf_key].copy_(_UpperCamelCase     )




@torch.no_grad()
def 			__lowercase	(						_UpperCamelCase,    _UpperCamelCase,    _UpperCamelCase,    _UpperCamelCase     )    ->Dict:







      """simple docstring"""

      lowercase			:   str					   =		64
      lowercase			:   str					   =		Tokenizer(_UpperCamelCase     )
      lowercase			:   Optional[int]					   =		align.Align('''efficientnet-b7''',    '''bert-base''',    640,    _UpperCamelCase,    tok.get_vocab_size()     )
      original_model.compile()
      original_model.load_weights(_UpperCamelCase     )

      lowercase			:   Optional[Any]					   =		original_model.trainable_variables
      lowercase			:   Tuple					   =		original_model.non_trainable_variables
      lowercase			:   Optional[int]					   =		{param.name: param.numpy() for param in tf_params}
      for param in tf_non_train_params:
            lowercase			:   List[str]					   =		param.numpy()
      lowercase			:   str					   =		list(tf_params.keys()     )

      # Load HuggingFace model
      lowercase			:   Tuple					   =		get_align_config()
      lowercase			:   List[Any]					   =		AlignModel(_UpperCamelCase     ).eval()
      lowercase			:   Union[str, Any]					   =		hf_model.state_dict()

      # Create src-to-dst parameter name mapping dictionary
      print('''Converting parameters...'''     )
      lowercase			:   List[str]					   =		rename_keys(_UpperCamelCase     )
      replace_params(_UpperCamelCase,    _UpperCamelCase,    _UpperCamelCase     )

      # Initialize processor
      lowercase			:   Union[str, Any]					   =		get_processor()
      lowercase			:   Optional[Any]					   =		processor(
          images=prepare_img(),    text='''A picture of a cat''',    padding='''max_length''',    max_length=64,    return_tensors='''pt'''     )

      # HF model inference
      hf_model.eval()
      with torch.no_grad():
            lowercase			:   Dict					   =		hf_model(**_UpperCamelCase     )

      lowercase			:   int					   =		outputs.image_embeds.detach().numpy()
      lowercase			:   Any					   =		outputs.text_embeds.detach().numpy()

      # Original model inference
      lowercase			:   Any					   =		False
      lowercase			:   Tuple					   =		EfficientNetImageProcessor(
          do_center_crop=_UpperCamelCase,    do_rescale=_UpperCamelCase,    do_normalize=_UpperCamelCase,    include_top=_UpperCamelCase,    resample=Image.BILINEAR,    )
      lowercase			:   str					   =		tf_image_processor(images=prepare_img(),    return_tensors='''tf''',    data_format='''channels_last'''     )['''pixel_values''']
      lowercase			:   int					   =		tok(tf.constant(['''A picture of a cat''']     )     )

      lowercase			:   Optional[int]					   =		original_model.image_encoder(_UpperCamelCase,    training=_UpperCamelCase     )
      lowercase			:   Optional[int]					   =		original_model.text_encoder(_UpperCamelCase,    training=_UpperCamelCase     )

      lowercase			:   Union[str, Any]					   =		tf.nn.la_normalize(_UpperCamelCase,    axis=-1     )
      lowercase			:   List[str]					   =		tf.nn.la_normalize(_UpperCamelCase,    axis=-1     )

      # Check whether original and HF model outputs match  -> np.allclose
      if not np.allclose(_UpperCamelCase,    _UpperCamelCase,    atol=1e-3     ):
            raise ValueError('''The predicted image features are not the same.'''     )
      if not np.allclose(_UpperCamelCase,    _UpperCamelCase,    atol=1e-3     ):
            raise ValueError('''The predicted text features are not the same.'''     )
      print('''Model outputs match!'''     )

      if save_model:
            # Create folder to save model
            if not os.path.isdir(_UpperCamelCase     ):
                  os.mkdir(_UpperCamelCase     )
            # Save converted model and image processor
            hf_model.save_pretrained(_UpperCamelCase     )
            processor.save_pretrained(_UpperCamelCase     )

      if push_to_hub:
            # Push model and image processor to hub
            print('''Pushing converted ALIGN to the hub...'''     )
            processor.push_to_hub('''align-base'''     )
            hf_model.push_to_hub('''align-base'''     )


if __name__ == "__main__":
  __a   				=      argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '''--checkpoint_path''',
      default='''./weights/model-weights''',
      type=str,
      help='''Path to the pretrained TF ALIGN checkpoint.''',
  )
  parser.add_argument(
      '''--pytorch_dump_folder_path''',
      default='''hf_model''',
      type=str,
      help='''Path to the output PyTorch model directory.''',
  )
  parser.add_argument('''--save_model''', action='''store_true''', help='''Save model to local''')
  parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image processor to the hub''')

  __a   				=      parser.parse_args()
  convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)

