




from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


SCREAMING_SNAKE_CASE__       :      Tuple  =   logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__       :      Dict  =   {
    'microsoft/resnet-50': 'https://huggingface.co/microsoft/resnet-50/blob/main/config.json',
}




class 		UpperCamelCase__	(lowerCAmelCase__    ,				lowerCAmelCase__     ):

    '''simple docstring'''



    lowerCamelCase_      :	List[str]	     =      """resnet"""
    lowerCamelCase_      :	List[Any]	     =      ["""basic""", """bottleneck"""]







    def __init__(    self     ,  UpperCamelCase__=3     ,  UpperCamelCase__=64     ,  UpperCamelCase__=[256, 512, 1024, 2048]     ,  UpperCamelCase__=[3, 4, 6, 3]     ,  UpperCamelCase__="bottleneck"     ,  UpperCamelCase__="relu"     ,  UpperCamelCase__=False     ,  UpperCamelCase__=None     ,  UpperCamelCase__=None     ,  **UpperCamelCase__     ,  )     ->      str:
           super().__init__(**UpperCamelCase__     )
           if layer_type not in self.layer_types:
                  raise ValueError(F'''layer_type={layer_type} is not one of {",".join(self.layer_types     )}'''     )
           lowerCamelCase   :      Optional[Any]			 =					num_channels
           lowerCamelCase   :      List[str]			 =					embedding_size
           lowerCamelCase   :      Any			 =					hidden_sizes
           lowerCamelCase   :      List[str]			 =					depths
           lowerCamelCase   :      Optional[int]			 =					layer_type
           lowerCamelCase   :      Any			 =					hidden_act
           lowerCamelCase   :      Optional[Any]			 =					downsample_in_first_stage
           lowerCamelCase   :      List[str]			 =					["stem"] + [F'''stage{idx}''' for idx in range(1     ,  len(UpperCamelCase__     ) + 1     )]
           lowerCamelCase ,      lowerCamelCase   :      Tuple			 =					get_aligned_output_features_output_indices(
               out_features=UpperCamelCase__     ,  out_indices=UpperCamelCase__     ,  stage_names=self.stage_names     )




class 		UpperCamelCase__	(lowerCAmelCase__     ):

    '''simple docstring'''



    lowerCamelCase_      :	int	     =      version.parse("""1.11"""     )







    @property
    def     _lowercase     (    self     )     ->      Mapping[str, Mapping[int, str]]:
           return OrderedDict(
               [
                   ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
               ]     )







    @property
    def     _lowercase     (    self     )     ->      float:
           return 1e-3














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=SCREAMING_SNAKE_CASE							,   )







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






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







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   int						):


      '''simple docstring'''






      import apache_beam as beam

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







class 			__a  (    datasets.BeamBasedBuilder       ):







 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






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







 def 					__lowercase					(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   List[Any]						):


      '''simple docstring'''






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







 def 					__lowercase					(		self					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Tuple							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






      import apache_beam as beam

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







def 		SCREAMING_SNAKE_CASE							(					) ->	List[str]:
     return [(i, {"content": content}) for i, content in enumerate(["foo", "bar", "foobar"]			)]







def 		SCREAMING_SNAKE_CASE							(					) ->	Optional[Any]:
     return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["foo", "bar", "foobar"]			)]





class 			__a  (    A__       ):







 @require_beam
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






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







 @require_beam
 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      import apache_beam as beam

      UpperCamelCase__			:   Any       				=  beam.io.parquetio.WriteToParquet

      UpperCamelCase__			:   Tuple       				=  len(get_test_dummy_examples()						)
      with tempfile.TemporaryDirectory() as tmp_cache_dir:
           UpperCamelCase__			:   str       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE							,   beam_runner="DirectRunner"						)
           with patch("apache_beam.io.parquetio.WriteToParquet"						) as write_parquet_mock:
                UpperCamelCase__			:   Optional[Any]       				=  partial(SCREAMING_SNAKE_CASE							,   num_shards=2						)
                builder.download_and_prepare()
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       SCREAMING_SNAKE_CASE							,   builder.name							,   "default"							,   "0.0.0"							,   F'{builder.name}-train-00000-of-00002.arrow'						)						)						)
           self.assertTrue(
               os.path.exists(
                   os.path.join(
                       SCREAMING_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"						)}						)						)
           UpperCamelCase__			:   List[Any]       				=  builder.as_dataset()
           self.assertEqual(dset["train"].num_rows							,   SCREAMING_SNAKE_CASE						)
           self.assertEqual(dset["train"].info.splits["train"].num_examples							,   SCREAMING_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(SCREAMING_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:
           UpperCamelCase__			:   int       				=  DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE						)
           self.assertRaises(datasets.builder.MissingBeamOptions							,   builder.download_and_prepare						)







 @require_beam
 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






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