




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    ):







		def        lowercase       (    self				)      ->      List[str]:


							'''simple docstring'''

							return datasets.DatasetInfo(
							    features=datasets.Features({"content": datasets.Value("string"				)}				) ,						supervised_keys=__snake_case ,						)







		def        lowercase       (    self ,						__snake_case ,						__snake_case				)      ->      Tuple:


							'''simple docstring'''

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







		def        lowercase       (    self ,						__snake_case ,						__snake_case				)      ->      int:


							'''simple docstring'''

							import apache_beam as beam

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






class        _a(    datasets.BeamBasedBuilder    ):







		def        lowercase       (    self				)      ->      List[Any]:


							'''simple docstring'''

							return datasets.DatasetInfo(
							    features=datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"				)}				)}				) ,						supervised_keys=__snake_case ,						)







		def        lowercase       (    self ,						__snake_case ,						__snake_case				)      ->      Tuple:


							'''simple docstring'''

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







		def        lowercase       (    self ,						__snake_case ,						__snake_case				)      ->      Tuple:


							'''simple docstring'''

							import apache_beam as beam

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







def    A				(							):
					return [(i, {"content": content}) for i, content in enumerate(["foo", "bar", "foobar"]				)]







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






class        _a(    __A    ):







		@require_beam
		def        lowercase       (    self				)      ->      Union[str, Any]:


							'''simple docstring'''

							_snake_case       :					List[Any]				       =   len(get_test_dummy_examples()				)
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
												_snake_case       :					Optional[int]				       =   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"				)}				)				)
												_snake_case       :					str				       =   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        lowercase       (    self				)      ->      List[Any]:


							'''simple docstring'''

							import apache_beam as beam

							_snake_case       :					Dict				       =   beam.io.parquetio.WriteToParquet

							_snake_case       :					List[Any]				       =   len(get_test_dummy_examples()				)
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
												_snake_case       :					List[Any]				       =   DummyBeamDataset(cache_dir=__snake_case ,						beam_runner="DirectRunner"				)
												with patch("apache_beam.io.parquetio.WriteToParquet"				) as write_parquet_mock:
																	_snake_case       :					Union[str, Any]				       =   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"				)}				)				)
												_snake_case       :					Dict				       =   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        lowercase       (    self				)      ->      List[Any]:


							'''simple docstring'''

							with tempfile.TemporaryDirectory() as tmp_cache_dir:
												_snake_case       :					str				       =   DummyBeamDataset(cache_dir=__snake_case				)
												self.assertRaises(datasets.builder.MissingBeamOptions ,						builder.download_and_prepare				)







		@require_beam
		def        lowercase       (    self				)      ->      Dict:


							'''simple docstring'''

							_snake_case       :					Optional[int]				       =   len(get_test_nested_examples()				)
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
												_snake_case       :					Optional[Any]				       =   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"				)}				)}				)				)
												_snake_case       :					Union[str, Any]				       =   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


__lowerCAmelCase :str     =    {
    'configuration_nllb_moe': [
        'NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'NllbMoeConfig',
    ]
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__lowerCAmelCase :Optional[int]     =    [
							    'NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'NllbMoeForConditionalGeneration',
							    'NllbMoeModel',
							    'NllbMoePreTrainedModel',
							    'NllbMoeTop2Router',
							    'NllbMoeSparseMLP',
							]


if TYPE_CHECKING:
							from .configuration_nllb_moe import (
							    NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    NllbMoeConfig,
							)

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_nllb_moe import (
														    NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST,
														    NllbMoeForConditionalGeneration,
														    NllbMoeModel,
														    NllbMoePreTrainedModel,
														    NllbMoeSparseMLP,
														    NllbMoeTopaRouter,
														)


else:
							import sys

							__lowerCAmelCase :int     =    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)