


"""simple docstring"""




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__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__=None						):
 assert torch_layer.weight.shape == weight.shape, f'{torch_layer} layer.weight does not match'
 __lowerCamelCase		:							List[Any]		 =						nn.Parameter(__a						)
 if bias is not None:
  assert torch_layer.bias.shape == bias.shape, f'{torch_layer} layer.bias does not match'
  __lowerCamelCase		:							Tuple		 =						nn.Parameter(__a						)
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							int		 =						np.asarray(weights[0]						)
 __lowerCamelCase		:							Optional[Any]		 =						np.asarray(weights[1]						)
 __lowerCamelCase		:							Tuple		 =						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__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							List[Any]		 =						np.asarray(weights[0]						)
 __lowerCamelCase		:							Optional[Any]		 =						np.asarray(weights[1]						)
 __lowerCamelCase		:							List[str]		 =						np.asarray(weights[2]						)
 __lowerCamelCase		:							List[Any]		 =						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__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							Dict		 =						weights[0][0][0]
 __lowerCamelCase		:							Tuple		 =						np.asarray(layer_norm_a[0]						)
 __lowerCamelCase		:							Dict		 =						np.asarray(layer_norm_a[1]						)
 set_param(
     torch_block.attention.layer_norm      ,	torch.tensor(__a						)      ,	torch.tensor(__a						)      ,	)

 # lsh weights + output
 __lowerCamelCase		:							List[str]		 =						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
 __lowerCamelCase		:							str		 =						weights[2][0][1][2]

 # Chunked Feed Forward
 if len(__a						) == 4:
  __lowerCamelCase		:							Optional[Any]		 =						intermediate_weights[2]

 # layernorm 2
 __lowerCamelCase		:							str		 =						np.asarray(intermediate_weights[0][0]						)
 __lowerCamelCase		:							List[str]		 =						np.asarray(intermediate_weights[0][1]						)
 set_param(
     torch_block.feed_forward.layer_norm      ,	torch.tensor(__a						)      ,	torch.tensor(__a						)      ,	)

 # intermediate dense
 __lowerCamelCase		:							Tuple		 =						np.asarray(intermediate_weights[1][0]						)
 __lowerCamelCase		:							Dict		 =						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
 __lowerCamelCase		:							Union[str, Any]		 =						np.asarray(intermediate_weights[4][0]						)
 __lowerCamelCase		:							int		 =						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__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							Union[str, Any]		 =						torch_model.reformer

 # word embeds
 __lowerCamelCase		:							List[str]		 =						np.asarray(weights[1]						)
 set_param(
     torch_model_reformer.embeddings.word_embeddings      ,	torch.tensor(__a						)      ,	)

 if isinstance(weights[3]      ,	__a						):
  __lowerCamelCase		:							str		 =						torch_model_reformer.embeddings.position_embeddings
  for emb_idx in range(len(position_embeddings.weights						)						):
   __lowerCamelCase		:							Union[str, Any]		 =						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'
   __lowerCamelCase		:							int		 =						nn.Parameter(torch.tensor(__a						)						)

 __lowerCamelCase		:							Optional[Any]		 =						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						):
  __lowerCamelCase		:							Dict		 =						trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
  set_block_weights_in_torch(__a      ,	__a      ,	__a						)

 # output layer norm
 __lowerCamelCase		:							str		 =						np.asarray(weights[7][0]						)
 __lowerCamelCase		:							Any		 =						np.asarray(weights[7][1]						)
 set_param(
     torch_model_reformer.encoder.layer_norm      ,	torch.tensor(__a						)      ,	torch.tensor(__a						)      ,	)

 # output embeddings
 __lowerCamelCase		:							Dict		 =						np.asarray(weights[9][0]						)
 __lowerCamelCase		:							List[str]		 =						np.asarray(weights[9][1]						)
 set_param(
     torch_model.lm_head.decoder      ,	torch.tensor(__a						).transpose(0      ,	1						).contiguous()      ,	torch.tensor(__a						)      ,	)
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 __lowerCamelCase		:							Tuple		 =						ReformerConfig.from_json_file(__a						)
 print(f'Building PyTorch model from configuration: {config}'						)
 __lowerCamelCase		:							int		 =						ReformerModelWithLMHead(__a						)

 with open(__a      ,	'rb'						) as f:
  __lowerCamelCase		:							List[str]		 =						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__":
       lowercase_	   =     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.'
       )
       lowercase_	   =     parser.parse_args()
       convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)





import datasets


lowercase_	   =     '\\n@InProceedings{conneau2018xnli,\n  author = "Conneau, Alexis\n                 and Rinott, Ruty\n                 and Lample, Guillaume\n                 and Williams, Adina\n                 and Bowman, Samuel R.\n                 and Schwenk, Holger\n                 and Stoyanov, Veselin",\n  title = "XNLI: Evaluating Cross-lingual Sentence Representations",\n  booktitle = "Proceedings of the 2018 Conference on Empirical Methods\n               in Natural Language Processing",\n  year = "2018",\n  publisher = "Association for Computational Linguistics",\n  location = "Brussels, Belgium",\n}\n'

lowercase_	   =     '\\nXNLI is a subset of a few thousand examples from MNLI which has been translated\ninto a 14 different languages (some low-ish resource). As with MNLI, the goal is\nto predict textual entailment (does sentence A imply/contradict/neither sentence\nB) and is a classification task (given two sentences, predict one of three\nlabels).\n'

lowercase_	   =     '\nComputes XNLI score which is just simple accuracy.\nArgs:\n    predictions: Predicted labels.\n    references: Ground truth labels.\nReturns:\n    \'accuracy\': accuracy\nExamples:\n\n    >>> predictions = [0, 1]\n    >>> references = [0, 1]\n    >>> xnli_metric = datasets.load_metric("xnli")\n    >>> results = xnli_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'accuracy\': 1.0}\n'
def 				UpperCamelCase__					(  SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__						):
 return (preds == labels).mean()



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




      '''simple docstring'''


      def 					_snake_case     (    self:	Union[str, Any]    ):
       return datasets.MetricInfo(
           description=_DESCRIPTION						,				citation=_CITATION						,				inputs_description=_KWARGS_DESCRIPTION						,				features=datasets.Features(
               {
                   'predictions': datasets.Value('int64' if self.config_name != 'sts-b' else 'float32'    ),
                   'references': datasets.Value('int64' if self.config_name != 'sts-b' else 'float32'    ),
               }    )						,				codebase_urls=[]						,				reference_urls=[]						,				format='numpy'						,				)




      def 					_snake_case     (    self:	int						,				a:	Optional[Any]						,				a:	Optional[Any]    ):
       return {"accuracy": simple_accuracy(a						,				a    )}

