


"""simple docstring"""







import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
       import tensorflow as tf
       from tensorflow.python.framework.errors_impl import ResourceExhaustedError

       from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
       import pyanvml.pyanvml as nvml

lowerCamelCase_   :   Union[str, Any]									=  logging.get_logger(__name__)
def _A						(						lowercase       ,    lowercase		):







    """simple docstring"""
    def run_func(lowercase		):
        @wraps(lowercase		)
        def run_in_eager_mode(*lowercase       ,    **lowercase		):
            return func(*lowercase       ,    **lowercase		)

        @wraps(lowercase		)
        @tf.function(experimental_compile=lowercase		)
        def run_in_graph_mode(*lowercase       ,    **lowercase		):
            return func(*lowercase       ,    **lowercase		)

        if do_eager_mode is True:
            if use_xla is not False:
                raise ValueError(
                    '''Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.'''		)
            return run_in_eager_mode
        else:
            return run_in_graph_mode

    return run_func
def _A						(						lowercase       ,    lowercase       ,    lowercase		):







    """simple docstring"""
    a 		=random.Random()
    a 		=[rng.randint(0       ,    vocab_size - 1		) for i in range(batch_size * sequence_length		)]
    return tf.constant(lowercase       ,    shape=(batch_size, sequence_length)       ,    dtype=tf.intaa		)






class       __A				(       _SCREAMING_SNAKE_CASE			):






 """simple docstring"""




 __lowerCAmelCase 					=  42
 __lowerCAmelCase 					=  42
 __lowerCAmelCase 					=  "TensorFlow"


 @property
 def       SCREAMING_SNAKE_CASE							(		self						)					->     Optional[Any]:
     return tf.__version__


 def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A						)					->     float:
     # initialize GPU on separate process
     a 		=self.args.strategy
     if strategy is None:
         raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''						)
     a 		=self._prepare_inference_func(__A ,    __A ,    __A						)
     return self._measure_speed(_inference						)


 def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A						)					->     float:
     a 		=self.args.strategy
     if strategy is None:
         raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''						)
     a 		=self._prepare_train_func(__A ,    __A ,    __A						)
     return self._measure_speed(_train						)


 def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A						)					->     [Memory, Optional[MemorySummary]]:
     # initialize GPU on separate process
     if self.args.is_gpu:
         tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx] ,    __A						)
     a 		=self.args.strategy
     if strategy is None:
         raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''						)
     a 		=self._prepare_inference_func(__A ,    __A ,    __A						)
     return self._measure_memory(_inference						)


 def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A						)					->     [Memory, Optional[MemorySummary]]:
     if self.args.is_gpu:
         tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx] ,    __A						)
     a 		=self.args.strategy
     if strategy is None:
         raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''						)

     a 		=self._prepare_train_func(__A ,    __A ,    __A						)
     return self._measure_memory(_train						)


 def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A						)					->     Callable[[], None]:
     a 		=self.config_dict[model_name]

     if self.args.fpaa:
         raise NotImplementedError('''Mixed precision is currently not supported.'''						)

     a 		=(
         hasattr(__A ,    '''architectures'''						)
         and isinstance(config.architectures ,    __A						)
         and len(config.architectures						) > 0
     )
     if not self.args.only_pretrain_model and has_model_class_in_config:
         try:
             a 		='''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
             a 		=__import__('''transformers''' ,    fromlist=[model_class]						)
             a 		=getattr(__A ,    __A						)
             a 		=model_cls(__A						)
         except ImportError:
             raise ImportError(
                 f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                 ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''						)
     else:
         a 		=TF_MODEL_MAPPING[config.__class__](__A						)

     # encoder-decoder has vocab size saved differently
     a 		=config.vocab_size if hasattr(__A ,    '''vocab_size'''						) else config.encoder.vocab_size
     a 		=random_input_ids(__A ,    __A ,    __A						)

     @run_with_tf_optimizations(self.args.eager_mode ,    self.args.use_xla						)
     def encoder_decoder_forward():
         return model(__A ,    decoder_input_ids=__A ,    training=__A						)

     @run_with_tf_optimizations(self.args.eager_mode ,    self.args.use_xla						)
     def encoder_forward():
         return model(__A ,    training=__A						)

     a 		=encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

     return _inference


 def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A						)					->     Callable[[], None]:
     a 		=self.config_dict[model_name]

     if self.args.eager_mode is not False:
         raise ValueError('''Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.'''						)

     if self.args.fpaa:
         raise NotImplementedError('''Mixed precision is currently not supported.'''						)

     a 		=(
         hasattr(__A ,    '''architectures'''						)
         and isinstance(config.architectures ,    __A						)
         and len(config.architectures						) > 0
     )
     if not self.args.only_pretrain_model and has_model_class_in_config:
         try:
             a 		='''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
             a 		=__import__('''transformers''' ,    fromlist=[model_class]						)
             a 		=getattr(__A ,    __A						)
             a 		=model_cls(__A						)
         except ImportError:
             raise ImportError(
                 f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                 ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''						)
     else:
         a 		=TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__A						)

     # encoder-decoder has vocab size saved differently
     a 		=config.vocab_size if hasattr(__A ,    '''vocab_size'''						) else config.encoder.vocab_size
     a 		=random_input_ids(__A ,    __A ,    __A						)

     @run_with_tf_optimizations(self.args.eager_mode ,    self.args.use_xla						)
     def encoder_decoder_train():
         a 		=model(__A ,    decoder_input_ids=__A ,    labels=__A ,    training=__A						)[0]
         a 		=tf.gradients(__A ,    model.trainable_variables						)
         return gradients

     @run_with_tf_optimizations(self.args.eager_mode ,    self.args.use_xla						)
     def encoder_train():
         a 		=model(__A ,    labels=__A ,    training=__A						)[0]
         a 		=tf.gradients(__A ,    model.trainable_variables						)
         return gradients

     a 		=encoder_decoder_train if config.is_encoder_decoder else encoder_train

     return _train


 def       SCREAMING_SNAKE_CASE							(		self ,    __A						)					->     float:
     with self.args.strategy.scope():
         try:
             if self.args.is_tpu or self.args.use_xla:
                 # run additional 10 times to stabilize compilation for tpu
                 logger.info('''Do inference on TPU. Running model 5 times to stabilize compilation'''						)
                 timeit.repeat(__A ,    repeat=1 ,    number=5						)

             # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
             a 		=timeit.repeat(
                 __A ,    repeat=self.args.repeat ,    number=10 ,    )

             return min(__A						) / 10.0
         except ResourceExhaustedError as e:
             self.print_fn(f'''Doesn\'t fit on GPU. {e}'''						)





 def       SCREAMING_SNAKE_CASE							(		self ,    __A						)					->     [Memory, MemorySummary]:
     logger.info(
         '''Note that TensorFlow allocates more memory than '''
         '''it might need to speed up computation. '''
         '''The memory reported here corresponds to the memory '''
         '''reported by `nvidia-smi`, which can vary depending '''
         '''on total available memory on the GPU that is used.'''						)
     with self.args.strategy.scope():
         try:
             if self.args.trace_memory_line_by_line:
                 if not self.args.eager_mode:
                     raise ValueError(
                         '''`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory'''
                         ''' consumption line by line.'''						)
                 a 		=start_memory_tracing('''transformers'''						)

             if self.args.is_tpu:
                 # tpu
                 raise NotImplementedError(
                     '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking'''
                     ''' with `args.memory=False`'''						)
             elif self.args.is_gpu:
                 # gpu
                 if not is_pyanvml_available():
                     logger.warning(
                         '''py3nvml not installed, we won\'t log GPU memory usage. '''
                         '''Install py3nvml (pip install py3nvml) to log information about GPU.'''						)
                     a 		='''N/A'''
                 else:
                     logger.info(
                         '''Measuring total GPU usage on GPU device. Make sure to not have additional processes'''
                         ''' running on the same GPU.'''						)
                     # init nvml
                     nvml.nvmlInit()
                     func()
                     a 		=nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx						)
                     a 		=nvml.nvmlDeviceGetMemoryInfo(__A						)
                     a 		=meminfo.used
                     a 		=Memory(__A						)
                     # shutdown nvml
                     nvml.nvmlShutdown()
             else:
                 # cpu
                 if self.args.trace_memory_line_by_line:
                     logger.info(
                         '''When enabling line by line tracing, the max peak memory for CPU is inaccurate in'''
                         ''' TensorFlow.'''						)
                     a 		=None
                 else:
                     a 		=measure_peak_memory_cpu(__A						)
                     a 		=Memory(__A						) if isinstance(__A ,    __A						) else memory_bytes
             if self.args.trace_memory_line_by_line:
                 a 		=stop_memory_tracing(__A						)
                 if memory is None:
                     a 		=summary.total
             else:
                 a 		=None

             return memory, summary
         except ResourceExhaustedError as e:
             self.print_fn(f'''Doesn\'t fit on GPU. {e}'''						)
             return "N/A", None





"""simple docstring"""




import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from tax import checkpoints

from transformers import MTaConfig, UMTaEncoderModel, UMTaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()



def     a__ (				snake_case__       , snake_case__       , snake_case__				)						->					List[Any]:
			return params[F'{prefix}/{prefix}/relpos_bias/rel_embedding'][:, i, :]



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__="attention"				)						->					List[Any]:
			lowerCamelCase			      =					lowerCamelCase			      =					np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/key/kernel'][:, i, :, :]				)
			lowerCamelCase			      =					k_tmp.reshape(k_tmp.shape[0]       , k_tmp.shape[1] * k_tmp.shape[2]				)
			lowerCamelCase			      =					np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/out/kernel'][:, i, :, :]				)
			lowerCamelCase			      =					o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]       , o_tmp.shape[2]				)
			lowerCamelCase			      =					np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/query/kernel'][:, i, :, :]				)
			lowerCamelCase			      =					q_tmp.reshape(q_tmp.shape[0]       , q_tmp.shape[1] * q_tmp.shape[2]				)
			lowerCamelCase			      =					np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/value/kernel'][:, i, :, :]				)
			lowerCamelCase			      =					v_tmp.reshape(v_tmp.shape[0]       , v_tmp.shape[1] * v_tmp.shape[2]				)
			return k, o, q, v



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__=False				)						->					List[str]:
			if split_mlp_wi:
						lowerCamelCase			      =					params[F'{prefix}/{prefix}/mlp/wi_0/kernel'][:, i, :]
						lowerCamelCase			      =					params[F'{prefix}/{prefix}/mlp/wi_1/kernel'][:, i, :]
						lowerCamelCase			      =					(wi_a, wi_a)
			else:
						lowerCamelCase			      =					params[F'{prefix}/{prefix}/mlp/wi/kernel'][:, i, :]

			lowerCamelCase			      =					params[F'{prefix}/{prefix}/mlp/wo/kernel'][:, i, :]
			return wi, wo



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__				)						->					Tuple:
			return params[F'{prefix}/{prefix}/{layer_name}/scale'][:, i]



def     a__ (				snake_case__       , *, snake_case__       , snake_case__       , snake_case__ = False				)						->					Dict:
			lowerCamelCase			      =					traverse_util.flatten_dict(variables["""target"""]				)
			lowerCamelCase			      =					{"""/""".join(snake_case__				): v for k, v in old.items()}

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

			lowerCamelCase			      =					collections.OrderedDict()

			# Shared embeddings.
			lowerCamelCase			      =					old["""token_embedder/embedding"""]

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

						# Block i, layer 1 (MLP).
						lowerCamelCase			      =					tax_layer_norm_lookup(snake_case__       , snake_case__       , """encoder"""       , """pre_mlp_layer_norm"""				)
						lowerCamelCase      ,   lowerCamelCase			      =					tax_mlp_lookup(snake_case__       , snake_case__       , """encoder"""       , snake_case__				)
						lowerCamelCase			      =					layer_norm
						if split_mlp_wi:
									lowerCamelCase			      =					wi[0].T
									lowerCamelCase			      =					wi[1].T
						else:
									lowerCamelCase			      =					wi.T
						lowerCamelCase			      =					wo.T
						if scalable_attention:
									# convert the rel_embedding of each layer
									lowerCamelCase			      =					tax_relpos_bias_lookup(
									    snake_case__       , snake_case__       , """encoder"""				).T

			lowerCamelCase			      =					old["""encoder/encoder_norm/scale"""]

			if not scalable_attention:
						lowerCamelCase			      =					tax_relpos_bias_lookup(
						    snake_case__       , 0       , """encoder"""				).T
						lowerCamelCase			      =					tax_relpos_bias_lookup(
						    snake_case__       , 0       , """decoder"""				).T

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

									# Block i, layer 1 (Cross Attention).
									lowerCamelCase			      =					tax_layer_norm_lookup(snake_case__       , snake_case__       , """decoder"""       , """pre_cross_attention_layer_norm"""				)
									lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					tax_attention_lookup(snake_case__       , snake_case__       , """decoder"""       , """encoder_decoder_attention"""				)
									lowerCamelCase			      =					layer_norm
									lowerCamelCase			      =					k.T
									lowerCamelCase			      =					o.T
									lowerCamelCase			      =					q.T
									lowerCamelCase			      =					v.T

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

									if scalable_attention:
												# convert the rel_embedding of each layer
												lowerCamelCase			      =					tax_relpos_bias_lookup(snake_case__       , snake_case__       , """decoder"""				).T

						lowerCamelCase			      =					old["""decoder/decoder_norm/scale"""]

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

			return new



def     a__ (				snake_case__       , snake_case__				)						->					Optional[int]:
			lowerCamelCase			      =					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:
						lowerCamelCase			      =					state_dict["""shared.weight"""]

			if not is_encoder_only:
						if "decoder.embed_tokens.weight" not in state_dict:
									lowerCamelCase			      =					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."""				)
									lowerCamelCase			      =					state_dict["""shared.weight"""]

			return state_dict



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__       , snake_case__				)						->					List[Any]:
			lowerCamelCase			      =					checkpoints.load_tax_checkpoint(snake_case__				)
			lowerCamelCase			      =					convert_tax_to_pytorch(
			    snake_case__       , num_layers=config.num_layers       , is_encoder_only=snake_case__       , scalable_attention=snake_case__				)
			lowerCamelCase			      =					make_state_dict(snake_case__       , snake_case__				)
			model.load_state_dict(snake_case__       , strict=snake_case__				)



def     a__ (				snake_case__       , snake_case__       , snake_case__       , snake_case__ = False       , snake_case__ = False       , )						->					str:
			lowerCamelCase			      =					MTaConfig.from_json_file(snake_case__				)
			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:
						lowerCamelCase			      =					UMTaEncoderModel(snake_case__				)
			else:
						lowerCamelCase			      =					UMTaForConditionalGeneration(snake_case__				)

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

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

			# Verify that we can load the checkpoint.
			model.from_pretrained(snake_case__				)
			print("""Done"""				)


if __name__ == "__main__":
					lowerCAmelCase							:			Optional[int]     							=							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
					)
					parser.add_argument(
					    """--scalable_attention""",
					    action="""store_true""",
					    help="""Whether the model uses scaled attention (umt5 model)""",
					    default=False,
					)
					lowerCAmelCase							:			int     							=							parser.parse_args()
					convert_tax_checkpoint_to_pytorch(
					    args.tax_checkpoint_path,
					    args.config_file,
					    args.pytorch_dump_path,
					    args.is_encoder_only,
					    args.scalable_attention,
					)



