

import argparse

from . import (
    BertConfig,
    BertGenerationConfig,
    BertGenerationDecoder,
    BertGenerationEncoder,
    load_tf_weights_in_bert_generation,
    logging,
)


logging.set_verbosity_info()





def 	snake_case_      (__A     :						Tuple					,   __A     :						Union[str, Any]					,   __A     :						str					,   __A     :						int					,   __A     :						List[Any]     )			->	Optional[Any]:
   # Initialise PyTorch model
   __lowerCAmelCase  :			Optional[Any]					    =						BertConfig.from_pretrained(
       """bert-large-cased"""					,   vocab_size=__A					,   max_position_embeddings=5_1_2					,   is_decoder=__A					,   add_cross_attention=__A					,   )
   __lowerCAmelCase  :			Optional[int]					    =						bert_config.to_dict()
   del bert_config_dict["type_vocab_size"]
   __lowerCAmelCase  :			Optional[Any]					    =						BertGenerationConfig(**__A     )
   if is_encoder:
      __lowerCAmelCase  :			List[Any]					    =						BertGenerationEncoder(__A     )
   else:
      __lowerCAmelCase  :			int					    =						BertGenerationDecoder(__A     )
   print(f'''Building PyTorch model from configuration: {config}'''     )

   # Load weights from tf checkpoint
   load_tf_weights_in_bert_generation(
       __A					,   __A					,   model_class="""bert"""					,   is_encoder_named_decoder=__A					,   is_encoder=__A					,   )

   # Save pytorch-model
   print(f'''Save PyTorch model and config to {pytorch_dump_path}'''     )
   model.save_pretrained(__A     )


if __name__ == "__main__":
   __UpperCAmelCase      = argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--tf_hub_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
   )
   parser.add_argument(
       """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
   )
   parser.add_argument(
       """--is_encoder_named_decoder""",
       action="""store_true""",
       help="""If decoder has to be renamed to encoder in PyTorch model.""",
   )
   parser.add_argument("""--is_encoder""", action="""store_true""", help="""If model is an encoder.""")
   parser.add_argument("""--vocab_size""", default=50_358, type=int, help="""Vocab size of model""")
   __UpperCAmelCase      = parser.parse_args()
   convert_tf_checkpoint_to_pytorch(
       args.tf_hub_path,
       args.pytorch_dump_path,
       args.is_encoder_named_decoder,
       args.vocab_size,
       is_encoder=args.is_encoder,
   )





import importlib.metadata
import subprocess
import sys





def 	snake_case_      (__A     :						bool = False     )			->	str:
   __lowerCAmelCase  :			Any					    =						False
   if "IPython" in sys.modules:
      __lowerCAmelCase  :			Optional[int]					    =						"""google.colab""" in str(sys.modules["""IPython"""].get_ipython()     )

   if in_colab:
      if upgrade:
         __lowerCAmelCase  :			int					    =						["""pip""", """install""", """-U""", """torch"""]
         subprocess.run(__A					,   check=__A     )
      # get the current version of torch
      __lowerCAmelCase  :			List[Any]					    =						importlib.metadata.version("""torch"""     )
      __lowerCAmelCase  :			List[Any]					    =						torch_version[: torch_version.rindex("""."""     )]
      __lowerCAmelCase  :			int					    =						f'''https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl'''
      __lowerCAmelCase  :			List[Any]					    =						["""pip""", """install""", xla_wheel]
      subprocess.run(__A					,   check=__A     )
   else:
      raise RuntimeError("""`install_xla` utility works only on google colab."""     )


