




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()
lowercase__    :					List[str]					=						logging.get_logger(__name__)


def 						lowerCamelCase__   (		_A							):







	'''simple docstring'''





	snake_case_		   =					tf.image.resize(lowercase__					,     (346, 346)							)
	snake_case_		   =					tf.image.crop_to_bounding_box(lowercase__					,     (346 - 289) // 2					,     (346 - 289) // 2					,     289					,     289							)
	return image


def 						lowerCamelCase__   (		):







	'''simple docstring'''





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

	snake_case_		   =					BertConfig()
	snake_case_		   =					AlignConfig.from_text_vision_configs(
	    text_config=lowercase__					,     vision_config=lowercase__					,     projection_dim=640							)
	return config


def 						lowerCamelCase__   (		):







	'''simple docstring'''





	snake_case_		   =					"http://images.cocodataset.org/val2017/000000039769.jpg"
	snake_case_		   =					Image.open(requests.get(lowercase__					,     stream=lowercase__							).raw							)
	return im


def 						lowerCamelCase__   (		):







	'''simple docstring'''





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


def 						lowerCamelCase__   (		_A							):







	'''simple docstring'''





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

	snake_case_		   =					[]
	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:
		snake_case_		   =					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")							)

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

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


def 						lowerCamelCase__   (		_A					,     _A					,     _A							):







	'''simple docstring'''





	list(hf_params.keys()							)

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

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

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


@torch.no_grad()
def 						lowerCamelCase__   (		_A					,     _A					,     _A					,     _A							):







	'''simple docstring'''





	snake_case_		   =					64
	snake_case_		   =					Tokenizer(lowercase__							)
	snake_case_		   =					align.Align("efficientnet-b7"					,     "bert-base"					,     640					,     lowercase__					,     tok.get_vocab_size()							)
	original_model.compile()
	original_model.load_weights(lowercase__							)

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

	# Load HuggingFace model
	snake_case_		   =					get_align_config()
	snake_case_		   =					AlignModel(lowercase__							).eval()
	snake_case_		   =					hf_model.state_dict()

	# Create src-to-dst parameter name mapping dictionary
	print("Converting parameters..."							)
	snake_case_		   =					rename_keys(lowercase__							)
	replace_params(lowercase__					,     lowercase__					,     lowercase__							)

	# Initialize processor
	snake_case_		   =					get_processor()
	snake_case_		   =					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():
		snake_case_		   =					hf_model(**lowercase__							)

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

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

	snake_case_		   =					original_model.image_encoder(lowercase__					,     training=lowercase__							)
	snake_case_		   =					original_model.text_encoder(lowercase__					,     training=lowercase__							)

	snake_case_		   =					tf.nn.la_normalize(lowercase__					,     axis=-1							)
	snake_case_		   =					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__":
		lowercase__    :					str					=						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")

		lowercase__    :					str					=						parser.parse_args()
		convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)






import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


lowercase__    :					str					=						datasets.utils.logging.get_logger(__name__)
@dataclass
class 		UpperCAmelCase       (    datasets.BuilderConfig      ):







						'''simple docstring'''





						lowerCAmelCase_							=       None







class 		UpperCAmelCase       (    datasets.ArrowBasedBuilder      ):







						'''simple docstring'''





						lowerCAmelCase_							=       ArrowConfig






						def 	snake_case__						(						self							:					Dict						):






							"""simple docstring"""






							return datasets.DatasetInfo(features=self.config.features						)






						def 	snake_case__						(						self							:					Tuple	,     __lowercase							:					Union[str, Any]						):






							"""simple docstring"""






							if not self.config.data_files:
								raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}"						)
							snake_case_		   =					dl_manager.download_and_extract(self.config.data_files						)
							if isinstance(__lowercase	,     (str, list, tuple)						):
								snake_case_		   =					data_files
								if isinstance(__lowercase	,     __lowercase						):
									snake_case_		   =					[files]
								# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
								snake_case_		   =					[dl_manager.iter_files(__lowercase						) for file in files]
								return [datasets.SplitGenerator(name=datasets.Split.TRAIN	,     gen_kwargs={"files": files}						)]
							snake_case_		   =					[]
							for split_name, files in data_files.items():
								if isinstance(__lowercase	,     __lowercase						):
									snake_case_		   =					[files]
								# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
								snake_case_		   =					[dl_manager.iter_files(__lowercase						) for file in files]
								# Infer features is they are stoed in the arrow schema
								if self.info.features is None:
									for file in itertools.chain.from_iterable(__lowercase						):
										with open(__lowercase	,     "rb"						) as f:
											snake_case_		   =					datasets.Features.from_arrow_schema(pa.ipc.open_stream(__lowercase						).schema						)
										break
								splits.append(datasets.SplitGenerator(name=__lowercase	,     gen_kwargs={"files": files}						)						)
							return splits






						def 	snake_case__						(						self							:					str	,     __lowercase							:					pa.Table						):






							"""simple docstring"""






							if self.info.features is not None:
								# more expensive cast to support nested features with keys in a different order
								# allows str <-> int/float or str to Audio for example
								snake_case_		   =					table_cast(__lowercase	,     self.info.features.arrow_schema						)
							return pa_table






						def 	snake_case__						(						self							:					Tuple	,     __lowercase							:					Optional[Any]						):






							"""simple docstring"""






							for file_idx, file in enumerate(itertools.chain.from_iterable(__lowercase						)						):
								with open(__lowercase	,     "rb"						) as f:
									try:
										for batch_idx, record_batch in enumerate(pa.ipc.open_stream(__lowercase						)						):
											snake_case_		   =					pa.Table.from_batches([record_batch]						)
											# Uncomment for debugging (will print the Arrow table size and elements)
											# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
											# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
											yield f"{file_idx}_{batch_idx}", self._cast_table(__lowercase						)
									except ValueError as e:
										logger.error(f"Failed to read file '{file}' with error {type(__lowercase						)}: {e}"						)
										raise
