
'''simple docstring'''



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  _a					(datasets.BeamBasedBuilder):





       """simple docstring"""




       def 							UpperCamelCase					(				self   )     ->      List[str]:
          return datasets.DatasetInfo(
              features=datasets.Features({"""content""": datasets.Value("""string"""   )}   )							,				supervised_keys=snake_case__							,				)






       def 							UpperCamelCase					(				self							,				A__							,				A__   )     ->      str:
          return [datasets.SplitGenerator(name=datasets.Split.TRAIN							,				gen_kwargs={"""examples""": get_test_dummy_examples()}   )]






       def 							UpperCamelCase					(				self							,				A__							,				A__   )     ->      Optional[int]:
          import apache_beam as beam

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






class  _a					(datasets.BeamBasedBuilder):





       """simple docstring"""




       def 							UpperCamelCase					(				self   )     ->      List[Any]:
          return datasets.DatasetInfo(
              features=datasets.Features({"""a""": datasets.Sequence({"""b""": datasets.Value("""string"""   )}   )}   )							,				supervised_keys=snake_case__							,				)






       def 							UpperCamelCase					(				self							,				A__							,				A__   )     ->      Optional[int]:
          return [
              datasets.SplitGenerator(name=datasets.Split.TRAIN							,				gen_kwargs={"""examples""": get_test_nested_examples()}   )
          ]






       def 							UpperCamelCase					(				self							,				A__							,				A__   )     ->      List[str]:
          import apache_beam as beam

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






def  lowerCAmelCase_						(					)  ->  Dict:




   """simple docstring"""
   return [(i, {"content": content}) for i, content in enumerate(["""foo""", """bar""", """foobar"""]    )]
def  lowerCAmelCase_						(					)  ->  Optional[int]:




   """simple docstring"""
   return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["""foo""", """bar""", """foobar"""]    )]







class  _a					(SCREAMING_SNAKE_CASE__):





       """simple docstring"""




       @require_beam
       def 							UpperCamelCase					(				self   )     ->      Dict:
          _SCREAMING_SNAKE_CASE  			=					len(get_test_dummy_examples()   )
          with tempfile.TemporaryDirectory() as tmp_cache_dir:
             _SCREAMING_SNAKE_CASE  			=					DummyBeamDataset(cache_dir=snake_case__							,				beam_runner="""DirectRunner"""   )
             builder.download_and_prepare()
             self.assertTrue(
                 os.path.exists(
                     os.path.join(snake_case__							,				builder.name							,				"""default"""							,				"""0.0.0"""							,				F"{builder.name}-train.arrow"   )   )   )
             self.assertDictEqual(builder.info.features							,				datasets.Features({"""content""": datasets.Value("""string"""   )}   )   )
             _SCREAMING_SNAKE_CASE  			=					builder.as_dataset()
             self.assertEqual(dset["""train"""].num_rows							,				snake_case__   )
             self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples							,				snake_case__   )
             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(snake_case__							,				builder.name							,				"""default"""							,				"""0.0.0"""							,				"""dataset_info.json"""   )   )   )
             del dset






       @require_beam
       def 							UpperCamelCase					(				self   )     ->      int:
          import apache_beam as beam

          _SCREAMING_SNAKE_CASE  			=					beam.io.parquetio.WriteToParquet

          _SCREAMING_SNAKE_CASE  			=					len(get_test_dummy_examples()   )
          with tempfile.TemporaryDirectory() as tmp_cache_dir:
             _SCREAMING_SNAKE_CASE  			=					DummyBeamDataset(cache_dir=snake_case__							,				beam_runner="""DirectRunner"""   )
             with patch("""apache_beam.io.parquetio.WriteToParquet"""   ) as write_parquet_mock:
                _SCREAMING_SNAKE_CASE  			=					partial(snake_case__							,				num_shards=2   )
                builder.download_and_prepare()
             self.assertTrue(
                 os.path.exists(
                     os.path.join(
                         snake_case__							,				builder.name							,				"""default"""							,				"""0.0.0"""							,				F"{builder.name}-train-00000-of-00002.arrow"   )   )   )
             self.assertTrue(
                 os.path.exists(
                     os.path.join(
                         snake_case__							,				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"""   )}   )   )
             _SCREAMING_SNAKE_CASE  			=					builder.as_dataset()
             self.assertEqual(dset["""train"""].num_rows							,				snake_case__   )
             self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples							,				snake_case__   )
             # 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(snake_case__							,				builder.name							,				"""default"""							,				"""0.0.0"""							,				"""dataset_info.json"""   )   )   )
             del dset






       @require_beam
       def 							UpperCamelCase					(				self   )     ->      Any:

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






       @require_beam
       def 							UpperCamelCase					(				self   )     ->      Optional[Any]:
          _SCREAMING_SNAKE_CASE  			=					len(get_test_nested_examples()   )
          with tempfile.TemporaryDirectory() as tmp_cache_dir:
             _SCREAMING_SNAKE_CASE  			=					NestedBeamDataset(cache_dir=snake_case__							,				beam_runner="""DirectRunner"""   )
             builder.download_and_prepare()
             self.assertTrue(
                 os.path.exists(
                     os.path.join(snake_case__							,				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"""   )}   )}   )   )
             _SCREAMING_SNAKE_CASE  			=					builder.as_dataset()
             self.assertEqual(dset["""train"""].num_rows							,				snake_case__   )
             self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples							,				snake_case__   )
             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(snake_case__							,				builder.name							,				"""default"""							,				"""0.0.0"""							,				"""dataset_info.json"""   )   )   )
             del dset












from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


SCREAMING_SNAKE_CASE:     Any									= {
    '''configuration_pegasus_x''': ['''PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''PegasusXConfig'''],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						SCREAMING_SNAKE_CASE:     List[Any]									= [
						    '''PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''PegasusXForConditionalGeneration''',
						    '''PegasusXModel''',
						    '''PegasusXPreTrainedModel''',
						]


if TYPE_CHECKING:
						from .configuration_pegasus_x import PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusXConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_pegasus_x import (
												    PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST,
												    PegasusXForConditionalGeneration,
												    PegasusXModel,
												    PegasusXPreTrainedModel,
												)


else:
						import sys

						SCREAMING_SNAKE_CASE:     List[Any]									= _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)