





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  :					Optional[int]          =  logging.get_logger(__name__)





def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							tf.image.resize(lowercase							,      (346, 346)							)
				__lowercase			      =							tf.image.crop_to_bounding_box(lowercase							,      (346 - 289) // 2							,      (346 - 289) // 2							,      289							,      289							)
				return image





def     UpperCAmelCase						(   ):







				"""simple docstring"""



				__lowercase			      =							EfficientNetConfig.from_pretrained('''google/efficientnet-b7'''							)
				__lowercase			      =							289
				__lowercase			      =							640
				__lowercase			      =							{'''0''': '''LABEL_0''', '''1''': '''LABEL_1'''}
				__lowercase			      =							{'''LABEL_0''': 0, '''LABEL_1''': 1}
				__lowercase			      =							[]

				__lowercase			      =							BertConfig()
				__lowercase			      =							AlignConfig.from_text_vision_configs(
				    text_config=lowercase							,      vision_config=lowercase							,      projection_dim=640							)
				return config





def     UpperCAmelCase						(   ):







				"""simple docstring"""



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





def     UpperCAmelCase						(   ):







				"""simple docstring"""



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





def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



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

				__lowercase			      =							[]
				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			      =							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			      =							{}
				for item in rename_keys:
								if item[0] in original_param_names:
												__lowercase			      =							'''vision_model.''' + item[1]

    # BERT text encoder
				__lowercase			      =							[]
				__lowercase			      =							'''tf_bert_model/bert'''
				__lowercase			      =							'''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			      =							item[1]
				return key_mapping





def     UpperCAmelCase						(   lowercase							,      lowercase							,      lowercase							):







				"""simple docstring"""



				list(hf_params.keys()							)

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

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

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





@torch.no_grad()
def     UpperCAmelCase						(   lowercase							,      lowercase							,      lowercase							,      lowercase							):







				"""simple docstring"""



				__lowercase			      =							64
				__lowercase			      =							Tokenizer(lowercase							)
				__lowercase			      =							align.Align('''efficientnet-b7'''							,      '''bert-base'''							,      640							,      lowercase							,      tok.get_vocab_size()							)
				original_model.compile()
				original_model.load_weights(lowercase							)

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

				# Load HuggingFace model
				__lowercase			      =							get_align_config()
				__lowercase			      =							AlignModel(lowercase							).eval()
				__lowercase			      =							hf_model.state_dict()

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

				# Initialize processor
				__lowercase			      =							get_processor()
				__lowercase			      =							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			      =							hf_model(**lowercase							)

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

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

				__lowercase			      =							original_model.image_encoder(lowercase							,      training=lowercase							)
				__lowercase			      =							original_model.text_encoder(lowercase							,      training=lowercase							)

				__lowercase			      =							tf.nn.la_normalize(lowercase							,      axis=-1							)
				__lowercase			      =							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__":
		__a  :					Optional[int]          =  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  :					Dict          =  parser.parse_args()
		convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)
import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__a  :					int          =  5_0_0_0_0_0

__a   ,						__a  :					Any          =  os.path.split(__file__)
__a  :					Dict          =  os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.select(range(0							,      len(lowercase							)							,      2							)							)





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.sort('''numbers'''							)





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.shuffle()





@get_duration
def     UpperCAmelCase						(   lowercase							):







				"""simple docstring"""



				__lowercase			      =							dataset.train_test_split(0.1							)





@get_duration
def     UpperCAmelCase						(   lowercase							,      lowercase=10							):







				"""simple docstring"""



				for shard_id in range(lowercase							):
								__lowercase			      =							dataset.shard(lowercase							,      lowercase							)





def     UpperCAmelCase						(   ):







				"""simple docstring"""



				__lowercase			      =							{'''num examples''': SPEED_TEST_N_EXAMPLES}
				__lowercase			      =							(select, sort, shuffle, train_test_split, shard)
				with tempfile.TemporaryDirectory() as tmp_dir:
								print('''generating dataset'''							)
								__lowercase			      =							datasets.Features({'''text''': datasets.Value('''string'''							), '''numbers''': datasets.Value('''float32'''							)}							)
								__lowercase			      =							generate_example_dataset(
								    os.path.join(lowercase							,      '''dataset.arrow'''							)							,      lowercase							,      num_examples=lowercase							)
								print('''Functions'''							)
								for func in functions:
												print(func.__name__							)
												__lowercase			      =							func(lowercase							)

				with open(lowercase							,      '''wb'''							) as f:
								f.write(json.dumps(lowercase							).encode('''utf-8'''							)							)


if __name__ == "__main__":  # useful to run the profiler
		benchmark_indices_mapping()