






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		(					lowercase__  :  str  )		-> Any:  # picklable for multiprocessing
	return x.sum()




def 				A		(					lowercase__  :  Dict  )		-> List[Any]:  # picklable for multiprocessing
	return i + 1

@dataclass
class        lowerCAmelCase_				:







				"""simple docstring"""




				_snake_case				:   int
				_snake_case				:   str

class        lowerCAmelCase_				(      lowercase	):







				"""simple docstring"""




				def 				__a   ( self     :Union[str, Any]     ):
					UpperCamelCase__ :List[str]						=      {}
					UpperCamelCase__ :Dict						=      []
					UpperCamelCase__ :List[Any]						=      1
					UpperCamelCase__ :Optional[Any]						=      [1, 2]
					UpperCamelCase__ :Dict						=      {"""a""": 1, """b""": 2}
					UpperCamelCase__ :Dict						=      {"""a""": [1, 2], """b""": [3, 4]}
					UpperCamelCase__ :List[Any]						=      {"""a""": {"""1""": 1}, """b""": 2}
					UpperCamelCase__ :Optional[Any]						=      {"""a""": 1, """b""": 2, """c""": 3, """d""": 4}
					UpperCamelCase__ :Union[str, Any]						=      {}
					UpperCamelCase__ :Tuple						=      []
					UpperCamelCase__ :Dict						=      2
					UpperCamelCase__ :Optional[Any]						=      [2, 3]
					UpperCamelCase__ :Dict						=      {"""a""": 2, """b""": 3}
					UpperCamelCase__ :int						=      {"""a""": [2, 3], """b""": [4, 5]}
					UpperCamelCase__ :Optional[int]						=      {"""a""": {"""1""": 2}, """b""": 3}
					UpperCamelCase__ :Union[str, Any]						=      {"""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__     )

					UpperCamelCase__ :Dict						=      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__     )

					UpperCamelCase__ :int						=      {"""a""": np.eye(2     ), """b""": np.zeros(3     ), """c""": np.ones(2     )}
					UpperCamelCase__ :int						=      {"""a""": 2, """b""": 0, """c""": 2}
					UpperCamelCase__ :Optional[int]						=      {
					    """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 lowerCamelCase__     : x + 1  ,     lowerCamelCase__  ,     num_proc=lowerCamelCase__     )







				def 				__a   ( self     :Tuple     ):
					UpperCamelCase__ :List[str]						=      {"""a""": 1, """b""": 2}
					UpperCamelCase__ :Tuple						=      {"""a""": 3, """b""": 4}
					UpperCamelCase__ :Optional[Any]						=      {"""a""": 5, """b""": 6}
					UpperCamelCase__ :Union[str, Any]						=      sorted([("""a""", (1, 3, 5)), ("""b""", (2, 4, 6))]     )
					self.assertEqual(sorted(zip_dict(lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__     )     )  ,     lowerCamelCase__     )







				def 				__a   ( self     :Optional[Any]     ):

					class        lowerCAmelCase_				:







									"""simple docstring"""




									_snake_case				:   List[str]        =			"""bar"""

					UpperCamelCase__ :List[str]						=      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		(					lowercase__  :  str ,		lowercase__  :  Any ,		lowercase__  :  Tuple  )		-> Optional[int]:
	with patch("""datasets.utils.py_utils._single_map_nested"""  ) as mock_single_map_nested, patch(
	    """datasets.parallel.parallel.Pool"""  ) as mock_multiprocessing_pool:
		UpperCamelCase__ :Optional[int]						=      {f"""{i}""": i for i in range(lowercase__  )}
		UpperCamelCase__ :Any						=      map_nested(lambda lowercase__  : x + 10 ,		lowercase__ ,		num_proc=lowercase__ ,		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        lowerCAmelCase_				(      lowercase	):







				"""simple docstring"""




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

					UpperCamelCase__ :List[Any]						=      layers.Dense(2     )

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

					with temp_seed(42  ,     set_tensorflow=lowerCamelCase__     ):
						UpperCamelCase__ :List[Any]						=      gen_random_output()
					with temp_seed(42  ,     set_tensorflow=lowerCamelCase__     ):
						UpperCamelCase__ :Tuple						=      gen_random_output()
					UpperCamelCase__ :Any						=      gen_random_output()

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







				@require_torch
				def 				__a   ( self     :Union[str, Any]     ):
					import torch

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

					with temp_seed(42  ,     set_pytorch=lowerCamelCase__     ):
						UpperCamelCase__ :Optional[int]						=      gen_random_output()
					with temp_seed(42  ,     set_pytorch=lowerCamelCase__     ):
						UpperCamelCase__ :Union[str, Any]						=      gen_random_output()
					UpperCamelCase__ :Optional[Any]						=      gen_random_output()

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







				def 				__a   ( self     :List[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__ :str						=      gen_random_output()
					UpperCamelCase__ :Any						=      gen_random_output()

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




@pytest.mark.parametrize("""input_data""" ,		[{}]  )
def 				A		(					lowercase__  :  Dict  )		-> Union[str, Any]:
	UpperCamelCase__ :Tuple						=      NestedDataStructure(lowercase__  ).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		(					lowercase__  :  str ,		lowercase__  :  Union[str, Any]  )		-> List[Any]:
	UpperCamelCase__ :str						=      NestedDataStructure(lowercase__  ).flatten()
	assert output == expected_output




def 				A		(					)		-> Optional[Any]:
	UpperCamelCase__ :Optional[int]						=      A(x=1 ,		y="""foobar"""  )
	UpperCamelCase__ :Dict						=      {"""x""": 1, """y""": """foobar"""}
	assert asdict(lowercase__  ) == expected_output

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

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




def 				A		(					lowercase__  :  str  )		-> List[str]:
	return text.split()




def 				A		(					lowercase__  :  Any  )		-> str:
	yield (time.time(), content)
	time.sleep(2  )
	yield (time.time(), content)




def 				A		(					)		-> str:
	with Pool(2  ) as pool:
		UpperCamelCase__ :List[Any]						=      list(iflatmap_unordered(lowercase__ ,		_split_text ,		kwargs_iterable=[{"""text""": """hello there"""}] * 10  )  )
		assert out.count("""hello"""  ) == 10
		assert out.count("""there"""  ) == 10
		assert len(lowercase__  ) == 20

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

	# check that we get items as fast as possible
	with Pool(2  ) as pool:
		UpperCamelCase__ :Tuple						=      []
		for yield_time, content in iflatmap_unordered(
		    lowercase__ ,		_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(lowercase__  )
		assert out.count("""a"""  ) == 2
		assert out.count("""b"""  ) == 2
		assert len(lowercase__  ) == 4





import importlib
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
import transformers.models.auto
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.bert.configuration_bert import BertConfig
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / '''utils'''))

from test_module.custom_configuration import CustomConfig  # noqa E402


a__		     =     get_tests_dir('''fixtures/dummy-config.json''')


class    SCREAMING_SNAKE_CASE_				(       unittest.TestCase   ):



	"""simple docstring"""


	def 					lowerCamelCase__     (					self		:						str	)	->       Dict:







						"""simple docstring"""







						__UpperCamelCase					:	Optional[Any]             =      0






	def 					lowerCamelCase__     (					self		:						Union[str, Any]	)	->       List[str]:







						"""simple docstring"""







						self.assertIsNotNone(transformers.models.auto.__spec__	)
						self.assertIsNotNone(importlib.util.find_spec("""transformers.models.auto"""	)	)






	def 					lowerCamelCase__     (					self		:						Optional[Any]	)	->       Dict:







						"""simple docstring"""







						__UpperCamelCase					:	int             =      AutoConfig.from_pretrained("""bert-base-uncased"""	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						List[Any]	)	->       Optional[Any]:







						"""simple docstring"""







						__UpperCamelCase					:	Dict             =      AutoConfig.from_pretrained(lowerCAmelCase	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Any	)	->       Optional[int]:







						"""simple docstring"""







						__UpperCamelCase					:	Optional[int]             =      AutoConfig.from_pretrained(lowerCAmelCase	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Tuple	)	->       Any:







						"""simple docstring"""







						__UpperCamelCase					:	List[Any]             =      AutoConfig.for_model("""roberta"""	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						int	)	->       List[Any]:







						"""simple docstring"""







						with tempfile.TemporaryDirectory() as tmp_dir:
											# This model name contains bert and roberta, but roberta ends up being picked.
											__UpperCamelCase					:	Union[str, Any]             =      os.path.join(lowerCAmelCase					,     """fake-roberta"""	)
											os.makedirs(lowerCAmelCase					,     exist_ok=lowerCAmelCase	)
											with open(os.path.join(lowerCAmelCase					,     """config.json"""	)					,     """w"""	) as f:
																f.write(json.dumps({}	)	)
											__UpperCamelCase					:	Dict             =      AutoConfig.from_pretrained(lowerCAmelCase	)
											self.assertEqual(type(lowerCAmelCase	)					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						int	)	->       str:







						"""simple docstring"""







						try:
											AutoConfig.register("""custom"""					,     lowerCAmelCase	)
											# Wrong model type will raise an error
											with self.assertRaises(lowerCAmelCase	):
																AutoConfig.register("""model"""					,     lowerCAmelCase	)
											# Trying to register something existing in the Transformers library will raise an error
											with self.assertRaises(lowerCAmelCase	):
																AutoConfig.register("""bert"""					,     lowerCAmelCase	)

											# Now that the config is registered, it can be used as any other config with the auto-API
											__UpperCamelCase					:	Optional[Any]             =      CustomConfig()
											with tempfile.TemporaryDirectory() as tmp_dir:
																config.save_pretrained(lowerCAmelCase	)
																__UpperCamelCase					:	List[str]             =      AutoConfig.from_pretrained(lowerCAmelCase	)
																self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)

						finally:
											if "custom" in CONFIG_MAPPING._extra_content:
																del CONFIG_MAPPING._extra_content["custom"]






	def 					lowerCamelCase__     (					self		:						int	)	->       Optional[Any]:







						"""simple docstring"""







						with self.assertRaisesRegex(
						    lowerCAmelCase					,     """bert-base is not a local folder and is not a valid model identifier"""	):
											__UpperCamelCase					:	Tuple             =      AutoConfig.from_pretrained("""bert-base"""	)






	def 					lowerCamelCase__     (					self		:						Dict	)	->       Any:







						"""simple docstring"""







						with self.assertRaisesRegex(
						    lowerCAmelCase					,     R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""	):
											__UpperCamelCase					:	Optional[Any]             =      AutoConfig.from_pretrained(lowerCAmelCase					,     revision="""aaaaaa"""	)






	def 					lowerCamelCase__     (					self		:						str	)	->       List[str]:







						"""simple docstring"""







						with self.assertRaisesRegex(
						    lowerCAmelCase					,     """hf-internal-testing/no-config-test-repo does not appear to have a file named config.json."""					,     ):
											__UpperCamelCase					:	Union[str, Any]             =      AutoConfig.from_pretrained("""hf-internal-testing/no-config-test-repo"""	)






	def 					lowerCamelCase__     (					self		:						Optional[int]	)	->       Optional[int]:







						"""simple docstring"""







						with self.assertRaises(lowerCAmelCase	):
											__UpperCamelCase					:	str             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""	)
						# If remote code is disabled, we can't load this config.
						with self.assertRaises(lowerCAmelCase	):
											__UpperCamelCase					:	Tuple             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)

						__UpperCamelCase					:	List[str]             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)
						self.assertEqual(config.__class__.__name__					,     """NewModelConfig"""	)

						# Test config can be reloaded.
						with tempfile.TemporaryDirectory() as tmp_dir:
											config.save_pretrained(lowerCAmelCase	)
											__UpperCamelCase					:	str             =      AutoConfig.from_pretrained(lowerCAmelCase					,     trust_remote_code=lowerCAmelCase	)
						self.assertEqual(reloaded_config.__class__.__name__					,     """NewModelConfig"""	)






	def 					lowerCamelCase__     (					self		:						Optional[int]	)	->       str:







						"""simple docstring"""







						class    SCREAMING_SNAKE_CASE_				(       _UpperCamelCase   ):



							"""simple docstring"""


							__magic_name__   :       int      	=		'new-model'

						try:
											AutoConfig.register("""new-model"""					,     lowerCAmelCase	)
											# If remote code is not set, the default is to use local
											__UpperCamelCase					:	Tuple             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""	)
											self.assertEqual(config.__class__.__name__					,     """NewModelConfigLocal"""	)

											# If remote code is disabled, we load the local one.
											__UpperCamelCase					:	Any             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)
											self.assertEqual(config.__class__.__name__					,     """NewModelConfigLocal"""	)

											# If remote is enabled, we load from the Hub
											__UpperCamelCase					:	List[Any]             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)
											self.assertEqual(config.__class__.__name__					,     """NewModelConfig"""	)

						finally:
											if "new-model" in CONFIG_MAPPING._extra_content:
																del CONFIG_MAPPING._extra_content["new-model"]






