


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 __UpperCAmelCase       (		__A ,						__A ,						__A			)						->					Optional[Any]:



							'''simple docstring'''





							return params[F"""{prefix}/{prefix}/relpos_bias/rel_embedding"""][:, i, :]





def __UpperCAmelCase       (		__A ,						__A ,						__A ,						__A="attention"			)						->					List[Any]:



							'''simple docstring'''





							UpperCAmelCase__										=     UpperCAmelCase__										=     np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/key/kernel"""][:, i, :, :]			)
							UpperCAmelCase__										=     k_tmp.reshape(k_tmp.shape[0] ,						k_tmp.shape[1] * k_tmp.shape[2]			)
							UpperCAmelCase__										=     np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/out/kernel"""][:, i, :, :]			)
							UpperCAmelCase__										=     o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1] ,						o_tmp.shape[2]			)
							UpperCAmelCase__										=     np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/query/kernel"""][:, i, :, :]			)
							UpperCAmelCase__										=     q_tmp.reshape(q_tmp.shape[0] ,						q_tmp.shape[1] * q_tmp.shape[2]			)
							UpperCAmelCase__										=     np.ascontiguousarray(params[F"""{prefix}/{prefix}/{layer_name}/value/kernel"""][:, i, :, :]			)
							UpperCAmelCase__										=     v_tmp.reshape(v_tmp.shape[0] ,						v_tmp.shape[1] * v_tmp.shape[2]			)
							return k, o, q, v





def __UpperCAmelCase       (		__A ,						__A ,						__A ,						__A=False			)						->					Tuple:



							'''simple docstring'''





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

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





def __UpperCAmelCase       (		__A ,						__A ,						__A ,						__A			)						->					Any:



							'''simple docstring'''





							return params[F"""{prefix}/{prefix}/{layer_name}/scale"""][:, i]





def __UpperCAmelCase       (		__A ,						*, __A ,						__A ,						__A = False			)						->					Tuple:



							'''simple docstring'''





							UpperCAmelCase__										=     traverse_util.flatten_dict(variables["target"]			)
							UpperCAmelCase__										=     {"/".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
							UpperCAmelCase__										=     "encoder/encoder/mlp/wi_0/kernel" in old
							print("Split MLP:" ,						__A			)

							UpperCAmelCase__										=     collections.OrderedDict()

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

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

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

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

							if not scalable_attention:
														UpperCAmelCase__										=     tax_relpos_bias_lookup(
														    __A ,						0 ,						"encoder"			).T
														UpperCAmelCase__										=     tax_relpos_bias_lookup(
														    __A ,						0 ,						"decoder"			).T

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

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

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

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

														UpperCAmelCase__										=     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:
																					UpperCAmelCase__										=     old["decoder/logits_dense/kernel"].T

							return new





def __UpperCAmelCase       (		__A ,						__A			)						->					Optional[Any]:



							'''simple docstring'''





							UpperCAmelCase__										=     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:
														UpperCAmelCase__										=     state_dict["shared.weight"]

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

							return state_dict





def __UpperCAmelCase       (		__A ,						__A ,						__A ,						__A ,						__A			)						->					Dict:



							'''simple docstring'''





							UpperCAmelCase__										=     checkpoints.load_tax_checkpoint(__A			)
							UpperCAmelCase__										=     convert_tax_to_pytorch(
							    __A ,						num_layers=config.num_layers ,						is_encoder_only=__A ,						scalable_attention=__A			)
							UpperCAmelCase__										=     make_state_dict(__A ,						__A			)
							model.load_state_dict(__A ,						strict=__A			)





def __UpperCAmelCase       (		__A ,						__A ,						__A ,						__A = False ,						__A = False ,						)						->					List[Any]:



							'''simple docstring'''





							UpperCAmelCase__										=     MTaConfig.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:
														UpperCAmelCase__										=     UMTaEncoderModel(__A			)
							else:
														UpperCAmelCase__										=     UMTaForConditionalGeneration(__A			)

							# Load weights from tf checkpoint
							load_tax_weights_in_ta(__A ,						__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__":
				A      = 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,
				)
				A      = 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,
				)









import copy
import json
import os
import tempfile

from transformers import is_torch_available

from .test_configuration_utils import config_common_kwargs


class     lowercase__  (					__SCREAMING_SNAKE_CASE		):






  def __init__(					self      :					Dict		,						_lowercase      :					int		,						_lowercase      :					List[Any]=None		,						_lowercase      :					str=True		,						_lowercase      :					str=None		,						**_lowercase      :					Optional[int] ):




         """simple docstring"""







         UpperCAmelCase__										=     parent
         UpperCAmelCase__										=     config_class
         UpperCAmelCase__										=     has_text_modality
         UpperCAmelCase__										=     kwargs
         UpperCAmelCase__										=     common_properties






  def 			_UpperCAmelCase  (					self      :					Dict ):




         """simple docstring"""







         UpperCAmelCase__										=     self.config_class(**self.inputs_dict )
         UpperCAmelCase__										=     (
             ["hidden_size", "num_attention_heads", "num_hidden_layers"]
             if self.common_properties is None
             else self.common_properties
         )

         # Add common fields for text models
         if self.has_text_modality:
                common_properties.extend(["vocab_size"] )

         # Test that config has the common properties as getters
         for prop in common_properties:
                self.parent.assertTrue(hasattr(_lowercase		,						_lowercase )		,						msg=F"""`{prop}` does not exist""" )

         # Test that config has the common properties as setter
         for idx, name in enumerate(_lowercase ):
                try:
                       setattr(_lowercase		,						_lowercase		,						_lowercase )
                       self.parent.assertEqual(
                           getattr(_lowercase		,						_lowercase )		,						_lowercase		,						msg=F"""`{name} value {idx} expected, but was {getattr(_lowercase		,						_lowercase )}""" )
                except NotImplementedError:
                       # Some models might not be able to implement setters for common_properties
                       # In that case, a NotImplementedError is raised
                       pass

        # Test if config class can be called with Config(prop_name=..)
         for idx, name in enumerate(_lowercase ):
                try:
                       UpperCAmelCase__										=     self.config_class(**{name: idx} )
                       self.parent.assertEqual(
                           getattr(_lowercase		,						_lowercase )		,						_lowercase		,						msg=F"""`{name} value {idx} expected, but was {getattr(_lowercase		,						_lowercase )}""" )
                except NotImplementedError:
                       # Some models might not be able to implement setters for common_properties
                       # In that case, a NotImplementedError is raised
                       pass






  def 			_UpperCAmelCase  (					self      :					int ):




         """simple docstring"""







         UpperCAmelCase__										=     self.config_class(**self.inputs_dict )
         UpperCAmelCase__										=     json.loads(config.to_json_string() )
         for key, value in self.inputs_dict.items():
                self.parent.assertEqual(obj[key]		,						_lowercase )






  def 			_UpperCAmelCase  (					self      :					Optional[int] ):




         """simple docstring"""







         UpperCAmelCase__										=     self.config_class(**self.inputs_dict )

         with tempfile.TemporaryDirectory() as tmpdirname:
                UpperCAmelCase__										=     os.path.join(_lowercase		,						"config.json" )
                config_first.to_json_file(_lowercase )
                UpperCAmelCase__										=     self.config_class.from_json_file(_lowercase )

         self.parent.assertEqual(config_second.to_dict()		,						config_first.to_dict() )






  def 			_UpperCAmelCase  (					self      :					Optional[int] ):




         """simple docstring"""







         UpperCAmelCase__										=     self.config_class(**self.inputs_dict )

         with tempfile.TemporaryDirectory() as tmpdirname:
                config_first.save_pretrained(_lowercase )
                UpperCAmelCase__										=     self.config_class.from_pretrained(_lowercase )

         self.parent.assertEqual(config_second.to_dict()		,						config_first.to_dict() )






  def 			_UpperCAmelCase  (					self      :					List[str] ):




         """simple docstring"""







         UpperCAmelCase__										=     self.config_class(**self.inputs_dict )

         UpperCAmelCase__										=     "test"
         with tempfile.TemporaryDirectory() as tmpdirname:
                UpperCAmelCase__										=     os.path.join(_lowercase		,						_lowercase )
                config_first.save_pretrained(_lowercase )
                UpperCAmelCase__										=     self.config_class.from_pretrained(_lowercase		,						subfolder=_lowercase )

         self.parent.assertEqual(config_second.to_dict()		,						config_first.to_dict() )






  def 			_UpperCAmelCase  (					self      :					List[str] ):




         """simple docstring"""







         UpperCAmelCase__										=     self.config_class(**self.inputs_dict		,						num_labels=5 )
         self.parent.assertEqual(len(config.idalabel )		,						5 )
         self.parent.assertEqual(len(config.labelaid )		,						5 )

         UpperCAmelCase__										=     3
         self.parent.assertEqual(len(config.idalabel )		,						3 )
         self.parent.assertEqual(len(config.labelaid )		,						3 )






  def 			_UpperCAmelCase  (					self      :					Any ):




         """simple docstring"""







         if self.config_class.is_composition:
                return
         UpperCAmelCase__										=     self.config_class()
         self.parent.assertIsNotNone(_lowercase )






  def 			_UpperCAmelCase  (					self      :					Any ):




         """simple docstring"""







         UpperCAmelCase__										=     copy.deepcopy(_lowercase )
         UpperCAmelCase__										=     self.config_class(**_lowercase )
         UpperCAmelCase__										=     []
         for key, value in config_common_kwargs.items():
                if key == "torch_dtype":
                       if not is_torch_available():
                              continue
                       else:
                              import torch

                              if config.torch_dtype != torch.floataa:
                                     wrong_values.append(("torch_dtype", config.torch_dtype, torch.floataa) )
                elif getattr(_lowercase		,						_lowercase ) != value:
                       wrong_values.append((key, getattr(_lowercase		,						_lowercase ), value) )

         if len(_lowercase ) > 0:
                UpperCAmelCase__										=     "\n".join([F"""- {v[0]}: got {v[1]} instead of {v[2]}""" for v in wrong_values] )
                raise ValueError(F"""The following keys were not properly set in the config:\n{errors}""" )






  def 			_UpperCAmelCase  (					self      :					Any ):




         """simple docstring"""







         self.create_and_test_config_common_properties()
         self.create_and_test_config_to_json_string()
         self.create_and_test_config_to_json_file()
         self.create_and_test_config_from_and_save_pretrained()
         self.create_and_test_config_from_and_save_pretrained_subfolder()
         self.create_and_test_config_with_num_labels()
         self.check_config_can_be_init_without_params()
         self.check_config_arguments_init()





