


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()
snake_case__  				=			logging.get_logger(__name__)
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:Tuple					)							->	Dict:
						__magic_name__					:				Dict				=   tf.image.resize(lowercase__   ,    (3_4_6, 3_4_6)					)
						__magic_name__					:				Tuple				=   tf.image.crop_to_bounding_box(lowercase__   ,    (3_4_6 - 2_8_9) // 2   ,    (3_4_6 - 2_8_9) // 2   ,    2_8_9   ,    2_8_9					)
						return image
def 			__SCREAMING_SNAKE_CASE							(					)							->	Any:
						__magic_name__					:				List[str]				=   EfficientNetConfig.from_pretrained('google/efficientnet-b7'					)
						__magic_name__					:				List[str]				=   2_8_9
						__magic_name__					:				Optional[int]				=   6_4_0
						__magic_name__					:				Union[str, Any]				=   {'0': 'LABEL_0', '1': 'LABEL_1'}
						__magic_name__					:				List[str]				=   {'LABEL_0': 0, 'LABEL_1': 1}
						__magic_name__					:				Any				=   []

						__magic_name__					:				Optional[Any]				=   BertConfig()
						__magic_name__					:				Optional[int]				=   AlignConfig.from_text_vision_configs(
						    text_config=lowercase__   ,    vision_config=lowercase__   ,    projection_dim=6_4_0					)
						return config
def 			__SCREAMING_SNAKE_CASE							(					)							->	Tuple:
						__magic_name__					:				Optional[Any]				=   'http://images.cocodataset.org/val2017/000000039769.jpg'
						__magic_name__					:				Optional[Any]				=   Image.open(requests.get(lowercase__   ,    stream=lowercase__					).raw					)
						return im
def 			__SCREAMING_SNAKE_CASE							(					)							->	int:
						__magic_name__					:				Union[str, Any]				=   EfficientNetImageProcessor(
						    do_center_crop=lowercase__   ,    rescale_factor=1 / 1_2_7.5   ,    rescale_offset=lowercase__   ,    do_normalize=lowercase__   ,    include_top=lowercase__   ,    resample=Image.BILINEAR   ,    )
						__magic_name__					:				List[Any]				=   BertTokenizer.from_pretrained('bert-base-uncased'					)
						__magic_name__					:				int				=   6_4
						__magic_name__					:				Optional[int]				=   AlignProcessor(image_processor=lowercase__   ,    tokenizer=lowercase__					)
						return processor
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:Optional[Any]					)							->	Tuple:
						# EfficientNet image encoder
						__magic_name__					:				Optional[int]				=   [v.split('_'					)[0].split('block'					)[1] for v in original_param_names if v.startswith('block'					)]
						__magic_name__					:				List[str]				=   list(set(lowercase__					)					)
						__magic_name__					:				List[str]				=   sorted(lowercase__					)
						__magic_name__					:				Tuple				=   len(lowercase__					)
						__magic_name__					:				Union[str, Any]				=   {b: str(lowercase__					) for b, i in zip(lowercase__   ,    range(lowercase__					)					)}

						__magic_name__					:				str				=   []
						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:
												__magic_name__					:				Any				=   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")					)

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

    # BERT text encoder
						__magic_name__					:				Any				=   []
						__magic_name__					:				Any				=   'tf_bert_model/bert'
						__magic_name__					:				Dict				=   '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:
																		__magic_name__					:				int				=   item[1]
						return key_mapping
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:List[str]   ,    lowercase__	:Dict   ,    lowercase__	:Optional[Any]					)							->	Optional[int]:
						list(hf_params.keys()					)

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

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

												# Replace HF parameters with original TF model parameters
												hf_params[hf_key].copy_(lowercase__					)
@torch.no_grad()
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:Optional[int]   ,    lowercase__	:List[Any]   ,    lowercase__	:Any   ,    lowercase__	:Any					)							->	List[str]:
						__magic_name__					:				Union[str, Any]				=   6_4
						__magic_name__					:				Optional[int]				=   Tokenizer(lowercase__					)
						__magic_name__					:				Optional[Any]				=   align.Align('efficientnet-b7'   ,    'bert-base'   ,    6_4_0   ,    lowercase__   ,    tok.get_vocab_size()					)
						original_model.compile()
						original_model.load_weights(lowercase__					)

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

						# Load HuggingFace model
						__magic_name__					:				Dict				=   get_align_config()
						__magic_name__					:				Tuple				=   AlignModel(lowercase__					).eval()
						__magic_name__					:				Optional[int]				=   hf_model.state_dict()

						# Create src-to-dst parameter name mapping dictionary
						print('Converting parameters...'					)
						__magic_name__					:				Optional[Any]				=   rename_keys(lowercase__					)
						replace_params(lowercase__   ,    lowercase__   ,    lowercase__					)

						# Initialize processor
						__magic_name__					:				List[str]				=   get_processor()
						__magic_name__					:				List[str]				=   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():
												__magic_name__					:				Tuple				=   hf_model(**lowercase__					)

						__magic_name__					:				Optional[int]				=   outputs.image_embeds.detach().numpy()
						__magic_name__					:				List[Any]				=   outputs.text_embeds.detach().numpy()

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

						__magic_name__					:				Optional[int]				=   original_model.image_encoder(lowercase__   ,    training=lowercase__					)
						__magic_name__					:				List[Any]				=   original_model.text_encoder(lowercase__   ,    training=lowercase__					)

						__magic_name__					:				Union[str, Any]				=   tf.nn.la_normalize(lowercase__   ,    axis=-1					)
						__magic_name__					:				Any				=   tf.nn.la_normalize(lowercase__   ,    axis=-1					)

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

						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__":
			snake_case__  				=			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')

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






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


snake_case__  				=			{'configuration_qdqbert': ['QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'QDQBertConfig']}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			snake_case__  				=			[
			    'QDQBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
			    'QDQBertForMaskedLM',
			    'QDQBertForMultipleChoice',
			    'QDQBertForNextSentencePrediction',
			    'QDQBertForQuestionAnswering',
			    'QDQBertForSequenceClassification',
			    'QDQBertForTokenClassification',
			    'QDQBertLayer',
			    'QDQBertLMHeadModel',
			    'QDQBertModel',
			    'QDQBertPreTrainedModel',
			    'load_tf_weights_in_qdqbert',
			]


if TYPE_CHECKING:
			from .configuration_qdqbert import QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, QDQBertConfig

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_qdqbert import (
						    QDQBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    QDQBertForMaskedLM,
						    QDQBertForMultipleChoice,
						    QDQBertForNextSentencePrediction,
						    QDQBertForQuestionAnswering,
						    QDQBertForSequenceClassification,
						    QDQBertForTokenClassification,
						    QDQBertLayer,
						    QDQBertLMHeadModel,
						    QDQBertModel,
						    QDQBertPreTrainedModel,
						    load_tf_weights_in_qdqbert,
						)


else:
			import sys

			snake_case__  				=			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


