




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_							(						__a						,						__a						,						__a			)							->      List[Any]:





							"""simple docstring"""





							return params[f'{prefix}/{prefix}/relpos_bias/rel_embedding'][:, i, :]
def     lowerCAmelCase_							(						__a						,						__a						,						__a						,						__a="attention"			)							->      Optional[Any]:





							"""simple docstring"""





							SCREAMING_SNAKE_CASE      :			str						  =np.ascontiguousarray(params[f'{prefix}/{prefix}/{layer_name}/key/kernel'][:, i, :, :]			)
							SCREAMING_SNAKE_CASE      :			Optional[int]						  =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      :			Any						  =o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]						,						o_tmp.shape[2]			)
							SCREAMING_SNAKE_CASE      :			Optional[Any]						  =np.ascontiguousarray(params[f'{prefix}/{prefix}/{layer_name}/query/kernel'][:, i, :, :]			)
							SCREAMING_SNAKE_CASE      :			List[str]						  =q_tmp.reshape(q_tmp.shape[0]						,						q_tmp.shape[1] * q_tmp.shape[2]			)
							SCREAMING_SNAKE_CASE      :			str						  =np.ascontiguousarray(params[f'{prefix}/{prefix}/{layer_name}/value/kernel'][:, i, :, :]			)
							SCREAMING_SNAKE_CASE      :			str						  =v_tmp.reshape(v_tmp.shape[0]						,						v_tmp.shape[1] * v_tmp.shape[2]			)
							return k, o, q, v
def     lowerCAmelCase_							(						__a						,						__a						,						__a						,						__a=False			)							->      Union[str, Any]:





							"""simple docstring"""





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

							SCREAMING_SNAKE_CASE      :			Union[str, Any]						  =params[f'{prefix}/{prefix}/mlp/wo/kernel'][:, i, :]
							return wi, wo
def     lowerCAmelCase_							(						__a						,						__a						,						__a						,						__a			)							->      Tuple:





							"""simple docstring"""





							return params[f'{prefix}/{prefix}/{layer_name}/scale'][:, i]
def     lowerCAmelCase_							(						__a						,						*, __a						,						__a						,						__a = False			)							->      str:





							"""simple docstring"""





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

							SCREAMING_SNAKE_CASE      :			Any						  =collections.OrderedDict()

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

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

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

							SCREAMING_SNAKE_CASE      :			Any						  =old["encoder/encoder_norm/scale"]

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

																					# Block i, layer 1 (Cross Attention).
																					SCREAMING_SNAKE_CASE      :			str						  =tax_layer_norm_lookup(_lowercase						,						_lowercase						,						'''decoder'''						,						'''pre_cross_attention_layer_norm'''			)
																					SCREAMING_SNAKE_CASE      :			List[str]						  =tax_attention_lookup(_lowercase						,						_lowercase						,						'''decoder'''						,						'''encoder_decoder_attention'''			)
																					SCREAMING_SNAKE_CASE      :			str						  =layer_norm
																					SCREAMING_SNAKE_CASE      :			List[Any]						  =k.T
																					SCREAMING_SNAKE_CASE      :			str						  =o.T
																					SCREAMING_SNAKE_CASE      :			Any						  =q.T
																					SCREAMING_SNAKE_CASE      :			Optional[Any]						  =v.T

																					# Block i, layer 2 (MLP).
																					SCREAMING_SNAKE_CASE      :			List[Any]						  =tax_layer_norm_lookup(_lowercase						,						_lowercase						,						'''decoder'''						,						'''pre_mlp_layer_norm'''			)
																					SCREAMING_SNAKE_CASE      :			str						  =tax_mlp_lookup(_lowercase						,						_lowercase						,						'''decoder'''						,						_lowercase			)
																					SCREAMING_SNAKE_CASE      :			Optional[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      :			Any						  =wi.T
																					SCREAMING_SNAKE_CASE      :			int						  =wo.T

																					if scalable_attention:
																												# convert the rel_embedding of each layer
																												SCREAMING_SNAKE_CASE      :			List[Any]						  =tax_relpos_bias_lookup(_lowercase						,						_lowercase						,						'''decoder'''			).T

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

							return new
def     lowerCAmelCase_							(						__a						,						__a			)							->      Optional[Any]:





							"""simple docstring"""





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

							if not is_encoder_only:
														if "decoder.embed_tokens.weight" not in state_dict:
																					SCREAMING_SNAKE_CASE      :			Union[str, Any]						  =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      :			str						  =state_dict["shared.weight"]

							return state_dict
def     lowerCAmelCase_							(						__a						,						__a						,						__a						,						__a						,						__a			)							->      Optional[Any]:





							"""simple docstring"""





							SCREAMING_SNAKE_CASE      :			List[str]						  =checkpoints.load_tax_checkpoint(_lowercase			)
							SCREAMING_SNAKE_CASE      :			int						  =convert_tax_to_pytorch(
							    _lowercase						,						num_layers=config.num_layers						,						is_encoder_only=_lowercase						,						scalable_attention=_lowercase			)
							SCREAMING_SNAKE_CASE      :			Tuple						  =make_state_dict(_lowercase						,						_lowercase			)
							model.load_state_dict(_lowercase						,						strict=_lowercase			)
def     lowerCAmelCase_							(						__a						,						__a						,						__a						,						__a = False						,						__a = False						,						)							->      Optional[int]:





							"""simple docstring"""





							SCREAMING_SNAKE_CASE      :			Optional[int]						  =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:
														SCREAMING_SNAKE_CASE      :			Any						  =UMTaEncoderModel(_lowercase			)
							else:
														SCREAMING_SNAKE_CASE      :			Any						  =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				     =						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,
							)













"""simple docstring"""




import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


a		:					List[Any]	      =     logging.get_logger(__name__)

a		:					List[Any]	      =     {
    '''microsoft/unispeech-sat-base-100h-libri-ft''': (
        '''https://huggingface.co/microsoft/unispeech-sat-base-100h-libri-ft/resolve/main/config.json'''
    ),
    # See all UniSpeechSat models at https://huggingface.co/models?filter=unispeech_sat
}


class 			__UpperCamelCase							(		a__		):
							lowerCamelCase  :   Tuple    ="""unispeech-sat"""






							def __init__(  self ,   lowerCAmelCase__=32 ,   lowerCAmelCase__=768 ,   lowerCAmelCase__=12 ,   lowerCAmelCase__=12 ,   lowerCAmelCase__=3072 ,   lowerCAmelCase__="gelu" ,   lowerCAmelCase__=0.1 ,   lowerCAmelCase__=0.1 ,   lowerCAmelCase__=0.1 ,   lowerCAmelCase__=0.0 ,   lowerCAmelCase__=0.0 ,   lowerCAmelCase__=0.1 ,   lowerCAmelCase__=0.1 ,   lowerCAmelCase__=0.02 ,   lowerCAmelCase__=1E-5 ,   lowerCAmelCase__="group" ,   lowerCAmelCase__="gelu" ,   lowerCAmelCase__=(512, 512, 512, 512, 512, 512, 512) ,   lowerCAmelCase__=(5, 2, 2, 2, 2, 2, 2) ,   lowerCAmelCase__=(10, 3, 3, 3, 3, 2, 2) ,   lowerCAmelCase__=False ,   lowerCAmelCase__=128 ,   lowerCAmelCase__=16 ,   lowerCAmelCase__=False ,   lowerCAmelCase__=True ,   lowerCAmelCase__=0.05 ,   lowerCAmelCase__=10 ,   lowerCAmelCase__=2 ,   lowerCAmelCase__=0.0 ,   lowerCAmelCase__=10 ,   lowerCAmelCase__=0 ,   lowerCAmelCase__=320 ,   lowerCAmelCase__=2 ,   lowerCAmelCase__=0.1 ,   lowerCAmelCase__=100 ,   lowerCAmelCase__=256 ,   lowerCAmelCase__=256 ,   lowerCAmelCase__=0.1 ,   lowerCAmelCase__="mean" ,   lowerCAmelCase__=False ,   lowerCAmelCase__=False ,   lowerCAmelCase__=256 ,   lowerCAmelCase__=(512, 512, 512, 512, 1500) ,   lowerCAmelCase__=(5, 3, 3, 1, 1) ,   lowerCAmelCase__=(1, 2, 3, 1, 1) ,   lowerCAmelCase__=512 ,   lowerCAmelCase__=0 ,   lowerCAmelCase__=1 ,   lowerCAmelCase__=2 ,   lowerCAmelCase__=504 ,   **lowerCAmelCase__ ,   )      ->					List[str]:
								super().__init__(**lowerCAmelCase__ ,   pad_token_id=lowerCAmelCase__ ,   bos_token_id=lowerCAmelCase__ ,   eos_token_id=lowerCAmelCase__ )
								a :						Any      =							hidden_size
								a :						Tuple      =							feat_extract_norm
								a :						Tuple      =							feat_extract_activation
								a :						Dict      =							list(lowerCAmelCase__ )
								a :						int      =							list(lowerCAmelCase__ )
								a :						Optional[Any]      =							list(lowerCAmelCase__ )
								a :						int      =							conv_bias
								a :						str      =							num_conv_pos_embeddings
								a :						Dict      =							num_conv_pos_embedding_groups
								a :						Optional[int]      =							len(self.conv_dim )
								a :						int      =							num_hidden_layers
								a :						Any      =							intermediate_size
								a :						Any      =							hidden_act
								a :						List[Any]      =							num_attention_heads
								a :						Any      =							hidden_dropout
								a :						Union[str, Any]      =							attention_dropout
								a :						Tuple      =							activation_dropout
								a :						Dict      =							feat_proj_dropout
								a :						Optional[Any]      =							final_dropout
								a :						Union[str, Any]      =							layerdrop
								a :						str      =							layer_norm_eps
								a :						Optional[int]      =							initializer_range
								a :						Optional[int]      =							vocab_size
								a :						str      =							num_clusters
								a :						Any      =							do_stable_layer_norm
								a :						str      =							use_weighted_layer_sum

								if (
								    (len(self.conv_stride ) != self.num_feat_extract_layers)
								    or (len(self.conv_kernel ) != self.num_feat_extract_layers)
								    or (len(self.conv_dim ) != self.num_feat_extract_layers)
								):
									raise ValueError(
									    "Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="
									    " `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="
									    f""" {len(self.conv_dim )}`, `len(config.conv_stride) = {len(self.conv_stride )}`,"""
									    f""" `len(config.conv_kernel) = {len(self.conv_kernel )}`.""" )

								# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
								a :						List[Any]      =							apply_spec_augment
								a :						int      =							mask_time_prob
								a :						Optional[int]      =							mask_time_length
								a :						Dict      =							mask_time_min_masks
								a :						Optional[int]      =							mask_feature_prob
								a :						List[str]      =							mask_feature_length
								a :						Dict      =							mask_feature_min_masks

								# parameters for pretraining with codevector quantized representations
								a :						List[str]      =							num_codevectors_per_group
								a :						List[Any]      =							num_codevector_groups
								a :						Tuple      =							contrastive_logits_temperature
								a :						int      =							feat_quantizer_dropout
								a :						Optional[Any]      =							num_negatives
								a :						Optional[int]      =							codevector_dim
								a :						Tuple      =							proj_codevector_dim
								a :						Optional[int]      =							diversity_loss_weight

								# ctc loss
								a :						Dict      =							ctc_loss_reduction
								a :						Optional[Any]      =							ctc_zero_infinity

								# SequenceClassification-specific parameter. Feel free to ignore for other classes.
								a :						Dict      =							classifier_proj_size

								# XVector-specific parameters. Feel free to ignore for other classes.
								a :						Union[str, Any]      =							list(lowerCAmelCase__ )
								a :						List[str]      =							list(lowerCAmelCase__ )
								a :						Tuple      =							list(lowerCAmelCase__ )
								a :						Optional[int]      =							xvector_output_dim






							@property
							def        __a  (  self )      ->					Any:
								return functools.reduce(operator.mul ,   self.conv_stride ,   1 )


