

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()
lowercase_									=   logging.get_logger(__name__)
def      a      (			A__	:   Optional[Any]      ,   A__	:   str      ,   A__	:   Optional[Any]     )    ->      Any:





      """simple docstring"""



      _lowercase	       =os.path.abspath(A__     )
      logger.info(F'''Converting TensorFlow checkpoint from {tf_path}'''     )
      # Load weights from TF model
      _lowercase	       =tf.train.list_variables(A__     )
      _lowercase	       =[]
      _lowercase	       =[]
      _lowercase	       =[]
      for full_name, shape in init_vars:
            # logger.info(f"Loading TF weight {name} with shape {shape}")
            _lowercase	       =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'
                  _lowercase	       =name[1:]
            # figure out how many levels deep the name is
            _lowercase	       =0
            for _name in name:
                  if _name.startswith('layer_with_weights'     ):
                        depth += 1
                  else:
                        break
            layer_depth.append(A__     )
            # read data
            _lowercase	       =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__     )     )})'''     )
      _lowercase	       =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__     ):
            _lowercase	       =full_name.split('/'     )
            _lowercase	       =model
            _lowercase	       =[]
            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'     ):
                        _lowercase	       =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']     )
                              _lowercase	       =getattr(A__      ,   'embeddings'     )
                              _lowercase	       =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     )]     )
                              _lowercase	       =getattr(A__      ,   'encoder'     )
                              _lowercase	       =getattr(A__      ,   'layer'     )
                              _lowercase	       =pointer[layer_num - 4]
                        elif layer_num == config.num_hidden_layers + 4:
                              # pooler layer
                              trace.extend(['pooler', 'dense']     )
                              _lowercase	       =getattr(A__      ,   'pooler'     )
                              _lowercase	       =getattr(A__      ,   'dense'     )
                  elif m_name == "embeddings":
                        trace.append('embeddings'     )
                        _lowercase	       =getattr(A__      ,   'embeddings'     )
                        if layer_num == 0:
                              trace.append('word_embeddings'     )
                              _lowercase	       =getattr(A__      ,   'word_embeddings'     )
                        elif layer_num == 1:
                              trace.append('position_embeddings'     )
                              _lowercase	       =getattr(A__      ,   'position_embeddings'     )
                        elif layer_num == 2:
                              trace.append('token_type_embeddings'     )
                              _lowercase	       =getattr(A__      ,   'token_type_embeddings'     )
                        else:
                              raise ValueError(F'''Unknown embedding layer with name {full_name}'''     )
                        trace.append('weight'     )
                        _lowercase	       =getattr(A__      ,   'weight'     )
                  elif m_name == "_attention_layer":
                        # self-attention layer
                        trace.extend(['attention', 'self']     )
                        _lowercase	       =getattr(A__      ,   'attention'     )
                        _lowercase	       =getattr(A__      ,   'self'     )
                  elif m_name == "_attention_layer_norm":
                        # output attention norm
                        trace.extend(['attention', 'output', 'LayerNorm']     )
                        _lowercase	       =getattr(A__      ,   'attention'     )
                        _lowercase	       =getattr(A__      ,   'output'     )
                        _lowercase	       =getattr(A__      ,   'LayerNorm'     )
                  elif m_name == "_attention_output_dense":
                        # output attention dense
                        trace.extend(['attention', 'output', 'dense']     )
                        _lowercase	       =getattr(A__      ,   'attention'     )
                        _lowercase	       =getattr(A__      ,   'output'     )
                        _lowercase	       =getattr(A__      ,   'dense'     )
                  elif m_name == "_output_dense":
                        # output dense
                        trace.extend(['output', 'dense']     )
                        _lowercase	       =getattr(A__      ,   'output'     )
                        _lowercase	       =getattr(A__      ,   'dense'     )
                  elif m_name == "_output_layer_norm":
                        # output dense
                        trace.extend(['output', 'LayerNorm']     )
                        _lowercase	       =getattr(A__      ,   'output'     )
                        _lowercase	       =getattr(A__      ,   'LayerNorm'     )
                  elif m_name == "_key_dense":
                        # attention key
                        trace.append('key'     )
                        _lowercase	       =getattr(A__      ,   'key'     )
                  elif m_name == "_query_dense":
                        # attention query
                        trace.append('query'     )
                        _lowercase	       =getattr(A__      ,   'query'     )
                  elif m_name == "_value_dense":
                        # attention value
                        trace.append('value'     )
                        _lowercase	       =getattr(A__      ,   'value'     )
                  elif m_name == "_intermediate_dense":
                        # attention intermediate dense
                        trace.extend(['intermediate', 'dense']     )
                        _lowercase	       =getattr(A__      ,   'intermediate'     )
                        _lowercase	       =getattr(A__      ,   'dense'     )
                  elif m_name == "_output_layer_norm":
                        # output layer norm
                        trace.append('output'     )
                        _lowercase	       =getattr(A__      ,   'output'     )
                  # weights & biases
                  elif m_name in ["bias", "beta"]:
                        trace.append('bias'     )
                        _lowercase	       =getattr(A__      ,   'bias'     )
                  elif m_name in ["kernel", "gamma"]:
                        trace.append('weight'     )
                        _lowercase	       =getattr(A__      ,   'weight'     )
                  else:
                        logger.warning(F'''Ignored {m_name}'''     )
        # for certain layers reshape is necessary
            _lowercase	       ='.'.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__     ):
                  _lowercase	       =array.reshape(pointer.data.shape     )
            if "kernel" in full_name:
                  _lowercase	       =array.transpose()
            if pointer.shape == array.shape:
                  _lowercase	       =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__	:   Optional[Any]      ,   A__	:   List[Any]      ,   A__	:   List[str]     )    ->      Optional[int]:





      """simple docstring"""



      logger.info(F'''Loading model based on config from {config_path}...'''     )
      _lowercase	       =BertConfig.from_json_file(A__     )
      _lowercase	       =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__":
  lowercase_									=   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).',
  )
  lowercase_									=   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.

from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


lowercase_									=   'Create a default config file for Accelerate with only a few flags set.'
def      a      (			A__	:   Optional[Any]="no"      ,   A__	:   str = default_json_config_file      ,   A__	:   bool = False     )    ->      Optional[int]:





      """simple docstring"""



      _lowercase	       =Path(A__     )
      path.parent.mkdir(parents=A__      ,   exist_ok=A__     )
      if path.exists():
            print(
                F'''Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.'''     )
            return False
      _lowercase	       =mixed_precision.lower()
      if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
            raise ValueError(
                F'''`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}'''     )
      _lowercase	       ={
          'compute_environment': 'LOCAL_MACHINE',
          'mixed_precision': mixed_precision,
      }
      if torch.cuda.is_available():
            _lowercase	       =torch.cuda.device_count()
            _lowercase	       =num_gpus
            _lowercase	       =False
            if num_gpus > 1:
                  _lowercase	       ='MULTI_GPU'
            else:
                  _lowercase	       ='NO'
      elif is_xpu_available() and use_xpu:
            _lowercase	       =torch.xpu.device_count()
            _lowercase	       =num_xpus
            _lowercase	       =False
            if num_xpus > 1:
                  _lowercase	       ='MULTI_XPU'
            else:
                  _lowercase	       ='NO'
      elif is_npu_available():
            _lowercase	       =torch.npu.device_count()
            _lowercase	       =num_npus
            _lowercase	       =False
            if num_npus > 1:
                  _lowercase	       ='MULTI_NPU'
            else:
                  _lowercase	       ='NO'
      else:
            _lowercase	       =0
            _lowercase	       =True
            _lowercase	       =1
            _lowercase	       ='NO'
      _lowercase	       =ClusterConfig(**A__     )
      config.to_json_file(A__     )
      return path
def      a      (			A__	:   Dict      ,   A__	:   Optional[Any]     )    ->      List[Any]:





      """simple docstring"""



      _lowercase	       =parser.add_parser('default'      ,   parents=A__      ,   help=A__      ,   formatter_class=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\'.'
          )      ,   dest='save_location'      ,   )

      parser.add_argument(
          '--mixed_precision'      ,   choices=['no', 'fp16', 'bf16']      ,   type=A__      ,   help='Whether or not to use mixed precision training. '
          'Choose between FP16 and BF16 (bfloat16) training. '
          'BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.'      ,   default='no'      ,   )
      parser.set_defaults(func=A__     )
      return parser







def      a      (			A__	:   List[str]     )    ->      Any:





      """simple docstring"""



      _lowercase	       =write_basic_config(args.mixed_precision      ,   args.save_location     )
      if config_file:
            print(F'''accelerate configuration saved at {config_file}'''     )


