



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     (					_UpperCAmelCase    :       List[Any]   )    ->			Dict:  # picklable for multiprocessing







			'''simple docstring'''
			return x.sum()





def 							A     (					_UpperCAmelCase    :       Any   )    ->			Union[str, Any]:  # picklable for multiprocessing







			'''simple docstring'''
			return i + 1







@dataclass
class       __lowerCAmelCase   :
			UpperCamelCase        =					42
			UpperCamelCase        =					42







class       __lowerCAmelCase   (      A						):



			def    _lowerCamelCase	(						self       :  Tuple)       ->   Union[str, Any]:



						"""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    _lowerCamelCase	(						self       :  Optional[int])       ->   int:



						"""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    _lowerCamelCase	(						self       :  Optional[int])       ->   Optional[int]:



						"""simple docstring"""






						class       __lowerCAmelCase   :
									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 							A     (					_UpperCAmelCase    :       Optional[Any]							,       _UpperCAmelCase    :       Union[str, Any]							,       _UpperCAmelCase    :       Dict   )    ->			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			=       {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       __lowerCAmelCase   (      A						):



			@require_tf
			def    _lowerCamelCase	(						self       :  List[str])       ->   Union[str, Any]:



						"""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(42			,		set_tensorflow=A):
									_UpperCAmelCase			=       gen_random_output()
						with temp_seed(42			,		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    _lowerCamelCase	(						self       :  int)       ->   str:



						"""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(42			,		set_pytorch=A):
									_UpperCAmelCase			=       gen_random_output()
						with temp_seed(42			,		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    _lowerCamelCase	(						self       :  Dict)       ->   Optional[Any]:



						"""simple docstring"""






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

						with temp_seed(42):
									_UpperCAmelCase			=       gen_random_output()
						with temp_seed(42):
									_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 							A     (					_UpperCAmelCase    :       Optional[Any]   )    ->			Optional[int]:







			'''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 							A     (					_UpperCAmelCase    :       Dict							,       _UpperCAmelCase    :       List[str]   )    ->			List[Any]:







			'''simple docstring'''
			_UpperCAmelCase			=       NestedDataStructure(_UpperCAmelCase   ).flatten()
			assert output == expected_output





def 							A     (					)    ->			Optional[int]:







			'''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 							A     (					_UpperCAmelCase    :       str   )    ->			int:







			'''simple docstring'''
			return text.split()





def 							A     (					_UpperCAmelCase    :       List[Any]   )    ->			int:







			'''simple docstring'''
			yield (time.time(), content)
			time.sleep(2   )
			yield (time.time(), content)





def 							A     (					)    ->			Optional[int]:







			'''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


UpperCAmelCase__       =	logging.get_logger(__name__)

UpperCAmelCase__       =	{
    "facebook/xmod-base": "https://huggingface.co/facebook/xmod-base/resolve/main/config.json",
    "facebook/xmod-large-prenorm": "https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json",
    "facebook/xmod-base-13-125k": "https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json",
    "facebook/xmod-base-30-125k": "https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json",
    "facebook/xmod-base-30-195k": "https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json",
    "facebook/xmod-base-60-125k": "https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json",
    "facebook/xmod-base-60-265k": "https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json",
    "facebook/xmod-base-75-125k": "https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json",
    "facebook/xmod-base-75-269k": "https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json",
}







class       __lowerCAmelCase   (      A						):
			UpperCamelCase        =					'''xmod'''



			def __init__(						self       :  Optional[int]			,		A       :  int=3_05_22			,		A       :  Tuple=7_68			,		A       :  Optional[Any]=12			,		A       :  Tuple=12			,		A       :  str=30_72			,		A       :  List[str]="gelu"			,		A       :  Any=0.1			,		A       :  int=0.1			,		A       :  Dict=5_12			,		A       :  Optional[Any]=2			,		A       :  Optional[Any]=0.0_2			,		A       :  List[Any]=1E-12			,		A       :  int=1			,		A       :  Tuple=0			,		A       :  Optional[Any]=2			,		A       :  int="absolute"			,		A       :  Union[str, Any]=True			,		A       :  List[Any]=None			,		A       :  Optional[Any]=False			,		A       :  List[str]=2			,		A       :  int=False			,		A       :  str=True			,		A       :  Optional[Any]=True			,		A       :  Tuple=("en_XX",)			,		A       :  Optional[int]=None			,		**A       :  List[str]			,		)       ->   Dict:



						"""simple docstring"""






						super().__init__(pad_token_id=A			,		bos_token_id=A			,		eos_token_id=A			,		**A)

						_UpperCAmelCase			=       vocab_size
						_UpperCAmelCase			=       hidden_size
						_UpperCAmelCase			=       num_hidden_layers
						_UpperCAmelCase			=       num_attention_heads
						_UpperCAmelCase			=       hidden_act
						_UpperCAmelCase			=       intermediate_size
						_UpperCAmelCase			=       hidden_dropout_prob
						_UpperCAmelCase			=       attention_probs_dropout_prob
						_UpperCAmelCase			=       max_position_embeddings
						_UpperCAmelCase			=       type_vocab_size
						_UpperCAmelCase			=       initializer_range
						_UpperCAmelCase			=       layer_norm_eps
						_UpperCAmelCase			=       position_embedding_type
						_UpperCAmelCase			=       use_cache
						_UpperCAmelCase			=       classifier_dropout
						_UpperCAmelCase			=       pre_norm
						_UpperCAmelCase			=       adapter_reduction_factor
						_UpperCAmelCase			=       adapter_layer_norm
						_UpperCAmelCase			=       adapter_reuse_layer_norm
						_UpperCAmelCase			=       ln_before_adapter
						_UpperCAmelCase			=       list(A)
						_UpperCAmelCase			=       default_language







class       __lowerCAmelCase   (      A						):



			@property
			def    _lowerCamelCase	(						self       :  Optional[Any])       ->   Mapping[str, Mapping[int, str]]:



						"""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),
						    ])
