


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








'''simple docstring'''






from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType

class 				a__(   lowerCamelCase__      ):
    lowercase__					=       ["""image_processor""", """tokenizer"""]
    lowercase__					=       """BridgeTowerImageProcessor"""
    lowercase__					=       ("""RobertaTokenizer""", """RobertaTokenizerFast""")







    def __init__( self					:    int		,     __snake_case					:    str		,     __snake_case					:    List[str]      ):
      super().__init__(__snake_case		,     __snake_case      )







    def __call__( self					:    int		,     __snake_case					:    Optional[Any]		,     __snake_case					:    Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None		,     __snake_case					:    bool = True		,     __snake_case					:    Union[bool, str, PaddingStrategy] = False		,     __snake_case					:    Union[bool, str, TruncationStrategy] = None		,     __snake_case					:    Optional[int] = None		,     __snake_case					:    int = 0		,     __snake_case					:    Optional[int] = None		,     __snake_case					:    Optional[bool] = None		,     __snake_case					:    Optional[bool] = None		,     __snake_case					:    bool = False		,     __snake_case					:    bool = False		,     __snake_case					:    bool = False		,     __snake_case					:    bool = False		,     __snake_case					:    bool = True		,     __snake_case					:    Optional[Union[str, TensorType]] = None		,     **__snake_case					:    List[Any]		,     ):
      a		:   Optional[int]	  =	self.tokenizer(
          text=__snake_case		,     add_special_tokens=__snake_case		,     padding=__snake_case		,     truncation=__snake_case		,     max_length=__snake_case		,     stride=__snake_case		,     pad_to_multiple_of=__snake_case		,     return_token_type_ids=__snake_case		,     return_attention_mask=__snake_case		,     return_overflowing_tokens=__snake_case		,     return_special_tokens_mask=__snake_case		,     return_offsets_mapping=__snake_case		,     return_length=__snake_case		,     verbose=__snake_case		,     return_tensors=__snake_case		,     **__snake_case		,     )
      # add pixel_values + pixel_mask
      a		:   List[str]	  =	self.image_processor(
          __snake_case		,     return_tensors=__snake_case		,     do_normalize=__snake_case		,     do_center_crop=__snake_case		,     **__snake_case      )
      encoding.update(__snake_case      )

      return encoding







    def      lowercase_			( self					:    int		,     *__snake_case					:    List[str]		,     **__snake_case					:    List[str]      ):
      return self.tokenizer.batch_decode(*__snake_case		,     **__snake_case      )







    def      lowercase_			( self					:    List[str]		,     *__snake_case					:    Tuple		,     **__snake_case					:    Union[str, Any]      ):
      return self.tokenizer.decode(*__snake_case		,     **__snake_case      )







    @property
    def      lowercase_			( self					:    Optional[Any]      ):
      a		:   Optional[Any]	  =	self.tokenizer.model_input_names
      a		:   Union[str, Any]	  =	self.image_processor.model_input_names
      return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names      )      )