
import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
__lowerCAmelCase     :	int       			=  logging.get_logger(__name__)

def a__       (			A_,	A_,	A_   ):






  '''simple docstring'''





  __magic_name__								=			os.path.abspath(A_   )
  logger.info(f'''Converting TensorFlow checkpoint from {tf_path}'''   )
  # Load weights from TF model
  __magic_name__								=			tf.train.list_variables(A_   )
  __magic_name__								=			[]
  __magic_name__								=			[]
  __magic_name__								=			[]
  for full_name, shape in init_vars:
    # logger.info(f"Loading TF weight {name} with shape {shape}")
    __magic_name__								=			full_name.split("""/"""   )
    if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
      logger.info(f'''Skipping non-model layer {full_name}'''   )
      continue
    if "optimizer" in full_name:
      logger.info(f'''Skipping optimization layer {full_name}'''   )
      continue
    if name[0] == "model":
      # ignore initial 'model'
      __magic_name__								=			name[1:]
    # figure out how many levels deep the name is
    __magic_name__								=			0
    for _name in name:
      if _name.startswith("""layer_with_weights"""   ):
        depth += 1
      else:
        break
    layer_depth.append(A_   )
    # read data
    __magic_name__								=			tf.train.load_variable(A_,	A_   )
    names.append("""/""".join(A_   )   )
    arrays.append(A_   )
  logger.info(f'''Read a total of {len(A_   ):,} layers'''   )

  # Sanity check
  if len(set(A_   )   ) != 1:
    raise ValueError(f'''Found layer names with different depths (layer depth {list(set(A_   )   )})'''   )
  __magic_name__								=			list(set(A_   )   )[0]
  if layer_depth != 1:
    raise ValueError(
        """The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP"""
        """ heads."""   )

  # convert layers
  logger.info("""Converting weights..."""   )
  for full_name, array in zip(A_,	A_   ):
    __magic_name__								=			full_name.split("""/"""   )
    __magic_name__								=			model
    __magic_name__								=			[]
    for i, m_name in enumerate(A_   ):
      if m_name == ".ATTRIBUTES":
        # variable names end with .ATTRIBUTES/VARIABLE_VALUE
        break
      if m_name.startswith("""layer_with_weights"""   ):
        __magic_name__								=			int(m_name.split("""-"""   )[-1]   )
        if layer_num <= 2:
          # embedding layers
          # layer_num 0: word_embeddings
          # layer_num 1: position_embeddings
          # layer_num 2: token_type_embeddings
          continue
        elif layer_num == 3:
          # embedding LayerNorm
          trace.extend(["""embeddings""", """LayerNorm"""]   )
          __magic_name__								=			getattr(A_,	"""embeddings"""   )
          __magic_name__								=			getattr(A_,	"""LayerNorm"""   )
        elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
          # encoder layers
          trace.extend(["""encoder""", """layer""", str(layer_num - 4   )]   )
          __magic_name__								=			getattr(A_,	"""encoder"""   )
          __magic_name__								=			getattr(A_,	"""layer"""   )
          __magic_name__								=			pointer[layer_num - 4]
        elif layer_num == config.num_hidden_layers + 4:
          # pooler layer
          trace.extend(["""pooler""", """dense"""]   )
          __magic_name__								=			getattr(A_,	"""pooler"""   )
          __magic_name__								=			getattr(A_,	"""dense"""   )
      elif m_name == "embeddings":
        trace.append("""embeddings"""   )
        __magic_name__								=			getattr(A_,	"""embeddings"""   )
        if layer_num == 0:
          trace.append("""word_embeddings"""   )
          __magic_name__								=			getattr(A_,	"""word_embeddings"""   )
        elif layer_num == 1:
          trace.append("""position_embeddings"""   )
          __magic_name__								=			getattr(A_,	"""position_embeddings"""   )
        elif layer_num == 2:
          trace.append("""token_type_embeddings"""   )
          __magic_name__								=			getattr(A_,	"""token_type_embeddings"""   )
        else:
          raise ValueError(f'''Unknown embedding layer with name {full_name}'''   )
        trace.append("""weight"""   )
        __magic_name__								=			getattr(A_,	"""weight"""   )
      elif m_name == "_attention_layer":
        # self-attention layer
        trace.extend(["""attention""", """self"""]   )
        __magic_name__								=			getattr(A_,	"""attention"""   )
        __magic_name__								=			getattr(A_,	"""self"""   )
      elif m_name == "_attention_layer_norm":
        # output attention norm
        trace.extend(["""attention""", """output""", """LayerNorm"""]   )
        __magic_name__								=			getattr(A_,	"""attention"""   )
        __magic_name__								=			getattr(A_,	"""output"""   )
        __magic_name__								=			getattr(A_,	"""LayerNorm"""   )
      elif m_name == "_attention_output_dense":
        # output attention dense
        trace.extend(["""attention""", """output""", """dense"""]   )
        __magic_name__								=			getattr(A_,	"""attention"""   )
        __magic_name__								=			getattr(A_,	"""output"""   )
        __magic_name__								=			getattr(A_,	"""dense"""   )
      elif m_name == "_output_dense":
        # output dense
        trace.extend(["""output""", """dense"""]   )
        __magic_name__								=			getattr(A_,	"""output"""   )
        __magic_name__								=			getattr(A_,	"""dense"""   )
      elif m_name == "_output_layer_norm":
        # output dense
        trace.extend(["""output""", """LayerNorm"""]   )
        __magic_name__								=			getattr(A_,	"""output"""   )
        __magic_name__								=			getattr(A_,	"""LayerNorm"""   )
      elif m_name == "_key_dense":
        # attention key
        trace.append("""key"""   )
        __magic_name__								=			getattr(A_,	"""key"""   )
      elif m_name == "_query_dense":
        # attention query
        trace.append("""query"""   )
        __magic_name__								=			getattr(A_,	"""query"""   )
      elif m_name == "_value_dense":
        # attention value
        trace.append("""value"""   )
        __magic_name__								=			getattr(A_,	"""value"""   )
      elif m_name == "_intermediate_dense":
        # attention intermediate dense
        trace.extend(["""intermediate""", """dense"""]   )
        __magic_name__								=			getattr(A_,	"""intermediate"""   )
        __magic_name__								=			getattr(A_,	"""dense"""   )
      elif m_name == "_output_layer_norm":
        # output layer norm
        trace.append("""output"""   )
        __magic_name__								=			getattr(A_,	"""output"""   )
      # weights & biases
      elif m_name in ["bias", "beta"]:
        trace.append("""bias"""   )
        __magic_name__								=			getattr(A_,	"""bias"""   )
      elif m_name in ["kernel", "gamma"]:
        trace.append("""weight"""   )
        __magic_name__								=			getattr(A_,	"""weight"""   )
      else:
        logger.warning(f'''Ignored {m_name}'''   )
        # for certain layers reshape is necessary
    __magic_name__								=			""".""".join(A_   )
    if re.match(R"""(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)""",	A_   ) or re.match(
        R"""(\S+)\.attention\.output\.dense\.weight""",	A_   ):
      __magic_name__								=			array.reshape(pointer.data.shape   )
    if "kernel" in full_name:
      __magic_name__								=			array.transpose()
    if pointer.shape == array.shape:
      __magic_name__								=			torch.from_numpy(A_   )
    else:
      raise ValueError(
          f'''Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:'''
          f''' {array.shape}'''   )
    logger.info(f'''Successfully set variable {full_name} to PyTorch layer {trace}'''   )
  return model

def a__       (			A_,	A_,	A_   ):






  '''simple docstring'''





  logger.info(f'''Loading model based on config from {config_path}...'''   )
  __magic_name__								=			BertConfig.from_json_file(A_   )
  __magic_name__								=			BertModel(A_   )

  # Load weights from checkpoint
  logger.info(f'''Loading weights from checkpoint {tf_checkpoint_path}...'''   )
  load_tfa_weights_in_bert(A_,	A_,	A_   )

  # Save pytorch-model
  logger.info(f'''Saving PyTorch model to {pytorch_dump_path}...'''   )
  torch.save(model.state_dict(),	A_   )


if __name__ == "__main__":
    __lowerCAmelCase     :	Dict       			=  argparse.ArgumentParser()
    parser.add_argument(
        '--tf_checkpoint_path', type=str, required=True, help='Path to the TensorFlow 2.x checkpoint path.'
    )
    parser.add_argument(
        '--bert_config_file',
        type=str,
        required=True,
        help='The config json file corresponding to the BERT model. This specifies the model architecture.',
    )
    parser.add_argument(
        '--pytorch_dump_path',
        type=str,
        required=True,
        help='Path to the output PyTorch model (must include filename).',
    )
    __lowerCAmelCase     :	Tuple       			=  parser.parse_args()
    convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)



# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os

from accelerate.utils import ComputeEnvironment

from .cluster import get_cluster_input
from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
from .sagemaker import get_sagemaker_input


__lowerCAmelCase     :	Optional[Any]       			=  'Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine'

def a__       (			):






  '''simple docstring'''





  __magic_name__								=			_ask_options(
      """In which compute environment are you running?""",	["""This machine""", """AWS (Amazon SageMaker)"""],	_convert_compute_environment,	)
  if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
    __magic_name__								=			get_sagemaker_input()
  else:
    __magic_name__								=			get_cluster_input()
  return config

def a__       (			A_=None   ):






  '''simple docstring'''





  if subparsers is not None:
    __magic_name__								=			subparsers.add_parser("""config""",	description=A_   )
  else:
    __magic_name__								=			argparse.ArgumentParser("""Accelerate config command""",	description=A_   )

  parser.add_argument(
      """--config_file""",	default=A_,	help=(
          """The path to use to store the config file. Will default to a file named default_config.yaml in the cache """
          """location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have """
          """such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed """
          """with 'huggingface'."""
      ),	)

  if subparsers is not None:
    parser.set_defaults(func=A_   )
  return parser

def a__       (			A_   ):






  '''simple docstring'''





  __magic_name__								=			get_user_input()
  if args.config_file is not None:
    __magic_name__								=			args.config_file
  else:
    if not os.path.isdir(A_   ):
      os.makedirs(A_   )
    __magic_name__								=			default_yaml_config_file

  if config_file.endswith(""".json"""   ):
    config.to_json_file(A_   )
  else:
    config.to_yaml_file(A_   )
  print(f'''accelerate configuration saved at {config_file}'''   )

def a__       (			):






  '''simple docstring'''





  __magic_name__								=			config_command_parser()
  __magic_name__								=			parser.parse_args()
  config_command(A_   )


if __name__ == "__main__":
    main()

