


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


     """simple docstring"""


     def 	snake_case_  (	self    ):
        return datasets.DatasetInfo(
            features=datasets.Features({"""content""": datasets.Value("""string"""    )}    )  ,       supervised_keys=__A  ,       )


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


     def 	snake_case_  (	self  ,       __A  ,       __A    ):
        import apache_beam as beam

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







class     __UpperCAmelCase							(       datasets.BeamBasedBuilder							):


     """simple docstring"""


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


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


     def 	snake_case_  (	self  ,       __A  ,       __A    ):
        import apache_beam as beam

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







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     __UpperCAmelCase							(       __A							):


     """simple docstring"""


     @require_beam
     def 	snake_case_  (	self    ):
        __a					=  len(get_test_dummy_examples()    )
        with tempfile.TemporaryDirectory() as tmp_cache_dir:
           __a					=  DummyBeamDataset(cache_dir=__A  ,       beam_runner="""DirectRunner"""    )
           builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(__A  ,       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  ,       __A    )
           self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples  ,       __A    )
           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(__A  ,       builder.name  ,       """default"""  ,       """0.0.0"""  ,       """dataset_info.json"""    )    )    )
           del dset


     @require_beam
     def 	snake_case_  (	self    ):
        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=__A  ,       beam_runner="""DirectRunner"""    )
           with patch("""apache_beam.io.parquetio.WriteToParquet"""    ) as write_parquet_mock:
              __a					=  partial(__A  ,       num_shards=2    )
              builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       __A  ,       builder.name  ,       """default"""  ,       """0.0.0"""  ,       f'''{builder.name}-train-00000-of-00002.arrow'''    )    )    )
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       __A  ,       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  ,       __A    )
           self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples  ,       __A    )
           # 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(__A  ,       builder.name  ,       """default"""  ,       """0.0.0"""  ,       """dataset_info.json"""    )    )    )
           del dset


     @require_beam
     def 	snake_case_  (	self    ):
        with tempfile.TemporaryDirectory() as tmp_cache_dir:
           __a					=  DummyBeamDataset(cache_dir=__A    )
           self.assertRaises(datasets.builder.MissingBeamOptions  ,       builder.download_and_prepare    )


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







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


SCREAMING_SNAKE_CASE							    =							{
    'configuration_tapas': ['TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TapasConfig'],
    'tokenization_tapas': ['TapasTokenizer'],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  SCREAMING_SNAKE_CASE							    =							[
      'TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
      'TapasForMaskedLM',
      'TapasForQuestionAnswering',
      'TapasForSequenceClassification',
      'TapasModel',
      'TapasPreTrainedModel',
      'load_tf_weights_in_tapas',
  ]
try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  SCREAMING_SNAKE_CASE							    =							[
      'TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
      'TFTapasForMaskedLM',
      'TFTapasForQuestionAnswering',
      'TFTapasForSequenceClassification',
      'TFTapasModel',
      'TFTapasPreTrainedModel',
  ]


if TYPE_CHECKING:
  from .configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
  from .tokenization_tapas import TapasTokenizer

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tapas import (
        TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
        TapasForMaskedLM,
        TapasForQuestionAnswering,
        TapasForSequenceClassification,
        TapasModel,
        TapasPreTrainedModel,
        load_tf_weights_in_tapas,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_tapas import (
        TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFTapasForMaskedLM,
        TFTapasForQuestionAnswering,
        TFTapasForSequenceClassification,
        TFTapasModel,
        TFTapasPreTrainedModel,
    )


else:
  import sys

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



