




import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam

class 		__lowerCAmelCase	(	datasets.BeamBasedBuilder ):
							def       lowerCamelCase__							(						self						:List[Any]				):







								'''simple docstring'''







								return datasets.DatasetInfo(
								    features=datasets.Features({"""content""": datasets.Value("""string"""				)}				)		,      supervised_keys=__magic_name__		,      )
							def       lowerCamelCase__							(						self						:Optional[Any]		,      __magic_name__						:int		,      __magic_name__						:Tuple				):







								'''simple docstring'''







								return [datasets.SplitGenerator(name=datasets.Split.TRAIN		,      gen_kwargs={"""examples""": get_test_dummy_examples()}				)]





							def       lowerCamelCase__							(						self						:Any		,      __magic_name__						:Dict		,      __magic_name__						:Union[str, Any]				):







								'''simple docstring'''







								import apache_beam as beam

								return pipeline | "Load Examples" >> beam.Create(__magic_name__				)







class 		__lowerCAmelCase	(	datasets.BeamBasedBuilder ):
							def       lowerCamelCase__							(						self						:str				):







								'''simple docstring'''







								return datasets.DatasetInfo(
								    features=datasets.Features({"""a""": datasets.Sequence({"""b""": datasets.Value("""string"""				)}				)}				)		,      supervised_keys=__magic_name__		,      )
							def       lowerCamelCase__							(						self						:Optional[int]		,      __magic_name__						:str		,      __magic_name__						:Optional[Any]				):







								'''simple docstring'''







								return [
								    datasets.SplitGenerator(name=datasets.Split.TRAIN		,      gen_kwargs={"""examples""": get_test_nested_examples()}				)
								]





							def       lowerCamelCase__							(						self						:Union[str, Any]		,      __magic_name__						:Dict		,      __magic_name__						:int				):







								'''simple docstring'''







								import apache_beam as beam

								return pipeline | "Load Examples" >> beam.Create(__magic_name__				)







def 		__A		(	)      ->       Any:
	return [(i, {"content": content}) for i, content in enumerate(["""foo""", """bar""", """foobar"""]    )]





def 		__A		(	)      ->       List[Any]:
	return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["""foo""", """bar""", """foobar"""]    )]



class 		__lowerCAmelCase	(	__magic_name__ ):
							@require_beam
							def       lowerCamelCase__							(						self						:str				):







								'''simple docstring'''







								a										=							len(get_test_dummy_examples()				)
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
									a										=							DummyBeamDataset(cache_dir=__magic_name__		,      beam_runner="""DirectRunner"""				)
									builder.download_and_prepare()
									self.assertTrue(
									    os.path.exists(
									        os.path.join(__magic_name__		,      builder.name		,      """default"""		,      """0.0.0"""		,      F'{builder.name}-train.arrow'				)				)				)
									self.assertDictEqual(builder.info.features		,      datasets.Features({"""content""": datasets.Value("""string"""				)}				)				)
									a										=							builder.as_dataset()
									self.assertEqual(dset["""train"""].num_rows		,      __magic_name__				)
									self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples		,      __magic_name__				)
									self.assertDictEqual(dset["""train"""][0]		,      get_test_dummy_examples()[0][1]				)
									self.assertDictEqual(
									    dset["""train"""][expected_num_examples - 1]		,      get_test_dummy_examples()[expected_num_examples - 1][1]				)
									self.assertTrue(
									    os.path.exists(os.path.join(__magic_name__		,      builder.name		,      """default"""		,      """0.0.0"""		,      """dataset_info.json"""				)				)				)
									del dset
							@require_beam
							def       lowerCamelCase__							(						self						:Optional[int]				):







								'''simple docstring'''







								import apache_beam as beam

								a										=							beam.io.parquetio.WriteToParquet

								a										=							len(get_test_dummy_examples()				)
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
									a										=							DummyBeamDataset(cache_dir=__magic_name__		,      beam_runner="""DirectRunner"""				)
									with patch("""apache_beam.io.parquetio.WriteToParquet"""				) as write_parquet_mock:
										a										=							partial(__magic_name__		,      num_shards=2				)
										builder.download_and_prepare()
									self.assertTrue(
									    os.path.exists(
									        os.path.join(
									            __magic_name__		,      builder.name		,      """default"""		,      """0.0.0"""		,      F'{builder.name}-train-00000-of-00002.arrow'				)				)				)
									self.assertTrue(
									    os.path.exists(
									        os.path.join(
									            __magic_name__		,      builder.name		,      """default"""		,      """0.0.0"""		,      F'{builder.name}-train-00000-of-00002.arrow'				)				)				)
									self.assertDictEqual(builder.info.features		,      datasets.Features({"""content""": datasets.Value("""string"""				)}				)				)
									a										=							builder.as_dataset()
									self.assertEqual(dset["""train"""].num_rows		,      __magic_name__				)
									self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples		,      __magic_name__				)
									# Order is not preserved when sharding, so we just check that all the elements are there
									self.assertListEqual(sorted(dset["""train"""]["""content"""]				)		,      sorted(["""foo""", """bar""", """foobar"""]				)				)
									self.assertTrue(
									    os.path.exists(os.path.join(__magic_name__		,      builder.name		,      """default"""		,      """0.0.0"""		,      """dataset_info.json"""				)				)				)
									del dset
							@require_beam
							def       lowerCamelCase__							(						self						:Dict				):







								'''simple docstring'''







								with tempfile.TemporaryDirectory() as tmp_cache_dir:
									a										=							DummyBeamDataset(cache_dir=__magic_name__				)
									self.assertRaises(datasets.builder.MissingBeamOptions		,      builder.download_and_prepare				)





							@require_beam
							def       lowerCamelCase__							(						self						:Optional[Any]				):







								'''simple docstring'''







								a										=							len(get_test_nested_examples()				)
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
									a										=							NestedBeamDataset(cache_dir=__magic_name__		,      beam_runner="""DirectRunner"""				)
									builder.download_and_prepare()
									self.assertTrue(
									    os.path.exists(
									        os.path.join(__magic_name__		,      builder.name		,      """default"""		,      """0.0.0"""		,      F'{builder.name}-train.arrow'				)				)				)
									self.assertDictEqual(
									    builder.info.features		,      datasets.Features({"""a""": datasets.Sequence({"""b""": datasets.Value("""string"""				)}				)}				)				)
									a										=							builder.as_dataset()
									self.assertEqual(dset["""train"""].num_rows		,      __magic_name__				)
									self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples		,      __magic_name__				)
									self.assertDictEqual(dset["""train"""][0]		,      get_test_nested_examples()[0][1]				)
									self.assertDictEqual(
									    dset["""train"""][expected_num_examples - 1]		,      get_test_nested_examples()[expected_num_examples - 1][1]				)
									self.assertTrue(
									    os.path.exists(os.path.join(__magic_name__		,      builder.name		,      """default"""		,      """0.0.0"""		,      """dataset_info.json"""				)				)				)
									del dset







import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp





def 		__A		(	)      ->       Any:
	a										=							ArgumentParser(
	    description=(
	        """PyTorch TPU distributed training launch """
	        """helper utility that will spawn up """
	        """multiple distributed processes"""
	    )    )

	# Optional arguments for the launch helper
	parser.add_argument("""--num_cores"""							,			type=__lowerCamelCase							,			default=1							,			help="""Number of TPU cores to use (1 or 8)."""    )

	# positional
	parser.add_argument(
	    """training_script"""							,			type=__lowerCamelCase							,			help=(
	        """The full path to the single TPU training """
	        """program/script to be launched in parallel, """
	        """followed by all the arguments for the """
	        """training script"""
	    )							,			)

	# rest from the training program
	parser.add_argument("""training_script_args"""							,			nargs=__lowerCamelCase    )

	return parser.parse_args()





def 		__A		(	)      ->       Union[str, Any]:
	a										=							parse_args()

	# Import training_script as a module.
	a										=							Path(args.training_script    )
	sys.path.append(str(script_fpath.parent.resolve()    )    )
	a										=							script_fpath.stem
	a										=							importlib.import_module(__lowerCamelCase    )

	# Patch sys.argv
	a										=							[args.training_script] + args.training_script_args + ["""--tpu_num_cores""", str(args.num_cores    )]
	xmp.spawn(mod._mp_fn							,			args=()							,			nprocs=args.num_cores    )


if __name__ == "__main__":
		main()

