





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 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				Any:  # picklable for multiprocessing







			'''simple docstring'''

			return x.sum()





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				Optional[Any]:  # picklable for multiprocessing







			'''simple docstring'''

			return i + 1

@dataclass
class 	A_       :
	UpperCAmelCase__					=			42
	UpperCAmelCase__					=			42

class 	A_       (a_			):
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				{}
				UpperCAmelCase   		=				[]
				UpperCAmelCase   		=				1
				UpperCAmelCase   		=				[1, 2]
				UpperCAmelCase   		=				{'''a''': 1, '''b''': 2}
				UpperCAmelCase   		=				{'''a''': [1, 2], '''b''': [3, 4]}
				UpperCAmelCase   		=				{'''a''': {'''1''': 1}, '''b''': 2}
				UpperCAmelCase   		=				{'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
				UpperCAmelCase   		=				{}
				UpperCAmelCase   		=				[]
				UpperCAmelCase   		=				2
				UpperCAmelCase   		=				[2, 3]
				UpperCAmelCase   		=				{'''a''': 2, '''b''': 3}
				UpperCAmelCase   		=				{'''a''': [2, 3], '''b''': [4, 5]}
				UpperCAmelCase   		=				{'''a''': {'''1''': 2}, '''b''': 3}
				UpperCAmelCase   		=				{'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}
				self.assertEqual(map_nested(_A							,					_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A		)							,					_A		)

				UpperCAmelCase   		=				2
				self.assertEqual(map_nested(_A							,					_A							,					num_proc=_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A							,					num_proc=_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A							,					num_proc=_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A							,					num_proc=_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A							,					num_proc=_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A							,					num_proc=_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A							,					num_proc=_A		)							,					_A		)
				self.assertEqual(map_nested(_A							,					_A							,					num_proc=_A		)							,					_A		)

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






				'''simple docstring'''







				UpperCAmelCase   		=				{'''a''': 1, '''b''': 2}
				UpperCAmelCase   		=				{'''a''': 3, '''b''': 4}
				UpperCAmelCase   		=				{'''a''': 5, '''b''': 6}
				UpperCAmelCase   		=				sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))]		)
				self.assertEqual(sorted(zip_dict(_A							,					_A							,					_A		)		)							,					_A		)







	def      _lowercase				(    self		):






				'''simple docstring'''







				class 	A_       :
					UpperCAmelCase__					=			'''bar'''
				UpperCAmelCase   		=				Foo()
				self.assertEqual(foo.my_attr							,					'''bar'''		)
				with temporary_assignment(_A							,					'''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 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__ )  ->				Tuple:







			'''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   		=				{F"""{i}""": i for i in range(UpperCamelCase__ )}
						UpperCAmelCase   		=				map_nested(lambda UpperCamelCase__ : 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 	A_       (a_			):
	@require_tf
	def      _lowercase				(    self		):






				'''simple docstring'''







				import tensorflow as tf
				from tensorflow.keras import layers

				UpperCAmelCase   		=				layers.Dense(2		)

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

				with temp_seed(4_2							,					set_tensorflow=_A		):
							UpperCAmelCase   		=				gen_random_output()
				with temp_seed(4_2							,					set_tensorflow=_A		):
							UpperCAmelCase   		=				gen_random_output()
				UpperCAmelCase   		=				gen_random_output()

				np.testing.assert_equal(_A							,					_A		)
				self.assertGreater(np.abs(outa - outa		).sum()							,					0		)
	@require_torch
	def      _lowercase				(    self		):






				'''simple docstring'''







				import torch

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

				with temp_seed(4_2							,					set_pytorch=_A		):
							UpperCAmelCase   		=				gen_random_output()
				with temp_seed(4_2							,					set_pytorch=_A		):
							UpperCAmelCase   		=				gen_random_output()
				UpperCAmelCase   		=				gen_random_output()

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







	def      _lowercase				(    self		):






				'''simple docstring'''







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

				with temp_seed(4_2		):
							UpperCAmelCase   		=				gen_random_output()
				with temp_seed(4_2		):
							UpperCAmelCase   		=				gen_random_output()
				UpperCAmelCase   		=				gen_random_output()

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





@pytest.mark.parametrize('''input_data'''			,			[{}] )
def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				List[str]:







			'''simple docstring'''

			UpperCAmelCase   		=				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 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__ )  ->				List[Any]:







			'''simple docstring'''

			UpperCAmelCase   		=				NestedDataStructure(UpperCamelCase__ ).flatten()
			assert output == expected_output





def 						__SCREAMING_SNAKE_CASE  (				)  ->				str:







			'''simple docstring'''

			UpperCAmelCase   		=				A(x=1			,			y='''foobar''' )
			UpperCAmelCase   		=				{'''x''': 1, '''y''': '''foobar'''}
			assert asdict(UpperCamelCase__ ) == expected_output

			UpperCAmelCase   		=				{'''a''': {'''b''': A(x=10			,			y='''foo''' )}, '''c''': [A(x=20			,			y='''bar''' )]}
			UpperCAmelCase   		=				{'''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 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				Union[str, Any]:







			'''simple docstring'''

			return text.split()





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				Dict:







			'''simple docstring'''

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






def 						__SCREAMING_SNAKE_CASE  (				)  ->				Tuple:







			'''simple docstring'''

			with Pool(2 ) as pool:
						UpperCAmelCase   		=				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   		=				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   		=				[]
						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 collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__A				:    int       =   logging.get_logger(__name__)

__A				:    Tuple       =   {
    "google/bigbird-roberta-base": "https://huggingface.co/google/bigbird-roberta-base/resolve/main/config.json",
    "google/bigbird-roberta-large": "https://huggingface.co/google/bigbird-roberta-large/resolve/main/config.json",
    "google/bigbird-base-trivia-itc": "https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/config.json",
    # See all BigBird models at https://huggingface.co/models?filter=big_bird
}

class 	A_       (a_			):
	UpperCAmelCase__					=			'''big_bird'''
	def __init__(    self							,					_A=5_0_3_5_8							,					_A=7_6_8							,					_A=1_2							,					_A=1_2							,					_A=3_0_7_2							,					_A="gelu_new"							,					_A=0.1							,					_A=0.1							,					_A=4_0_9_6							,					_A=2							,					_A=0.02							,					_A=1E-12							,					_A=True							,					_A=0							,					_A=1							,					_A=2							,					_A=6_6							,					_A="block_sparse"							,					_A=True							,					_A=False							,					_A=6_4							,					_A=3							,					_A=None							,					**_A							,					):






				'''simple docstring'''







				super().__init__(
				    pad_token_id=_A							,					bos_token_id=_A							,					eos_token_id=_A							,					sep_token_id=_A							,					**_A							,					)

				UpperCAmelCase   		=				vocab_size
				UpperCAmelCase   		=				max_position_embeddings
				UpperCAmelCase   		=				hidden_size
				UpperCAmelCase   		=				num_hidden_layers
				UpperCAmelCase   		=				num_attention_heads
				UpperCAmelCase   		=				intermediate_size
				UpperCAmelCase   		=				hidden_act
				UpperCAmelCase   		=				hidden_dropout_prob
				UpperCAmelCase   		=				attention_probs_dropout_prob
				UpperCAmelCase   		=				initializer_range
				UpperCAmelCase   		=				type_vocab_size
				UpperCAmelCase   		=				layer_norm_eps
				UpperCAmelCase   		=				use_cache

				UpperCAmelCase   		=				rescale_embeddings
				UpperCAmelCase   		=				attention_type
				UpperCAmelCase   		=				use_bias
				UpperCAmelCase   		=				block_size
				UpperCAmelCase   		=				num_random_blocks
				UpperCAmelCase   		=				classifier_dropout

class 	A_       (a_			):
	@property
	def      _lowercase				(    self		):






				'''simple docstring'''







				if self.task == "multiple-choice":
							UpperCAmelCase   		=				{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
				else:
							UpperCAmelCase   		=				{0: '''batch''', 1: '''sequence'''}
				return OrderedDict(
				    [
				        ('''input_ids''', dynamic_axis),
				        ('''attention_mask''', dynamic_axis),
				    ]		)



