



'''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       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase		)      ->			Dict:
							return params[F'''{prefix}/{prefix}/relpos_bias/rel_embedding'''][:, i, :]



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



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase	,				_lowercase=False		)      ->			Dict:
							if split_mlp_wi:
														UpperCAmelCase     :		List[Any]							     =		params[F'''{prefix}/{prefix}/mlp/wi_0/kernel'''][:, i, :]
														UpperCAmelCase     :		List[Any]							     =		params[F'''{prefix}/{prefix}/mlp/wi_1/kernel'''][:, i, :]
														UpperCAmelCase     :		List[str]							     =		(wi_a, wi_a)
							else:
														UpperCAmelCase     :		List[str]							     =		params[F'''{prefix}/{prefix}/mlp/wi/kernel'''][:, i, :]

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



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase	,				_lowercase		)      ->			int:
							return params[F'''{prefix}/{prefix}/{layer_name}/scale'''][:, i]



def       __lowerCamelCase  (					_lowercase	,				*, _lowercase	,				_lowercase	,				_lowercase = False		)      ->			List[str]:
							UpperCAmelCase     :		str							     =		traverse_util.flatten_dict(variables["""target"""]		)
							UpperCAmelCase     :		Dict							     =		{"""/""".join(_lowercase		): v for k, v in old.items()}

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

							UpperCAmelCase     :		int							     =		collections.OrderedDict()

							# Shared embeddings.
							UpperCAmelCase     :		Any							     =		old["""token_embedder/embedding"""]

							# Encoder.
							for i in range(_lowercase		):
														# Block i, layer 0 (Self Attention).
														UpperCAmelCase     :		str							     =		tax_layer_norm_lookup(_lowercase	,				_lowercase	,				"""encoder"""	,				"""pre_attention_layer_norm"""		)
														UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Tuple							     =		tax_attention_lookup(_lowercase	,				_lowercase	,				"""encoder"""	,				"""attention"""		)
														UpperCAmelCase     :		Optional[Any]							     =		layer_norm
														UpperCAmelCase     :		Union[str, Any]							     =		k.T
														UpperCAmelCase     :		Union[str, Any]							     =		o.T
														UpperCAmelCase     :		Tuple							     =		q.T
														UpperCAmelCase     :		str							     =		v.T

														# Block i, layer 1 (MLP).
														UpperCAmelCase     :		Optional[int]							     =		tax_layer_norm_lookup(_lowercase	,				_lowercase	,				"""encoder"""	,				"""pre_mlp_layer_norm"""		)
														UpperCAmelCase     ,     UpperCAmelCase     :		List[Any]							     =		tax_mlp_lookup(_lowercase	,				_lowercase	,				"""encoder"""	,				_lowercase		)
														UpperCAmelCase     :		Any							     =		layer_norm
														if split_mlp_wi:
																					UpperCAmelCase     :		Tuple							     =		wi[0].T
																					UpperCAmelCase     :		List[Any]							     =		wi[1].T
														else:
																					UpperCAmelCase     :		List[Any]							     =		wi.T
														UpperCAmelCase     :		Tuple							     =		wo.T
														if scalable_attention:
																					# convert the rel_embedding of each layer
																					UpperCAmelCase     :		Union[str, Any]							     =		tax_relpos_bias_lookup(
																					    _lowercase	,				_lowercase	,				"""encoder"""		).T

							UpperCAmelCase     :		List[Any]							     =		old["""encoder/encoder_norm/scale"""]

							if not scalable_attention:
														UpperCAmelCase     :		Optional[int]							     =		tax_relpos_bias_lookup(
														    _lowercase	,				0	,				"""encoder"""		).T
														UpperCAmelCase     :		List[Any]							     =		tax_relpos_bias_lookup(
														    _lowercase	,				0	,				"""decoder"""		).T

							if not is_encoder_only:
														# Decoder.
														for i in range(_lowercase		):
																					# Block i, layer 0 (Self Attention).
																					UpperCAmelCase     :		List[Any]							     =		tax_layer_norm_lookup(_lowercase	,				_lowercase	,				"""decoder"""	,				"""pre_self_attention_layer_norm"""		)
																					UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Dict							     =		tax_attention_lookup(_lowercase	,				_lowercase	,				"""decoder"""	,				"""self_attention"""		)
																					UpperCAmelCase     :		List[Any]							     =		layer_norm
																					UpperCAmelCase     :		str							     =		k.T
																					UpperCAmelCase     :		Dict							     =		o.T
																					UpperCAmelCase     :		str							     =		q.T
																					UpperCAmelCase     :		Optional[int]							     =		v.T

																					# Block i, layer 1 (Cross Attention).
																					UpperCAmelCase     :		Union[str, Any]							     =		tax_layer_norm_lookup(_lowercase	,				_lowercase	,				"""decoder"""	,				"""pre_cross_attention_layer_norm"""		)
																					UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Any							     =		tax_attention_lookup(_lowercase	,				_lowercase	,				"""decoder"""	,				"""encoder_decoder_attention"""		)
																					UpperCAmelCase     :		Optional[int]							     =		layer_norm
																					UpperCAmelCase     :		Union[str, Any]							     =		k.T
																					UpperCAmelCase     :		Tuple							     =		o.T
																					UpperCAmelCase     :		int							     =		q.T
																					UpperCAmelCase     :		Optional[int]							     =		v.T

																					# Block i, layer 2 (MLP).
																					UpperCAmelCase     :		Any							     =		tax_layer_norm_lookup(_lowercase	,				_lowercase	,				"""decoder"""	,				"""pre_mlp_layer_norm"""		)
																					UpperCAmelCase     ,     UpperCAmelCase     :		List[str]							     =		tax_mlp_lookup(_lowercase	,				_lowercase	,				"""decoder"""	,				_lowercase		)
																					UpperCAmelCase     :		List[str]							     =		layer_norm
																					if split_mlp_wi:
																												UpperCAmelCase     :		int							     =		wi[0].T
																												UpperCAmelCase     :		Optional[int]							     =		wi[1].T
																					else:
																												UpperCAmelCase     :		Union[str, Any]							     =		wi.T
																					UpperCAmelCase     :		Tuple							     =		wo.T

																					if scalable_attention:
																												# convert the rel_embedding of each layer
																												UpperCAmelCase     :		Tuple							     =		tax_relpos_bias_lookup(_lowercase	,				_lowercase	,				"""decoder"""		).T

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

							return new



def       __lowerCamelCase  (					_lowercase	,				_lowercase		)      ->			Tuple:
							UpperCAmelCase     :		Union[str, Any]							     =		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     :		Union[str, Any]							     =		state_dict["""shared.weight"""]

							if not is_encoder_only:
														if "decoder.embed_tokens.weight" not in state_dict:
																					UpperCAmelCase     :		Optional[int]							     =		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     :		Any							     =		state_dict["""shared.weight"""]

							return state_dict



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase	,				_lowercase	,				_lowercase		)      ->			Any:
							UpperCAmelCase     :		Union[str, Any]							     =		checkpoints.load_tax_checkpoint(_lowercase		)
							UpperCAmelCase     :		Optional[int]							     =		convert_tax_to_pytorch(
							    _lowercase	,				num_layers=config.num_layers	,				is_encoder_only=_lowercase	,				scalable_attention=_lowercase		)
							UpperCAmelCase     :		str							     =		make_state_dict(_lowercase	,				_lowercase		)
							model.load_state_dict(_lowercase	,				strict=_lowercase		)



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase	,				_lowercase = False	,				_lowercase = False	,				)      ->			Optional[int]:
							UpperCAmelCase     :		Optional[Any]							     =		MTaConfig.from_json_file(_lowercase		)
							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     :		Optional[int]							     =		UMTaEncoderModel(_lowercase		)
							else:
														UpperCAmelCase     :		Dict							     =		UMTaForConditionalGeneration(_lowercase		)

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

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

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


if __name__ == "__main__":
							a   :			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,
							)
							a   :			List[str]					  =    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,
							)






'''simple docstring'''




import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
							from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
							from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
							from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


a   :			str					  =    DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

a   :			int					  =    """main"""
# Default branch name
a   :			Any					  =    """f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"""
# One particular commit (not the top of `main`)
a   :			str					  =    """aaaaaaa"""
# This commit does not exist, so we should 404.

a   :			int					  =    """d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"""
# Sha-1 of config.json on the top of `main`, for checking purposes
a   :			Any					  =    """4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"""



@contextlib.contextmanager
def       __lowerCamelCase  (					)      ->			List[str]:
							print("""Welcome!"""		)
							yield
							print("""Bye!"""		)



@contextlib.contextmanager
def       __lowerCamelCase  (					)      ->			Optional[int]:
							print("""Bonjour!"""		)
							yield
							print("""Au revoir!"""		)

class 			UpperCamelCase_ (							unittest.TestCase			):





	def  _lowercase(  self    )			->			List[Any]:
								# If the spec is missing, importlib would not be able to import the module dynamically.
								assert transformers.__spec__ is not None
								assert importlib.util.find_spec("""transformers"""    ) is not None




class 			UpperCamelCase_ (							unittest.TestCase			):





	@unittest.mock.patch("""sys.stdout"""		,		new_callable=io.StringIO    )
	def  _lowercase(  self		,		A    )			->			Tuple:
								with ContextManagers([]    ):
															print("""Transformers are awesome!"""    )
								# The print statement adds a new line at the end of the output
								self.assertEqual(mock_stdout.getvalue()		,		"""Transformers are awesome!\n"""    )





	@unittest.mock.patch("""sys.stdout"""		,		new_callable=io.StringIO    )
	def  _lowercase(  self		,		A    )			->			Dict:
								with ContextManagers([context_en()]    ):
															print("""Transformers are awesome!"""    )
								# The output should be wrapped with an English welcome and goodbye
								self.assertEqual(mock_stdout.getvalue()		,		"""Welcome!\nTransformers are awesome!\nBye!\n"""    )





	@unittest.mock.patch("""sys.stdout"""		,		new_callable=io.StringIO    )
	def  _lowercase(  self		,		A    )			->			Union[str, Any]:
								with ContextManagers([context_fr(), context_en()]    ):
															print("""Transformers are awesome!"""    )
								# The output should be wrapped with an English and French welcome and goodbye
								self.assertEqual(mock_stdout.getvalue()		,		"""Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n"""    )





	@require_torch
	def  _lowercase(  self    )			->			Optional[int]:
								self.assertEqual(find_labels(A    )		,		["""labels"""]    )
								self.assertEqual(find_labels(A    )		,		["""labels""", """next_sentence_label"""]    )
								self.assertEqual(find_labels(A    )		,		["""start_positions""", """end_positions"""]    )

								class 			UpperCamelCase_ (							__magic_name__			):
									pass




								self.assertEqual(find_labels(A    )		,		["""labels"""]    )





	@require_tf
	def  _lowercase(  self    )			->			int:
								self.assertEqual(find_labels(A    )		,		["""labels"""]    )
								self.assertEqual(find_labels(A    )		,		["""labels""", """next_sentence_label"""]    )
								self.assertEqual(find_labels(A    )		,		["""start_positions""", """end_positions"""]    )

								class 			UpperCamelCase_ (							__magic_name__			):
									pass




								self.assertEqual(find_labels(A    )		,		["""labels"""]    )






	@require_flax
	def  _lowercase(  self    )			->			Any:
								# Flax models don't have labels
								self.assertEqual(find_labels(A    )		,		[]    )
								self.assertEqual(find_labels(A    )		,		[]    )
								self.assertEqual(find_labels(A    )		,		[]    )

								class 			UpperCamelCase_ (							__magic_name__			):
									pass




								self.assertEqual(find_labels(A    )		,		[]    )

