





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_       (							UpperCamelCase_     )   ->						Union[str, Any]:
       UpperCamelCase_      				=     tf.image.resize(UpperCamelCase_    ,						(346, 346)     )
       UpperCamelCase_      				=     tf.image.crop_to_bounding_box(UpperCamelCase_    ,						(346 - 289) // 2    ,						(346 - 289) // 2    ,						289    ,						289     )
       return image




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

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




def        lowerCAmelCase_       (							)   ->						Any:
       UpperCamelCase_      				=     "http://images.cocodataset.org/val2017/000000039769.jpg"
       UpperCamelCase_      				=     Image.open(requests.get(UpperCamelCase_    ,						stream=UpperCamelCase_     ).raw     )
       return im




def        lowerCAmelCase_       (							)   ->						Tuple:
       UpperCamelCase_      				=     EfficientNetImageProcessor(
           do_center_crop=UpperCamelCase_    ,						rescale_factor=1 / 1_27.5    ,						rescale_offset=UpperCamelCase_    ,						do_normalize=UpperCamelCase_    ,						include_top=UpperCamelCase_    ,						resample=Image.BILINEAR    ,						)
       UpperCamelCase_      				=     BertTokenizer.from_pretrained("bert-base-uncased"     )
       UpperCamelCase_      				=     64
       UpperCamelCase_      				=     AlignProcessor(image_processor=UpperCamelCase_    ,						tokenizer=UpperCamelCase_     )
       return processor




def        lowerCAmelCase_       (							UpperCamelCase_     )   ->						Dict:
       # EfficientNet image encoder
       UpperCamelCase_      				=     [v.split("_"     )[0].split("block"     )[1] for v in original_param_names if v.startswith("block"     )]
       UpperCamelCase_      				=     list(set(UpperCamelCase_     )     )
       UpperCamelCase_      				=     sorted(UpperCamelCase_     )
       UpperCamelCase_      				=     len(UpperCamelCase_     )
       UpperCamelCase_      				=     {b: str(UpperCamelCase_     ) for b, i in zip(UpperCamelCase_    ,						range(UpperCamelCase_     )     )}

       UpperCamelCase_      				=     []
       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:
              UpperCamelCase_      				=     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''')     )

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

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




def        lowerCAmelCase_       (							UpperCamelCase_    ,						UpperCamelCase_    ,						UpperCamelCase_     )   ->						Dict:
       list(hf_params.keys()     )

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

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

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




@torch.no_grad()
def        lowerCAmelCase_       (							UpperCamelCase_    ,						UpperCamelCase_    ,						UpperCamelCase_    ,						UpperCamelCase_     )   ->						Optional[int]:
       UpperCamelCase_      				=     64
       UpperCamelCase_      				=     Tokenizer(UpperCamelCase_     )
       UpperCamelCase_      				=     align.Align("efficientnet-b7"    ,						"bert-base"    ,						640    ,						UpperCamelCase_    ,						tok.get_vocab_size()     )
       original_model.compile()
       original_model.load_weights(UpperCamelCase_     )

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

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

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

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

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

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

       UpperCamelCase_      				=     original_model.image_encoder(UpperCamelCase_    ,						training=UpperCamelCase_     )
       UpperCamelCase_      				=     original_model.text_encoder(UpperCamelCase_    ,						training=UpperCamelCase_     )

       UpperCamelCase_      				=     tf.nn.la_normalize(UpperCamelCase_    ,						axis=-1     )
       UpperCamelCase_      				=     tf.nn.la_normalize(UpperCamelCase_    ,						axis=-1     )

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

       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)








def        lowerCAmelCase_       (							UpperCamelCase_    ,						UpperCamelCase_     )   ->						list:
       UpperCamelCase_      				=     sorted(UpperCamelCase_    ,						key=lambda UpperCamelCase_     : value[2]    ,						reverse=UpperCamelCase_     )

       # Create a list of size equal to the maximum deadline
       # and initialize it with -1
       UpperCamelCase_      				=     max(UpperCamelCase_    ,						key=lambda UpperCamelCase_     : value[1]     )[1]
       UpperCamelCase_      				=     [-1] * max_deadline

       # Finding the maximum profit and the count of jobs
       UpperCamelCase_      				=     0
       UpperCamelCase_      				=     0
       for job in jobs:
              # Find a free time slot for this job
              # (Note that we start from the last possible slot)
              for i in range(job[1] - 1    ,						-1    ,						-1     ):
                     if time_slots[i] == -1:
                            UpperCamelCase_      				=     job[0]
                            count += 1
                            max_profit += job[2]
                            break
       return [count, max_profit]


if __name__ == "__main__":
     import doctest

     doctest.testmod()

