





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






							def 					__a     (     self   )			->				Any:
								return datasets.DatasetInfo(
								    features=datasets.Features({"content": datasets.Value("string"   )}   )   ,      supervised_keys=lowerCAmelCase__   ,      )






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






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				List[Any]:
								import apache_beam as beam

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


class  __UpperCamelCase       (		datasets.BeamBasedBuilder		):






							def 					__a     (     self   )			->				str:
								return datasets.DatasetInfo(
								    features=datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"   )}   )}   )   ,      supervised_keys=lowerCAmelCase__   ,      )






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






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				str:
								import apache_beam as beam

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






def    _SCREAMING_SNAKE_CASE	(			)				->List[str]:






	'''simple docstring'''






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






def    _SCREAMING_SNAKE_CASE	(			)				->Tuple:






	'''simple docstring'''






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







class  __UpperCamelCase       (		a__		):






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






							@require_beam
							def 					__a     (     self   )			->				Union[str, Any]:
								import apache_beam as beam

								a   :	int			=       beam.io.parquetio.WriteToParquet

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






							@require_beam
							def 					__a     (     self   )			->				Optional[Any]:
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
									a   :	int			=       DummyBeamDataset(cache_dir=lowerCAmelCase__   )
									self.assertRaises(datasets.builder.MissingBeamOptions   ,      builder.download_and_prepare   )






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









"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


if is_sentencepiece_available():
		from ..ta.tokenization_ta import TaTokenizer
else:
		from ...utils.dummy_sentencepiece_objects import TaTokenizer

a		:				List[Any]  			=    TaTokenizer

if is_tokenizers_available():
		from ..ta.tokenization_ta_fast import TaTokenizerFast
else:
		from ...utils.dummy_tokenizers_objects import TaTokenizerFast

a		:				Union[str, Any]  			=    TaTokenizerFast

a		:				Union[str, Any]  			=    {'''configuration_mt5''': ['''MT5Config''', '''MT5OnnxConfig''']}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		a		:				str  			=    [
		    '''MT5EncoderModel''',
		    '''MT5ForConditionalGeneration''',
		    '''MT5ForQuestionAnswering''',
		    '''MT5Model''',
		    '''MT5PreTrainedModel''',
		    '''MT5Stack''',
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		a		:				List[str]  			=    ['''TFMT5EncoderModel''', '''TFMT5ForConditionalGeneration''', '''TFMT5Model''']

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		a		:				Optional[Any]  			=    ['''FlaxMT5EncoderModel''', '''FlaxMT5ForConditionalGeneration''', '''FlaxMT5Model''']


if TYPE_CHECKING:
		from .configuration_mta import MTaConfig, MTaOnnxConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_mta import (
				    MTaEncoderModel,
				    MTaForConditionalGeneration,
				    MTaForQuestionAnswering,
				    MTaModel,
				    MTaPreTrainedModel,
				    MTaStack,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_mta import TFMTaEncoderModel, TFMTaForConditionalGeneration, TFMTaModel

		try:
				if not is_flax_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_flax_mta import FlaxMTaEncoderModel, FlaxMTaForConditionalGeneration, FlaxMTaModel

else:
		import sys

		a		:				List[Any]  			=    _LazyModule(
		    __name__,
		    globals()['''__file__'''],
		    _import_structure,
		    extra_objects={'''MT5Tokenizer''': MTaTokenizer, '''MT5TokenizerFast''': MTaTokenizerFast},
		    module_spec=__spec__,
		)


