






from scipy.stats import pearsonr

import datasets


a__       =     "\nPearson correlation coefficient and p-value for testing non-correlation.\nThe Pearson correlation coefficient measures the linear relationship between two datasets. The calculation of the p-value relies on the assumption that each dataset is normally distributed. Like other correlation coefficients, this one varies between -1 and +1 with 0 implying no correlation. Correlations of -1 or +1 imply an exact linear relationship. Positive correlations imply that as x increases, so does y. Negative correlations imply that as x increases, y decreases.\nThe p-value roughly indicates the probability of an uncorrelated system producing datasets that have a Pearson correlation at least as extreme as the one computed from these datasets.\n"


a__       =     "\nArgs:\n    predictions (`list` of `int`): Predicted class labels, as returned by a model.\n    references (`list` of `int`): Ground truth labels.\n    return_pvalue (`boolean`): If `True`, returns the p-value, along with the correlation coefficient. If `False`, returns only the correlation coefficient. Defaults to `False`.\n\nReturns:\n    pearsonr (`float`): Pearson correlation coefficient. Minimum possible value is -1. Maximum possible value is 1. Values of 1 and -1 indicate exact linear positive and negative relationships, respectively. A value of 0 implies no correlation.\n    p-value (`float`): P-value, which roughly indicates the probability of an The p-value roughly indicates the probability of an uncorrelated system producing datasets that have a Pearson correlation at least as extreme as the one computed from these datasets. Minimum possible value is 0. Maximum possible value is 1. Higher values indicate higher probabilities.\n\nExamples:\n\n    Example 1-A simple example using only predictions and references.\n        >>> pearsonr_metric = datasets.load_metric(\"pearsonr\")\n        >>> results = pearsonr_metric.compute(predictions=[10, 9, 2.5, 6, 4], references=[1, 2, 3, 4, 5])\n        >>> print(round(results['pearsonr'], 2))\n        -0.74\n\n    Example 2-The same as Example 1, but that also returns the `p-value`.\n        >>> pearsonr_metric = datasets.load_metric(\"pearsonr\")\n        >>> results = pearsonr_metric.compute(predictions=[10, 9, 2.5, 6, 4], references=[1, 2, 3, 4, 5], return_pvalue=True)\n        >>> print(sorted(list(results.keys())))\n        ['p-value', 'pearsonr']\n        >>> print(round(results['pearsonr'], 2))\n        -0.74\n        >>> print(round(results['p-value'], 2))\n        0.15\n"


a__       =     "\n@article{2020SciPy-NMeth,\nauthor  = {Virtanen, Pauli and Gommers, Ralf and Oliphant, Travis E. and\n      Haberland, Matt and Reddy, Tyler and Cournapeau, David and\n      Burovski, Evgeni and Peterson, Pearu and Weckesser, Warren and\n      Bright, Jonathan and {van der Walt}, St{\'e}fan J. and\n      Brett, Matthew and Wilson, Joshua and Millman, K. Jarrod and\n      Mayorov, Nikolay and Nelson, Andrew R. J. and Jones, Eric and\n      Kern, Robert and Larson, Eric and Carey, C J and\n      Polat, Ilhan and Feng, Yu and Moore, Eric W. and\n      {VanderPlas}, Jake and Laxalde, Denis and Perktold, Josef and\n      Cimrman, Robert and Henriksen, Ian and Quintero, E. A. and\n      Harris, Charles R. and Archibald, Anne M. and\n      Ribeiro, Antonio H. and Pedregosa, Fabian and\n      {van Mulbregt}, Paul and {SciPy 1.0 Contributors}},\ntitle   = {{{SciPy} 1.0: Fundamental Algorithms for Scientific\n      Computing in Python}},\njournal = {Nature Methods},\nyear    = {2020},\nvolume  = {17},\npages   = {261--272},\nadsurl  = {https://rdcu.be/b08Wh},\ndoi = {10.1038/s41592-019-0686-2},\n}\n"


@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION     ,      _KWARGS_DESCRIPTION      )
class    _lowerCAmelCase			(		datasets.Metric      ):



				"""simple docstring"""



				def 					__magic_name__							(      self    :    int):






											'''simple docstring'''







											return datasets.MetricInfo(
											    description=_DESCRIPTION			,		citation=_CITATION			,		inputs_description=_KWARGS_DESCRIPTION			,		features=datasets.Features(
											        {
											            """predictions""": datasets.Value("""float"""),
											            """references""": datasets.Value("""float"""),
											        })			,		reference_urls=["""https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.pearsonr.html"""]			,		)



				def 					__magic_name__							(      self    :    Any			,		UpperCamelCase__    :    str			,		UpperCamelCase__    :    Tuple			,		UpperCamelCase__    :    List[str]=False):






											'''simple docstring'''







											if return_pvalue:
																		snake_case__	  =       pearsonr(_UpperCamelCase			,		_UpperCamelCase)
																		return {"pearsonr": results[0], "p-value": results[1]}
											else:
																		return {"pearsonr": float(pearsonr(_UpperCamelCase			,		_UpperCamelCase)[0])}








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)
