






from ...processing_utils import ProcessorMixin


class    _lowerCAmelCase			(		lowercase_      ):



				"""simple docstring"""



				_lowercase				:   str      =  '''WhisperFeatureExtractor'''
				_lowercase				:   int      =  '''WhisperTokenizer'''



				def __init__(      self    :    Dict			,		UpperCamelCase__    :    List[str]			,		UpperCamelCase__    :    Tuple):






											'''simple docstring'''







											super().__init__(UpperCamelCase__			,		UpperCamelCase__)
											snake_case__	  =       self.feature_extractor
											snake_case__	  =       False



				def 					__magic_name__							(      self    :    Tuple			,		UpperCamelCase__    :    Tuple=None			,		UpperCamelCase__    :    Optional[int]=None			,		UpperCamelCase__    :    Tuple=True):






											'''simple docstring'''







											return self.tokenizer.get_decoder_prompt_ids(task=UpperCamelCase__			,		language=UpperCamelCase__			,		no_timestamps=UpperCamelCase__)



				def __call__(      self    :    int			,		*UpperCamelCase__    :    Dict			,		**UpperCamelCase__    :    Union[str, Any]):






											'''simple docstring'''







											if self._in_target_context_manager:
																		return self.current_processor(*UpperCamelCase__			,		**UpperCamelCase__)

											snake_case__	  =       kwargs.pop("""audio"""			,		UpperCamelCase__)
											snake_case__	  =       kwargs.pop("""sampling_rate"""			,		UpperCamelCase__)
											snake_case__	  =       kwargs.pop("""text"""			,		UpperCamelCase__)
											if len(UpperCamelCase__) > 0:
																		snake_case__	  =       args[0]
																		snake_case__	  =       args[1:]

											if audio is None and text is None:
																		raise ValueError("""You need to specify either an `audio` or `text` input to process.""")

											if audio is not None:
																		snake_case__	  =       self.feature_extractor(UpperCamelCase__			,		*UpperCamelCase__			,		sampling_rate=UpperCamelCase__			,		**UpperCamelCase__)
											if text is not None:
																		snake_case__	  =       self.tokenizer(UpperCamelCase__			,		**UpperCamelCase__)

											if text is None:
																		return inputs

											elif audio is None:
																		return encodings
											else:
																		snake_case__	  =       encodings["""input_ids"""]
																		return inputs



				def 					__magic_name__							(      self    :    Any			,		*UpperCamelCase__    :    Dict			,		**UpperCamelCase__    :    Union[str, Any]):






											'''simple docstring'''







											return self.tokenizer.batch_decode(*UpperCamelCase__			,		**UpperCamelCase__)



				def 					__magic_name__							(      self    :    Tuple			,		*UpperCamelCase__    :    Dict			,		**UpperCamelCase__    :    str):






											'''simple docstring'''







											return self.tokenizer.decode(*UpperCamelCase__			,		**UpperCamelCase__)



				def 					__magic_name__							(      self    :    Optional[Any]			,		UpperCamelCase__    :    str			,		UpperCamelCase__    :    Dict="np"):






											'''simple docstring'''







											return self.tokenizer.get_prompt_ids(UpperCamelCase__			,		return_tensors=UpperCamelCase__)








import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()

def _UpperCAmelCase   (	a						:	Optional[Any]  ,					a						:	Dict  ,					a						:	Union[str, Any]=None ):
							# set parameter of one layer
							assert torch_layer.weight.shape == weight.shape, F'''{torch_layer} layer.weight does not match'''
							snake_case__	  =       nn.Parameter(a )
							if bias is not None:
														assert torch_layer.bias.shape == bias.shape, F'''{torch_layer} layer.bias does not match'''
														snake_case__	  =       nn.Parameter(a )

def _UpperCAmelCase   (	a						:	List[Any]  ,					a						:	List[str]  ,					a						:	List[Any] ):
							# set torch weights for 1-to-1 comparison
							snake_case__	  =       np.asarray(weights[0] )
							snake_case__	  =       np.asarray(weights[1] )
							snake_case__	  =       np.asarray(weights[2] )

							set_param(
							    torch_layer.self_attention.query_key  ,					torch.tensor(a ).transpose(1  ,					2 ).contiguous().view(-1  ,					a )  ,					)
							set_param(
							    torch_layer.self_attention.value  ,					torch.tensor(a ).transpose(1  ,					2 ).contiguous().view(-1  ,					a )  ,					)
							set_param(
							    torch_layer.output.dense  ,					torch.tensor(a ).view(-1  ,					a ).contiguous().transpose(0  ,					1 )  ,					)

def _UpperCAmelCase   (	a						:	Optional[int]  ,					a						:	Dict  ,					a						:	Optional[int] ):
							# set torch weights for 1-to-1 comparison
							snake_case__	  =       np.asarray(weights[0] )
							snake_case__	  =       np.asarray(weights[1] )
							snake_case__	  =       np.asarray(weights[2] )
							snake_case__	  =       np.asarray(weights[3] )

							set_param(
							    torch_layer.self_attention.query  ,					torch.tensor(a ).transpose(1  ,					2 ).contiguous().view(-1  ,					a )  ,					)
							set_param(
							    torch_layer.self_attention.key  ,					torch.tensor(a ).transpose(1  ,					2 ).contiguous().view(-1  ,					a )  ,					)
							set_param(
							    torch_layer.self_attention.value  ,					torch.tensor(a ).transpose(1  ,					2 ).contiguous().view(-1  ,					a )  ,					)
							set_param(
							    torch_layer.output.dense  ,					torch.tensor(a ).view(-1  ,					a ).contiguous().transpose(0  ,					1 )  ,					)

def _UpperCAmelCase   (	a						:	List[Any]  ,					a						:	Tuple  ,					a						:	List[Any] ):
							# layernorm 1
							snake_case__	  =       weights[0][0][0]
							snake_case__	  =       np.asarray(layer_norm_a[0] )
							snake_case__	  =       np.asarray(layer_norm_a[1] )
							set_param(
							    torch_block.attention.layer_norm  ,					torch.tensor(a )  ,					torch.tensor(a )  ,					)

							# lsh weights + output
							snake_case__	  =       weights[0][1]
							if len(a ) < 4:
														set_layer_weights_in_torch_lsh(a  ,					torch_block.attention  ,					a )
							else:
														set_layer_weights_in_torch_local(a  ,					torch_block.attention  ,					a )

							# intermediate weighs
							snake_case__	  =       weights[2][0][1][2]

							# Chunked Feed Forward
							if len(a ) == 4:
														snake_case__	  =       intermediate_weights[2]

							# layernorm 2
							snake_case__	  =       np.asarray(intermediate_weights[0][0] )
							snake_case__	  =       np.asarray(intermediate_weights[0][1] )
							set_param(
							    torch_block.feed_forward.layer_norm  ,					torch.tensor(a )  ,					torch.tensor(a )  ,					)

							# intermediate dense
							snake_case__	  =       np.asarray(intermediate_weights[1][0] )
							snake_case__	  =       np.asarray(intermediate_weights[1][1] )
							set_param(
							    torch_block.feed_forward.dense.dense  ,					torch.tensor(a ).transpose(0  ,					1 ).contiguous()  ,					torch.tensor(a )  ,					)

							# intermediate out
							snake_case__	  =       np.asarray(intermediate_weights[4][0] )
							snake_case__	  =       np.asarray(intermediate_weights[4][1] )
							set_param(
							    torch_block.feed_forward.output.dense  ,					torch.tensor(a ).transpose(0  ,					1 ).contiguous()  ,					torch.tensor(a )  ,					)

def _UpperCAmelCase   (	a						:	List[Any]  ,					a						:	Dict  ,					a						:	Union[str, Any] ):
							# reformer model
							snake_case__	  =       torch_model.reformer

							# word embeds
							snake_case__	  =       np.asarray(weights[1] )
							set_param(
							    torch_model_reformer.embeddings.word_embeddings  ,					torch.tensor(a )  ,					)

							if isinstance(weights[3]  ,					a ):
														snake_case__	  =       torch_model_reformer.embeddings.position_embeddings
														for emb_idx in range(len(position_embeddings.weights ) ):
																					snake_case__	  =       np.asarray(weights[3][emb_idx][0] )
																					assert (
																					    position_embeddings.weights[emb_idx].shape == emb_weights.shape
																					), F'''{position_embeddings[emb_idx]} emb does not match'''
																					snake_case__	  =       nn.Parameter(torch.tensor(a ) )

							snake_case__	  =       weights[5]
							assert len(torch_model_reformer.encoder.layers ) * 4 == len(
							    a ), "HF and trax model do not have the same number of layers"
							for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers ):
														snake_case__	  =       trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
														set_block_weights_in_torch(a  ,					a  ,					a )

							# output layer norm
							snake_case__	  =       np.asarray(weights[7][0] )
							snake_case__	  =       np.asarray(weights[7][1] )
							set_param(
							    torch_model_reformer.encoder.layer_norm  ,					torch.tensor(a )  ,					torch.tensor(a )  ,					)

							# output embeddings
							snake_case__	  =       np.asarray(weights[9][0] )
							snake_case__	  =       np.asarray(weights[9][1] )
							set_param(
							    torch_model.lm_head.decoder  ,					torch.tensor(a ).transpose(0  ,					1 ).contiguous()  ,					torch.tensor(a )  ,					)

def _UpperCAmelCase   (	a						:	Any  ,					a						:	str  ,					a						:	str ):
							# Initialise PyTorch model
							snake_case__	  =       ReformerConfig.from_json_file(a )
							print(F'''Building PyTorch model from configuration: {config}''' )
							snake_case__	  =       ReformerModelWithLMHead(a )

							with open(a  ,					"""rb""" ) as f:
														snake_case__	  =       pickle.load(a )["""weights"""]

							set_model_weights_in_torch(a  ,					a  ,					config.hidden_size )

							# Save pytorch-model
							print(F'''Save PyTorch model to {pytorch_dump_path}''' )
							torch.save(model.state_dict()  ,					a )


if __name__ == "__main__":
	a__       =     argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    """--trax_model_pkl_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
	)
	parser.add_argument(
	    """--config_file""",
	    default=None,
	    type=str,
	    required=True,
	    help=(
	        """The config json file corresponding to the pre-trained Reformer model. \n"""
	        """This specifies the model architecture."""
	    ),
	)
	parser.add_argument(
	    """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
	)
	a__       =     parser.parse_args()
	convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)
