

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



def    snake_case_(_UpperCamelCase  )							->			Optional[Any]:



  """simple docstring"""






  _snake_case		=							tf.image.resize(SCREAMING_SNAKE_CASE__				, (346, 346)  )
  _snake_case		=							tf.image.crop_to_bounding_box(SCREAMING_SNAKE_CASE__				, (346 - 289) // 2				, (346 - 289) // 2				, 289				, 289  )
  return image



def    snake_case_()							->			int:



  """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=SCREAMING_SNAKE_CASE__				, vision_config=SCREAMING_SNAKE_CASE__				, projection_dim=640  )
  return config



def    snake_case_()							->			List[Any]:



  """simple docstring"""






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



def    snake_case_()							->			int:



  """simple docstring"""






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



def    snake_case_(_UpperCamelCase  )							->			List[Any]:



  """simple docstring"""






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

  _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    snake_case_(_UpperCamelCase				, _UpperCamelCase				, _UpperCamelCase  )							->			Optional[int]:



  """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(SCREAMING_SNAKE_CASE__  ).permute(3				, 2				, 0				, 1  )
    elif "embeddings" in key:
      _snake_case		=							torch.from_numpy(SCREAMING_SNAKE_CASE__  )
    elif "depthwise_kernel" in key:
      _snake_case		=							torch.from_numpy(SCREAMING_SNAKE_CASE__  ).permute(2				, 3				, 0				, 1  )
    elif "kernel" in key:
      _snake_case		=							torch.from_numpy(np.transpose(SCREAMING_SNAKE_CASE__  )  )
    elif "temperature" in key:
      _snake_case		=							value
    elif "bn/gamma" or "bn/beta" in key:
      _snake_case		=							torch.from_numpy(np.transpose(SCREAMING_SNAKE_CASE__  )  ).squeeze()
    else:
      _snake_case		=							torch.from_numpy(SCREAMING_SNAKE_CASE__  )

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



@torch.no_grad()
def    snake_case_(_UpperCamelCase				, _UpperCamelCase				, _UpperCamelCase				, _UpperCamelCase  )							->			int:



  """simple docstring"""






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

  _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(SCREAMING_SNAKE_CASE__  ).eval()
  _snake_case		=							hf_model.state_dict()

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

  # 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(**SCREAMING_SNAKE_CASE__  )

  _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=SCREAMING_SNAKE_CASE__				, do_rescale=SCREAMING_SNAKE_CASE__				, do_normalize=SCREAMING_SNAKE_CASE__				, include_top=SCREAMING_SNAKE_CASE__				, 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(SCREAMING_SNAKE_CASE__				, training=SCREAMING_SNAKE_CASE__  )
  _snake_case		=							original_model.text_encoder(SCREAMING_SNAKE_CASE__				, training=SCREAMING_SNAKE_CASE__  )

  _snake_case		=							tf.nn.la_normalize(SCREAMING_SNAKE_CASE__				, axis=-1  )
  _snake_case		=							tf.nn.la_normalize(SCREAMING_SNAKE_CASE__				, axis=-1  )

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

  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    	=			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    	=			parser.parse_args()
   convert_align_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)




import random
import sys

from . import cryptomath_module as cryptomath

__A    	=			(
    r''' !"#$%&\'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_`'''
    r'''abcdefghijklmnopqrstuvwxyz{|}~'''
)



def    snake_case_(_UpperCamelCase				, _UpperCamelCase				, _UpperCamelCase  )							->			None:



  """simple docstring"""






  if mode == "encrypt":
    if key_a == 1:
      sys.exit(
          '''The affine cipher becomes weak when key '''
          '''A is set to 1. Choose different key'''  )
    if key_b == 0:
      sys.exit(
          '''The affine cipher becomes weak when key '''
          '''B is set to 0. Choose different key'''  )
  if key_a < 0 or key_b < 0 or key_b > len(_UpperCamelCase  ) - 1:
    sys.exit(
        '''Key A must be greater than 0 and key B must '''
        F"""be between 0 and {len(_UpperCamelCase  ) - 1}."""  )
  if cryptomath.gcd(_UpperCamelCase				, len(_UpperCamelCase  )  ) != 1:
    sys.exit(
        F"""Key A {key_a} and the symbol set size {len(_UpperCamelCase  )} """
        '''are not relatively prime. Choose a different key.'''  )



def    snake_case_(_UpperCamelCase				, _UpperCamelCase  )							->			str:



  """simple docstring"""






  _snake_case,  _snake_case		=							divmod(_UpperCamelCase				, len(_UpperCamelCase  )  )
  check_keys(_UpperCamelCase				, _UpperCamelCase				, '''encrypt'''  )
  _snake_case		=							''''''
  for symbol in message:
    if symbol in SYMBOLS:
      _snake_case		=							SYMBOLS.find(_UpperCamelCase  )
      cipher_text += SYMBOLS[(sym_index * key_a + key_b) % len(_UpperCamelCase  )]
    else:
      cipher_text += symbol
  return cipher_text



def    snake_case_(_UpperCamelCase				, _UpperCamelCase  )							->			str:



  """simple docstring"""






  _snake_case,  _snake_case		=							divmod(_UpperCamelCase				, len(_UpperCamelCase  )  )
  check_keys(_UpperCamelCase				, _UpperCamelCase				, '''decrypt'''  )
  _snake_case		=							''''''
  _snake_case		=							cryptomath.find_mod_inverse(_UpperCamelCase				, len(_UpperCamelCase  )  )
  for symbol in message:
    if symbol in SYMBOLS:
      _snake_case		=							SYMBOLS.find(_UpperCamelCase  )
      plain_text += SYMBOLS[
          (sym_index - key_b) * mod_inverse_of_key_a % len(_UpperCamelCase  )
      ]
    else:
      plain_text += symbol
  return plain_text



def    snake_case_()							->			int:



  """simple docstring"""






  while True:
    _snake_case		=							random.randint(2				, len(_UpperCamelCase  )  )
    _snake_case		=							random.randint(2				, len(_UpperCamelCase  )  )
    if cryptomath.gcd(_UpperCamelCase				, len(_UpperCamelCase  )  ) == 1 and key_b % len(_UpperCamelCase  ) != 0:
      return key_b * len(_UpperCamelCase  ) + key_b



def    snake_case_()							->			None:



  """simple docstring"""






  _snake_case		=							input('''Enter message: '''  ).strip()
  _snake_case		=							int(input('''Enter key [2000 - 9000]: '''  ).strip()  )
  _snake_case		=							input('''Encrypt/Decrypt [E/D]: '''  ).strip().lower()

  if mode.startswith('''e'''  ):
    _snake_case		=							'''encrypt'''
    _snake_case		=							encrypt_message(_UpperCamelCase				, _UpperCamelCase  )
  elif mode.startswith('''d'''  ):
    _snake_case		=							'''decrypt'''
    _snake_case		=							decrypt_message(_UpperCamelCase				, _UpperCamelCase  )
  print(F"""\n{mode.title()}ed text: \n{translated}"""  )


if __name__ == "__main__":
 import doctest

 doctest.testmod()
 # main()

