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
__UpperCamelCase	: List[Any]    		=  data_utils.TransfoXLTokenizer
__UpperCamelCase	: Dict    		=  data_utils.TransfoXLCorpus
__UpperCamelCase	: Optional[int]    		=  data_utils
__UpperCamelCase	: int    		=  data_utils
def    A			(					_lowercase	,						_lowercase	,						_lowercase	,						_lowercase		):
						if transfo_xl_dataset_file:
												# Convert a pre-processed corpus (see original TensorFlow repo)
												with open(lowerCamelCase_	,						'''rb'''		) as fp:
																		SCREAMING_SNAKE_CASE							:					List[str]           =							pickle.load(lowerCamelCase_	,						encoding='''latin1'''		)
												# Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
												SCREAMING_SNAKE_CASE							:					str           =							pytorch_dump_folder_path + '/' + VOCAB_FILES_NAMES['pretrained_vocab_file']
												print(f"""Save vocabulary to {pytorch_vocab_dump_path}"""		)
												SCREAMING_SNAKE_CASE							:					Any           =							corpus.vocab.__dict__
												torch.save(lowerCamelCase_	,						lowerCamelCase_		)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							corpus.__dict__
												corpus_dict_no_vocab.pop('''vocab'''	,						lowerCamelCase_		)
												SCREAMING_SNAKE_CASE							:					Tuple           =							pytorch_dump_folder_path + '/' + CORPUS_NAME
												print(f"""Save dataset to {pytorch_dataset_dump_path}"""		)
												torch.save(lowerCamelCase_	,						lowerCamelCase_		)

						if tf_checkpoint_path:
												# Convert a pre-trained TensorFlow model
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							os.path.abspath(lowerCamelCase_		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							os.path.abspath(lowerCamelCase_		)

												print(f"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""		)
												# Initialise PyTorch model
												if transfo_xl_config_file == "":
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							TransfoXLConfig()
												else:
																		SCREAMING_SNAKE_CASE							:					str           =							TransfoXLConfig.from_json_file(lowerCamelCase_		)
												print(f"""Building PyTorch model from configuration: {config}"""		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							TransfoXLLMHeadModel(lowerCamelCase_		)

												SCREAMING_SNAKE_CASE							:					int           =							load_tf_weights_in_transfo_xl(lowerCamelCase_	,						lowerCamelCase_	,						lowerCamelCase_		)
												# Save pytorch-model
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							os.path.join(lowerCamelCase_	,						lowerCamelCase_		)
												SCREAMING_SNAKE_CASE							:					str           =							os.path.join(lowerCamelCase_	,						lowerCamelCase_		)
												print(f"""Save PyTorch model to {os.path.abspath(lowerCamelCase_		)}"""		)
												torch.save(model.state_dict()	,						lowerCamelCase_		)
												print(f"""Save configuration file to {os.path.abspath(lowerCamelCase_		)}"""		)
												with open(lowerCamelCase_	,						'''w'''	,						encoding='''utf-8'''		) as f:
																		f.write(config.to_json_string()		)


if __name__ == "__main__":
					__UpperCamelCase	: Optional[Any]    		=  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.',
					)
					__UpperCamelCase	: Union[str, Any]    		=  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,
					)






import argparse
import json
import logging
import os
import sys
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, get_gpu_count, slow


__UpperCamelCase	: str    		=  [
    os.path.join(os.path.dirname(__file__), dirname)
    for dirname in [
        'text-classification',
        'language-modeling',
        'summarization',
        'token-classification',
        'question-answering',
    ]
]
sys.path.extend(SRC_DIRS)


if SRC_DIRS is not None:
					import run_clm_flax
					import run_flax_glue
					import run_flax_ner
					import run_mlm_flax
					import run_qa
					import run_summarization_flax
					import run_ta_mlm_flax


logging.basicConfig(level=logging.DEBUG)

__UpperCamelCase	: int    		=  logging.getLogger()
def    A			(					):
						SCREAMING_SNAKE_CASE							:					str           =							argparse.ArgumentParser()
						parser.add_argument('''-f'''		)
						SCREAMING_SNAKE_CASE							:					List[str]           =							parser.parse_args()
						return args.f
def    A			(					_lowercase	,						_lowercase="eval"		):
						SCREAMING_SNAKE_CASE							:					Dict           =							os.path.join(_lowercase	,						f"""{split}_results.json"""		)
						if os.path.exists(_lowercase		):
												with open(_lowercase	,						'''r'''		) as f:
																		return json.load(_lowercase		)
						raise ValueError(f"""can't find {path}"""		)


__UpperCamelCase	: Optional[Any]    		=  logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)





class       lowercase__		(     UpperCamelCase_):


						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					int           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Tuple           =							f"""
            run_glue.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --eval_steps=2
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_flax_glue.main()
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.75		)


						@slow
						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					str           =							f"""
            run_clm_flax.py
            --model_name_or_path distilgpt2
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --block_size 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_clm_flax.main()
																		SCREAMING_SNAKE_CASE							:					Dict           =							get_results(UpperCamelCase__		)
																		self.assertLess(result['''eval_perplexity''']					,	100		)


						@slow
						def  __A ( self			:       Any		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_summarization.py
            --model_name_or_path t5-small
            --train_file tests/fixtures/tests_samples/xsum/sample.json
            --validation_file tests/fixtures/tests_samples/xsum/sample.json
            --test_file tests/fixtures/tests_samples/xsum/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=8
            --do_train
            --do_eval
            --do_predict
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --predict_with_generate
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_summarization_flax.main()
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							get_results(UpperCamelCase__					,	split='''test'''		)
																		self.assertGreaterEqual(result['''test_rouge1''']					,	10		)
																		self.assertGreaterEqual(result['''test_rouge2''']					,	2		)
																		self.assertGreaterEqual(result['''test_rougeL''']					,	7		)
																		self.assertGreaterEqual(result['''test_rougeLsum''']					,	7		)


						@slow
						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[Any]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Dict           =							f"""
            run_mlm.py
            --model_name_or_path distilroberta-base
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --logging_steps 2 --eval_steps 2
            --do_train
            --do_eval
            --num_train_epochs=1
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_mlm_flax.main()
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							get_results(UpperCamelCase__		)
																		self.assertLess(result['''eval_perplexity''']					,	42		)


						@slow
						def  __A ( self			:       List[Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_t5_mlm_flax.py
            --model_name_or_path t5-small
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_ta_mlm_flax.main()
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.42		)


						@slow
						def  __A ( self			:       Union[str, Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Tuple           =							7 if get_gpu_count() > 1 else 2

												SCREAMING_SNAKE_CASE							:					Dict           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Any           =							f"""
            run_flax_ner.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/conll/sample.json
            --validation_file tests/fixtures/tests_samples/conll/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
            --warmup_steps=2
            --learning_rate=2e-4
            --logging_steps 2 --eval_steps 2
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=2
            --num_train_epochs={epochs}
            --seed 7
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_flax_ner.main()
																		SCREAMING_SNAKE_CASE							:					List[str]           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_accuracy''']					,	0.75		)
																		self.assertGreaterEqual(result['''eval_f1''']					,	0.3		)



						@slow
						def  __A ( self			:       List[str]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							self.get_auto_remove_tmp_dir()
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							f"""
            run_qa.py
            --model_name_or_path bert-base-uncased
            --version_2_with_negative
            --train_file tests/fixtures/tests_samples/SQUAD/sample.json
            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=2
            --do_train
            --do_eval
            --logging_steps 2 --eval_steps 2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

												with patch.object(UpperCamelCase__					,	'''argv'''					,	UpperCamelCase__		):
																		run_qa.main()
																		SCREAMING_SNAKE_CASE							:					str           =							get_results(UpperCamelCase__		)
																		self.assertGreaterEqual(result['''eval_f1''']					,	30		)
																		self.assertGreaterEqual(result['''eval_exact''']					,	30		)





