






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__				:		Dict    ,    lowerCAmelCase__				:		Optional[Any]    ,    lowerCAmelCase__				:		Dict      ):





				'''simple docstring'''







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



def 					_lowerCamelCase(			lowerCAmelCase__				:		str    ,    lowerCAmelCase__				:		int    ,    lowerCAmelCase__				:		List[str]    ,    lowerCAmelCase__				:		str="attention"      ):





				'''simple docstring'''







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



def 					_lowerCamelCase(			lowerCAmelCase__				:		Dict    ,    lowerCAmelCase__				:		int    ,    lowerCAmelCase__				:		List[str]    ,    lowerCAmelCase__				:		Any=False      ):





				'''simple docstring'''







				if split_mlp_wi:
								SCREAMING_SNAKE_CASE_   :	Dict													=   params[F'''{prefix}/{prefix}/mlp/wi_0/kernel'''][:, i, :]
								SCREAMING_SNAKE_CASE_   :	Dict													=   params[F'''{prefix}/{prefix}/mlp/wi_1/kernel'''][:, i, :]
								SCREAMING_SNAKE_CASE_   :	str													=   (wi_a, wi_a)
				else:
								SCREAMING_SNAKE_CASE_   :	Any													=   params[F'''{prefix}/{prefix}/mlp/wi/kernel'''][:, i, :]

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



def 					_lowerCamelCase(			lowerCAmelCase__				:		List[Any]    ,    lowerCAmelCase__				:		Tuple    ,    lowerCAmelCase__				:		Any    ,    lowerCAmelCase__				:		List[str]      ):





				'''simple docstring'''







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



def 					_lowerCamelCase(			lowerCAmelCase__				:		dict    ,    *, lowerCAmelCase__				:		int    ,    lowerCAmelCase__				:		bool    ,    lowerCAmelCase__				:		bool = False      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   traverse_util.flatten_dict(variables['target']      )
				SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   {'/'.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
				SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   'encoder/encoder/mlp/wi_0/kernel' in old
				print('Split MLP:'    ,    lowerCAmelCase__      )

				SCREAMING_SNAKE_CASE_   :	Any													=   collections.OrderedDict()

				# Shared embeddings.
				SCREAMING_SNAKE_CASE_   :	Optional[int]													=   old['token_embedder/embedding']

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

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

				SCREAMING_SNAKE_CASE_   :	Dict													=   old['encoder/encoder_norm/scale']

				if not scalable_attention:
								SCREAMING_SNAKE_CASE_   :	Any													=   tax_relpos_bias_lookup(
								    lowerCAmelCase__    ,    0    ,    'encoder'      ).T
								SCREAMING_SNAKE_CASE_   :	str													=   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).
												SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   tax_layer_norm_lookup(lowerCAmelCase__    ,    lowerCAmelCase__    ,    'decoder'    ,    'pre_self_attention_layer_norm'      )
												SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_   :	Any													=   tax_attention_lookup(lowerCAmelCase__    ,    lowerCAmelCase__    ,    'decoder'    ,    'self_attention'      )
												SCREAMING_SNAKE_CASE_   :	List[str]													=   layer_norm
												SCREAMING_SNAKE_CASE_   :	Tuple													=   k.T
												SCREAMING_SNAKE_CASE_   :	Tuple													=   o.T
												SCREAMING_SNAKE_CASE_   :	int													=   q.T
												SCREAMING_SNAKE_CASE_   :	str													=   v.T

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

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

												if scalable_attention:
																# convert the rel_embedding of each layer
																SCREAMING_SNAKE_CASE_   :	Tuple													=   tax_relpos_bias_lookup(lowerCAmelCase__    ,    lowerCAmelCase__    ,    'decoder'      ).T

								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   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:
												SCREAMING_SNAKE_CASE_   :	Optional[int]													=   old['decoder/logits_dense/kernel'].T

				return new



def 					_lowerCamelCase(			lowerCAmelCase__				:		Any    ,    lowerCAmelCase__				:		bool      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Optional[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:
								SCREAMING_SNAKE_CASE_   :	List[Any]													=   state_dict['shared.weight']

				if not is_encoder_only:
								if "decoder.embed_tokens.weight" not in state_dict:
												SCREAMING_SNAKE_CASE_   :	List[str]													=   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.'      )
												SCREAMING_SNAKE_CASE_   :	Optional[int]													=   state_dict['shared.weight']

				return state_dict



def 					_lowerCamelCase(			lowerCAmelCase__				:		List[Any]    ,    lowerCAmelCase__				:		List[str]    ,    lowerCAmelCase__				:		Dict    ,    lowerCAmelCase__				:		List[Any]    ,    lowerCAmelCase__				:		Optional[Any]      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	List[str]													=   checkpoints.load_tax_checkpoint(lowerCAmelCase__      )
				SCREAMING_SNAKE_CASE_   :	Dict													=   convert_tax_to_pytorch(
				    lowerCAmelCase__    ,    num_layers=config.num_layers    ,    is_encoder_only=lowerCAmelCase__    ,    scalable_attention=lowerCAmelCase__      )
				SCREAMING_SNAKE_CASE_   :	Optional[int]													=   make_state_dict(lowerCAmelCase__    ,    lowerCAmelCase__      )
				model.load_state_dict(lowerCAmelCase__    ,    strict=lowerCAmelCase__      )



def 					_lowerCamelCase(			lowerCAmelCase__				:		int    ,    lowerCAmelCase__				:		Optional[Any]    ,    lowerCAmelCase__				:		str    ,    lowerCAmelCase__				:		bool = False    ,    lowerCAmelCase__				:		bool = False    ,    ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	str													=   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:
								SCREAMING_SNAKE_CASE_   :	str													=   UMTaEncoderModel(lowerCAmelCase__      )
				else:
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   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__":
				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 importlib
import inspect
import os
import re


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_config_docstrings.py
A        =			'src/transformers'


# This is to make sure the transformers module imported is the one in the repo.
A        =			importlib.util.spec_from_file_location(
    'transformers',
    os.path.join(PATH_TO_TRANSFORMERS, '__init__.py'),
    submodule_search_locations=[PATH_TO_TRANSFORMERS],
)
A        =			spec.loader.load_module()

A        =			transformers.models.auto.configuration_auto.CONFIG_MAPPING

# Regex pattern used to find the checkpoint mentioned in the docstring of `config_class`.
# For example, `[bert-base-uncased](https://huggingface.co/bert-base-uncased)`
A        =			re.compile('\[(.+?)\]\((https://huggingface\.co/.+?)\)')


A        =			{
    'CLIPConfigMixin',
    'DecisionTransformerConfigMixin',
    'EncoderDecoderConfigMixin',
    'RagConfigMixin',
    'SpeechEncoderDecoderConfigMixin',
    'VisionEncoderDecoderConfigMixin',
    'VisionTextDualEncoderConfigMixin',
}



def 					_lowerCamelCase(			):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Optional[int]													=   []

				for config_class in list(CONFIG_MAPPING.values()      ):
								SCREAMING_SNAKE_CASE_   :	int													=   False

								# source code of `config_class`
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   inspect.getsource(lowerCAmelCase__      )
								SCREAMING_SNAKE_CASE_   :	List[Any]													=   _re_checkpoint.findall(lowerCAmelCase__      )

								for checkpoint in checkpoints:
												# Each `checkpoint` is a tuple of a checkpoint name and a checkpoint link.
												# For example, `('bert-base-uncased', 'https://huggingface.co/bert-base-uncased')`
												SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_   :	Tuple													=   checkpoint

												# verify the checkpoint name corresponds to the checkpoint link
												SCREAMING_SNAKE_CASE_   :	List[str]													=   F'''https://huggingface.co/{ckpt_name}'''
												if ckpt_link == ckpt_link_from_name:
																SCREAMING_SNAKE_CASE_   :	Optional[int]													=   True
																break

								SCREAMING_SNAKE_CASE_   :	Tuple													=   config_class.__name__
								if not checkpoint_found and name not in CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK:
												configs_without_checkpoint.append(lowerCAmelCase__      )

				if len(lowerCAmelCase__      ) > 0:
								SCREAMING_SNAKE_CASE_   :	str													=   '\n'.join(sorted(lowerCAmelCase__      )      )
								raise ValueError(F'''The following configurations don\'t contain any valid checkpoint:\n{message}'''      )


if __name__ == "__main__":
				check_config_docstrings_have_checkpoints()