




"""simple docstring"""

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			(    __snake_case  :					Union[str, Any]	):  # picklable for multiprocessing



				'''simple docstring'''





				return x.sum()



def 					a			(    __snake_case  :					List[str]	):  # picklable for multiprocessing



				'''simple docstring'''





				return i + 1



@dataclass
class 			_snake_case					:

				'''simple docstring'''




				UpperCamelCase__		  =42
				UpperCamelCase__		  =42



class 			_snake_case					(				__snake_case							):

				'''simple docstring'''




				def 					snake_case_			(				self   :	Any  ):
								UpperCAmelCase_       :Union[str, Any]         =   {}
								UpperCAmelCase_       :List[str]         =   []
								UpperCAmelCase_       :Tuple         =   1
								UpperCAmelCase_       :Optional[int]         =   [1, 2]
								UpperCAmelCase_       :Tuple         =   {"a": 1, "b": 2}
								UpperCAmelCase_       :int         =   {"a": [1, 2], "b": [3, 4]}
								UpperCAmelCase_       :Union[str, Any]         =   {"a": {"1": 1}, "b": 2}
								UpperCAmelCase_       :Optional[int]         =   {"a": 1, "b": 2, "c": 3, "d": 4}
								UpperCAmelCase_       :List[Any]         =   {}
								UpperCAmelCase_       :str         =   []
								UpperCAmelCase_       :Optional[int]         =   2
								UpperCAmelCase_       :List[str]         =   [2, 3]
								UpperCAmelCase_       :Optional[int]         =   {"a": 2, "b": 3}
								UpperCAmelCase_       :Optional[int]         =   {"a": [2, 3], "b": [4, 5]}
								UpperCAmelCase_       :int         =   {"a": {"1": 2}, "b": 3}
								UpperCAmelCase_       :Any         =   {"a": 2, "b": 3, "c": 4, "d": 5}
								self.assertEqual(map_nested(snake_case       ,				snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case  )       ,				snake_case  )

								UpperCAmelCase_       :str         =   2
								self.assertEqual(map_nested(snake_case       ,				snake_case       ,				num_proc=snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case       ,				num_proc=snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case       ,				num_proc=snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case       ,				num_proc=snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case       ,				num_proc=snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case       ,				num_proc=snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case       ,				num_proc=snake_case  )       ,				snake_case  )
								self.assertEqual(map_nested(snake_case       ,				snake_case       ,				num_proc=snake_case  )       ,				snake_case  )

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





				def 					snake_case_			(				self   :	Optional[int]  ):
								UpperCAmelCase_       :Optional[Any]         =   {"a": 1, "b": 2}
								UpperCAmelCase_       :List[str]         =   {"a": 3, "b": 4}
								UpperCAmelCase_       :Tuple         =   {"a": 5, "b": 6}
								UpperCAmelCase_       :List[Any]         =   sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))]  )
								self.assertEqual(sorted(zip_dict(snake_case       ,				snake_case       ,				snake_case  )  )       ,				snake_case  )





				def 					snake_case_			(				self   :	str  ):



								class 			_snake_case					:

												'''simple docstring'''




												UpperCamelCase__		  ="""bar"""


								UpperCAmelCase_       :Any         =   Foo()
								self.assertEqual(foo.my_attr       ,				'''bar'''  )
								with temporary_assignment(snake_case       ,				'''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			(    __snake_case  :					Dict,     __snake_case  :					Union[str, Any],     __snake_case  :					Optional[Any]	):



				'''simple docstring'''





				with patch('''datasets.utils.py_utils._single_map_nested'''	) as mock_single_map_nested, patch(
				    '''datasets.parallel.parallel.Pool'''	) as mock_multiprocessing_pool:
								UpperCAmelCase_       :List[Any]         =   {f'{i}': i for i in range(_UpperCAmelCase	)}
								UpperCAmelCase_       :str         =   map_nested(lambda __snake_case	: 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 			_snake_case					(				__snake_case							):

				'''simple docstring'''




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

								UpperCAmelCase_       :Any         =   layers.Dense(2  )

								def gen_random_output():
												UpperCAmelCase_       :List[Any]         =   tf.random.uniform((1, 3)  )
												return model(snake_case  ).numpy()

								with temp_seed(42       ,				set_tensorflow=snake_case  ):
												UpperCAmelCase_       :Union[str, Any]         =   gen_random_output()
								with temp_seed(42       ,				set_tensorflow=snake_case  ):
												UpperCAmelCase_       :List[str]         =   gen_random_output()
								UpperCAmelCase_       :Tuple         =   gen_random_output()

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





				@require_torch
				def 					snake_case_			(				self   :	Tuple  ):
								import torch

								def gen_random_output():
												UpperCAmelCase_       :List[Any]         =   torch.nn.Linear(3       ,				2  )
												UpperCAmelCase_       :int         =   torch.rand(1       ,				3  )
												return model(snake_case  ).detach().numpy()

								with temp_seed(42       ,				set_pytorch=snake_case  ):
												UpperCAmelCase_       :int         =   gen_random_output()
								with temp_seed(42       ,				set_pytorch=snake_case  ):
												UpperCAmelCase_       :Optional[Any]         =   gen_random_output()
								UpperCAmelCase_       :Any         =   gen_random_output()

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





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

								with temp_seed(42  ):
												UpperCAmelCase_       :Optional[int]         =   gen_random_output()
								with temp_seed(42  ):
												UpperCAmelCase_       :Optional[int]         =   gen_random_output()
								UpperCAmelCase_       :Any         =   gen_random_output()

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



@pytest.mark.parametrize('''input_data''',     [{}]	)
def 					a			(    __snake_case  :					str	):



				'''simple docstring'''





				UpperCAmelCase_       :Union[str, Any]         =   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			(    __snake_case  :					Any,     __snake_case  :					Any	):



				'''simple docstring'''





				UpperCAmelCase_       :Optional[Any]         =   NestedDataStructure(_UpperCAmelCase	).flatten()
				assert output == expected_output



def 					a			(    ):



				'''simple docstring'''





				UpperCAmelCase_       :Optional[int]         =   A(x=1,     y='''foobar'''	)
				UpperCAmelCase_       :Union[str, Any]         =   {"x": 1, "y": "foobar"}
				assert asdict(_UpperCAmelCase	) == expected_output

				UpperCAmelCase_       :Optional[Any]         =   {"a": {"b": A(x=10,     y='''foo'''	)}, "c": [A(x=20,     y='''bar'''	)]}
				UpperCAmelCase_       :Optional[Any]         =   {"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			(    __snake_case  :					str	):



				'''simple docstring'''





				return text.split()



def 					a			(    __snake_case  :					Tuple	):



				'''simple docstring'''





				yield (time.time(), content)
				time.sleep(2	)
				yield (time.time(), content)





def 					a			(    ):



				'''simple docstring'''





				with Pool(2	) as pool:
								UpperCAmelCase_       :Dict         =   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:
								UpperCAmelCase_       :Optional[int]         =   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:
								UpperCAmelCase_       :Tuple         =   []
								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 TYPE_CHECKING

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


UpperCamelCase_						:     int	       =							{
    """configuration_transfo_xl""": ["""TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP""", """TransfoXLConfig"""],
    """tokenization_transfo_xl""": ["""TransfoXLCorpus""", """TransfoXLTokenizer"""],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCamelCase_						:     Any	       =							[
							    """TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """AdaptiveEmbedding""",
							    """TransfoXLForSequenceClassification""",
							    """TransfoXLLMHeadModel""",
							    """TransfoXLModel""",
							    """TransfoXLPreTrainedModel""",
							    """load_tf_weights_in_transfo_xl""",
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCamelCase_						:     Tuple	       =							[
							    """TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """TFAdaptiveEmbedding""",
							    """TFTransfoXLForSequenceClassification""",
							    """TFTransfoXLLMHeadModel""",
							    """TFTransfoXLMainLayer""",
							    """TFTransfoXLModel""",
							    """TFTransfoXLPreTrainedModel""",
							]


if TYPE_CHECKING:
							from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
							from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_transfo_xl import (
														    TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
														    AdaptiveEmbedding,
														    TransfoXLForSequenceClassification,
														    TransfoXLLMHeadModel,
														    TransfoXLModel,
														    TransfoXLPreTrainedModel,
														    load_tf_weights_in_transfo_xl,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_transfo_xl import (
														    TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TFAdaptiveEmbedding,
														    TFTransfoXLForSequenceClassification,
														    TFTransfoXLLMHeadModel,
														    TFTransfoXLMainLayer,
														    TFTransfoXLModel,
														    TFTransfoXLPreTrainedModel,
														)

else:
							import sys

							UpperCamelCase_						:     Optional[int]	       =							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




