


'''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_				(	UpperCamelCase_ ):  # picklable for multiprocessing
							return x.sum()







def       a_				(	UpperCamelCase_ ):  # picklable for multiprocessing
							return i + 1






@dataclass
class 					__lowerCAmelCase   :

						"""simple docstring"""






						_UpperCAmelCase		:       int
						_UpperCAmelCase		:       str






class 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						def 				_UpperCAmelCase       (  self							:      Optional[Any]	):
													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 lowerCAmelCase	: x + 1 ,  lowerCAmelCase ,  num_proc=lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      Dict	):
													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 				_UpperCAmelCase       (  self							:      List[Any]	):






													class 					__lowerCAmelCase   :

																			"""simple docstring"""






																			_UpperCAmelCase		:       List[str]   						="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),
        (1_6, 1_6, 1_6),
        (1_6, 1_7, 1_6),
        (1_7, 1_6, 1_6),
    ]						,  )
def       a_				(	UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
							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 UpperCamelCase_ : x + 1_0						,  UpperCamelCase_						,  num_proc=UpperCamelCase_						,  parallel_min_length=1_6 )
														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 				_UpperCAmelCase       (  self							:      Tuple	):
													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 				_UpperCAmelCase       (  self							:      Any	):
													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 				_UpperCAmelCase       (  self							:      Optional[int]	):
													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_				(	UpperCamelCase_ ):
							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_				(	UpperCamelCase_						,  UpperCamelCase_ ):
							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=1_0						,  y="foo" )}, "c": [A(x=2_0						,  y="bar" )]}
							A_			    =		{"a": {"b": {"x": 1_0, "y": "foo"}}, "c": [{"x": 2_0, "y": "bar"}]}
							assert asdict(UpperCamelCase_ ) == expected_output

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







def       a_				(	UpperCamelCase_ ):
							return text.split()







def       a_				(	UpperCamelCase_ ):
							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"}] * 1_0 ) )
														assert out.count("hello" ) == 1_0
														assert out.count("there" ) == 1_0
														assert len(UpperCamelCase_ ) == 2_0

							# 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"}] * 1_0 ) )
														assert out.count("hello" ) == 1_0
														assert out.count("there" ) == 1_0
														assert len(UpperCamelCase_ ) == 2_0

							# 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









'''simple docstring'''
from ...configuration_utils import PretrainedConfig
from ...utils import logging


__SCREAMING_SNAKE_CASE					:				Tuple									=		logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE					:				Union[str, Any]									=		{
    '''transfo-xl-wt103''': '''https://huggingface.co/transfo-xl-wt103/resolve/main/config.json''',
}






class 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						_UpperCAmelCase		:       Optional[Any]   						="transfo-xl"
						_UpperCAmelCase		:       str   						=["mems"]
						_UpperCAmelCase		:       Optional[int]   						={
						    "n_token": "vocab_size",
						    "hidden_size": "d_model",
						    "num_attention_heads": "n_head",
						    "num_hidden_layers": "n_layer",
						}






						def __init__(  self							:      str ,  lowerCAmelCase							:      Tuple=26_77_35 ,  lowerCAmelCase							:      List[Any]=[2_00_00, 4_00_00, 20_00_00] ,  lowerCAmelCase							:      Dict=10_24 ,  lowerCAmelCase							:      List[str]=10_24 ,  lowerCAmelCase							:      int=16 ,  lowerCAmelCase							:      Optional[Any]=64 ,  lowerCAmelCase							:      str=40_96 ,  lowerCAmelCase							:      Optional[int]=4 ,  lowerCAmelCase							:      Any=False ,  lowerCAmelCase							:      List[str]=18 ,  lowerCAmelCase							:      Any=16_00 ,  lowerCAmelCase							:      List[str]=10_00 ,  lowerCAmelCase							:      int=True ,  lowerCAmelCase							:      List[str]=True ,  lowerCAmelCase							:      List[str]=0 ,  lowerCAmelCase							:      Tuple=-1 ,  lowerCAmelCase							:      Optional[int]=True ,  lowerCAmelCase							:      List[str]=0.1 ,  lowerCAmelCase							:      int=0.0 ,  lowerCAmelCase							:      Tuple=True ,  lowerCAmelCase							:      List[str]="normal" ,  lowerCAmelCase							:      Tuple=0.0_1 ,  lowerCAmelCase							:      Optional[int]=0.0_1 ,  lowerCAmelCase							:      List[Any]=0.0_2 ,  lowerCAmelCase							:      Dict=1e-5 ,  lowerCAmelCase							:      Optional[int]=0 ,  **lowerCAmelCase							:      List[str] ,  ):
													A_			    =		vocab_size
													A_			    =		[]
													self.cutoffs.extend(lowerCAmelCase	)
													if proj_share_all_but_first:
																				A_			    =		[False] + [True] * len(self.cutoffs	)
													else:
																				A_			    =		[False] + [False] * len(self.cutoffs	)
													A_			    =		d_model
													A_			    =		d_embed
													A_			    =		d_head
													A_			    =		d_inner
													A_			    =		div_val
													A_			    =		pre_lnorm
													A_			    =		n_layer
													A_			    =		n_head
													A_			    =		mem_len
													A_			    =		same_length
													A_			    =		attn_type
													A_			    =		clamp_len
													A_			    =		sample_softmax
													A_			    =		adaptive
													A_			    =		dropout
													A_			    =		dropatt
													A_			    =		untie_r
													A_			    =		init
													A_			    =		init_range
													A_			    =		proj_init_std
													A_			    =		init_std
													A_			    =		layer_norm_epsilon
													super().__init__(eos_token_id=lowerCAmelCase ,  **lowerCAmelCase	)






						@property
						def 				_UpperCAmelCase       (  self							:      str	):
													# Message copied from Transformer-XL documentation
													logger.info(F"The model {self.model_type} is one of the few models that has no sequence length limit."	)
													return -1






						@max_position_embeddings.setter
						def 				_UpperCAmelCase       (  self							:      Tuple ,  lowerCAmelCase							:      Dict	):
													# Message copied from Transformer-XL documentation
													raise NotImplementedError(
													    F"The model {self.model_type} is one of the few models that has no sequence length limit."	)





