


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		(  __lowerCamelCase	:    Union[str, Any]				,   __lowerCamelCase	:    Optional[Any]				,   __lowerCamelCase	:    Tuple	)  ->						Any:
					return params[f"""{prefix}/{prefix}/relpos_bias/rel_embedding"""][:, i, :]






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Optional[int]				,   __lowerCamelCase	:    Optional[Any]				,   __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    int="attention"	)  ->						Dict:
					__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 					__lowerCAmelCase		(  __lowerCamelCase	:    Any				,   __lowerCamelCase	:    int				,   __lowerCamelCase	:    Union[str, Any]				,   __lowerCamelCase	:    Tuple=False	)  ->						List[Any]:
					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 					__lowerCAmelCase		(  __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    str				,   __lowerCamelCase	:    Optional[Any]				,   __lowerCamelCase	:    str	)  ->						List[Any]:
					return params[f"""{prefix}/{prefix}/{layer_name}/scale"""][:, i]






def 					__lowerCAmelCase		(  __lowerCamelCase	:    dict				,   *, __lowerCamelCase	:    int				,   __lowerCamelCase	:    bool				,   __lowerCamelCase	:    bool = False	)  ->						Optional[Any]:
					__lowerCAmelCase	   =traverse_util.flatten_dict(variables["""target"""]	)
					__lowerCAmelCase	   ={"""/""".join(__lowerCamelCase	): 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:"""				,   __lowerCamelCase	)

					__lowerCAmelCase	   =collections.OrderedDict()

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

					# Encoder.
					for i in range(__lowerCamelCase	):
										# Block i, layer 0 (Self Attention).
										__lowerCAmelCase	   =tax_layer_norm_lookup(__lowerCamelCase				,   __lowerCamelCase				,   """encoder"""				,   """pre_attention_layer_norm"""	)
										__lowerCAmelCase   ,	__lowerCAmelCase   ,	__lowerCAmelCase   ,	__lowerCAmelCase	   =tax_attention_lookup(__lowerCamelCase				,   __lowerCamelCase				,   """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(__lowerCamelCase				,   __lowerCamelCase				,   """encoder"""				,   """pre_mlp_layer_norm"""	)
										__lowerCAmelCase   ,	__lowerCAmelCase	   =tax_mlp_lookup(__lowerCamelCase				,   __lowerCamelCase				,   """encoder"""				,   __lowerCamelCase	)
										__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(
															    __lowerCamelCase				,   __lowerCamelCase				,   """encoder"""	).T

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

					if not scalable_attention:
										__lowerCAmelCase	   =tax_relpos_bias_lookup(
										    __lowerCamelCase				,   0				,   """encoder"""	).T
										__lowerCAmelCase	   =tax_relpos_bias_lookup(
										    __lowerCamelCase				,   0				,   """decoder"""	).T

					if not is_encoder_only:
										# Decoder.
										for i in range(__lowerCamelCase	):
															# Block i, layer 0 (Self Attention).
															__lowerCAmelCase	   =tax_layer_norm_lookup(__lowerCamelCase				,   __lowerCamelCase				,   """decoder"""				,   """pre_self_attention_layer_norm"""	)
															__lowerCAmelCase   ,	__lowerCAmelCase   ,	__lowerCAmelCase   ,	__lowerCAmelCase	   =tax_attention_lookup(__lowerCamelCase				,   __lowerCamelCase				,   """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(__lowerCamelCase				,   __lowerCamelCase				,   """decoder"""				,   """pre_cross_attention_layer_norm"""	)
															__lowerCAmelCase   ,	__lowerCAmelCase   ,	__lowerCAmelCase   ,	__lowerCAmelCase	   =tax_attention_lookup(__lowerCamelCase				,   __lowerCamelCase				,   """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(__lowerCamelCase				,   __lowerCamelCase				,   """decoder"""				,   """pre_mlp_layer_norm"""	)
															__lowerCAmelCase   ,	__lowerCAmelCase	   =tax_mlp_lookup(__lowerCamelCase				,   __lowerCamelCase				,   """decoder"""				,   __lowerCamelCase	)
															__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(__lowerCamelCase				,   __lowerCamelCase				,   """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 					__lowerCAmelCase		(  __lowerCamelCase	:    Union[str, Any]				,   __lowerCamelCase	:    bool	)  ->						str:
					__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 					__lowerCAmelCase		(  __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    int				,   __lowerCamelCase	:    List[Any]				,   __lowerCamelCase	:    Union[str, Any]				,   __lowerCamelCase	:    Tuple	)  ->						Any:
					__lowerCAmelCase	   =checkpoints.load_tax_checkpoint(__lowerCamelCase	)
					__lowerCAmelCase	   =convert_tax_to_pytorch(
					    __lowerCamelCase				,   num_layers=config.num_layers				,   is_encoder_only=__lowerCamelCase				,   scalable_attention=__lowerCamelCase	)
					__lowerCAmelCase	   =make_state_dict(__lowerCamelCase				,   __lowerCamelCase	)
					model.load_state_dict(__lowerCamelCase				,   strict=__lowerCamelCase	)






def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    Optional[Any]				,   __lowerCamelCase	:    Optional[int]				,   __lowerCamelCase	:    bool = False				,   __lowerCamelCase	:    bool = False				,   )  ->						List[str]:
					__lowerCAmelCase	   =MTaConfig.from_json_file(__lowerCamelCase	)
					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(__lowerCamelCase	)
					else:
										__lowerCAmelCase	   =UMTaForConditionalGeneration(__lowerCamelCase	)

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

					# Save pytorch-model
					print(f"""Save PyTorch model to {pytorch_dump_path}"""	)
					model.save_pretrained(__lowerCamelCase	)

					# Verify that we can load the checkpoint.
					model.from_pretrained(__lowerCamelCase	)
					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
				)
				parser.add_argument(
				    '''--scalable_attention''',
				    action='''store_true''',
				    help='''Whether the model uses scaled attention (umt5 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,
				    args.scalable_attention,
				)









import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
				from sagemaker import Session, TrainingJobAnalytics
				from sagemaker.huggingface import HuggingFace






@pytest.mark.skipif(
    literal_eval(os.getenv("TEST_SAGEMAKER"	,    "False"				)				) is not True	,    reason="Skipping test because should only be run when releasing minor transformers version"	,    )
@pytest.mark.usefixtures("sm_env"				)
@parameterized_class(
    [
        {
            "framework": "pytorch",
            "script": "run_glue.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 650, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "pytorch",
            "script": "run_ddp.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "tensorflow",
            "script": "run_tf_dist.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.6, "eval_loss": 0.7},
        },
    ]				)
class 						__a (     unittest.TestCase				):



	def 		UpperCamelCase							(						self							:    Dict)->							Any:
						if self.framework == "pytorch":
											subprocess.run(
											    F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()				,	encoding="""utf-8"""				,	check=snake_case_				,	)
						assert hasattr(self				,	"""env""")



	def 		UpperCamelCase							(						self							:    Dict				,	snake_case_							:    str)->							Dict:
						__lowerCAmelCase	   =F"""{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"""
						# distributed data settings
						__lowerCAmelCase	   ={"""smdistributed""": {"""dataparallel""": {"""enabled""": True}}} if self.script != """run_ddp.py""" else None

						# creates estimator
						return HuggingFace(
						    entry_point=self.script				,	source_dir=self.env.test_path				,	role=self.env.role				,	image_uri=self.env.image_uri				,	base_job_name=snake_case_				,	instance_count=snake_case_				,	instance_type=self.instance_type				,	debugger_hook_config=snake_case_				,	hyperparameters={**self.env.distributed_hyperparameters, """model_name_or_path""": self.model_name_or_path}				,	metric_definitions=self.env.metric_definitions				,	distribution=snake_case_				,	py_version="""py36"""				,	)



	def 		UpperCamelCase							(						self							:    int				,	snake_case_							:    List[Any])->							int:
						TrainingJobAnalytics(snake_case_).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv""")



	@parameterized.expand([(2,)])
	def 		UpperCamelCase							(						self							:    Any				,	snake_case_							:    Dict)->							Optional[Any]:
						# create estimator
						__lowerCAmelCase	   =self.create_estimator(snake_case_)

						# run training
						estimator.fit()

						# result dataframe
						__lowerCAmelCase	   =TrainingJobAnalytics(estimator.latest_training_job.name).dataframe()

						# extract kpis
						__lowerCAmelCase	   =list(result_metrics_df[result_metrics_df.metric_name == """eval_accuracy"""]["""value"""])
						__lowerCAmelCase	   =list(result_metrics_df[result_metrics_df.metric_name == """eval_loss"""]["""value"""])
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						__lowerCAmelCase	   =(
						    Session().describe_training_job(estimator.latest_training_job.name).get("""TrainingTimeInSeconds"""				,	99_99_99)
						)

						# assert kpis
						assert train_runtime <= self.results["train_runtime"]
						assert all(t >= self.results["""eval_accuracy"""] for t in eval_accuracy)
						assert all(t <= self.results["""eval_loss"""] for t in eval_loss)

						# dump tests result into json file to share in PR
						with open(F"""{estimator.latest_training_job.name}.json"""				,	"""w""") as outfile:
											json.dump({"""train_time""": train_runtime, """eval_accuracy""": eval_accuracy, """eval_loss""": eval_loss}				,	snake_case_)




