




'''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
_snake_case : Optional[Any]  				=						data_utils.TransfoXLTokenizer
_snake_case : int  				=						data_utils.TransfoXLCorpus
_snake_case : Union[str, Any]  				=						data_utils
_snake_case : List[Any]  				=						data_utils
def 						snake_case_       (UpperCamelCase			:    Dict					,		UpperCamelCase			:    Union[str, Any]					,		UpperCamelCase			:    Any					,		UpperCamelCase			:    Optional[int] ):





							'''simple docstring'''


							if transfo_xl_dataset_file:
														# Convert a pre-processed corpus (see original TensorFlow repo)
														with open(UpperCamelCase					,		'''rb''' ) as fp:
																					_a      	=							pickle.load(UpperCamelCase					,		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(UpperCamelCase					,		UpperCamelCase )

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

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

														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(UpperCamelCase )
														print(f'Building PyTorch model from configuration: {config}' )
														_a      	=							TransfoXLLMHeadModel(UpperCamelCase )

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


if __name__ == "__main__":
				_snake_case : int  				=						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.',
				)
				_snake_case : int  				=						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,
				)







'''simple docstring'''


import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
				import optax
				from flax.training.common_utils import onehot

				from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
				from transformers.models.ta.modeling_flax_ta import shift_tokens_right

@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class      A       (						unittest.TestCase       ):







							@slow
							def   __lowerCAmelCase     (    self			:    List[Any]   )     ->				Union[str, Any]:


														"""simple docstring"""


														_a      	=							FlaxMTaForConditionalGeneration.from_pretrained('''google/mt5-small'''   )
														_a      	=							AutoTokenizer.from_pretrained('''google/mt5-small'''   )

														_a      	=							tokenizer('''Hello there'''				,     return_tensors='''np'''   ).input_ids
														_a      	=							tokenizer('''Hi I am'''				,     return_tensors='''np'''   ).input_ids

														_a      	=							shift_tokens_right(lowerCAmelCase_				,     model.config.pad_token_id				,     model.config.decoder_start_token_id   )

														_a      	=							model(lowerCAmelCase_				,     decoder_input_ids=lowerCAmelCase_   ).logits
														_a      	=							optax.softmax_cross_entropy(lowerCAmelCase_				,     onehot(lowerCAmelCase_				,     logits.shape[-1]   )   ).mean()

														_a      	=							-(labels.shape[-1] * loss.item())

														_a      	=							-8_4.9_1_2_7
														self.assertTrue(abs(mtf_score - EXPECTED_SCORE   ) < 1e-4   )

