


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()
lowerCAmelCase__						:Any         =							logging.get_logger(__name__)

def  lowerCAmelCase__				(    a__:	List[str] )						->	Tuple:


							'''simple docstring'''





							_UpperCAmelCase					=    tf.image.resize(a__    ,						(3_4_6, 3_4_6) )
							_UpperCAmelCase					=    tf.image.crop_to_bounding_box(a__    ,						(3_4_6 - 2_8_9) // 2    ,						(3_4_6 - 2_8_9) // 2    ,						2_8_9    ,						2_8_9 )
							return image

def  lowerCAmelCase__				(    )						->	Optional[int]:


							'''simple docstring'''





							_UpperCAmelCase					=    EfficientNetConfig.from_pretrained('google/efficientnet-b7' )
							_UpperCAmelCase					=    2_8_9
							_UpperCAmelCase					=    6_4_0
							_UpperCAmelCase					=    {'0': 'LABEL_0', '1': 'LABEL_1'}
							_UpperCAmelCase					=    {'LABEL_0': 0, 'LABEL_1': 1}
							_UpperCAmelCase					=    []

							_UpperCAmelCase					=    BertConfig()
							_UpperCAmelCase					=    AlignConfig.from_text_vision_configs(
							    text_config=a__    ,						vision_config=a__    ,						projection_dim=6_4_0 )
							return config

def  lowerCAmelCase__				(    )						->	List[str]:


							'''simple docstring'''





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

def  lowerCAmelCase__				(    )						->	int:


							'''simple docstring'''





							_UpperCAmelCase					=    EfficientNetImageProcessor(
							    do_center_crop=a__    ,						rescale_factor=1 / 127.5    ,						rescale_offset=a__    ,						do_normalize=a__    ,						include_top=a__    ,						resample=Image.BILINEAR    ,						)
							_UpperCAmelCase					=    BertTokenizer.from_pretrained('bert-base-uncased' )
							_UpperCAmelCase					=    6_4
							_UpperCAmelCase					=    AlignProcessor(image_processor=a__    ,						tokenizer=a__ )
							return processor

def  lowerCAmelCase__				(    a__:	Union[str, Any] )						->	List[Any]:


							'''simple docstring'''





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

							_UpperCAmelCase					=    []
							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:
														_UpperCAmelCase					=    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''') )

							_UpperCAmelCase					=    {}
							for item in rename_keys:
														if item[0] in original_param_names:
																					_UpperCAmelCase					=    'vision_model.' + item[1]

    # BERT text encoder
							_UpperCAmelCase					=    []
							_UpperCAmelCase					=    'tf_bert_model/bert'
							_UpperCAmelCase					=    'text_model'
							for i in range(1_2 ):
														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:
																					_UpperCAmelCase					=    item[1]
							return key_mapping

def  lowerCAmelCase__				(    a__:	Any    ,						a__:	int    ,						a__:	Optional[int] )						->	Optional[int]:


							'''simple docstring'''





							list(hf_params.keys() )

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

														_UpperCAmelCase					=    key_mapping[key]
														if "_conv" in key and "kernel" in key:
																					_UpperCAmelCase					=    torch.from_numpy(a__ ).permute(3    ,						2    ,						0    ,						1 )
														elif "embeddings" in key:
																					_UpperCAmelCase					=    torch.from_numpy(a__ )
														elif "depthwise_kernel" in key:
																					_UpperCAmelCase					=    torch.from_numpy(a__ ).permute(2    ,						3    ,						0    ,						1 )
														elif "kernel" in key:
																					_UpperCAmelCase					=    torch.from_numpy(np.transpose(a__ ) )
														elif "temperature" in key:
																					_UpperCAmelCase					=    value
														elif "bn/gamma" or "bn/beta" in key:
																					_UpperCAmelCase					=    torch.from_numpy(np.transpose(a__ ) ).squeeze()
														else:
																					_UpperCAmelCase					=    torch.from_numpy(a__ )

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

@torch.no_grad()
def  lowerCAmelCase__				(    a__:	List[Any]    ,						a__:	Union[str, Any]    ,						a__:	Dict    ,						a__:	Dict )						->	Optional[Any]:


							'''simple docstring'''





							_UpperCAmelCase					=    6_4
							_UpperCAmelCase					=    Tokenizer(a__ )
							_UpperCAmelCase					=    align.Align('efficientnet-b7'    ,						'bert-base'    ,						6_4_0    ,						a__    ,						tok.get_vocab_size() )
							original_model.compile()
							original_model.load_weights(a__ )

							_UpperCAmelCase					=    original_model.trainable_variables
							_UpperCAmelCase					=    original_model.non_trainable_variables
							_UpperCAmelCase					=    {param.name: param.numpy() for param in tf_params}
							for param in tf_non_train_params:
														_UpperCAmelCase					=    param.numpy()
							_UpperCAmelCase					=    list(tf_params.keys() )

							# Load HuggingFace model
							_UpperCAmelCase					=    get_align_config()
							_UpperCAmelCase					=    AlignModel(a__ ).eval()
							_UpperCAmelCase					=    hf_model.state_dict()

							# Create src-to-dst parameter name mapping dictionary
							print('Converting parameters...' )
							_UpperCAmelCase					=    rename_keys(a__ )
							replace_params(a__    ,						a__    ,						a__ )

							# Initialize processor
							_UpperCAmelCase					=    get_processor()
							_UpperCAmelCase					=    processor(
							    images=prepare_img()    ,						text='A picture of a cat'    ,						padding='max_length'    ,						max_length=6_4    ,						return_tensors='pt' )

							# HF model inference
							hf_model.eval()
							with torch.no_grad():
														_UpperCAmelCase					=    hf_model(**a__ )

							_UpperCAmelCase					=    outputs.image_embeds.detach().numpy()
							_UpperCAmelCase					=    outputs.text_embeds.detach().numpy()

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

							_UpperCAmelCase					=    original_model.image_encoder(a__    ,						training=a__ )
							_UpperCAmelCase					=    original_model.text_encoder(a__    ,						training=a__ )

							_UpperCAmelCase					=    tf.nn.la_normalize(a__    ,						axis=-1 )
							_UpperCAmelCase					=    tf.nn.la_normalize(a__    ,						axis=-1 )

							# Check whether original and HF model outputs match  -> np.allclose
							if not np.allclose(a__    ,						a__    ,						atol=1e-3 ):
														raise ValueError('The predicted image features are not the same.' )
							if not np.allclose(a__    ,						a__    ,						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(a__ ):
																					os.mkdir(a__ )
														# Save converted model and image processor
														hf_model.save_pretrained(a__ )
														processor.save_pretrained(a__ )

							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__":
		lowerCAmelCase__						:List[Any]         =							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''')

		lowerCAmelCase__						:Union[str, Any]         =							parser.parse_args()
		convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)




import unittest

from transformers import load_tool
from transformers.tools.agent_types import AGENT_TYPE_MAPPING

from .test_tools_common import ToolTesterMixin, output_types

class 	__a   (       unittest.TestCase					,   UpperCAmelCase    ):
		def 			UpperCAmelCase__       (    self		)				-> str:




									"""simple docstring"""





									_UpperCAmelCase					=    load_tool('translation'		)
									self.tool.setup()
									_UpperCAmelCase					=    load_tool('translation'    ,					remote=_SCREAMING_SNAKE_CASE		)
		def 			UpperCAmelCase__       (    self		)				-> Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase					=    self.tool('Hey, what\'s up?'    ,					src_lang='English'    ,					tgt_lang='French'		)
									self.assertEqual(_SCREAMING_SNAKE_CASE    ,					'- Hé, comment ça va?'		)
		def 			UpperCAmelCase__       (    self		)				-> Tuple:




									"""simple docstring"""





									_UpperCAmelCase					=    self.remote_tool('Hey, what\'s up?'    ,					src_lang='English'    ,					tgt_lang='French'		)
									self.assertEqual(_SCREAMING_SNAKE_CASE    ,					'- Hé, comment ça va?'		)
		def 			UpperCAmelCase__       (    self		)				-> Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase					=    self.tool(text='Hey, what\'s up?'    ,					src_lang='English'    ,					tgt_lang='French'		)
									self.assertEqual(_SCREAMING_SNAKE_CASE    ,					'- Hé, comment ça va?'		)
		def 			UpperCAmelCase__       (    self		)				-> int:




									"""simple docstring"""





									_UpperCAmelCase					=    self.remote_tool(text='Hey, what\'s up?'    ,					src_lang='English'    ,					tgt_lang='French'		)
									self.assertEqual(_SCREAMING_SNAKE_CASE    ,					'- Hé, comment ça va?'		)
		def 			UpperCAmelCase__       (    self		)				-> Tuple:




									"""simple docstring"""





									_UpperCAmelCase					=    ['Hey, what\'s up?', 'English', 'Spanish']
									_UpperCAmelCase					=    self.tool(*_SCREAMING_SNAKE_CASE		)

									# There is a single output
									if len(self.tool.outputs		) == 1:
																_UpperCAmelCase					=    [outputs]

									self.assertListEqual(output_types(_SCREAMING_SNAKE_CASE		)    ,					self.tool.outputs		)
		def 			UpperCAmelCase__       (    self		)				-> Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    ['Hey, what\'s up?', 'English', 'Spanish']
									_UpperCAmelCase					=    self.tool(*_SCREAMING_SNAKE_CASE		)

									if not isinstance(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		):
																_UpperCAmelCase					=    [outputs]

									self.assertEqual(len(_SCREAMING_SNAKE_CASE		)    ,					len(self.tool.outputs		)		)

									for output, output_type in zip(_SCREAMING_SNAKE_CASE    ,					self.tool.outputs		):
																_UpperCAmelCase					=    AGENT_TYPE_MAPPING[output_type]
																self.assertTrue(isinstance(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)		)






		def 			UpperCAmelCase__       (    self		)				-> Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase					=    ['Hey, what\'s up?', 'English', 'Spanish']

									_UpperCAmelCase					=    []

									for _input, input_type in zip(_SCREAMING_SNAKE_CASE    ,					self.tool.inputs		):
																if isinstance(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		):
																							_inputs.append([AGENT_TYPE_MAPPING[_input_type](_input		) for _input_type in input_type]		)
																else:
																							_inputs.append(AGENT_TYPE_MAPPING[input_type](_input		)		)

        # Should not raise an error
									_UpperCAmelCase					=    self.tool(*_SCREAMING_SNAKE_CASE		)

									if not isinstance(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		):
																_UpperCAmelCase					=    [outputs]

									self.assertEqual(len(_SCREAMING_SNAKE_CASE		)    ,					len(self.tool.outputs		)		)
