
"""simple docstring"""


import unittest

from transformers import BertGenerationConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import BertGenerationDecoder, BertGenerationEncoder







class 				lowerCAmelCase_				:



						'''simple docstring'''


						def __init__( self     :      Optional[int]  ,A_     :      Any  ,A_     :      int=13  ,A_     :      Optional[int]=7  ,A_     :      List[Any]=True  ,A_     :      str=True  ,A_     :      Optional[Any]=99  ,A_     :      Optional[Any]=32  ,A_     :      Optional[int]=5  ,A_     :      str=4  ,A_     :      int=37  ,A_     :      Optional[Any]="gelu"  ,A_     :      int=0.1  ,A_     :      Dict=0.1  ,A_     :      Any=50  ,A_     :      Tuple=0.02  ,A_     :      Tuple=True  ,A_     :      Optional[Any]=None  ,)     ->   int:
							A						   =   parent
							A						   =   batch_size
							A						   =   seq_length
							A						   =   is_training
							A						   =   use_input_mask
							A						   =   vocab_size
							A						   =   hidden_size
							A						   =   num_hidden_layers
							A						   =   num_attention_heads
							A						   =   intermediate_size
							A						   =   hidden_act
							A						   =   hidden_dropout_prob
							A						   =   attention_probs_dropout_prob
							A						   =   max_position_embeddings
							A						   =   initializer_range
							A						   =   use_labels
							A						   =   scope


						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   Optional[Any]:
							A						   =   ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size      )

							A						   =   None
							if self.use_input_mask:
								A						   =   random_attention_mask([self.batch_size, self.seq_length]      )

							if self.use_labels:
								A						   =   ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size      )

							A						   =   self.get_config()

							return config, input_ids, input_mask, token_labels


						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   List[str]:
							return BertGenerationConfig(
							    vocab_size=self.vocab_size  ,hidden_size=self.hidden_size  ,num_hidden_layers=self.num_hidden_layers  ,num_attention_heads=self.num_attention_heads  ,intermediate_size=self.intermediate_size  ,hidden_act=self.hidden_act  ,hidden_dropout_prob=self.hidden_dropout_prob  ,attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,max_position_embeddings=self.max_position_embeddings  ,is_decoder=A_  ,initializer_range=self.initializer_range  ,)


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[int]      )     ->   Optional[int]:
							(
							    (
							    A      
							)   ,      (
							    A      
							)   ,      (
							    A      
							)   ,      (
							    A      
							)   ,            
							)						   =   self.prepare_config_and_inputs()

							A						   =   True
							A						   =   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
							A						   =   ids_tensor([self.batch_size, self.seq_length]  ,vocab_size=2      )

							return (
							    config,
							    input_ids,
							    input_mask,
							    token_labels,
							    encoder_hidden_states,
							    encoder_attention_mask,
							)


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Union[str, Any]  ,A_     :      int  ,A_     :      Dict  ,A_     :      Any  ,A_     :      Union[str, Any]  ,**A_     :      int  ,)     ->   List[Any]:
							A						   =   BertGenerationEncoder(config=A_      )
							model.to(A_      )
							model.eval()
							A						   =   model(A_  ,attention_mask=A_      )
							A						   =   model(A_      )
							self.parent.assertEqual(result.last_hidden_state.shape  ,(self.batch_size, self.seq_length, self.hidden_size)      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[int]  ,A_     :      str  ,A_     :      Union[str, Any]  ,A_     :      int  ,A_     :      Union[str, Any]  ,A_     :      Dict  ,A_     :      int  ,**A_     :      Union[str, Any]  ,)     ->   int:
							A						   =   True
							A						   =   BertGenerationEncoder(config=A_      )
							model.to(A_      )
							model.eval()
							A						   =   model(
							    A_  ,attention_mask=A_  ,encoder_hidden_states=A_  ,encoder_attention_mask=A_  ,)
							A						   =   model(
							    A_  ,attention_mask=A_  ,encoder_hidden_states=A_  ,)
							self.parent.assertEqual(result.last_hidden_state.shape  ,(self.batch_size, self.seq_length, self.hidden_size)      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      str  ,A_     :      Union[str, Any]  ,A_     :      List[Any]  ,A_     :      Optional[Any]  ,A_     :      Any  ,A_     :      Optional[int]  ,A_     :      Dict  ,**A_     :      Tuple  ,)     ->   int:
							A						   =   True
							A						   =   True
							A						   =   BertGenerationDecoder(config=A_      ).to(A_      ).eval()

							# first forward pass
							A						   =   model(
							    A_  ,attention_mask=A_  ,encoder_hidden_states=A_  ,encoder_attention_mask=A_  ,use_cache=A_  ,)
							A						   =   outputs.past_key_values

							# create hypothetical multiple next token and extent to next_input_ids
							A						   =   ids_tensor((self.batch_size, 3)  ,config.vocab_size      )
							A						   =   ids_tensor((self.batch_size, 3)  ,vocab_size=2      )

							# append to next input_ids and
							A						   =   torch.cat([input_ids, next_tokens]  ,dim=-1      )
							A						   =   torch.cat([input_mask, next_mask]  ,dim=-1      )

							A						   =   model(
							    A_  ,attention_mask=A_  ,encoder_hidden_states=A_  ,encoder_attention_mask=A_  ,output_hidden_states=A_  ,)['hidden_states'][0]
							A						   =   model(
							    A_  ,attention_mask=A_  ,encoder_hidden_states=A_  ,encoder_attention_mask=A_  ,past_key_values=A_  ,output_hidden_states=A_  ,)['hidden_states'][0]

							# select random slice
							A						   =   ids_tensor((1,)  ,output_from_past.shape[-1]      ).item()
							A						   =   output_from_no_past[:, -3:, random_slice_idx].detach()
							A						   =   output_from_past[:, :, random_slice_idx].detach()

							self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]      )

							# test that outputs are equal for slice
							self.parent.assertTrue(torch.allclose(A_  ,A_  ,atol=1e-3      )      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      str  ,A_     :      Optional[int]  ,A_     :      int  ,A_     :      int  ,A_     :      List[str]  ,*A_     :      int  ,)     ->   List[str]:
							A						   =   BertGenerationDecoder(A_      )
							model.to(A_      )
							model.eval()
							A						   =   model(A_  ,attention_mask=A_  ,labels=A_      )
							self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.vocab_size)      )




						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   Dict:
							A   ,      A   ,      A   ,      A						   =   self.prepare_config_and_inputs()
							A						   =   {'input_ids': input_ids, 'attention_mask': input_mask}
							return config, inputs_dict







@require_torch
class 				lowerCAmelCase_				(      _lowercase		,					_lowercase		,					_lowercase		,					unittest.TestCase    ):



						'''simple docstring'''


						_lowerCamelCase:    str      =    (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
						_lowerCamelCase:    Optional[int]      =    (BertGenerationDecoder,) if is_torch_available() else ()
						_lowerCamelCase:    Union[str, Any]      =    (
						    {'''feature-extraction''': BertGenerationEncoder, '''text-generation''': BertGenerationDecoder}
						    if is_torch_available()
						    else {}
						)


						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   Optional[int]:
							A						   =   BertGenerationEncoderTester(self      )
							A						   =   ConfigTester(self  ,config_class=A_  ,hidden_size=37      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   Optional[int]:
							self.config_tester.run_common_tests()


						def 						_SCREAMING_SNAKE_CASE      ( self     :      str      )     ->   Optional[Any]:
							A						   =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   int:
							A   ,      A   ,      A   ,      A						   =   self.model_tester.prepare_config_and_inputs()
							A						   =   'bert'
							self.model_tester.create_and_check_model(A_  ,A_  ,A_  ,A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   Optional[Any]:
							A						   =   self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_model_as_decoder(*A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   str:
							A						   =   self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_decoder_model_past_large_inputs(*A_      )


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]      )     ->   Dict:
							# This regression test was failing with PyTorch < 1.3
							(
							    (
							    A      
							)   ,      (
							    A      
							)   ,      (
							    A      
							)   ,      (
							    A      
							)   ,      (
							    A      
							)   ,      (
							    A      
							)   ,            
							)						   =   self.model_tester.prepare_config_and_inputs_for_decoder()

							A						   =   None

							self.model_tester.create_and_check_model_as_decoder(
							    A_  ,A_  ,A_  ,A_  ,A_  ,A_  ,)


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   List[str]:
							A						   =   self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_for_causal_lm(*A_      )




						@slow
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Any      )     ->   List[str]:
							A						   =   BertGenerationEncoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'      )
							self.assertIsNotNone(A_      )







@require_torch
class 				lowerCAmelCase_				(      unittest.TestCase    ):



						'''simple docstring'''


						@slow
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   Any:
							A						   =   BertGenerationEncoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'      )
							A						   =   torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]      )
							with torch.no_grad():
								A						   =   model(A_      )[0]
							A						   =   torch.Size([1, 8, 1024]      )
							self.assertEqual(output.shape  ,A_      )
							A						   =   torch.tensor(
							    [[[0.17_75, 0.00_83, -0.03_21], [1.60_02, 0.12_87, 0.39_12], [2.14_73, 0.57_91, 0.60_66]]]      )
							self.assertTrue(torch.allclose(output[:, :3, :3]  ,A_  ,atol=1e-4      )      )







@require_torch
class 				lowerCAmelCase_				(      unittest.TestCase    ):



						'''simple docstring'''


						@slow
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   Tuple:
							A						   =   BertGenerationDecoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'      )
							A						   =   torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]      )
							with torch.no_grad():
								A						   =   model(A_      )[0]
							A						   =   torch.Size([1, 8, 5_0358]      )
							self.assertEqual(output.shape  ,A_      )
							A						   =   torch.tensor(
							    [[[-0.57_88, -2.59_94, -3.70_54], [0.04_38, 4.79_97, 1.87_95], [1.58_62, 6.64_09, 4.46_38]]]      )
							self.assertTrue(torch.allclose(output[:, :3, :3]  ,A_  ,atol=1e-4      )      )

"""simple docstring"""


import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
_lowercase							   =				data_utils.TransfoXLTokenizer
_lowercase							   =				data_utils.TransfoXLCorpus
_lowercase							   =				data_utils
_lowercase							   =				data_utils






def       _snake_case				(		snake_case__   :     Union[str, Any]	,      snake_case__   :     Dict	,      snake_case__   :     Tuple	,      snake_case__   :     int     ):
	if transfo_xl_dataset_file:
		# Convert a pre-processed corpus (see original TensorFlow repo)
		with open(snake_case__	,      'rb'     ) as fp:
			A						   =   pickle.load(snake_case__	,      encoding='latin1'     )
		# Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
		A						   =   pytorch_dump_folder_path + '/' + VOCAB_FILES_NAMES['pretrained_vocab_file']
		print(F'Save vocabulary to {pytorch_vocab_dump_path}'     )
		A						   =   corpus.vocab.__dict__
		torch.save(snake_case__	,      snake_case__     )

		A						   =   corpus.__dict__
		corpus_dict_no_vocab.pop('vocab'	,      snake_case__     )
		A						   =   pytorch_dump_folder_path + '/' + CORPUS_NAME
		print(F'Save dataset to {pytorch_dataset_dump_path}'     )
		torch.save(snake_case__	,      snake_case__     )

	if tf_checkpoint_path:
		# Convert a pre-trained TensorFlow model
		A						   =   os.path.abspath(snake_case__     )
		A						   =   os.path.abspath(snake_case__     )

		print(F'Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.'     )
		# Initialise PyTorch model
		if transfo_xl_config_file == "":
			A						   =   TransfoXLConfig()
		else:
			A						   =   TransfoXLConfig.from_json_file(snake_case__     )
		print(F'Building PyTorch model from configuration: {config}'     )
		A						   =   TransfoXLLMHeadModel(snake_case__     )

		A						   =   load_tf_weights_in_transfo_xl(snake_case__	,      snake_case__	,      snake_case__     )
		# Save pytorch-model
		A						   =   os.path.join(snake_case__	,      snake_case__     )
		A						   =   os.path.join(snake_case__	,      snake_case__     )
		print(F'Save PyTorch model to {os.path.abspath(snake_case__     )}'     )
		torch.save(model.state_dict()	,      snake_case__     )
		print(F'Save configuration file to {os.path.abspath(snake_case__     )}'     )
		with open(snake_case__	,      'w'	,      encoding='utf-8'     ) as f:
			f.write(config.to_json_string()     )


if __name__ == "__main__":
		_lowercase							   =				argparse.ArgumentParser()
		parser.add_argument(
		    '''--pytorch_dump_folder_path''',
		    default=None,
		    type=str,
		    required=True,
		    help='''Path to the folder to store the PyTorch model or dataset/vocab.''',
		)
		parser.add_argument(
		    '''--tf_checkpoint_path''',
		    default='''''',
		    type=str,
		    help='''An optional path to a TensorFlow checkpoint path to be converted.''',
		)
		parser.add_argument(
		    '''--transfo_xl_config_file''',
		    default='''''',
		    type=str,
		    help=(
		        '''An optional config json file corresponding to the pre-trained BERT model. \n'''
		        '''This specifies the model architecture.'''
		    ),
		)
		parser.add_argument(
		    '''--transfo_xl_dataset_file''',
		    default='''''',
		    type=str,
		    help='''An optional dataset file to be converted in a vocabulary.''',
		)
		_lowercase							   =				parser.parse_args()
		convert_transfo_xl_checkpoint_to_pytorch(
		    args.tf_checkpoint_path,
		    args.transfo_xl_config_file,
		    args.pytorch_dump_folder_path,
		    args.transfo_xl_dataset_file,
		)