

import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
def      a      (			A__	:   Union[str, Any]      ,   A__	:   int      ,   A__	:   Optional[Any]      ,   A__	:   Optional[int]="attention"     )    ->      Optional[int]:





      """simple docstring"""



      _lowercase	       =params[F'''{prefix}/layers_{i}/{layer_name}/key/kernel''']
      _lowercase	       =params[F'''{prefix}/layers_{i}/{layer_name}/out/kernel''']
      _lowercase	       =params[F'''{prefix}/layers_{i}/{layer_name}/query/kernel''']
      _lowercase	       =params[F'''{prefix}/layers_{i}/{layer_name}/value/kernel''']
      return k, o, q, v
def      a      (			A__	:   Any      ,   A__	:   List[str]      ,   A__	:   str      ,   A__	:   Optional[Any]=False     )    ->      Optional[int]:





      """simple docstring"""



      if split_mlp_wi:
            _lowercase	       =params[F'''{prefix}/layers_{i}/mlp/wi_0/kernel''']
            _lowercase	       =params[F'''{prefix}/layers_{i}/mlp/wi_1/kernel''']
            _lowercase	       =(wi_a, wi_a)
      else:
            _lowercase	       =params[F'''{prefix}/layers_{i}/mlp/wi/kernel''']

      _lowercase	       =params[F'''{prefix}/layers_{i}/mlp/wo/kernel''']
      return wi, wo
def      a      (			A__	:   Dict      ,   A__	:   Tuple      ,   A__	:   Union[str, Any]      ,   A__	:   Tuple     )    ->      Any:





      """simple docstring"""



      return params[F'''{prefix}/layers_{i}/{layer_name}/scale''']
def      a      (			A__	:   dict      ,   *, A__	:   int      ,   A__	:   bool     )    ->      str:





      """simple docstring"""



      _lowercase	       =traverse_util.flatten_dict(variables['target']     )
      _lowercase	       ={'/'.join(A__     ): v for k, v in old.items()}

      # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
      _lowercase	       ='encoder/layers_0/mlp/wi_0/kernel' in old
      print('Split MLP:'      ,   A__     )

      _lowercase	       =collections.OrderedDict()

      # Shared embeddings.
      _lowercase	       =old['token_embedder/embedding']

      # Encoder.
      for i in range(A__     ):
            # Block i, layer 0 (Self Attention).
            _lowercase	       =tax_layer_norm_lookup(A__      ,   A__      ,   'encoder'      ,   'pre_attention_layer_norm'     )
            _lowercase       ,       _lowercase       ,       _lowercase       ,       _lowercase	       =tax_attention_lookup(A__      ,   A__      ,   'encoder'      ,   'attention'     )
            _lowercase	       =layer_norm
            _lowercase	       =k.T
            _lowercase	       =o.T
            _lowercase	       =q.T
            _lowercase	       =v.T

            # Block i, layer 1 (MLP).
            _lowercase	       =tax_layer_norm_lookup(A__      ,   A__      ,   'encoder'      ,   'pre_mlp_layer_norm'     )
            _lowercase       ,       _lowercase	       =tax_mlp_lookup(A__      ,   A__      ,   'encoder'      ,   A__     )
            _lowercase	       =layer_norm
            if split_mlp_wi:
                  _lowercase	       =wi[0].T
                  _lowercase	       =wi[1].T
            else:
                  _lowercase	       =wi.T
            _lowercase	       =wo.T

      _lowercase	       =old[
          'encoder/relpos_bias/rel_embedding'
      ].T
      _lowercase	       =old['encoder/encoder_norm/scale']

      if not is_encoder_only:
            # Decoder.
            for i in range(A__     ):
                  # Block i, layer 0 (Self Attention).
                  _lowercase	       =tax_layer_norm_lookup(A__      ,   A__      ,   'decoder'      ,   'pre_self_attention_layer_norm'     )
                  _lowercase       ,       _lowercase       ,       _lowercase       ,       _lowercase	       =tax_attention_lookup(A__      ,   A__      ,   'decoder'      ,   'self_attention'     )
                  _lowercase	       =layer_norm
                  _lowercase	       =k.T
                  _lowercase	       =o.T
                  _lowercase	       =q.T
                  _lowercase	       =v.T

                  # Block i, layer 1 (Cross Attention).
                  _lowercase	       =tax_layer_norm_lookup(A__      ,   A__      ,   'decoder'      ,   'pre_cross_attention_layer_norm'     )
                  _lowercase       ,       _lowercase       ,       _lowercase       ,       _lowercase	       =tax_attention_lookup(A__      ,   A__      ,   'decoder'      ,   'encoder_decoder_attention'     )
                  _lowercase	       =layer_norm
                  _lowercase	       =k.T
                  _lowercase	       =o.T
                  _lowercase	       =q.T
                  _lowercase	       =v.T

                  # Block i, layer 2 (MLP).
                  _lowercase	       =tax_layer_norm_lookup(A__      ,   A__      ,   'decoder'      ,   'pre_mlp_layer_norm'     )
                  _lowercase       ,       _lowercase	       =tax_mlp_lookup(A__      ,   A__      ,   'decoder'      ,   A__     )
                  _lowercase	       =layer_norm
                  if split_mlp_wi:
                        _lowercase	       =wi[0].T
                        _lowercase	       =wi[1].T
                  else:
                        _lowercase	       =wi.T
                  _lowercase	       =wo.T

            _lowercase	       =old['decoder/decoder_norm/scale']
            _lowercase	       =old[
                'decoder/relpos_bias/rel_embedding'
            ].T

            # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
            if "decoder/logits_dense/kernel" in old:
                  _lowercase	       =old['decoder/logits_dense/kernel'].T

      return new
def      a      (			A__	:   List[Any]      ,   A__	:   bool     )    ->      int:





      """simple docstring"""



      _lowercase	       =collections.OrderedDict([(k, torch.from_numpy(v.copy()     )) for (k, v) in converted_params.items()]     )

      # Add what is missing.
      if "encoder.embed_tokens.weight" not in state_dict:
            _lowercase	       =state_dict['shared.weight']

      if not is_encoder_only:
            if "decoder.embed_tokens.weight" not in state_dict:
                  _lowercase	       =state_dict['shared.weight']

            if "lm_head.weight" not in state_dict:  # For old 1.0 models.
                  print('Using shared word embeddings as lm_head.'     )
                  _lowercase	       =state_dict['shared.weight']

      return state_dict
def      a      (			A__	:   List[Any]      ,   A__	:   Optional[int]      ,   A__	:   List[str]      ,   A__	:   Dict     )    ->      Union[str, Any]:





      """simple docstring"""



      _lowercase	       =checkpoints.load_tax_checkpoint(A__     )
      _lowercase	       =convert_tax_to_pytorch(A__      ,   num_layers=config.num_layers      ,   is_encoder_only=A__     )
      _lowercase	       =make_state_dict(A__      ,   A__     )
      model.load_state_dict(A__      ,   strict=A__     )
def      a      (			A__	:   Optional[Any]      ,   A__	:   List[Any]      ,   A__	:   Optional[Any]      ,   A__	:   bool = False     )    ->      Optional[Any]:





      """simple docstring"""



      _lowercase	       =TaConfig.from_json_file(A__     )
      print(F'''Building PyTorch model from configuration: {config}'''     )
      # Non-v1.1 checkpoints could also use T5Model, but this works for all.
      # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
      if is_encoder_only:
            _lowercase	       =TaEncoderModel(A__     )
      else:
            _lowercase	       =TaForConditionalGeneration(A__     )

      # Load weights from tf checkpoint
      load_tax_weights_in_ta(A__      ,   A__      ,   A__      ,   A__     )

      # Save pytorch-model
      print(F'''Save PyTorch model to {pytorch_dump_path}'''     )
      model.save_pretrained(A__     )

      # Verify that we can load the checkpoint.
      model.from_pretrained(A__     )
      print('Done'     )


if __name__ == "__main__":
  lowercase_									=   argparse.ArgumentParser(description='Converts a native T5X checkpoint into a PyTorch checkpoint.')
  # Required parameters
  parser.add_argument(
      '--t5x_checkpoint_path', default=None, type=str, required=True, help='Path to the T5X checkpoint.'
  )
  parser.add_argument(
      '--config_file',
      default=None,
      type=str,
      required=True,
      help='The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.',
  )
  parser.add_argument(
      '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
  )
  parser.add_argument(
      '--is_encoder_only', action='store_true', help='Check if the model is encoder-decoder model', default=False
  )
  lowercase_									=   parser.parse_args()
  convert_tax_checkpoint_to_pytorch(
      args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
  )





# 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}'''     )


