





import argparse
import logging
import os
import sys
import tempfile
from pathlib import Path

import lightning_base
import pytest
import pytorch_lightning as pl
import torch
from convert_pl_checkpoint_to_hf import convert_pl_to_hf
from distillation import distill_main
from finetune import SummarizationModule, main
from huggingface_hub import list_models
from parameterized import parameterized
from run_eval import generate_summaries_or_translations
from torch import nn

from transformers import AutoConfig, AutoModelForSeqaSeqLM
from transformers.testing_utils import CaptureStderr, CaptureStdout, TestCasePlus, require_torch_gpu, slow
from utils import label_smoothed_nll_loss, lmap, load_json


logging.basicConfig(level=logging.DEBUG)

_lowerCamelCase							=     logging.getLogger()
_lowerCamelCase							=     torch.cuda.is_available()
_lowerCamelCase							=     {
    """max_tokens_per_batch""": None,
    """supervise_forward""": True,
    """normalize_hidden""": True,
    """label_smoothing""": 0.2,
    """eval_max_gen_length""": None,
    """eval_beams""": 1,
    """val_metric""": """loss""",
    """save_top_k""": 1,
    """adafactor""": True,
    """early_stopping_patience""": 2,
    """logger_name""": """default""",
    """length_penalty""": 0.5,
    """cache_dir""": """""",
    """task""": """summarization""",
    """num_workers""": 2,
    """alpha_hid""": 0,
    """freeze_embeds""": True,
    """enc_only""": False,
    """tgt_suffix""": """""",
    """resume_from_checkpoint""": None,
    """sortish_sampler""": True,
    """student_decoder_layers""": 1,
    """val_check_interval""": 1.0,
    """output_dir""": """""",
    """fp16""": False,  # TODO(SS): set this to CUDA_AVAILABLE if ci installs apex or start using native amp
    """no_teacher""": False,
    """fp16_opt_level""": """O1""",
    """gpus""": 1 if CUDA_AVAILABLE else 0,
    """n_tpu_cores""": 0,
    """max_grad_norm""": 1.0,
    """do_train""": True,
    """do_predict""": True,
    """accumulate_grad_batches""": 1,
    """server_ip""": """""",
    """server_port""": """""",
    """seed""": 42,
    """model_name_or_path""": """sshleifer/bart-tiny-random""",
    """config_name""": """""",
    """tokenizer_name""": """facebook/bart-large""",
    """do_lower_case""": False,
    """learning_rate""": 0.3,
    """lr_scheduler""": """linear""",
    """weight_decay""": 0.0,
    """adam_epsilon""": 1e-08,
    """warmup_steps""": 0,
    """max_epochs""": 1,
    """train_batch_size""": 2,
    """eval_batch_size""": 2,
    """max_source_length""": 12,
    """max_target_length""": 12,
    """val_max_target_length""": 12,
    """test_max_target_length""": 12,
    """fast_dev_run""": False,
    """no_cache""": False,
    """n_train""": -1,
    """n_val""": -1,
    """n_test""": -1,
    """student_encoder_layers""": 1,
    """freeze_encoder""": False,
    """auto_scale_batch_size""": False,
    """overwrite_output_dir""": False,
    """student""": None,
}






def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							List[Any]   ,				UpperCamelCase__:							Optional[Any]      ):
		SCREAMING_SNAKE_CASE__								= """\n""".join(UpperCamelCase__      )
		Path(UpperCamelCase__      ).open("""w"""      ).writelines(UpperCamelCase__      )


_lowerCamelCase							=     [""" Sam ate lunch today.""", """Sams lunch ingredients."""]
_lowerCamelCase							=     ["""A very interesting story about what I ate for lunch.""", """Avocado, celery, turkey, coffee"""]
_lowerCamelCase							=     """patrickvonplaten/t5-tiny-random"""
_lowerCamelCase							=     """sshleifer/t5-tinier-random"""
_lowerCamelCase							=     """sshleifer/bart-tiny-random"""
_lowerCamelCase							=     """sshleifer/tiny-mbart"""
_lowerCamelCase							=     """sshleifer/tiny-marian-en-de"""
_lowerCamelCase							=     """stas/tiny-wmt19-en-de"""


_lowerCamelCase							=     logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks






def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							Union[str, Any]      ):



		for split in ["train", "val", "test"]:
				_dump_articles(os.path.join(UpperCamelCase__   ,				f'''{split}.source'''      )   ,				UpperCamelCase__      )
				_dump_articles(os.path.join(UpperCamelCase__   ,				f'''{split}.target'''      )   ,				UpperCamelCase__      )
		return tmp_dir




class 					UpperCamelCase_    (       _UpperCAmelCase      ):


							@classmethod
							def        _snake_case					(							cls      :Dict						)				->				Dict:



									"""simple docstring"""



									logging.disable(logging.CRITICAL						)  # remove noisy download output from tracebacks
									return cls


							@slow
							@require_torch_gpu
							def        _snake_case					(							self      :Optional[Any]						)				->				str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= list_models()
									SCREAMING_SNAKE_CASE__								= """sshleifer"""
									SCREAMING_SNAKE_CASE__								= [x.modelId for x in model_list if x.modelId.startswith(lowerCAmelCase__						)]
									SCREAMING_SNAKE_CASE__								= ["""sshleifer/blenderbot-3B""", """sshleifer/blenderbot-90M"""]
									SCREAMING_SNAKE_CASE__								= []
									for m in model_ids:
											if m in allowed_to_be_broken:
													continue
											try:
													AutoConfig.from_pretrained(lowerCAmelCase__						)
											except Exception:
													failures.append(lowerCAmelCase__						)
									assert not failures, f'''The following models could not be loaded through AutoConfig: {failures}'''


							def        _snake_case					(							self      :Any						)				->				Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= {"""student_encoder_layers""": 2, """student_decoder_layers""": 1, """no_teacher""": True}
									self._test_distiller_cli(lowerCAmelCase__						)


							def        _snake_case					(							self      :List[Any]						)				->				Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= {
									    """student_encoder_layers""": 2,
									    """student_decoder_layers""": 1,
									    """max_epochs""": 4,
									    """val_check_interval""": 0.2_5,
									    """alpha_hid""": 2.0,
									    """model_name_or_path""": """IGNORE_THIS_IT_DOESNT_GET_USED""",
									}
									SCREAMING_SNAKE_CASE__								= self._test_distiller_cli(lowerCAmelCase__							,			check_contents=lowerCAmelCase__						)

									SCREAMING_SNAKE_CASE__								= list(Path(model.output_dir						).glob("""*.ckpt"""						)						)
									self.assertEqual(1							,			len(lowerCAmelCase__						)						)
									SCREAMING_SNAKE_CASE__								= list(Path(model.output_dir						).glob("""**/*.bin"""						)						)
									self.assertEqual(len(lowerCAmelCase__						)							,			2						)
									SCREAMING_SNAKE_CASE__								= lmap(str.strip							,			Path(model.hparams.data_dir						).joinpath("""test.source"""						).open().readlines()						)
									SCREAMING_SNAKE_CASE__								= tempfile.mktemp()  # XXX: not being cleaned up
									generate_summaries_or_translations(lowerCAmelCase__							,			lowerCAmelCase__							,			str(model.output_dir / """best_tfmr"""						)						)
									self.assertTrue(Path(lowerCAmelCase__						).exists()						)

									SCREAMING_SNAKE_CASE__								= self.get_auto_remove_tmp_dir()
									convert_pl_to_hf(ckpts[0]							,			transformer_ckpts[0].parent							,			lowerCAmelCase__						)
									assert os.path.exists(os.path.join(lowerCAmelCase__							,			"""pytorch_model.bin"""						)						)


							def        _snake_case					(							self      :int						)				->				int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= AutoModelForSeqaSeqLM.from_pretrained(lowerCAmelCase__						)
									SCREAMING_SNAKE_CASE__		,  SCREAMING_SNAKE_CASE__								= model.dummy_inputs["""input_ids"""], model.dummy_inputs["""attention_mask"""]
									SCREAMING_SNAKE_CASE__								= torch.tensor([[0, 4, 8, 2], [0, 8, 2, 1]]							,			dtype=torch.long							,			device=model.device						)
									SCREAMING_SNAKE_CASE__								= target_ids[:, :-1].contiguous()  # Why this line?
									SCREAMING_SNAKE_CASE__								= target_ids[:, 1:].clone()  # why clone?
									SCREAMING_SNAKE_CASE__								= model(
									    lowerCAmelCase__							,			attention_mask=lowerCAmelCase__							,			decoder_input_ids=lowerCAmelCase__							,			labels=lowerCAmelCase__							,			use_cache=lowerCAmelCase__						).loss

									SCREAMING_SNAKE_CASE__								= model(lowerCAmelCase__							,			attention_mask=lowerCAmelCase__							,			decoder_input_ids=lowerCAmelCase__							,			use_cache=lowerCAmelCase__						).logits

									SCREAMING_SNAKE_CASE__								= nn.functional.log_softmax(lowerCAmelCase__							,			dim=-1						)
									SCREAMING_SNAKE_CASE__		,  SCREAMING_SNAKE_CASE__								= label_smoothed_nll_loss(
									    lowerCAmelCase__							,			lowerCAmelCase__							,			0.1							,			ignore_index=model.config.pad_token_id						)
									with self.assertRaises(lowerCAmelCase__						):
											# TODO: understand why this breaks
											self.assertEqual(lowerCAmelCase__							,			lowerCAmelCase__						)


							def        _snake_case					(							self      :int						)				->				Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= {
									    """student_encoder_layers""": 2,
									    """student_decoder_layers""": 1,
									    """num_train_epochs""": 4,
									    """val_check_interval""": 0.2_5,
									    """alpha_hid""": 2.0,
									    """task""": """translation""",
									    """model_name_or_path""": """IGNORE_THIS_IT_DOESNT_GET_USED""",
									    """tokenizer_name""": MBART_TINY,
									    """teacher""": MBART_TINY,
									    """src_lang""": """en_XX""",
									    """tgt_lang""": """ro_RO""",
									}
									SCREAMING_SNAKE_CASE__								= self._test_distiller_cli(lowerCAmelCase__							,			check_contents=lowerCAmelCase__						)
									assert model.model.config.model_type == "mbart"

									SCREAMING_SNAKE_CASE__								= list(Path(model.output_dir						).glob("""*.ckpt"""						)						)
									self.assertEqual(1							,			len(lowerCAmelCase__						)						)
									SCREAMING_SNAKE_CASE__								= list(Path(model.output_dir						).glob("""**/*.bin"""						)						)
									SCREAMING_SNAKE_CASE__								= list(Path(model.output_dir						).glob("""best_tfmr/*"""						)						)
									assert len(lowerCAmelCase__						) > 2
									self.assertEqual(len(lowerCAmelCase__						)							,			2						)


							def        _snake_case					(							self      :Dict						)				->				Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= {
									    """student_encoder_layers""": 1,
									    """student_decoder_layers""": 1,
									    """alpha_hid""": 2.0,
									    """teacher""": T5_TINY,
									    """model_name_or_path""": T5_TINY,
									    """tokenizer_name""": T5_TINY,
									}
									self._test_distiller_cli(lowerCAmelCase__						)


							def        _snake_case					(							self      :int						)				->				List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= {
									    """teacher""": T5_TINY,
									    """student""": T5_TINIER,
									    """model_name_or_path""": T5_TINIER,
									    """tokenizer_name""": T5_TINIER,
									}
									self._test_distiller_cli(lowerCAmelCase__						)


							def        _snake_case					(							self      :str							,			__A      :Union[str, Any]							,			__A      :List[Any]=True						)				->				List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= {
									    """label_smoothing""": 0.0,
									    """early_stopping_patience""": -1,
									    """train_batch_size""": 1,
									    """eval_batch_size""": 2,
									    """max_epochs""": 2,
									    """alpha_mlm""": 0.2,
									    """alpha_ce""": 0.8,
									    """do_predict""": True,
									    """model_name_or_path""": """sshleifer/tinier_bart""",
									    """teacher""": CHEAP_ARGS["""model_name_or_path"""],
									    """val_check_interval""": 0.5,
									}
									default_updates.update(lowerCAmelCase__						)
									SCREAMING_SNAKE_CASE__								= CHEAP_ARGS.copy()
									SCREAMING_SNAKE_CASE__								= make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()						)
									SCREAMING_SNAKE_CASE__								= self.get_auto_remove_tmp_dir()

									args_d.update(data_dir=lowerCAmelCase__							,			output_dir=lowerCAmelCase__							,			**lowerCAmelCase__						)
									SCREAMING_SNAKE_CASE__								= distill_main(argparse.Namespace(**lowerCAmelCase__						)						)
									if not check_contents:
											return model
									SCREAMING_SNAKE_CASE__								= os.listdir(lowerCAmelCase__						)
									SCREAMING_SNAKE_CASE__								= {os.path.basename(lowerCAmelCase__						) for p in contents}
									SCREAMING_SNAKE_CASE__								= [p for p in contents if p.endswith("""ckpt"""						)]
									assert len(lowerCAmelCase__						) > 0

									self.assertIn("""test_generations.txt"""							,			lowerCAmelCase__						)
									self.assertIn("""test_results.txt"""							,			lowerCAmelCase__						)

									SCREAMING_SNAKE_CASE__								= load_json(model.metrics_save_path						)
									SCREAMING_SNAKE_CASE__								= metrics["""val"""][-1]
									self.assertGreaterEqual(last_step_stats["""val_avg_gen_time"""]							,			0.0_1						)
									self.assertGreaterEqual(1.0							,			last_step_stats["""val_avg_gen_time"""]						)
									self.assertIsInstance(last_step_stats[f'''val_avg_{model.val_metric}''']							,			lowerCAmelCase__						)
									SCREAMING_SNAKE_CASE__								= int(args_d["""max_epochs"""] * (1 / args_d["""val_check_interval"""]) + 1						)
									self.assertEqual(len(metrics["""val"""]						)							,			lowerCAmelCase__						)
									self.assertEqual(len(metrics["""test"""]						)							,			1						)
									return model




class 					UpperCamelCase_    (       _UpperCAmelCase      ):


							@parameterized.expand(
							    [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY]							,			)
							def        _snake_case					(							self      :Dict							,			__A      :List[str]						)				->				int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= CHEAP_ARGS.copy()
									SCREAMING_SNAKE_CASE__								= """translation""" if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else """summarization"""
									SCREAMING_SNAKE_CASE__								= 0.1 if task == """translation""" else 0

									SCREAMING_SNAKE_CASE__								= make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()						)
									SCREAMING_SNAKE_CASE__								= self.get_auto_remove_tmp_dir()
									args_d.update(
									    data_dir=lowerCAmelCase__							,			model_name_or_path=lowerCAmelCase__							,			tokenizer_name=lowerCAmelCase__							,			train_batch_size=2							,			eval_batch_size=2							,			output_dir=lowerCAmelCase__							,			do_predict=lowerCAmelCase__							,			task=lowerCAmelCase__							,			src_lang="""en_XX"""							,			tgt_lang="""ro_RO"""							,			freeze_encoder=lowerCAmelCase__							,			freeze_embeds=lowerCAmelCase__							,			)
									assert "n_train" in args_d
									SCREAMING_SNAKE_CASE__								= argparse.Namespace(**lowerCAmelCase__						)
									SCREAMING_SNAKE_CASE__								= main(lowerCAmelCase__						)

									SCREAMING_SNAKE_CASE__								= module.model.get_input_embeddings()
									assert not input_embeds.weight.requires_grad
									if model == T5_TINY:
											SCREAMING_SNAKE_CASE__								= module.model.lm_head
											assert not lm_head.weight.requires_grad
											assert (lm_head.weight == input_embeds.weight).all().item()
									elif model == FSMT_TINY:
											SCREAMING_SNAKE_CASE__								= module.model.model
											SCREAMING_SNAKE_CASE__								= fsmt.decoder.embed_positions
											assert not embed_pos.weight.requires_grad
											assert not fsmt.decoder.embed_tokens.weight.requires_grad
											# check that embeds are not the same
											assert fsmt.decoder.embed_tokens != fsmt.encoder.embed_tokens
									else:
											SCREAMING_SNAKE_CASE__								= module.model.model
											SCREAMING_SNAKE_CASE__								= bart.decoder.embed_positions
											assert not embed_pos.weight.requires_grad
											assert not bart.shared.weight.requires_grad
											# check that embeds are the same
											assert bart.decoder.embed_tokens == bart.encoder.embed_tokens
											assert bart.decoder.embed_tokens == bart.shared

									SCREAMING_SNAKE_CASE__								= load_json(module.output_dir / """text_batch.json"""						)
									assert isinstance(lowerCAmelCase__							,			lowerCAmelCase__						)
									assert len(lowerCAmelCase__						) >= 4


							def        _snake_case					(							self      :Union[str, Any]						)				->				List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= CHEAP_ARGS.copy()

									SCREAMING_SNAKE_CASE__								= """summarization"""
									SCREAMING_SNAKE_CASE__								= make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()						)

									args_d.update(
									    data_dir=lowerCAmelCase__							,			tokenizer_name=lowerCAmelCase__							,			train_batch_size=2							,			eval_batch_size=2							,			do_predict=lowerCAmelCase__							,			task=lowerCAmelCase__							,			src_lang="""en_XX"""							,			tgt_lang="""ro_RO"""							,			freeze_encoder=lowerCAmelCase__							,			freeze_embeds=lowerCAmelCase__							,			)

									# test models whose config includes the extra_model_args
									SCREAMING_SNAKE_CASE__								= BART_TINY
									SCREAMING_SNAKE_CASE__								= self.get_auto_remove_tmp_dir()
									SCREAMING_SNAKE_CASE__								= args_d.copy()
									args_da.update(
									    model_name_or_path=lowerCAmelCase__							,			output_dir=lowerCAmelCase__							,			)
									SCREAMING_SNAKE_CASE__								= ("""encoder_layerdrop""", """decoder_layerdrop""", """dropout""", """attention_dropout""")
									for p in extra_model_params:
											SCREAMING_SNAKE_CASE__								= 0.5
									SCREAMING_SNAKE_CASE__								= argparse.Namespace(**lowerCAmelCase__						)
									SCREAMING_SNAKE_CASE__								= main(lowerCAmelCase__						)
									for p in extra_model_params:
											assert getattr(model.config							,			lowerCAmelCase__						) == 0.5, f'''failed to override the model config for param {p}'''

									# test models whose config doesn't include the extra_model_args
									SCREAMING_SNAKE_CASE__								= T5_TINY
									SCREAMING_SNAKE_CASE__								= self.get_auto_remove_tmp_dir()
									SCREAMING_SNAKE_CASE__								= args_d.copy()
									args_da.update(
									    model_name_or_path=lowerCAmelCase__							,			output_dir=lowerCAmelCase__							,			)
									SCREAMING_SNAKE_CASE__								= """encoder_layerdrop"""
									SCREAMING_SNAKE_CASE__								= 0.5
									SCREAMING_SNAKE_CASE__								= argparse.Namespace(**lowerCAmelCase__						)
									with pytest.raises(lowerCAmelCase__						) as excinfo:
											SCREAMING_SNAKE_CASE__								= main(lowerCAmelCase__						)
									assert str(excinfo.value						) == f'''model config doesn\'t have a `{unsupported_param}` attribute'''


							def        _snake_case					(							self      :str						)				->				Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= CHEAP_ARGS.copy()

									SCREAMING_SNAKE_CASE__								= """summarization"""
									SCREAMING_SNAKE_CASE__								= make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()						)

									SCREAMING_SNAKE_CASE__								= BART_TINY
									SCREAMING_SNAKE_CASE__								= self.get_auto_remove_tmp_dir()

									args_d.update(
									    data_dir=lowerCAmelCase__							,			model_name_or_path=lowerCAmelCase__							,			output_dir=lowerCAmelCase__							,			tokenizer_name=lowerCAmelCase__							,			train_batch_size=2							,			eval_batch_size=2							,			do_predict=lowerCAmelCase__							,			task=lowerCAmelCase__							,			src_lang="""en_XX"""							,			tgt_lang="""ro_RO"""							,			freeze_encoder=lowerCAmelCase__							,			freeze_embeds=lowerCAmelCase__							,			)

									# emulate finetune.py
									SCREAMING_SNAKE_CASE__								= argparse.ArgumentParser()
									SCREAMING_SNAKE_CASE__								= pl.Trainer.add_argparse_args(lowerCAmelCase__						)
									SCREAMING_SNAKE_CASE__								= SummarizationModule.add_model_specific_args(lowerCAmelCase__							,			os.getcwd()						)
									SCREAMING_SNAKE_CASE__								= {"""--help""": True}

									# --help test
									with pytest.raises(lowerCAmelCase__						) as excinfo:
											with CaptureStdout() as cs:
													SCREAMING_SNAKE_CASE__								= parser.parse_args(lowerCAmelCase__						)
											assert False, "--help is expected to sys.exit"
									assert excinfo.type == SystemExit
									SCREAMING_SNAKE_CASE__								= lightning_base.arg_to_scheduler_metavar
									assert expected in cs.out, "--help is expected to list the supported schedulers"

									# --lr_scheduler=non_existing_scheduler test
									SCREAMING_SNAKE_CASE__								= """non_existing_scheduler"""
									SCREAMING_SNAKE_CASE__								= {f'''--lr_scheduler={unsupported_param}'''}
									with pytest.raises(lowerCAmelCase__						) as excinfo:
											with CaptureStderr() as cs:
													SCREAMING_SNAKE_CASE__								= parser.parse_args(lowerCAmelCase__						)
											assert False, "invalid argument is expected to sys.exit"
									assert excinfo.type == SystemExit
									SCREAMING_SNAKE_CASE__								= f'''invalid choice: \'{unsupported_param}\''''
									assert expected in cs.err, f'''should have bailed on invalid choice of scheduler {unsupported_param}'''

									# --lr_scheduler=existing_scheduler test
									SCREAMING_SNAKE_CASE__								= """cosine"""
									SCREAMING_SNAKE_CASE__								= args_d.copy()
									SCREAMING_SNAKE_CASE__								= supported_param
									SCREAMING_SNAKE_CASE__								= argparse.Namespace(**lowerCAmelCase__						)
									SCREAMING_SNAKE_CASE__								= main(lowerCAmelCase__						)
									assert (
									    getattr(model.hparams							,			"""lr_scheduler"""						) == supported_param
									), f'''lr_scheduler={supported_param} shouldn\'t fail'''






import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


_lowerCamelCase							=     500000

_lowerCamelCase			,    _lowerCamelCase							=     os.path.split(__file__)
_lowerCamelCase							=     os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset      ):
		SCREAMING_SNAKE_CASE__								= dataset.select(range(0   ,				len(UpperCamelCase__      )   ,				2      )      )






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset      ):
		SCREAMING_SNAKE_CASE__								= dataset.sort("""numbers"""      )






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset      ):
		SCREAMING_SNAKE_CASE__								= dataset.shuffle()






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset      ):
		SCREAMING_SNAKE_CASE__								= dataset.train_test_split(0.1      )






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset   ,				UpperCamelCase__:							Tuple=10      ):
		for shard_id in range(UpperCamelCase__      ):
				SCREAMING_SNAKE_CASE__								= dataset.shard(UpperCamelCase__   ,				UpperCamelCase__      )






def 			SCREAMING_SNAKE_CASE__							(    ):
		SCREAMING_SNAKE_CASE__								= {"""num examples""": SPEED_TEST_N_EXAMPLES}
		SCREAMING_SNAKE_CASE__								= (select, sort, shuffle, train_test_split, shard)
		with tempfile.TemporaryDirectory() as tmp_dir:
				print("""generating dataset"""      )
				SCREAMING_SNAKE_CASE__								= datasets.Features({"""text""": datasets.Value("""string"""      ), """numbers""": datasets.Value("""float32"""      )}      )
				SCREAMING_SNAKE_CASE__								= generate_example_dataset(
				    os.path.join(UpperCamelCase__   ,				"""dataset.arrow"""      )   ,				UpperCamelCase__   ,				num_examples=UpperCamelCase__      )
				print("""Functions"""      )
				for func in functions:
						print(func.__name__      )
						SCREAMING_SNAKE_CASE__								= func(UpperCamelCase__      )

		with open(UpperCamelCase__   ,				"""wb"""      ) as f:
				f.write(json.dumps(UpperCamelCase__      ).encode("""utf-8"""      )      )


if __name__ == "__main__":  # useful to run the profiler
	benchmark_indices_mapping()