


import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch




def  a   (lowerCAmelCase__   ):  # picklable for multiprocessing
   return x.sum()




def  a   (lowerCAmelCase__   ):  # picklable for multiprocessing
   return i + 1


@dataclass
class     __UpperCAmelCase							:


     """simple docstring"""
     _lowerCamelCase   =      42
     _lowerCamelCase   =      42







class     __UpperCAmelCase							(       __A							):


     """simple docstring"""


     def 	snake_case_  (	self    ):
        __a					=  {}
        __a					=  []
        __a					=  1
        __a					=  [1, 2]
        __a					=  {"""a""": 1, """b""": 2}
        __a					=  {"""a""": [1, 2], """b""": [3, 4]}
        __a					=  {"""a""": {"""1""": 1}, """b""": 2}
        __a					=  {"""a""": 1, """b""": 2, """c""": 3, """d""": 4}
        __a					=  {}
        __a					=  []
        __a					=  2
        __a					=  [2, 3]
        __a					=  {"""a""": 2, """b""": 3}
        __a					=  {"""a""": [2, 3], """b""": [4, 5]}
        __a					=  {"""a""": {"""1""": 2}, """b""": 3}
        __a					=  {"""a""": 2, """b""": 3, """c""": 4, """d""": 5}
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__    )  ,       lowerCamelCase__    )

        __a					=  2
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       num_proc=lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       num_proc=lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       num_proc=lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       num_proc=lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       num_proc=lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       num_proc=lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       num_proc=lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       num_proc=lowerCamelCase__    )  ,       lowerCamelCase__    )

        __a					=  {"""a""": np.eye(2    ), """b""": np.zeros(3    ), """c""": np.ones(2    )}
        __a					=  {"""a""": 2, """b""": 0, """c""": 2}
        __a					=  {
            """a""": np.eye(2    ).astype(lowerCamelCase__    ),
            """b""": np.zeros(3    ).astype(lowerCamelCase__    ),
            """c""": np.ones(2    ).astype(lowerCamelCase__    ),
        }
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       map_numpy=lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(
            {k: v.tolist() for k, v in map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       map_numpy=lowerCamelCase__    ).items()}  ,       {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}  ,       )
        self.assertEqual(map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       map_numpy=lowerCamelCase__  ,       num_proc=lowerCamelCase__    )  ,       lowerCamelCase__    )
        self.assertEqual(
            {k: v.tolist() for k, v in map_nested(lowerCamelCase__  ,       lowerCamelCase__  ,       map_numpy=lowerCamelCase__  ,       num_proc=lowerCamelCase__    ).items()}  ,       {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}  ,       )
        with self.assertRaises(lowerCamelCase__    ):  # can't pickle a local lambda
           map_nested(lambda __A    : x + 1  ,       lowerCamelCase__  ,       num_proc=lowerCamelCase__    )


     def 	snake_case_  (	self    ):
        __a					=  {"""a""": 1, """b""": 2}
        __a					=  {"""a""": 3, """b""": 4}
        __a					=  {"""a""": 5, """b""": 6}
        __a					=  sorted([("""a""", (1, 3, 5)), ("""b""", (2, 4, 6))]    )
        self.assertEqual(sorted(zip_dict(lowerCamelCase__  ,       lowerCamelCase__  ,       lowerCamelCase__    )    )  ,       lowerCamelCase__    )


     def 	snake_case_  (	self    ):


        class     __UpperCAmelCase							:


             """simple docstring"""
             _lowerCamelCase   =      """bar"""







        __a					=  Foo()
        self.assertEqual(foo.my_attr  ,       """bar"""    )
        with temporary_assignment(lowerCamelCase__  ,       """my_attr"""  ,       """BAR"""    ):
           self.assertEqual(foo.my_attr  ,       """BAR"""    )
        self.assertEqual(foo.my_attr  ,       """bar"""    )







@pytest.mark.parametrize(
    """iterable_length, num_proc, expected_num_proc"""		,				[
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]		,				)
def  a   (lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__   ):


   with patch("""datasets.utils.py_utils._single_map_nested"""   ) as mock_single_map_nested, patch(
       """datasets.parallel.parallel.Pool"""   ) as mock_multiprocessing_pool:
      __a					=  {f'''{i}''': i for i in range(UpperCamelCase__   )}
      __a					=  map_nested(lambda lowerCAmelCase__   : x + 10		,				UpperCamelCase__		,				num_proc=UpperCamelCase__		,				parallel_min_length=16   )
      if expected_num_proc == 1:
         assert mock_single_map_nested.called
         assert not mock_multiprocessing_pool.called
      else:
         assert not mock_single_map_nested.called
         assert mock_multiprocessing_pool.called
         assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc


class     __UpperCAmelCase							(       __A							):


     """simple docstring"""


     @require_tf
     def 	snake_case_  (	self    ):
        import tensorflow as tf
        from tensorflow.keras import layers

        __a					=  layers.Dense(2    )

        def gen_random_output():
           __a					=  tf.random.uniform((1, 3)    )
           return model(lowerCamelCase__    ).numpy()

        with temp_seed(42  ,       set_tensorflow=lowerCamelCase__    ):
           __a					=  gen_random_output()
        with temp_seed(42  ,       set_tensorflow=lowerCamelCase__    ):
           __a					=  gen_random_output()
        __a					=  gen_random_output()

        np.testing.assert_equal(lowerCamelCase__  ,       lowerCamelCase__    )
        self.assertGreater(np.abs(outa - outa    ).sum()  ,       0    )


     @require_torch
     def 	snake_case_  (	self    ):
        import torch

        def gen_random_output():
           __a					=  torch.nn.Linear(3  ,       2    )
           __a					=  torch.rand(1  ,       3    )
           return model(lowerCamelCase__    ).detach().numpy()

        with temp_seed(42  ,       set_pytorch=lowerCamelCase__    ):
           __a					=  gen_random_output()
        with temp_seed(42  ,       set_pytorch=lowerCamelCase__    ):
           __a					=  gen_random_output()
        __a					=  gen_random_output()

        np.testing.assert_equal(lowerCamelCase__  ,       lowerCamelCase__    )
        self.assertGreater(np.abs(outa - outa    ).sum()  ,       0    )


     def 	snake_case_  (	self    ):
        def gen_random_output():
           return np.random.rand(1  ,       3    )

        with temp_seed(42    ):
           __a					=  gen_random_output()
        with temp_seed(42    ):
           __a					=  gen_random_output()
        __a					=  gen_random_output()

        np.testing.assert_equal(lowerCamelCase__  ,       lowerCamelCase__    )
        self.assertGreater(np.abs(outa - outa    ).sum()  ,       0    )







@pytest.mark.parametrize("""input_data"""		,				[{}]   )
def  a   (lowerCAmelCase__   ):
   __a					=  NestedDataStructure(UpperCamelCase__   ).data
   assert output_data == input_data




@pytest.mark.parametrize(
    """data, expected_output"""		,				[
        ({}, []),
        ([], []),
        ("""foo""", ["""foo"""]),
        (["""foo""", """bar"""], ["""foo""", """bar"""]),
        ([["""foo""", """bar"""]], ["""foo""", """bar"""]),
        ([[["""foo"""], ["""bar"""]]], ["""foo""", """bar"""]),
        ([[["""foo"""], """bar"""]], ["""foo""", """bar"""]),
        ({"""a""": 1, """b""": 2}, [1, 2]),
        ({"""a""": [1, 2], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, [4]]}, [1, 2, 3, 4]),
        ({"""a""": {"""1""": 1}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": [2]}, [1, 2]),
    ]		,				)
def  a   (lowerCAmelCase__		,				lowerCAmelCase__   ):
   __a					=  NestedDataStructure(UpperCamelCase__   ).flatten()
   assert output == expected_output




def  a   ():
   __a					=  A(x=1		,				y="""foobar"""   )
   __a					=  {"""x""": 1, """y""": """foobar"""}
   assert asdict(UpperCamelCase__   ) == expected_output

   __a					=  {"""a""": {"""b""": A(x=10		,				y="""foo"""   )}, """c""": [A(x=20		,				y="""bar"""   )]}
   __a					=  {"""a""": {"""b""": {"""x""": 10, """y""": """foo"""}}, """c""": [{"""x""": 20, """y""": """bar"""}]}
   assert asdict(UpperCamelCase__   ) == expected_output

   with pytest.raises(UpperCamelCase__   ):
      asdict([1, A(x=10		,				y="""foo"""   )]   )




def  a   (lowerCAmelCase__   ):
   return text.split()




def  a   (lowerCAmelCase__   ):
   yield (time.time(), content)
   time.sleep(2   )
   yield (time.time(), content)




def  a   ():


   with Pool(2   ) as pool:
      __a					=  list(iflatmap_unordered(UpperCamelCase__		,				_split_text		,				kwargs_iterable=[{"""text""": """hello there"""}] * 10   )   )
      assert out.count("""hello"""   ) == 10
      assert out.count("""there"""   ) == 10
      assert len(UpperCamelCase__   ) == 20

   # check multiprocess from pathos (uses dill for pickling)
   with multiprocess.Pool(2   ) as pool:
      __a					=  list(iflatmap_unordered(UpperCamelCase__		,				_split_text		,				kwargs_iterable=[{"""text""": """hello there"""}] * 10   )   )
      assert out.count("""hello"""   ) == 10
      assert out.count("""there"""   ) == 10
      assert len(UpperCamelCase__   ) == 20

   # check that we get items as fast as possible
   with Pool(2   ) as pool:
      __a					=  []
      for yield_time, content in iflatmap_unordered(
          UpperCamelCase__		,				_aseconds_generator_of_aitems_with_timing		,				kwargs_iterable=[{"""content""": """a"""}, {"""content""": """b"""}]   ):
         assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
         out.append(UpperCamelCase__   )
      assert out.count("""a"""   ) == 2
      assert out.count("""b"""   ) == 2
      assert len(UpperCamelCase__   ) == 4









from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, logging


if is_torch_available():
					import torch


__A								=       logging.get_logger(__name__)







class        __lowerCAmelCase					(			__magic_name__					):

			"""simple docstring"""







			snake_case_  						=			['''pixel_values''']


			def __init__( self     ,		lowerCamelCase__ = True     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = PILImageResampling.BILINEAR     ,		lowerCamelCase__ = True     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = True     ,		lowerCamelCase__ = 1 / 255     ,		lowerCamelCase__ = True     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		**lowerCamelCase__     ,		)       -> None:



							'''simple docstring'''
							super().__init__(**lowerCamelCase__				)
							__lowerCamelCase									=       size if size is not None else {'shortest_edge': 256}
							__lowerCamelCase									=       get_size_dict(lowerCamelCase__     ,		default_to_square=lowerCamelCase__				)
							__lowerCamelCase									=       crop_size if crop_size is not None else {'height': 224, 'width': 224}
							__lowerCamelCase									=       get_size_dict(lowerCamelCase__     ,		param_name='crop_size'				)
							__lowerCamelCase									=       do_resize
							__lowerCamelCase									=       size
							__lowerCamelCase									=       resample
							__lowerCamelCase									=       do_center_crop
							__lowerCamelCase									=       crop_size
							__lowerCamelCase									=       do_rescale
							__lowerCamelCase									=       rescale_factor
							__lowerCamelCase									=       do_normalize
							__lowerCamelCase									=       image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
							__lowerCamelCase									=       image_std if image_std is not None else IMAGENET_STANDARD_STD


			def  lowercase_       ( self     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__ = PILImageResampling.BICUBIC     ,		lowerCamelCase__ = None     ,		**lowerCamelCase__     ,		)       -> np.ndarray:



							'''simple docstring'''
							__lowerCamelCase									=       get_size_dict(lowerCamelCase__     ,		default_to_square=lowerCamelCase__				)
							if "shortest_edge" not in size:
											raise ValueError(f"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""				)
							__lowerCamelCase									=       get_resize_output_image_size(lowerCamelCase__     ,		size=size['shortest_edge']     ,		default_to_square=lowerCamelCase__				)
							return resize(lowerCamelCase__     ,		size=lowerCamelCase__     ,		resample=lowerCamelCase__     ,		data_format=lowerCamelCase__     ,		**lowerCamelCase__				)


			def  lowercase_       ( self     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__ = None     ,		**lowerCamelCase__     ,		)       -> np.ndarray:



							'''simple docstring'''
							__lowerCamelCase									=       get_size_dict(lowerCamelCase__				)
							if "height" not in size or "width" not in size:
											raise ValueError(f"""The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}"""				)
							return center_crop(lowerCamelCase__     ,		size=(size['height'], size['width'])     ,		data_format=lowerCamelCase__     ,		**lowerCamelCase__				)


			def  lowercase_       ( self     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__ = None     ,		**lowerCamelCase__				)       -> np.ndarray:



							'''simple docstring'''
							return rescale(lowerCamelCase__     ,		scale=lowerCamelCase__     ,		data_format=lowerCamelCase__     ,		**lowerCamelCase__				)


			def  lowercase_       ( self     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__ = None     ,		**lowerCamelCase__     ,		)       -> np.ndarray:



							'''simple docstring'''
							return normalize(lowerCamelCase__     ,		mean=lowerCamelCase__     ,		std=lowerCamelCase__     ,		data_format=lowerCamelCase__     ,		**lowerCamelCase__				)


			def  lowercase_       ( self     ,		lowerCamelCase__     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = ChannelDimension.FIRST     ,		**lowerCamelCase__     ,		)       -> str:



							'''simple docstring'''
							__lowerCamelCase									=       do_resize if do_resize is not None else self.do_resize
							__lowerCamelCase									=       size if size is not None else self.size
							__lowerCamelCase									=       get_size_dict(lowerCamelCase__     ,		default_to_square=lowerCamelCase__				)
							__lowerCamelCase									=       resample if resample is not None else self.resample
							__lowerCamelCase									=       do_center_crop if do_center_crop is not None else self.do_center_crop
							__lowerCamelCase									=       crop_size if crop_size is not None else self.crop_size
							__lowerCamelCase									=       get_size_dict(lowerCamelCase__     ,		param_name='crop_size'				)
							__lowerCamelCase									=       do_rescale if do_rescale is not None else self.do_rescale
							__lowerCamelCase									=       rescale_factor if rescale_factor is not None else self.rescale_factor
							__lowerCamelCase									=       do_normalize if do_normalize is not None else self.do_normalize
							__lowerCamelCase									=       image_mean if image_mean is not None else self.image_mean
							__lowerCamelCase									=       image_std if image_std is not None else self.image_std

							__lowerCamelCase									=       make_list_of_images(lowerCamelCase__				)

							if not valid_images(lowerCamelCase__				):
											raise ValueError(
											    'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '
											    'torch.Tensor, tf.Tensor or jax.ndarray.'				)

							if do_resize and size is None:
											raise ValueError('Size must be specified if do_resize is True.'				)

							if do_center_crop and crop_size is None:
											raise ValueError('Crop size must be specified if do_center_crop is True.'				)

							if do_rescale and rescale_factor is None:
											raise ValueError('Rescale factor must be specified if do_rescale is True.'				)

							if do_normalize and (image_mean is None or image_std is None):
											raise ValueError('Image mean and std must be specified if do_normalize is True.'				)

							# All transformations expect numpy arrays.
							__lowerCamelCase									=       [to_numpy_array(lowerCamelCase__				) for image in images]

							if do_resize:
											__lowerCamelCase									=       [self.resize(image=lowerCamelCase__     ,		size=lowerCamelCase__     ,		resample=lowerCamelCase__				) for image in images]

							if do_center_crop:
											__lowerCamelCase									=       [self.center_crop(image=lowerCamelCase__     ,		size=lowerCamelCase__				) for image in images]

							if do_rescale:
											__lowerCamelCase									=       [self.rescale(image=lowerCamelCase__     ,		scale=lowerCamelCase__				) for image in images]

							if do_normalize:
											__lowerCamelCase									=       [self.normalize(image=lowerCamelCase__     ,		mean=lowerCamelCase__     ,		std=lowerCamelCase__				) for image in images]

							__lowerCamelCase									=       [to_channel_dimension_format(lowerCamelCase__     ,		lowerCamelCase__				) for image in images]

							__lowerCamelCase									=       {'pixel_values': images}
							return BatchFeature(data=lowerCamelCase__     ,		tensor_type=lowerCamelCase__				)






			def  lowercase_       ( self     ,		lowerCamelCase__     ,		lowerCamelCase__ = None				)       -> Optional[Any]:



							'''simple docstring'''
							__lowerCamelCase									=       outputs.logits

							# Resize logits and compute semantic segmentation maps
							if target_sizes is not None:
											if len(lowerCamelCase__				) != len(lowerCamelCase__				):
															raise ValueError(
															    'Make sure that you pass in as many target sizes as the batch dimension of the logits'				)

											if is_torch_tensor(lowerCamelCase__				):
															__lowerCamelCase									=       target_sizes.numpy()

											__lowerCamelCase									=       []

											for idx in range(len(lowerCamelCase__				)				):
															__lowerCamelCase									=       torch.nn.functional.interpolate(
															    logits[idx].unsqueeze(dim=0				)     ,		size=target_sizes[idx]     ,		mode='bilinear'     ,		align_corners=lowerCamelCase__				)
															__lowerCamelCase									=       resized_logits[0].argmax(dim=0				)
															semantic_segmentation.append(lowerCamelCase__				)
							else:
											__lowerCamelCase									=       logits.argmax(dim=1				)
											__lowerCamelCase									=       [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]				)]

							return semantic_segmentation


