
"""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     __lowerCAmelCase				(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase=None		):
	# set parameter of one layer
	assert torch_layer.weight.shape == weight.shape, F"{torch_layer} layer.weight does not match"
	__lowerCAmelCase			:						Union[str, Any]									=      nn.Parameter(_UpperCamelCase		)
	if bias is not None:
		assert torch_layer.bias.shape == bias.shape, F"{torch_layer} layer.bias does not match"
		__lowerCAmelCase			:						Optional[int]									=      nn.Parameter(_UpperCamelCase		)
def     __lowerCAmelCase				(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase		):
	# set torch weights for 1-to-1 comparison
	__lowerCAmelCase			:						Dict									=      np.asarray(weights[0]		)
	__lowerCAmelCase			:						Dict									=      np.asarray(weights[1]		)
	__lowerCAmelCase			:						Tuple									=      np.asarray(weights[2]		)

	set_param(
	    torch_layer.self_attention.query_key       ,  torch.tensor(_UpperCamelCase		).transpose(1       ,  2		).contiguous().view(-1       ,  _UpperCamelCase		)       ,  )
	set_param(
	    torch_layer.self_attention.value       ,  torch.tensor(_UpperCamelCase		).transpose(1       ,  2		).contiguous().view(-1       ,  _UpperCamelCase		)       ,  )
	set_param(
	    torch_layer.output.dense       ,  torch.tensor(_UpperCamelCase		).view(-1       ,  _UpperCamelCase		).contiguous().transpose(0       ,  1		)       ,  )
def     __lowerCAmelCase				(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase		):
	# set torch weights for 1-to-1 comparison
	__lowerCAmelCase			:						List[str]									=      np.asarray(weights[0]		)
	__lowerCAmelCase			:						List[str]									=      np.asarray(weights[1]		)
	__lowerCAmelCase			:						int									=      np.asarray(weights[2]		)
	__lowerCAmelCase			:						int									=      np.asarray(weights[3]		)

	set_param(
	    torch_layer.self_attention.query       ,  torch.tensor(_UpperCamelCase		).transpose(1       ,  2		).contiguous().view(-1       ,  _UpperCamelCase		)       ,  )
	set_param(
	    torch_layer.self_attention.key       ,  torch.tensor(_UpperCamelCase		).transpose(1       ,  2		).contiguous().view(-1       ,  _UpperCamelCase		)       ,  )
	set_param(
	    torch_layer.self_attention.value       ,  torch.tensor(_UpperCamelCase		).transpose(1       ,  2		).contiguous().view(-1       ,  _UpperCamelCase		)       ,  )
	set_param(
	    torch_layer.output.dense       ,  torch.tensor(_UpperCamelCase		).view(-1       ,  _UpperCamelCase		).contiguous().transpose(0       ,  1		)       ,  )
def     __lowerCAmelCase				(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase		):
	# layernorm 1
	__lowerCAmelCase			:						Optional[Any]									=      weights[0][0][0]
	__lowerCAmelCase			:						Optional[int]									=      np.asarray(layer_norm_a[0]		)
	__lowerCAmelCase			:						List[str]									=      np.asarray(layer_norm_a[1]		)
	set_param(
	    torch_block.attention.layer_norm       ,  torch.tensor(_UpperCamelCase		)       ,  torch.tensor(_UpperCamelCase		)       ,  )

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

	# intermediate weighs
	__lowerCAmelCase			:						Dict									=      weights[2][0][1][2]

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

	# layernorm 2
	__lowerCAmelCase			:						Optional[int]									=      np.asarray(intermediate_weights[0][0]		)
	__lowerCAmelCase			:						Dict									=      np.asarray(intermediate_weights[0][1]		)
	set_param(
	    torch_block.feed_forward.layer_norm       ,  torch.tensor(_UpperCamelCase		)       ,  torch.tensor(_UpperCamelCase		)       ,  )

	# intermediate dense
	__lowerCAmelCase			:						List[Any]									=      np.asarray(intermediate_weights[1][0]		)
	__lowerCAmelCase			:						List[str]									=      np.asarray(intermediate_weights[1][1]		)
	set_param(
	    torch_block.feed_forward.dense.dense       ,  torch.tensor(_UpperCamelCase		).transpose(0       ,  1		).contiguous()       ,  torch.tensor(_UpperCamelCase		)       ,  )

	# intermediate out
	__lowerCAmelCase			:						Union[str, Any]									=      np.asarray(intermediate_weights[4][0]		)
	__lowerCAmelCase			:						List[str]									=      np.asarray(intermediate_weights[4][1]		)
	set_param(
	    torch_block.feed_forward.output.dense       ,  torch.tensor(_UpperCamelCase		).transpose(0       ,  1		).contiguous()       ,  torch.tensor(_UpperCamelCase		)       ,  )
def     __lowerCAmelCase				(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase		):
	# reformer model
	__lowerCAmelCase			:						Union[str, Any]									=      torch_model.reformer

	# word embeds
	__lowerCAmelCase			:						Union[str, Any]									=      np.asarray(weights[1]		)
	set_param(
	    torch_model_reformer.embeddings.word_embeddings       ,  torch.tensor(_UpperCamelCase		)       ,  )

	if isinstance(weights[3]       ,  _UpperCamelCase		):
		__lowerCAmelCase			:						Dict									=      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(_UpperCamelCase		)		)

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

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

	# output embeddings
	__lowerCAmelCase			:						Optional[int]									=      np.asarray(weights[9][0]		)
	__lowerCAmelCase			:						str									=      np.asarray(weights[9][1]		)
	set_param(
	    torch_model.lm_head.decoder       ,  torch.tensor(_UpperCamelCase		).transpose(0       ,  1		).contiguous()       ,  torch.tensor(_UpperCamelCase		)       ,  )
def     __lowerCAmelCase				(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase		):
	# Initialise PyTorch model
	__lowerCAmelCase			:						Optional[Any]									=      ReformerConfig.from_json_file(_UpperCamelCase		)
	print(F"Building PyTorch model from configuration: {config}"		)
	__lowerCAmelCase			:						int									=      ReformerModelWithLMHead(_UpperCamelCase		)

	with open(_UpperCamelCase       ,  'rb'		) as f:
		__lowerCAmelCase			:						Tuple									=      pickle.load(_UpperCamelCase		)['weights']

	set_model_weights_in_torch(_UpperCamelCase       ,  _UpperCamelCase       ,  config.hidden_size		)

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


if __name__ == "__main__":
							lowerCamelCase__   					=		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."""
							)
							lowerCamelCase__   					=		parser.parse_args()
							convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)






'''simple docstring'''




import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer


class  SCREAMING_SNAKE_CASE    (a__							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       tempfile.mkdtemp()
          __A						: str      =       5

          # Realm tok
          __A						: Union[str, Any]      =       [
              '[UNK]',
              '[CLS]',
              '[SEP]',
              '[PAD]',
              '[MASK]',
              'test',
              'question',
              'this',
              'is',
              'the',
              'first',
              'second',
              'third',
              'fourth',
              'fifth',
              'record',
              'want',
              '##want',
              '##ed',
              'wa',
              'un',
              'runn',
              '##ing',
              ',',
              'low',
              'lowest',
          ]
          __A						: Union[str, Any]      =       os.path.join(self.tmpdirname						,	'realm_tokenizer')
          os.makedirs(_UpperCAmelCase						,	exist_ok=_UpperCAmelCase)
          __A						: Tuple      =       os.path.join(_UpperCAmelCase						,	VOCAB_FILES_NAMES['vocab_file'])
          with open(self.vocab_file						,	'w'						,	encoding='utf-8') as vocab_writer:
               vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]))

          __A						: int      =       os.path.join(self.tmpdirname						,	'realm_block_records')
          os.makedirs(_UpperCAmelCase						,	exist_ok=_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname						,	'realm_tokenizer'))


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          shutil.rmtree(self.tmpdirname)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Union[str, Any]      =       RealmConfig(num_block_records=self.num_block_records)
          return config


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Dict      =       Dataset.from_dict(
              {
                  'id': ['0', '1'],
                  'question': ['foo', 'bar'],
                  'answers': [['Foo', 'Bar'], ['Bar']],
              })
          return dataset


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       np.array(
              [
                  B'This is the first record',
                  B'This is the second record',
                  B'This is the third record',
                  B'This is the fourth record',
                  B'This is the fifth record',
                  B'This is a longer longer longer record',
              ]						,	dtype=_UpperCAmelCase						,	)
          return block_records


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[Any]      =       RealmRetriever(
              block_records=self.get_dummy_block_records()						,	tokenizer=self.get_tokenizer()						,	)
          return retriever


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       self.get_config()
          __A						: str      =       self.get_dummy_retriever()
          __A						: List[str]      =       retriever.tokenizer

          __A						: Dict      =       np.array([0, 3]						,	dtype='long')
          __A						: Dict      =       tokenizer(['Test question']).input_ids
          __A						: Optional[Any]      =       tokenizer(
              ['the fourth']						,	add_special_tokens=_UpperCAmelCase						,	return_token_type_ids=_UpperCAmelCase						,	return_attention_mask=_UpperCAmelCase						,	).input_ids
          __A						: str      =       config.reader_seq_len

          __A ,__A ,__A ,__A						: List[str]      =       retriever(
              _UpperCAmelCase						,	_UpperCAmelCase						,	answer_ids=_UpperCAmelCase						,	max_length=_UpperCAmelCase						,	return_tensors='np')

          self.assertEqual(len(_UpperCAmelCase)						,	2)
          self.assertEqual(len(_UpperCAmelCase)						,	2)
          self.assertEqual(len(_UpperCAmelCase)						,	2)
          self.assertEqual(concat_inputs.input_ids.shape						,	(2, 10))
          self.assertEqual(concat_inputs.attention_mask.shape						,	(2, 10))
          self.assertEqual(concat_inputs.token_type_ids.shape						,	(2, 10))
          self.assertEqual(concat_inputs.special_tokens_mask.shape						,	(2, 10))
          self.assertEqual(
              tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0])						,	['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'first', 'record', '[SEP]']						,	)
          self.assertEqual(
              tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1])						,	['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'fourth', 'record', '[SEP]']						,	)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       self.get_config()
          __A						: Any      =       self.get_dummy_retriever()
          __A						: str      =       retriever.tokenizer

          __A						: Dict      =       np.array([0, 3, 5]						,	dtype='long')
          __A						: Tuple      =       tokenizer(['Test question']).input_ids
          __A						: Union[str, Any]      =       tokenizer(
              ['the fourth', 'longer longer']						,	add_special_tokens=_UpperCAmelCase						,	return_token_type_ids=_UpperCAmelCase						,	return_attention_mask=_UpperCAmelCase						,	).input_ids
          __A						: Dict      =       config.reader_seq_len

          __A ,__A ,__A ,__A						: str      =       retriever(
              _UpperCAmelCase						,	_UpperCAmelCase						,	answer_ids=_UpperCAmelCase						,	max_length=_UpperCAmelCase						,	return_tensors='np')

          self.assertEqual([False, True, True]						,	_UpperCAmelCase)
          self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]						,	_UpperCAmelCase)
          self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]						,	_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[Any]      =       self.get_dummy_retriever()
          retriever.save_pretrained(os.path.join(self.tmpdirname						,	'realm_block_records'))

          # Test local path
          __A						: str      =       retriever.from_pretrained(os.path.join(self.tmpdirname						,	'realm_block_records'))
          self.assertEqual(retriever.block_records[0]						,	B'This is the first record')

          # Test mocked remote path
          with patch('transformers.models.realm.retrieval_realm.hf_hub_download') as mock_hf_hub_download:
               __A						: int      =       os.path.join(
                   os.path.join(self.tmpdirname						,	'realm_block_records')						,	_REALM_BLOCK_RECORDS_FILENAME)
               __A						: Tuple      =       RealmRetriever.from_pretrained('google/realm-cc-news-pretrained-openqa')

          self.assertEqual(retriever.block_records[0]						,	B'This is the first record')