





from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


UpperCamelCase__      =logging.get_logger(__name__)






def     lowerCamelCase__			(__lowerCamelCase	):
     if isinstance(__lowerCamelCase,       np.ndarray	):
          return list(tensor.shape	)

     _SCREAMING_SNAKE_CASE       :	Any						    =							tf.shape(__lowerCamelCase	)

     if tensor.shape == tf.TensorShape(__lowerCamelCase	):
          return dynamic

     _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							tensor.shape.as_list()

     return [dynamic[i] if s is None else s for i, s in enumerate(__lowerCamelCase	)]






def     lowerCamelCase__			(__lowerCamelCase,       __lowerCamelCase = None,       __lowerCamelCase = None	):
     return tf.nn.softmax(logits=logits + 1e-9,       axis=__lowerCamelCase,       name=__lowerCamelCase	)






def     lowerCamelCase__			(__lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase=1e-5,       __lowerCamelCase=-1	):
     # This is a very simplified functional layernorm, designed to duplicate
     # the functionality of PyTorch nn.functional.layer_norm when this is needed to port
     # models in Transformers.

     if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(__lowerCamelCase,       __lowerCamelCase	):
          raise NotImplementedError("Only 1D weight and bias tensors are supported for now, with only a single axis."	)

     # Get mean and variance on the axis to be normalized
     _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							tf.nn.moments(__lowerCamelCase,       axes=[axis],       keepdims=__lowerCamelCase	)

     if axis != -1:
          # Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
          # on every dimension except axis
          _SCREAMING_SNAKE_CASE       :	List[Any]						    =							[1] * inputs.shape.rank
          _SCREAMING_SNAKE_CASE       :	List[Any]						    =							shape_list(__lowerCamelCase	)[axis]
          _SCREAMING_SNAKE_CASE       :	str						    =							tf.reshape(__lowerCamelCase,       __lowerCamelCase	)
          _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							tf.reshape(__lowerCamelCase,       __lowerCamelCase	)

     # Compute layer normalization using the batch_normalization
     # function.
     _SCREAMING_SNAKE_CASE       :	Any						    =							tf.nn.batch_normalization(
         __lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase,       offset=__lowerCamelCase,       scale=__lowerCamelCase,       variance_epsilon=__lowerCamelCase,       )
     return outputs






def     lowerCamelCase__			(__lowerCamelCase,       __lowerCamelCase=0,       __lowerCamelCase=-1	):
     # Replicates the behavior of torch.flatten in TF

     # If end_dim or start_dim is negative, count them from the end
     if end_dim < 0:
          end_dim += input.shape.rank
     if start_dim < 0:
          start_dim += input.shape.rank

     if start_dim == end_dim:
          return input

     _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							tf.shape(__lowerCamelCase	)
     _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]	)
     _SCREAMING_SNAKE_CASE       :	Tuple						    =							tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]],       axis=0	)
     return tf.reshape(__lowerCamelCase,       __lowerCamelCase	)






def     lowerCamelCase__			(__lowerCamelCase	):
     if not isinstance(__lowerCamelCase,       tf.Tensor	):
          _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							tf.convert_to_tensor(__lowerCamelCase	)  # Catches stray NumPy inputs
     if encoder_attention_mask.shape.rank == 3:
          _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							encoder_attention_mask[:, None, :, :]
     if encoder_attention_mask.shape.rank == 2:
          _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							encoder_attention_mask[:, None, None, :]
     # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
     # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
     # /transformer/transformer_layers.py#L270
     # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
     # encoder_extended_attention_mask.transpose(-1, -2))
     _SCREAMING_SNAKE_CASE       :	List[str]						    =							(
         tf.cast(1,       encoder_attention_mask.dtype	) - encoder_extended_attention_mask
     ) * encoder_extended_attention_mask.dtype.min

     return encoder_extended_attention_mask






def     lowerCamelCase__			(__lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase = "input_ids"	):
     tf.debugging.assert_less(
         __lowerCamelCase,       tf.cast(__lowerCamelCase,       dtype=tensor.dtype	),       message=(
             f"""The maximum value of {tensor_name} ({tf.math.reduce_max(__lowerCamelCase	)}) must be smaller than the embedding """
             f"""layer's input dimension ({embed_dim}). The likely cause is some problem at tokenization time."""
         ),       )






def     lowerCamelCase__			(__lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase	):
     _SCREAMING_SNAKE_CASE       :	List[str]						    =							64512
     # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
     # because in that case even chunking the array would not make the saving
     # possible.
     _SCREAMING_SNAKE_CASE       :	List[Any]						    =							[x for x in data if len(__lowerCamelCase	) > HDF5_OBJECT_HEADER_LIMIT]

     # Expecting this to never be true.
     if bad_attributes:
          raise RuntimeError(
              "The following attributes cannot be saved to HDF5 file because "
              f"""they are larger than {HDF5_OBJECT_HEADER_LIMIT} """
              f"""bytes: {bad_attributes}"""	)

     _SCREAMING_SNAKE_CASE       :	Any						    =							np.asarray(__lowerCamelCase	)

     _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							1
     _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							np.array_split(__lowerCamelCase,       __lowerCamelCase	)

     # This will never loop forever thanks to the test above.
     while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data	):
          num_chunks += 1
          _SCREAMING_SNAKE_CASE       :	List[Any]						    =							np.array_split(__lowerCamelCase,       __lowerCamelCase	)

     if num_chunks > 1:
          for chunk_id, chunk_data in enumerate(__lowerCamelCase	):
               _SCREAMING_SNAKE_CASE       :	Dict						    =							chunk_data
     else:
          _SCREAMING_SNAKE_CASE       :	Tuple						    =							data






def     lowerCamelCase__			(__lowerCamelCase,       __lowerCamelCase	):
     if name in group.attrs:
          _SCREAMING_SNAKE_CASE       :	int						    =							[n.decode("utf8"	) if hasattr(__lowerCamelCase,       "decode"	) else n for n in group.attrs[name]]
     else:
          _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							[]
          _SCREAMING_SNAKE_CASE       :	Tuple						    =							0
          while "%s%d" % (name, chunk_id) in group.attrs:
               data.extend(
                   [n.decode("utf8"	) if hasattr(__lowerCamelCase,       "decode"	) else n for n in group.attrs["%s%d" % (name, chunk_id)]]	)
               chunk_id += 1
     return data






def     lowerCamelCase__			(__lowerCamelCase	):

     def _expand_single_ad_tensor(__lowerCamelCase	):
          if isinstance(__lowerCamelCase,       tf.Tensor	) and t.shape.rank == 1:
               return tf.expand_dims(__lowerCamelCase,       axis=-1	)
          return t

     return tf.nest.map_structure(_expand_single_ad_tensor,       __lowerCamelCase	)






import argparse

import torch

from transformers import GPTaConfig, GPTaModel, load_tf_weights_in_gpta
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()






def     lowerCamelCase__			(__lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase	):
     # Construct model
     if gpta_config_file == "":
          _SCREAMING_SNAKE_CASE       :	str						    =							GPTaConfig()
     else:
          _SCREAMING_SNAKE_CASE       :	int						    =							GPTaConfig.from_json_file(__lowerCamelCase	)
     _SCREAMING_SNAKE_CASE       :	Dict						    =							GPTaModel(__lowerCamelCase	)

     # Load weights from numpy
     load_tf_weights_in_gpta(__lowerCamelCase,       __lowerCamelCase,       __lowerCamelCase	)

     # Save pytorch-model
     _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							pytorch_dump_folder_path + "/" + WEIGHTS_NAME
     _SCREAMING_SNAKE_CASE       :	Tuple						    =							pytorch_dump_folder_path + "/" + CONFIG_NAME
     print(f"""Save PyTorch model to {pytorch_weights_dump_path}"""	)
     torch.save(model.state_dict(),       __lowerCamelCase	)
     print(f"""Save configuration file to {pytorch_config_dump_path}"""	)
     with open(__lowerCamelCase,       "w",       encoding="utf-8"	) as f:
          f.write(config.to_json_string()	)


if __name__ == "__main__":
       UpperCamelCase__      =argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           '--gpt2_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
       )
       parser.add_argument(
           '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
       )
       parser.add_argument(
           '--gpt2_config_file',
           default='',
           type=str,
           help=(
               'An optional config json file corresponding to the pre-trained OpenAI model. \n'
               'This specifies the model architecture.'
           ),
       )
       UpperCamelCase__      =parser.parse_args()
       convert_gpta_checkpoint_to_pytorch(args.gpta_checkpoint_path, args.gpta_config_file, args.pytorch_dump_folder_path)