





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()
_UpperCamelCase    =	logging.get_logger(__name__)







def  lowerCAmelCase__(		lowercase    :      Optional[int]			)	->      Optional[int]:
  __snake_case				:			str            =    tf.image.resize(lowercase				,       (346, 346)			)
  __snake_case				:			Optional[Any]            =    tf.image.crop_to_bounding_box(lowercase				,       (346 - 289) // 2				,       (346 - 289) // 2				,       289				,       289			)
  return image







def  lowerCAmelCase__(		)	->      Optional[int]:
  __snake_case				:			Any            =    EfficientNetConfig.from_pretrained("google/efficientnet-b7"			)
  __snake_case				:			List[str]            =    289
  __snake_case				:			List[str]            =    640
  __snake_case				:			int            =    {"0": "LABEL_0", "1": "LABEL_1"}
  __snake_case				:			Any            =    {"LABEL_0": 0, "LABEL_1": 1}
  __snake_case				:			Optional[int]            =    []

  __snake_case				:			str            =    BertConfig()
  __snake_case				:			Any            =    AlignConfig.from_text_vision_configs(
      text_config=lowercase				,       vision_config=lowercase				,       projection_dim=640			)
  return config







def  lowerCAmelCase__(		)	->      Optional[Any]:
  __snake_case				:			str            =    "http://images.cocodataset.org/val2017/000000039769.jpg"
  __snake_case				:			Optional[int]            =    Image.open(requests.get(lowercase				,       stream=lowercase			).raw			)
  return im







def  lowerCAmelCase__(		)	->      Optional[int]:
  __snake_case				:			Optional[int]            =    EfficientNetImageProcessor(
      do_center_crop=lowercase				,       rescale_factor=1 / 1_2_7.5				,       rescale_offset=lowercase				,       do_normalize=lowercase				,       include_top=lowercase				,       resample=Image.BILINEAR				,       )
  __snake_case				:			Any            =    BertTokenizer.from_pretrained("bert-base-uncased"			)
  __snake_case				:			Any            =    64
  __snake_case				:			Tuple            =    AlignProcessor(image_processor=lowercase				,       tokenizer=lowercase			)
  return processor







def  lowerCAmelCase__(		lowercase    :      Optional[Any]			)	->      int:
  # EfficientNet image encoder
  __snake_case				:			Any            =    [v.split("_"			)[0].split("block"			)[1] for v in original_param_names if v.startswith("block"			)]
  __snake_case				:			Tuple            =    list(set(lowercase			)			)
  __snake_case				:			Optional[int]            =    sorted(lowercase			)
  __snake_case				:			Union[str, Any]            =    len(lowercase			)
  __snake_case				:			Union[str, Any]            =    {b: str(lowercase			) for b, i in zip(lowercase				,       range(lowercase			)			)}

  __snake_case				:			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:
    __snake_case				:			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""")			)

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

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







def  lowerCAmelCase__(		lowercase    :      Optional[Any]				,       lowercase    :      str				,       lowercase    :      Tuple			)	->      List[Any]:
  list(hf_params.keys()			)

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

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

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







@torch.no_grad()
def  lowerCAmelCase__(		lowercase    :      List[Any]				,       lowercase    :      Dict				,       lowercase    :      Dict				,       lowercase    :      str			)	->      Tuple:
  __snake_case				:			Tuple            =    64
  __snake_case				:			Any            =    Tokenizer(lowercase			)
  __snake_case				:			Any            =    align.Align("efficientnet-b7"				,       "bert-base"				,       640				,       lowercase				,       tok.get_vocab_size()			)
  original_model.compile()
  original_model.load_weights(lowercase			)

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

  # Load HuggingFace model
  __snake_case				:			Any            =    get_align_config()
  __snake_case				:			int            =    AlignModel(lowercase			).eval()
  __snake_case				:			List[Any]            =    hf_model.state_dict()

  # Create src-to-dst parameter name mapping dictionary
  print("Converting parameters..."			)
  __snake_case				:			Union[str, Any]            =    rename_keys(lowercase			)
  replace_params(lowercase				,       lowercase				,       lowercase			)

  # Initialize processor
  __snake_case				:			Tuple            =    get_processor()
  __snake_case				:			int            =    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				:			Optional[Any]            =    hf_model(**lowercase			)

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

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

  __snake_case				:			int            =    original_model.image_encoder(lowercase				,       training=lowercase			)
  __snake_case				:			Dict            =    original_model.text_encoder(lowercase				,       training=lowercase			)

  __snake_case				:			Dict            =    tf.nn.la_normalize(lowercase				,       axis=-1			)
  __snake_case				:			List[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__":
     _UpperCamelCase    =	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''')

     _UpperCamelCase    =	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_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_UpperCamelCase    =	{
    '''configuration_dpr''': ['''DPR_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''DPRConfig'''],
    '''tokenization_dpr''': [
        '''DPRContextEncoderTokenizer''',
        '''DPRQuestionEncoderTokenizer''',
        '''DPRReaderOutput''',
        '''DPRReaderTokenizer''',
    ],
}


try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCamelCase    =	[
         '''DPRContextEncoderTokenizerFast''',
         '''DPRQuestionEncoderTokenizerFast''',
         '''DPRReaderTokenizerFast''',
     ]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCamelCase    =	[
         '''DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''DPRContextEncoder''',
         '''DPRPretrainedContextEncoder''',
         '''DPRPreTrainedModel''',
         '''DPRPretrainedQuestionEncoder''',
         '''DPRPretrainedReader''',
         '''DPRQuestionEncoder''',
         '''DPRReader''',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCamelCase    =	[
         '''TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TFDPRContextEncoder''',
         '''TFDPRPretrainedContextEncoder''',
         '''TFDPRPretrainedQuestionEncoder''',
         '''TFDPRPretrainedReader''',
         '''TFDPRQuestionEncoder''',
         '''TFDPRReader''',
     ]


if TYPE_CHECKING:
     from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
     from .tokenization_dpr import (
         DPRContextEncoderTokenizer,
         DPRQuestionEncoderTokenizer,
         DPRReaderOutput,
         DPRReaderTokenizer,
     )

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_dpr_fast import (
              DPRContextEncoderTokenizerFast,
              DPRQuestionEncoderTokenizerFast,
              DPRReaderTokenizerFast,
          )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_dpr import (
              DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPRContextEncoder,
              DPRPretrainedContextEncoder,
              DPRPreTrainedModel,
              DPRPretrainedQuestionEncoder,
              DPRPretrainedReader,
              DPRQuestionEncoder,
              DPRReader,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_dpr import (
              TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFDPRContextEncoder,
              TFDPRPretrainedContextEncoder,
              TFDPRPretrainedQuestionEncoder,
              TFDPRPretrainedReader,
              TFDPRQuestionEncoder,
              TFDPRReader,
          )

else:
     import sys

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

